[kierandidi]

What does this PR do?

Adds RFD4-Proteína, a flow-matching model from the Baker Lab and NVIDIA that designs protein sequences and all-atom structures together, conditioned on a binding target, enzyme motif, ligand or other molecular context.

• RFD4ProteinaPipeline with RFD4ProteinaTransformer (denoiser), RFD4ProteinaAutoencoder (vae) and two RFD4ProteinaScheduler instances, one for the backbone and one for the local latents
• Full generation, partial diffusion and inverse folding; precomputed starting latents; num_designs_per_prompt for several designs per input
• LoRA loading (RFD4ProteinaLoraLoaderMixin), group offloading, regional compilation and quantization support
• A conversion script for the native checkpoints and the native PEFT adapter
• Docs for the pipeline, both models and the scheduler, and tests for each

The pipeline takes prepared feature tensors. Annotated-CIF preparation and writing live in the separate rfd4-proteina-prep package, so the library has no structure-parsing dependency.

Weights (private release preview): https://huggingface.co/nvidia/RFD4-Proteina-Diffusers

Example output and how to read it

What goes in

An annotated CIF file says, for every residue and atom, what is given (kept fixed) and what the model should generate. prepare_cif from rfd4-proteina-prep turns it into the tensors the pipeline takes. The task is chosen by the annotations, not by a pipeline argument: fixing a target and leaving a chain free designs a binder, fixing a few motif atoms designs a scaffold, fixing the backbone designs a sequence.

What comes out

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).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", "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")  # open design.cif in PyMOL, ChimeraX or Mol*

Output field	Shape	Meaning
coords_nm	(batch, tokens, 37, 3)	Atom coordinates in nanometers. Every residue has 37 atom slots (the "atom37" layout); which slots exist for the generated residue type is in atom_mask.
residue_type	(batch, tokens)	Generated residue class from a 147-class vocabulary (amino acids, nucleotides, ...), or -1 for tokens the model did not generate.
emit_mask	(batch, tokens)	True for the tokens the model generated, false for tokens that were given as context.
atom_mask	(batch, tokens, 37)	True for the atom slots that exist for each token.

A token is a residue for proteins and nucleic acids and a single atom for ligands. write_cif combines these tensors with the input annotations into the output CIF, so the given context and the generated parts end up in one file.

How to read the figures

Each figure shows the C-alpha trace (one point per residue, joined along the chain) of the input on the left and the output on the right, in the same orientation. Grey is held fixed from the input, colour is generated by the model, and red dots are ligand or motif atoms. Where an input has no coordinates for something (a ligand identity, or nothing at all), the left panel says so.

What a reasonable result looks like

The checks below test two things: the output respects what was given, and the generated structure is physically sensible. All six are one sample with seed 42, BF16 and the 60-step recipe on one H200 (4 to 8 s each).

Example	Given	Generated	Checks
Protein binder	115-residue PD-L1 target with its coordinates	A 56-residue binder chain	Target unchanged (C-alpha RMSD 0.07 Å); 16 binder residues touch the target (heavy atoms within 4.5 Å)
Antibody	Antibody and antigen (280 residues) with CDR H3 masked	New residues in CDR H3 (five of positions 99 to 104 changed: W99D, G101H, D102F, G103R, F104S)	The other 275 residues are unchanged (RMSD 0.19 Å)
Enzyme	Seven motif atoms placed in space (zinc, part of orotate, part of an aspartate)	A 180-residue scaffold and the rest of the ligand	Zinc and orotate atoms preserved (0.06 Å); 10 residues within 4 Å of the ligand
Small-molecule binder	The ligand identity (SAM), no coordinates	A 150-residue protein and the ligand pose	17 residues within 4 Å of the ligand
Unconditional	A length of 250 tokens	A 250-residue protein	Radius of gyration 17.1 Å, close to the roughly 18 Å typical of a compact 250-residue protein
Inverse folding	A 59-residue backbone (PDB 2KL8)	A new sequence on that backbone	Backbone unchanged (0.0 Å); 34% of residues identical to the native sequence

For every example, C-alpha spacing within a chain is about 3.8 Å (the inverse-folding input has one unresolved gap), and no two residues that are not neighbours in sequence have C-alpha atoms closer than 3.9 Å.

These checks show that the output is well formed and follows the conditioning. They do not show that a design binds, folds or is catalytically active. As with other generative design models, designs are meant to be filtered afterwards, for example by refolding the generated sequence with a structure predictor and comparing it with the generated structure.

Verification

• Fast tests for the pipeline, transformer, autoencoder and scheduler pass; make style and make fix-copies are clean.
• With FP32 weights and TF32 off, velocities match the original implementation to about 1e-6, and outputs are identical before and after the review refactors on 21 inputs in all three modes.

Earlier review rounds from @sayakpaul and @yiyixuxu happened in a private staging PR and are addressed here.

[sayakpaul]

Thanks! I left some comments most of which are centered around not trying to pass dicts to promote the use of tuples (to keep things explicitly visible to the callers and readers).

[sayakpaul]

/diffusers-bot pytest tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py tests/models/autoencoders/test_models_autoencoder_rfd4_proteina.py tests/models/transformers/test_models_transformer_rfd4_proteina.py

[github-actions]

✅ pytest tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py tests/models/autoencoders/test_models_autoencoder_rfd4_proteina.py tests/models/transformers/test_models_transformer_rfd4_proteina.py passed on GPU — view logs.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[kierandidi]

Thanks, went through all of it and applied the same ideas across the rest of the PR, not just the lines you commented on (tensors instead of dicts, explicit names, only what the release needs). Details in the replies below.

[kierandidi]

@sayakpaul @yiyixuxu everything from the review is addressed now (replies are in the threads, and they're resolved). Do things look fine, or is there anything else you need?

[yiyixuxu]

thanks a lot for the PR!
we are very very impressed on how many changes got applied within such a short time frame.
super excited to welcome first protein model into diffusers! will merge once CI passes