RFD4-Proteína

RFD4-Proteína é um modelo de flow-matching do Baker Lab e da NVIDIA para o design biomolecular. Ele gera sequências de proteínas e estruturas de todos os átomos simultaneamente, condicionado ao contexto molecular para preservar: um alvo de ligação, um motivo enzimático, um ligante, hotspots ou estrutura secundária. Um transformer em pares e um autoencoder são integrados por dois RFD4ProteinaScheduler, um para a backbone e outro para os latentes locais.

O pipeline consome tensores de recursos preparados. Converter um arquivo CIF anotado nesses tensores, e os tensores gerados de volta em um arquivo CIF, é um pré-processamento específico do domínio que reside no pacote rfd4-proteina-prep.

Exemplo de uso

Instale o pacote de preparação em uma revisão fixa. Ele determina os tensores de recursos que o modelo vê, portanto, uma instalação sem versão fixa altera as estruturas geradas quando o pacote é atualizado.

pip install diffusers git+https://huggingface.co/nvidia/RFD4-Proteina-prep@c9959efdebd75ac7fe9129fbbf0e291a6e0487d5

import torch
from huggingface_hub import hf_hub_download
from diffusers import RFD4ProteinaPipeline
from rfd4_proteina_prep import prepare_cif, write_cif

pipe = RFD4ProteinaPipeline.from_pretrained("nvidia/RFD4-Proteina-Diffusers", dtype=torch.bfloat16).to("cuda")

# An annotated CIF specifies what to hold fixed and what to design: here a PD-L1 target and a 50-120 residue binder.
input_cif = hf_hub_download("nvidia/RFD4-Proteina-Diffusers", "examples/protein_binder/input.cif")
features, atoms = prepare_cif(input_cif, seed=42, device=pipe.device)

# The masks and molecule tensors the pipeline always needs are arguments; the annotation tensors go in one mapping.
required = ["padding_mask", "is_in_context_conditioning", "supports_local_latent", "token_type", "atom37_mask"]
conditioning = {name: value for name, value in features.items() if name not in required}

output = pipe(
    padding_mask=features["padding_mask"],
    is_in_context_conditioning=features["is_in_context_conditioning"],
    supports_local_latent=features["supports_local_latent"],
    token_type=features["token_type"],
    atom37_mask=features["atom37_mask"],
    conditioning_features=conditioning,
    generator=torch.Generator("cuda").manual_seed(42),
)
write_cif(output, features, atoms, "design.cif")

Dicas

O repositório do modelo inclui uma entrada anotada e um design gerado para cada tarefa em examples/: proteína binder, anticorpo, VHH, enzima, scaffolding de motivo, binder de molécula pequena, incondicional, inverse folding e difusão parcial.
Pass partial_t_bb_ca e partial_t_local_latents com tempos de início em [0, 1) para redesenhar uma entrada a partir de um início com ruído, ou inverse_folding=True com ambos os tempos de início em 0.0 para manter o backbone e projetar a sequência. Ambos codificam a entrada, portanto, exigem coords_nm limpos e rejeitam coordenadas que não são finitas, que geralmente são resíduos não resolvidos.
O número de etapas é o comprimento da configuração de pipeline recycle_depths, 60 para a receita publicada (180 passagens do tronco). Altere o agendamento passando uma lista diferente de recycle_depths ao construir o pipeline.

Para gerar vários designs de cada entrada, defina num_designs_per_prompt e passe um gerador distinto para cada design. A saída tem batch * num_designs_per_prompt linhas, e a linha r é o design r % num_designs_per_prompt da linha de entrada r // num_designs_per_prompt.

generators = [torch.Generator("cuda").manual_seed(seed) for seed in range(8)]
output = pipe(
    padding_mask=features["padding_mask"],
    is_in_context_conditioning=features["is_in_context_conditioning"],
    supports_local_latent=features["supports_local_latent"],
    token_type=features["token_type"],
    atom37_mask=features["atom37_mask"],
    conditioning_features=conditioning,
    num_designs_per_prompt=8,
    generator=generators,
)

