RFD4-Proteína

RFD4-Proteína is a flow-matching model from the Baker Lab and NVIDIA for biomolecular design. It generates protein sequences and all-atom structures together, conditioned on the molecular context to preserve: a binding target, an enzyme motif, a ligand, hotspots or secondary structure. A paired transformer and autoencoder are integrated by two RFD4ProteinaScheduler, one for the backbone and one for the local latents.

The pipeline consumes prepared feature tensors. Turning an annotated CIF file into those tensors, and the generated tensors back into a CIF file, is domain-specific preprocessing that lives in the rfd4-proteina-prep package.

Usage example

Install the preparation package at a fixed revision. It decides the feature tensors the model sees, so an unpinned install changes generated structures when the package is updated.

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

Tips

The model repository ships an annotated input and a generated design for each task under examples/: protein binder, antibody, VHH, enzyme, motif scaffolding, small-molecule binder, unconditional, inverse folding and partial diffusion.
Pass partial_t_bb_ca and partial_t_local_latents with start times in [0, 1) to redesign an input from a noised start, or inverse_folding=True with both start times at 0.0 to keep the backbone and design the sequence. Both encode the input, so they need clean coords_nm and reject coordinates that are not finite, which are usually unresolved residues.
The number of steps is the length of the recycle_depths pipeline config, 60 for the released recipe (180 trunk passes). Change the schedule by passing a different recycle_depths list when constructing the pipeline.

To generate several designs of each input, set num_designs_per_prompt and pass one distinct generator per design. The output has batch * num_designs_per_prompt rows, and row r is design r % num_designs_per_prompt of input row 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 reads the features of the row it writes, so repeat them the same way, with repeat_interleave rather than 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)

Load in torch.bfloat16 for speed. The sampler state and the feature embeddings stay in float32. A different dtype, kernel or batch size can end in a different design for the same seed.
A generated structure is not a validated design. Refold and score it for the intended task.

LoRA

The released model is a rank-32 LoRA fine-tune of the pretrained base, merged into the weights. The base and the adapter are on the pretrain-base revision of the model repository, and the pipeline loads the adapter with load_lora_weights through 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

class diffusers.RFD4ProteinaPipeline

< source >

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

Parameters

transformer (RFD4ProteinaTransformer) — Paired backbone and local-latent velocity denoiser.
vae (RFD4ProteinaAutoencoder) — The autoencoder paired with the denoiser checkpoint.
bb_ca_scheduler (RFD4ProteinaScheduler) — Integrates the backbone coordinates.
local_latents_scheduler (RFD4ProteinaScheduler) — Integrates the local latents.
recycle_depths (list[int], optional) — Trunk recycling passes for each denoising step. Its length is the default number of inference steps. Defaults to the released schedule of 60 steps and 180 passes.

Generate molecular structures and sequences from prepared molecular tensors.

The pipeline consumes the feature tensors produced by the RFD4 input preparation (see the rfd4-proteina-prep package for annotated CIF files) and returns decoded atom37 coordinates and residue identities. Coordinates are in nanometers.

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

Parameters

