RFD4-Proteína

RFD4-Proteína est un modèle de flow-matching du Baker Lab et de NVIDIA pour la conception biomoléculaire. Il génère conjointement des séquences de protéines et des structures toutes atomes, conditionnées par le contexte moléculaire afin de préserver : une cible de liaison, un motif enzymatique, un ligand, des hotspots ou la structure secondaire. Un transformeur apparié et un auto-encodeur sont intégrés par deux RFD4ProteinaScheduler, l'un pour l'ossature et l'autre pour les latents locaux.

Le pipeline consomme des tenseurs de caractéristiques préparés. La conversion d'un fichier CIF annoté en ces tenseurs, ainsi que la reconstitution d'un fichier CIF à partir des tenseurs générés, constitue un prétraitement spécifique au domaine, géré par le package rfd4-proteina-prep.

Exemple d'utilisation

Installez le package de préparation à une révision fixe. Il détermine les tenseurs de caractéristiques que le modèle reçoit, de sorte qu’une installation non verrouillée modifie les structures générées lorsque le package est mis à jour.

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")

Astuces

Le dépôt du modèle fournit une entrée annotée et un design généré pour chaque tâche sous examples/ : protéine liante, anticorps, VHH, enzyme, motif scaffolding, petit-molécule liante, non conditionné, repliement inverse et diffusion partielle.
Passez partial_t_bb_ca et partial_t_local_latents avec des temps de départ dans [0, 1) pour redessiner une entrée à partir d'un départ bruité, ou inverse_folding=True avec les deux temps de départ à 0.0 pour conserver le squelette et concevoir la séquence. Les deux encodent l'entrée, ils nécessitent donc des coords_nm propres et rejettent les coordonnées non finies, qui correspondent généralement à des résidus non résolus.
Le nombre d'étapes correspond à la longueur de la configuration de pipeline recycle_depths, soit 60 pour la recette publiée (180 passages du tronc). Modifiez le planning en transmettant une liste recycle_depths différente lors de la construction du pipeline.

Pour générer plusieurs conceptions pour chaque entrée, définissez num_designs_per_prompt et passez un générateur distinct par conception. La sortie comporte batch * num_designs_per_prompt lignes, et la ligne r est la conception r % num_designs_per_prompt de la ligne d'entrée 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 lit les caractéristiques de la ligne qu’il écrit, il faut donc les répéter de la même manière, avec repeat_interleave plutôt que 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)

Chargez en torch.bfloat16 pour plus de rapidité. L’état du sampler et les embeddings de caractéristiques restent en float32. Un dtype, un kernel ou une taille de lot différente peut aboutir à un design différent pour la même graine.
Une structure générée n'est pas un design validé. Refold et évaluez-la pour la tâche prévue.

LoRA

Le modèle publié est un fine-tuning LoRA de rang 32 de la base pré-entraînée, fusionné dans les poids. La base et l'adaptateur se trouvent sur la révision pretrain-base du dépôt du modèle, et le pipeline charge l'adaptateur avec load_lora_weights via 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 )

Paramètres

transformer (RFD4ProteinaTransformer) — Backbone apparié et débruiteur de vitesse à latence locale.
vae (RFD4ProteinaAutoencoder) — L'autoencodeur associé au point de contrôle du débruiteur.
bb_ca_scheduler (RFD4ProteinaScheduler) — Intègre les coordonnées du squelette.
local_latents_scheduler (RFD4ProteinaScheduler) — Intègre les latents locaux.
recycle_depths (list[int], facultatif) — Passes de recyclage du tronc pour chaque étape de débruitage. Sa longueur correspond au nombre d'étapes d'inférence par défaut. Valeur par défaut : le calendrier publié de 60 étapes et 180 passes.

Générer des structures et des séquences moléculaires à partir de tenseurs moléculaires préparés.

Le pipeline consomme les tenseurs de caractéristiques produits par la préparation des entrées RFD4 (voir le package rfd4-proteina-prep pour les fichiers CIF annotés) et renvoie les coordonnées atom37 décodées ainsi que les identités des résidus. Les coordonnées sont exprimées en nanomètres.

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

Paramètres