write_cif lê os recursos da linha que escreve, então repita-os da mesma forma, com repeat_interleave em vez de repeat:

repeated = {k: v.repeat_interleave(8, dim=0) if torch.is_tensor(v) else v for k, v in features.items()}
for row in range(8):
    write_cif(output, repeated, atoms, f"design_{row}.cif", row=row)

Carregue em torch.bfloat16 para ganhar velocidade. O estado do sampler e os embeddings de features permanecem em float32. Um dtype, kernel ou tamanho de lote diferente pode resultar em um design diferente para a mesma seed.
Uma estrutura gerada não é um design validado. Refold e pontue-a para a tarefa pretendida.

LoRA

O modelo publicado é um fine-tune LoRA de rank-32 da base pré-treinada, mesclado nos pesos. A base e o adaptador estão na revisão pretrain-base do repositório do modelo, e o pipeline carrega o adaptador com load_lora_weights por meio de RFD4ProteinaLoraLoaderMixin:

pipe = RFD4ProteinaPipeline.from_pretrained(
    "nvidia/RFD4-Proteina-Diffusers", revision="pretrain-base", dtype=torch.bfloat16
).to("cuda")
pipe.load_lora_weights(
    "nvidia/RFD4-Proteina-Diffusers",
    revision="pretrain-base",
    subfolder="lora",
    weight_name="pytorch_lora_weights.safetensors",
)

RFD4ProteinaPipeline

classe diffusers.RFD4ProteinaPipeline

< source >

( transformer: RFD4ProteinaTransformervae: RFD4ProteinaAutoencoderbb_ca_scheduler: RFD4ProteinaSchedulerlocal_latents_scheduler: RFD4ProteinaSchedulerrecycle_depths: list[int] | None = None )

Parâmetros

transformer (RFD4ProteinaTransformer) — Backbone em pares e denoiser de velocidade de latente local.
vae (RFD4ProteinaAutoencoder) — O autoencoder associado ao checkpoint do denoiser.
bb_ca_scheduler (RFD4ProteinaScheduler) — Integra as coordenadas da espinha dorsal.
local_latents_scheduler (RFD4ProteinaScheduler) — Integra os latentes locais.
recycle_depths (list[int], opcional) — Passagens de reciclagem do tronco para cada etapa de desruidamento. Seu comprimento é o número padrão de etapas de inferência. Padrão é o cronograma liberado de 60 etapas e 180 passagens.

Gere estruturas e sequências moleculares a partir de tensores moleculares preparados.

O pipeline consome os tensores de características produzidos pela preparação de entrada do RFD4 (veja o pacote rfd4-proteina-prep para arquivos CIF anotados) e retorna as coordenadas decodificadas de atom37 e as identidades dos resíduos. As coordenadas estão em nanômetros.

call

< source >

( padding_mask: Tensoris_in_context_conditioning: Tensorsupports_local_latent: Tensortoken_type: typing.Optional[torch.Tensor] = Noneatom37_mask: typing.Optional[torch.Tensor] = Nonecoords_nm: typing.Optional[torch.Tensor] = Noneconditioning_features: dict[str, typing.Any] | None = Nonenum_designs_per_prompt: int = 1generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonepartial_t_bb_ca: float | None = Nonepartial_t_local_latents: float | None = Noneinverse_folding: bool = Falselatents_bb_ca: typing.Optional[torch.Tensor] = Nonelatents_local_latents: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonecallback_on_step_end: collections.abc.Callable[[diffusers.pipelines.pipeline_utils.DiffusionPipeline, int, torch.Tensor, dict], dict | None] | None = Nonecallback_on_step_end_tensor_inputs: list = ['bb_ca', 'local_latents']output_type: str = 'pt'return_dict: bool = True ) → RFD4ProteinaPipelineOutput or tuple

Parâmetros