padding_mask (torch.Tensor) — Boolean (batch, tokens), true where a token is real rather than padding.
is_in_context_conditioning (torch.Tensor) — Boolean (batch, tokens), true for tokens that are given as context and left unchanged.
supports_local_latent (torch.Tensor) — Boolean (batch, tokens), true for tokens that carry a local latent. All three masks share one shape and device.
token_type (torch.Tensor, optional) — (batch, tokens) token-type codes.
atom37_mask (torch.Tensor, optional) — (batch, tokens, 37), true where an atom slot is used.
coords_nm (torch.Tensor, optional) — Clean (batch, tokens, 37, 3) atom37 coordinates in nanometers, needed by partial diffusion and inverse folding.
conditioning_features (dict[str, Any], optional) — The remaining topology and conditioning tensors under the keys the preprocessor emits. It must hold the sequence condition, condition_sequence_1_residue, and molecule_iid. Which of the rest are read is fixed by the checkpoint’s feature configuration. It is one mapping because the preprocessor emits about 30 optional tensors under fixed names, and a signature with that many arguments would tie the pipeline to the preprocessor’s naming.
num_designs_per_prompt (int, defaults to 1) — Number of designs to generate for each input row. The inputs are repeated along the batch axis, so the output has batch * num_designs_per_prompt rows and row r is design r % num_designs_per_prompt of input row r // num_designs_per_prompt.
generator (torch.Generator or list[torch.Generator], optional) — One distinct generator per output row, for reproducible designs. The generators advance and the global PyTorch random stream is left untouched. Omit it to draw from the global stream.
partial_t_bb_ca (float, optional) — Start time in [0, 1) for the backbone, to start from a noised version of the input instead of pure noise. Give it together with partial_t_local_latents; requires coords_nm and atom37_mask.
partial_t_local_latents (float, optional) — Start time in [0, 1) for the local latents; see partial_t_bb_ca.
inverse_folding (bool, defaults to False) — Preserve the input backbone and design the sequence. Requires zero for both partial start times.
latents_bb_ca (torch.Tensor, optional) — Pre-generated starting backbone state (batch * num_designs_per_prompt, tokens, 3), to replace the one this call would otherwise sample. The schedule is unchanged. Give it together with latents_local_latents.
latents_local_latents (torch.Tensor, optional) — Pre-generated starting local latents (batch * num_designs_per_prompt, tokens, latent_dim).
attention_kwargs (dict, optional) — Passed to the transformer; scale sets the LoRA scale.
callback_on_step_end (Callable, optional) — Called as callback_on_step_end(pipeline, step, timestep, callback_kwargs) after each step. callback_kwargs holds the tensors named in callback_on_step_end_tensor_inputs; a returned dict may replace "bb_ca" and "local_latents".
output_type (str, defaults to "pt") — "pt" returns torch tensors and "np" returns numpy arrays.
callback_on_step_end_tensor_inputs (list[str], defaults to ["bb_ca", "local_latents"]) — Tensors passed to the callback: "bb_ca", "local_latents", "last_clean_bb_ca", "last_clean_local_latents", "sequence_logits", "insertion_logits".
return_dict (bool, defaults to True) — Return a typed output, or a tuple of its non-None fields.

Returns

RFD4ProteinaPipelineOutput or tuple

Decoded atom37 tensors (coords_nm in nanometers, residue_type, emit_mask, atom_mask), the final noisy state, the last clean estimate, and the last sequence and insertion logits.

Generate structures and sequences from prepared molecular tensors.

Examples:

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

Validate molecular masks, per-design generators, partial diffusion and callbacks.

noise_backbone

< source >

( target_backbonebb_maskconditioninggeneratorsindex )

Noise the preserved backbone to the scheduled time of step index, for inverse folding.

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 )

Draw the starting state: pure noise, or the clean input noised to start_times (backbone, latents). clean_latents can be left out when the latent start time is zero.

The draws follow the native sampler’s order on every row generator: for a full generation the backbone noise and its stochastic offsets come before the latent noise, while partial diffusion draws both noises before the offsets.

RFD4ProteinaPipelineOutput

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

Parameters

coords_nm (torch.Tensor) — Decoded atom37 coordinates (B, N, 37, 3).
residue_type (torch.Tensor) — Residue vocabulary indices (B, N); -1 where no residue is emitted.
emit_mask (torch.Tensor) — Boolean mask (B, N) of the tokens that were designed.
atom_mask (torch.Tensor) — Boolean mask (B, N, 37) of the valid atoms.
bb_ca (torch.Tensor) — Final backbone state (B, N, 3) in nanometers.
local_latents (torch.Tensor) — Final local-latent state (B, N, D).
clean_bb_ca (torch.Tensor) — Backbone prediction of the clean state at the last denoising step.
clean_local_latents (torch.Tensor) — Local-latent prediction of the clean state at the last denoising step.
sequence_logits (torch.Tensor) — Final per-token residue logits (B, N, 147).
insertion_logits (torch.Tensor) — Final insertion predictions (B, N, N) for unindexed motifs.

Molecular inference results. Coordinates are in nanometers.

Update on GitHub

Diffusers

RFD4-Proteína

Usage example

Tips

LoRA

RFD4ProteinaPipeline

class diffusers.RFD4ProteinaPipeline

__call__

check_inputs

noise_backbone

prepare_latents

RFD4ProteinaPipelineOutput

class diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

call