padding_mask (torch.Tensor) — (batch, tokens) booléen, vrai là où le jeton est réel et non un jeton de bourrage.
is_in_context_conditioning (torch.Tensor) — Booléen (batch, tokens), vrai pour les jetons fournis en contexte et laissés inchangés.
supports_local_latent (torch.Tensor) — Booléen (batch, tokens), vrai pour les jetons qui portent un latent local. Les trois masques partagent la même forme et le même appareil.
token_type (torch.Tensor, facultatif) — (batch, tokens) codes de type de jeton.
atom37_mask (torch.Tensor, facultatif) — (batch, tokens, 37), vrai là où un emplacement d'atome est utilisé.
coords_nm (torch.Tensor, facultatif) — Coordonnées atom37 propres de forme (batch, tokens, 37, 3) en nanomètres, nécessaires à la diffusion partielle et au repliement inverse.
conditioning_features (dict[str, Any], facultatif) — Les tenseurs de topologie et de conditionnement restants, sous les clés émises par le préprocesseur. Il doit contenir la condition de séquence, condition_sequence_1_residue et molecule_iid. Le choix des autres tenseurs lus est fixé par la configuration des fonctionnalités du point de contrôle. C’est une seule correspondance car le préprocesseur émet environ 30 tenseurs facultatifs sous des noms fixes, et une signature avec autant d’arguments lierait le pipeline à la nomenclature du préprocesseur.
num_designs_per_prompt (int, par défaut 1) — Nombre de designs à générer pour chaque ligne d'entrée. Les entrées sont répétées le long de l'axe du lot, de sorte que la sortie comporte batch * num_designs_per_prompt lignes et que la ligne r est le design r % num_designs_per_prompt de la ligne d'entrée r // num_designs_per_prompt.
generator (torch.Generator ou list[torch.Generator], facultatif) — Un générateur distinct par ligne de sortie, pour des conceptions reproductibles. Les générateurs avancent et le flux aléatoire global de PyTorch est laissé intact. Omettez-le pour tirer des valeurs du flux global.
partial_t_bb_ca (float, facultatif) — Temps de départ dans [0, 1) pour le backbone, afin de démarrer à partir d'une version bruitée de l'entrée plutôt que d'un bruit pur. À fournir avec partial_t_local_latents ; nécessite coords_nm et atom37_mask.
partial_t_local_latents (float, facultatif) — Temps de départ dans [0, 1) pour les latents locaux ; voir partial_t_bb_ca.
inverse_folding (bool, par défaut False) — Conserve le squelette d’entrée et conçoit la séquence. Exige zéro pour les deux temps de début partiels.
latents_bb_ca (torch.Tensor, facultatif) — État initial pré-généré du backbone (batch * num_designs_per_prompt, tokens, 3), destiné à remplacer celui qui serait autrement échantillonné par cet appel. L’horaire reste inchangé. Fournissez-le conjointement avec latents_local_latents.
latents_local_latents (torch.Tensor, facultatif) — Latents locaux de départ pré-générés (batch * num_designs_per_prompt, tokens, latent_dim).
attention_kwargs (dict, facultatif) — Transmis au transformer ; scale définit l'échelle de LoRA.
callback_on_step_end (Callable, facultatif) — Appelé sous la forme callback_on_step_end(pipeline, step, timestep, callback_kwargs) après chaque étape. callback_kwargs contient les tenseurs nommés dans callback_on_step_end_tensor_inputs ; un dictionnaire renvoyé peut remplacer "bb_ca" et "local_latents".
output_type (str, par défaut "pt") — "pt" renvoie des tenseurs torch et "np" renvoie des tableaux numpy.
callback_on_step_end_tensor_inputs (list[str], par défaut ["bb_ca", "local_latents"]) — Tenseurs transmis au rappel : "bb_ca", "local_latents", "last_clean_bb_ca", "last_clean_local_latents", "sequence_logits", "insertion_logits".
return_dict (bool, par défaut True) — Retourne une sortie typée, ou un tuple de ses champs non nuls.

Retourne

RFD4ProteinaPipelineOutput ou tuple

Tenseurs atom37 décodés (coords_nm en nanomètres, residue_type, emit_mask, atom_mask), l'état final bruité, la dernière estimation propre, ainsi que les derniers logits de séquence et d'insertion.

Générer des structures et des séquences à partir de tenseurs moléculaires préparés.

Exemples :

>>> 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 )

Valider les masques moléculaires, les générateurs par conception, la diffusion partielle et les rappels.

noise_backbone

< source >

( target_backbonebb_maskconditioninggeneratorsindex )

Brouiller l'épine dorsale préservée jusqu'au temps planifié de l'étape index, pour le repliement inverse.

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 )

Tracez l’état initial : du bruit pur, ou l’entrée propre bruitée jusqu’à start_times (backbone, latents). clean_latents peut être omis lorsque l’heure de départ des latents est nulle.

Les tirages suivent l’ordre du sampler natif sur chaque générateur de ligne : pour une génération complète, le bruit du backbone et ses décalages stochastiques précèdent le bruit latent, tandis que la diffusion partielle tire les deux bruits avant les décalages.

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 )

Paramètres

coords_nm (torch.Tensor) — Coordonnées atom37 décodées (B, N, 37, 3).
residue_type (torch.Tensor) — Indices du vocabulaire des résidus (B, N) ; -1 là où aucun résidu n'est émis.
emit_mask (torch.Tensor) — Masque booléen (B, N) des jetons qui ont été conçus.
atom_mask (torch.Tensor) — Masque booléen (B, N, 37) des atomes valides.
bb_ca (torch.Tensor) — État final du backbone (B, N, 3) en nanomètres.
local_latents (torch.Tensor) — État final des latents locaux (B, N, D).
clean_bb_ca (torch.Tensor) — Prédiction du backbone de l'état propre à la dernière étape de débruitage.
clean_local_latents (torch.Tensor) — Prédiction des latents locaux de l'état propre à la dernière étape de débruitage.
sequence_logits (torch.Tensor) — Logits de résidu par jeton finaux (B, N, 147).
insertion_logits (torch.Tensor) — Prédictions finales d'insertion (B, N, N) pour les motifs non indexés.

Résultats d'inférence moléculaire. Les coordonnées sont en nanomètres.

Update on GitHub

Diffusers

RFD4-Proteína

Exemple d'utilisation

Astuces

LoRA

RFD4ProteinaPipeline

classe diffusers.RFD4ProteinaPipeline

__call__

check_inputs

noise_backbone

prepare_latents

RFD4ProteinaPipelineOutput

classe diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

call