RFD4-Proteína

RFD4-Proteína es un modelo de flow-matching del Baker Lab y NVIDIA para el diseño biomolecular. Genera secuencias de proteínas y estructuras de todos los átomos de forma conjunta, condicionadas al contexto molecular para preservar: una diana de unión, un motivo enzimático, un ligando, puntos calientes o estructura secundaria. Un transformador emparejado y un autoencoder se integran mediante dos RFD4ProteinaScheduler, uno para la cadena principal y otro para los latentes locales.

El pipeline consume tensores de características preparados. Convertir un archivo CIF anotado en esos tensores, y los tensores generados de vuelta en un archivo CIF, es un preprocesamiento específico del dominio que se encuentra en el paquete rfd4-proteina-prep.

Ejemplo de uso

Instale el paquete de preparación en una revisión fija. Este paquete determina los tensores de características que ve el modelo, por lo que una instalación sin fijar la versión cambia las estructuras generadas cuando el paquete se actualiza.

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

Consejos

El repositorio del modelo incluye una entrada anotada y un diseño generado para cada tarea bajo examples/: proteína binder, anticuerpo, VHH, enzima, andamiaje de motivo, binder de molécula pequeña, incondicional, plegado inverso y difusión parcial.
Pase partial_t_bb_ca y partial_t_local_latents con tiempos de inicio en [0, 1) para rediseñar una entrada desde un inicio con ruido, o inverse_folding=True con ambos tiempos de inicio en 0.0 para mantener el esqueleto y diseñar la secuencia. Ambos codifican la entrada, por lo que requieren coords_nm limpias y rechazan coordenadas que no son finitas, que suelen ser residuos no resueltos.
El número de pasos es la longitud de la configuración de pipeline recycle_depths, 60 para la receta publicada (180 pasadas de trunk). Cambia el programa pasando una lista diferente de recycle_depths al construir el pipeline.

Para generar varios diseños de cada entrada, establezca num_designs_per_prompt y pase un generador distinto por diseño. La salida tiene batch * num_designs_per_prompt filas, y la fila r es el diseño r % num_designs_per_prompt de la fila 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 lee las características de la fila que escribe, por lo que repítelas de la misma manera, con repeat_interleave en lugar 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)

Cargue en torch.bfloat16 para mayor velocidad. El estado del muestreador y los embeddings de características permanecen en float32. Un dtype, kernel o tamaño de lote diferente puede dar lugar a un diseño diferente para la misma semilla.
Una estructura generada no es un diseño validado. Repliegue y evalúela para la tarea prevista.

LoRA

El modelo publicado es un ajuste fino LoRA de rango 32 de la base preentrenada, fusionado en los pesos. La base y el adaptador están en la revisión pretrain-base del repositorio del modelo, y la línea de carga carga el adaptador con load_lora_weights a través 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

clase diffusers.RFD4ProteinaPipeline

< source >

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

Parámetros

transformer (RFD4ProteinaTransformer) — Denoisor de velocidad de latente local y backbone emparejado.
vae (RFD4ProteinaAutoencoder) — El autoencoder asociado al punto de control del desruidador.
bb_ca_scheduler (RFD4ProteinaScheduler) — Integra las coordenadas de la cadena principal.
local_latents_scheduler (RFD4ProteinaScheduler) — Integra los latentes locales.
recycle_depths (list[int], opcional) — Pasadas de reciclaje del tronco para cada paso de desruido. Su longitud es el número predeterminado de pasos de inferencia. Por defecto, utiliza el programa publicado de 60 pasos y 180 pasadas.

Genere estructuras y secuencias moleculares a partir de tensores moleculares preparados.