padding_mask (torch.Tensor) — Booleano (batch, tokens), verdadeiro onde o token é real e não padding.
is_in_context_conditioning (torch.Tensor) — Booleano (batch, tokens), verdadeiro para tokens fornecidos como contexto e mantidos inalterados.
supports_local_latent (torch.Tensor) — Booleano (batch, tokens), verdadeiro para tokens que carregam um latente local. As três máscaras compartilham uma forma e um dispositivo.
token_type (torch.Tensor, opcional) — (batch, tokens) códigos de tipo de token.
atom37_mask (torch.Tensor, opcional) — (batch, tokens, 37), verdadeiro onde um slot de átomo é utilizado.
coords_nm (torch.Tensor, opcional) — Coordenadas limpas (batch, tokens, 37, 3) de atom37 em nanômetros, necessárias para difusão parcial e dobramento inverso.
conditioning_features (dict[str, Any], opcional) — Os tensores restantes de topologia e condicionamento sob as chaves emitidas pelo pré-processador. Deve conter a condição de sequência, condition_sequence_1_residue e molecule_iid. Quais dos demais são lidos é fixado pela configuração de recursos do checkpoint. É um único mapeamento porque o pré-processador emite cerca de 30 tensores opcionais sob nomes fixos, e uma assinatura com tantos argumentos vincularia o pipeline à nomenclatura do pré-processador.
num_designs_per_prompt (int, padrão 1) — Número de designs a serem gerados para cada linha de entrada. As entradas são repetidas ao longo do eixo do lote, de modo que a saída possui batch * num_designs_per_prompt linhas e a linha r é o design r % num_designs_per_prompt da linha de entrada r // num_designs_per_prompt.
generator (torch.Generator ou list[torch.Generator], opcional) — Um gerador distinto para cada linha de saída, para designs reproduzíveis. Os geradores avançam e o fluxo aleatório global do PyTorch permanece intacto. Omita-o para sortear do fluxo global.
partial_t_bb_ca (float, opcional) — Tempo inicial em [0, 1) para o backbone, para começar a partir de uma versão com ruído da entrada em vez de ruído puro. Forneça junto com partial_t_local_latents; requer coords_nm e atom37_mask.
partial_t_local_latents (float, opcional) — Tempo inicial em [0, 1) para os latentes locais; veja partial_t_bb_ca.
inverse_folding (bool, padrão False) — Preserva o backbone de entrada e projeta a sequência. Exige zero para ambos os tempos parciais de início.
latents_bb_ca (torch.Tensor, opcional) — Estado inicial pré-gerado do backbone (batch * num_designs_per_prompt, tokens, 3), para substituir o que seria amostrado por esta chamada. O agendamento permanece inalterado. Forneça-o junto com latents_local_latents.
latents_local_latents (torch.Tensor, opcional) — Latentes locais iniciais pré-gerados (batch * num_designs_per_prompt, tokens, latent_dim).
attention_kwargs (dict, opcional) — Passado para o transformer; scale define a escala do LoRA.
callback_on_step_end (Callable, opcional) — Chamado como callback_on_step_end(pipeline, step, timestep, callback_kwargs) após cada etapa. callback_kwargs contém os tensores nomeados em callback_on_step_end_tensor_inputs; um dicionário retornado pode substituir "bb_ca" e "local_latents".
output_type (str, padrão "pt") — "pt" retorna tensores do torch e "np" retorna arrays do numpy.
callback_on_step_end_tensor_inputs (list[str], padrão ["bb_ca", "local_latents"]) — Tensores passados para o callback: "bb_ca", "local_latents", "last_clean_bb_ca", "last_clean_local_latents", "sequence_logits", "insertion_logits".
return_dict (bool, padrão True) — Retorna uma saída tipada ou uma tupla de seus campos não-None.

Retorna

RFD4ProteinaPipelineOutput ou tuple

Tensores atom37 decodificados (coords_nm em nanômetros, residue_type, emit_mask, atom_mask), o estado final com ruído, a última estimativa limpa e os últimos logits de sequência e inserção.

Gere estruturas e sequências a partir de tensores moleculares preparados.

Exemplos:

>>> import torch
>>> from diffusers import RFD4ProteinaPipeline
>>> from rfd4_proteina_prep import prepare_cif, write_cif

