nvidia/RFD4-Proteina-prep

Project page · Model · Diffusers weights · Training data · Code

rfd4-proteina-prep

Annotated-CIF input preparation and output export for the Diffusers RFD4ProteinaPipeline.

The pipeline itself only consumes prepared feature tensors. This package holds the domain-specific parts that turn an annotated CIF file (the RFD4 condition annotations written by AtomWorks Studio) into those tensors, and the generated tensors back into an annotated CIF. It builds on AtomWorks and has no dependency on the original RFD4 runtime.

Install and authenticate

Python 3.11 or newer is required. The repository and the model are public, so no authentication is needed for the downloads below; log in only if you hit rate limits, and keep tokens out of notebooks and shell history.

pip install huggingface_hub
hf auth login
hf auth whoami
hf download nvidia/RFD4-Proteina-prep --local-dir rfd4-prep
pip install ./rfd4-prep
pip install git+https://github.com/huggingface/diffusers

RFD4ProteinaPipeline was merged into Diffusers on 2026-10-20 (huggingface/diffusers#15088) and ships in the next release, so install Diffusers from GitHub until then. The model revision below is the one whose parameter names and shapes match that code. Older model revisions, including b21bec64, use different parameter names and do not load. The helper does not convert weights or silently repair compatibility. trust_remote_code=True loads Python from the helper repo; inspect block.py and pin a revision for reproducible use. It does not install dependencies automatically.

Hub custom block: download, prepare, generate, export

Both interfaces use the same installed rfd4_proteina_prep implementation. The custom block is a thin wrapper, not a second vendored copy of the chemistry code, and does not add AtomWorks to Diffusers itself.

import torch
from huggingface_hub import HfApi, hf_hub_download
from diffusers import ModularPipeline, RFD4ProteinaPipeline

helper_repo = "nvidia/RFD4-Proteina-prep"
# Resolve once so code and inputs come from one immutable snapshot.
helper_revision = HfApi().model_info(helper_repo).sha
model_revision = "92fc3ae50e336b2ec59c713c19cfd43e5d3e923b"
helper = ModularPipeline.from_pretrained(
helper_repo, revision=helper_revision, trust_remote_code=True
)
input_cif = hf_hub_download(helper_repo, "examples/inputs/protein_binder.cif", revision=helper_revision)
prepared = helper(cif_path=input_cif, seed=42, device="cuda:0")
features, atoms = prepared.get("features"), prepared.get("atom_array")

pipe = RFD4ProteinaPipeline.from_pretrained(
"nvidia/RFD4-Proteina-Diffusers", revision=model_revision, dtype=torch.bfloat16
).to("cuda:0")
# The tensors the pipeline takes as arguments; the annotations go in one mapping. `coords_nm` is one of
# them: partial diffusion and inverse folding read it by name, and a full run ignores it.
required = ["padding_mask", "is_in_context_conditioning", "supports_local_latent", "token_type",
"atom37_mask", "coords_nm"]
inputs = {name: features[name] for name in required}
conditioning = {name: value for name, value in features.items() if name not in required}
output = pipe(**inputs, conditioning_features=conditioning, generator=torch.Generator("cuda:0").manual_seed(42))
cif_path = helper(
pipeline_output=output, features=features, atom_array=atoms,
output_path="design.cif", output="cif_path",
)
print(cif_path)

Use a new output filename each time. Export refuses to overwrite an existing file. row selects a batch row when generating several designs of the same prepared input. Tensor coordinates are in nanometers; CIF coordinates are in angstroms. Keep the atom mapping together with its features, not a different example's mapping.

The regular Python API remains available without remote-code loading:

from rfd4_proteina_prep import prepare_cif, write_cif
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", "coords_nm"]
inputs = {name: features[name] for name in required}
conditioning = {name: value for name, value in features.items() if name not in required}
output = pipe(**inputs, conditioning_features=conditioning, generator=torch.Generator("cuda:0").manual_seed(42))
write_cif(output, features, atoms, "another-design.cif")

prepare_cif returns a single-example batch; seed controls expansion of flexible segments and preserves the Python, NumPy and CPU PyTorch global random states.

Downloadable non-antibody examples

The catalog includes all seven non-antibody task presets from the model's Studio gallery, plus the three web-editor examples. Antibody and VHH/nanobody inputs are excluded. Existing assets in the separate model repository are not deleted. Inputs are copied without changing their annotations. Source revisions, SHA-256 hashes, task-specific arguments and model/code pins are recorded in [examples/manifest.json](examples/manifest.json).

| Example | Input | Special invocation | | --- | --- | --- | | PD-L1 protein binder | [protein_binder.cif](examples/inputs/protein_binder.cif) | Default conditioned generation | | Catalytic-motif enzyme | [enzyme_design.cif](examples/inputs/enzyme_design.cif) | Default | | 1QJG motif scaffold | [motif_scaffolding.cif](examples/inputs/motif_scaffolding.cif) | Default | | SAM ligand binder | [small_molecule_binder.cif](examples/inputs/small_molecule_binder.cif) | Default | | Unconditional protein | [unconditional.cif](examples/inputs/unconditional.cif) | Default | | 2KL8 inverse folding | [inverse_folding.cif](examples/inputs/inverse_folding.cif) | inverse_folding=True, both partial_t values zero | | 2KL8 partial diffusion | [partial_diffusion.cif](examples/inputs/partial_diffusion.cif) | Start times derived from the checkpoint schedule; 40% remaining | | Studio PD-L1 | [ppi_pdl1.cif](examples/studio/ppi_pdl1.cif) | Default | | Studio PH2 active site | [mcsa_ph2.cif](examples/studio/mcsa_ph2.cif) | Default | | Studio IAI ligand binder | [disco_iai.cif](examples/studio/disco_iai.cif) | Default |

To check every input without downloading weights or using a GPU:

python rfd4-prep/examples/run_hub_examples.py --prepare-only --out prep-check

To generate all ten presets with the pinned...