El pipeline consume los tensores de características generados por la preparación de entrada de RFD4 (consulte el paquete rfd4-proteina-prep para archivos CIF anotados) y devuelve las coordenadas decodificadas de atom37 y las identidades de los residuos. Las coordenadas están en 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), verdadero donde un token es real y no relleno.
is_in_context_conditioning (torch.Tensor) — Booleano (batch, tokens), verdadero para los tokens que se proporcionan como contexto y se mantienen sin cambios.
supports_local_latent (torch.Tensor) — Booleano (batch, tokens), verdadero para los tokens que llevan un latente local. Las tres máscaras comparten una misma forma y dispositivo.
token_type (torch.Tensor, opcional) — (batch, tokens) códigos de tipo de token.
atom37_mask (torch.Tensor, opcional) — (batch, tokens, 37), verdadero donde se utiliza un hueco de átomo.
coords_nm (torch.Tensor, opcional) — Coordenadas limpias de atom37 en formato (batch, tokens, 37, 3) en nanómetros, necesarias para la difusión parcial y el plegado inverso.
conditioning_features (dict[str, Any], opcional) — Los tensores restantes de topología y condicionamiento bajo las claves que emite el preprocesador. Debe contener la condición de secuencia, condition_sequence_1_residue y molecule_iid. Cuáles de los demás se leen queda fijado por la configuración de características del checkpoint. Es un único mapeo porque el preprocesador emite unos 30 tensores opcionales bajo nombres fijos, y una firma con tantos argumentos ataría el pipeline a la nomenclatura del preprocesador.
num_designs_per_prompt (int, por defecto 1) — Número de diseños a generar para cada fila de entrada. Las entradas se repiten a lo largo del eje del lote, por lo que la salida tiene batch * num_designs_per_prompt filas y la fila r es el diseño r % num_designs_per_prompt de la fila de entrada r // num_designs_per_prompt.
generator (torch.Generator o list[torch.Generator], opcional) — Un generador distinto por cada fila de salida, para diseños reproducibles. Los generadores avanzan y el flujo aleatorio global de PyTorch se deja intacto. Omítelo para extraer del flujo global.
partial_t_bb_ca (float, opcional) — Tiempo de inicio en [0, 1) para el backbone, para comenzar desde una versión con ruido de la entrada en lugar de ruido puro. Proporciónelo junto con partial_t_local_latents; requiere coords_nm y atom37_mask.
partial_t_local_latents (float, opcional) — Tiempo de inicio en [0, 1) para los latentes locales; consulte partial_t_bb_ca.
inverse_folding (bool, por defecto False) — Conserva el backbone de entrada y diseña la secuencia. Requiere cero para ambos tiempos de inicio parciales.
latents_bb_ca (torch.Tensor, opcional) — Estado inicial pregenerado del backbone (batch * num_designs_per_prompt, tokens, 3), para reemplazar el que esta llamada muestrearía de otro modo. El calendario no cambia. Proporciónalo junto con latents_local_latents.
latents_local_latents (torch.Tensor, opcional) — Latentes locales de inicio pregenerados (batch * num_designs_per_prompt, tokens, latent_dim).
attention_kwargs (dict, opcional) — Se pasa al transformer; scale establece la escala de LoRA.
callback_on_step_end (Callable, opcional) — Se llama como callback_on_step_end(pipeline, step, timestep, callback_kwargs) después de cada paso. callback_kwargs contiene los tensores especificados en callback_on_step_end_tensor_inputs; un diccionario devuelto puede sustituir "bb_ca" y "local_latents".
output_type (str, por defecto "pt") — "pt" devuelve tensores de torch y "np" devuelve arreglos de numpy.
callback_on_step_end_tensor_inputs (list[str], por defecto ["bb_ca", "local_latents"]) — Tensores pasados al callback: "bb_ca", "local_latents", "last_clean_bb_ca", "last_clean_local_latents", "sequence_logits", "insertion_logits".
return_dict (bool, por defecto True) — Devuelve una salida tipada o una tupla de sus campos no nulos.

Devuelve

RFD4ProteinaPipelineOutput o tuple

Tensores atom37 decodificados (coords_nm en nanómetros, residue_type, emit_mask, atom_mask), el estado ruidoso final, la última estimación limpia y los últimos logits de secuencia e inserción.

Genere estructuras y secuencias a partir de tensores moleculares preparados.

Ejemplos:

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

Validar máscaras moleculares, generadores por diseño, difusión parcial y callbacks.

noise_backbone

< source >

( target_backbonebb_maskconditioninggeneratorsindex )

Ruido el backbone preservado hasta el tiempo programado del paso index, para el plegado 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 )

Dibuje el estado inicial: ruido puro, o la entrada limpia con ruido añadido hasta start_times (backbone, latentes). clean_latents puede omitirse cuando el tiempo de inicio del latente es cero.

Los sorteos siguen el orden del muestreador nativo en cada generador de filas: para una generación completa, el ruido del backbone y sus desplazamientos estocásticos aparecen antes que el ruido latente, mientras que la difusión parcial sortea ambos ruidos antes que los desplazamientos.

RFD4ProteinaPipelineOutput

clase 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 del vocabulario de residuos (B, N); -1 donde no se emite ningún residuo.
emit_mask (torch.Tensor) — Máscara booleana (B, N) de los tokens que fueron diseñados.
atom_mask (torch.Tensor) — Máscara booleana (B, N, 37) de los átomos válidos.
bb_ca (torch.Tensor) — Estado final del backbone (B, N, 3) en nanómetros.
local_latents (torch.Tensor) — Estado final de latentes locales (B, N, D).
clean_bb_ca (torch.Tensor) — Predicción del backbone del estado limpio en el último paso de desruido.
clean_local_latents (torch.Tensor) — Predicción del estado limpio en el último paso de desruido.
sequence_logits (torch.Tensor) — Logits residuales finales por token (B, N, 147).
insertion_logits (torch.Tensor) — Predicciones finales de inserción (B, N, N) para motivos sin índice.

Resultados de inferencia molecular. Las coordenadas están en nanómetros.

Update on GitHub

Diffusers

RFD4-Proteína

Ejemplo de uso

Consejos

LoRA

RFD4ProteinaPipeline

clase diffusers.RFD4ProteinaPipeline

__call__

check_inputs

noise_backbone

prepare_latents

RFD4ProteinaPipelineOutput

clase diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

call