>>> pipe = RFD4ProteinaPipeline.from_pretrained("nvidia/RFD4-Proteina-Diffusers", dtype=torch.bfloat16)
>>> pipe.to("cuda")

>>> features, atoms = prepare_cif("input.cif", seed=42, device=pipe.device)
>>> required = ["padding_mask", "is_in_context_conditioning", "supports_local_latent", "token_type"]
>>> required += ["atom37_mask"]
>>> output = pipe(
...     padding_mask=features["padding_mask"],
...     is_in_context_conditioning=features["is_in_context_conditioning"],
...     supports_local_latent=features["supports_local_latent"],
...     token_type=features["token_type"],
...     atom37_mask=features["atom37_mask"],
...     conditioning_features={name: value for name, value in features.items() if name not in required},
...     generator=torch.Generator("cuda").manual_seed(42),
... )
>>> write_cif(output, features, atoms, "design.cif")

check_inputs

< source >

( padding_maskis_in_context_conditioningsupports_local_latentatom37_maskcoords_nmconditioning_featuresnum_designs_per_promptgeneratorpartial_tinverse_foldinglatentscallback_on_step_end_tensor_inputsoutput_type )

Valide máscaras moleculares, geradores por design, difusão parcial e callbacks.

noise_backbone

< source >

( target_backbonebb_maskconditioninggeneratorsindex )

Adicione ruído ao backbone preservado até o tempo agendado do passo index, para o desdobramento inverso.

prepare_latents

< source >

( bb_mask: Tensorlatent_mask: Tensorconditioning: dictgenerators: list[torch.Generator] | Noneclean_bb: typing.Optional[torch.Tensor] = Noneclean_latents: typing.Optional[torch.Tensor] = Nonestart_times: tuple[float, float] | None = None )

Desenhe o estado inicial: ruído puro, ou a entrada limpa com ruído adicionado até start_times (backbone, latentes). clean_latents pode ser omitido quando o tempo inicial do latente for zero.

Os sorteios seguem a ordem do sampler nativo em cada gerador de linha: para uma geração completa, o ruído do backbone e seus deslocamentos estocásticos vêm antes do ruído latente, enquanto a difusão parcial sorteia ambos os ruídos antes dos deslocamentos.

RFD4ProteinaPipelineOutput

classe diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

< source >

( coords_nm: Tensorresidue_type: Tensoremit_mask: Tensoratom_mask: Tensorbb_ca: Tensorlocal_latents: Tensorclean_bb_ca: Tensorclean_local_latents: Tensorsequence_logits: Tensorinsertion_logits: Tensor )

Parâmetros

coords_nm (torch.Tensor) — Coordenadas decodificadas de atom37 (B, N, 37, 3).
residue_type (torch.Tensor) — Índices do vocabulário de resíduos (B, N); -1 onde nenhum resíduo é emitido.
emit_mask (torch.Tensor) — Máscara booleana (B, N) dos tokens que foram projetados.
atom_mask (torch.Tensor) — Máscara booleana (B, N, 37) dos átomos válidos.
bb_ca (torch.Tensor) — Estado final do backbone (B, N, 3) em nanômetros.
local_latents (torch.Tensor) — Estado final de latentes locais (B, N, D).
clean_bb_ca (torch.Tensor) — Previsão do backbone do estado limpo na última etapa de desruidamento.
clean_local_latents (torch.Tensor) — Previsão do estado limpo em latentes locais na última etapa de desruido.
sequence_logits (torch.Tensor) — Logits residuais finais por token (B, N, 147).
insertion_logits (torch.Tensor) — Previsões finais de inserção (B, N, N) para motivos não indexados.

Resultados de inferência molecular. As coordenadas estão em nanômetros.

Update on GitHub

Diffusers

RFD4-Proteína

Exemplo de uso

Dicas

LoRA

RFD4ProteinaPipeline

classe diffusers.RFD4ProteinaPipeline

__call__

check_inputs

noise_backbone

prepare_latents

RFD4ProteinaPipelineOutput

classe diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

call