From 6212a68a5dfe5cf4110a69381dfb845b20ac233c Mon Sep 17 00:00:00 2001 From: Kieran Didi Date: Tue, 20 Oct 2026 10:29:16 +0000 Subject: [PATCH 1/8] Add RFD4-Proteina Adds the RFD4-Proteina pipeline with its transformer, autoencoder and per-modality schedulers, a LoRA loader mixin, a checkpoint conversion script, tests and documentation. --- docs/source/en/_toctree.yml | 10 + docs/source/en/api/loaders/lora.md | 5 + .../api/models/rfd4_proteina_autoencoder.md | 38 + .../api/models/rfd4_proteina_transformer.md | 35 + docs/source/en/api/pipelines/rfd4_proteina.md | 94 ++ .../source/en/api/schedulers/rfd4_proteina.md | 29 + scripts/convert_rfd4_proteina_to_diffusers.py | 601 +++++++++ src/diffusers/__init__.py | 8 + src/diffusers/loaders/__init__.py | 4 +- src/diffusers/loaders/lora_pipeline.py | 189 +++ src/diffusers/models/__init__.py | 4 + src/diffusers/models/autoencoders/__init__.py | 1 + .../autoencoders/autoencoder_rfd4_proteina.py | 690 ++++++++++ src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_rfd4_proteina.py | 1189 +++++++++++++++++ src/diffusers/pipelines/__init__.py | 2 + .../pipelines/rfd4_proteina/__init__.py | 35 + .../rfd4_proteina/pipeline_output.py | 41 + .../rfd4_proteina/pipeline_rfd4_proteina.py | 564 ++++++++ src/diffusers/schedulers/__init__.py | 2 + .../schedulers/scheduling_rfd4_proteina.py | 276 ++++ src/diffusers/utils/dummy_pt_objects.py | 60 + .../test_models_autoencoder_rfd4_proteina.py | 101 ++ .../test_models_transformer_rfd4_proteina.py | 143 ++ tests/pipelines/rfd4_proteina/__init__.py | 0 .../test_pipeline_rfd4_proteina.py | 407 ++++++ .../test_scheduler_rfd4_proteina.py | 157 +++ 27 files changed, 4685 insertions(+), 1 deletion(-) create mode 100644 docs/source/en/api/models/rfd4_proteina_autoencoder.md create mode 100644 docs/source/en/api/models/rfd4_proteina_transformer.md create mode 100644 docs/source/en/api/pipelines/rfd4_proteina.md create mode 100644 docs/source/en/api/schedulers/rfd4_proteina.md create mode 100644 scripts/convert_rfd4_proteina_to_diffusers.py create mode 100644 src/diffusers/models/autoencoders/autoencoder_rfd4_proteina.py create mode 100644 src/diffusers/models/transformers/transformer_rfd4_proteina.py create mode 100644 src/diffusers/pipelines/rfd4_proteina/__init__.py create mode 100644 src/diffusers/pipelines/rfd4_proteina/pipeline_output.py create mode 100644 src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py create mode 100644 src/diffusers/schedulers/scheduling_rfd4_proteina.py create mode 100644 tests/models/autoencoders/test_models_autoencoder_rfd4_proteina.py create mode 100644 tests/models/transformers/test_models_transformer_rfd4_proteina.py create mode 100644 tests/pipelines/rfd4_proteina/__init__.py create mode 100644 tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py create mode 100644 tests/schedulers/test_scheduler_rfd4_proteina.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 266fc60f79b9..38b1de7a83d2 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -413,6 +413,10 @@ title: QwenImage21Transformer2DModel - local: api/models/qwenimage_transformer2d title: QwenImageTransformer2DModel + - local: api/models/rfd4_proteina_autoencoder + title: RFD4ProteinaAutoencoder + - local: api/models/rfd4_proteina_transformer + title: RFD4ProteinaTransformer - local: api/models/sana_transformer2d title: SanaTransformer2DModel - local: api/models/sana_video_transformer3d @@ -710,6 +714,10 @@ - local: api/pipelines/z_image title: Z-Image title: Image + - sections: + - local: api/pipelines/rfd4_proteina + title: RFD4-Proteína + title: Protein - sections: - local: api/pipelines/diffusion_gemma title: DiffusionGemma @@ -840,6 +848,8 @@ title: PNDMScheduler - local: api/schedulers/repaint title: RePaintScheduler + - local: api/schedulers/rfd4_proteina + title: RFD4ProteinaScheduler - local: api/schedulers/score_sde_ve title: ScoreSdeVeScheduler - local: api/schedulers/tcd diff --git a/docs/source/en/api/loaders/lora.md b/docs/source/en/api/loaders/lora.md index 4c8eb20b299a..1a4fded1fa10 100644 --- a/docs/source/en/api/loaders/lora.md +++ b/docs/source/en/api/loaders/lora.md @@ -39,11 +39,16 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi - [`ErnieImageLoraLoaderMixin`] provides similar functions for [Ernie-Image](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ernie_image). - [`LTX2LoraLoaderMixin`] provides similar functions for [Flux2](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ltx2). - [`MiniMaxH3LoraLoaderMixin`] provides similar functions for [MiniMax-H3](https://huggingface.co/docs/diffusers/main/en/api/pipelines/minimax_h3). +- [`RFD4ProteinaLoraLoaderMixin`] provides similar functions for [RFD4-Proteína](https://huggingface.co/docs/diffusers/main/en/api/pipelines/rfd4_proteina). - [`LoraBaseMixin`] provides a base class with several utility methods to fuse, unfuse, unload, LoRAs and more. > [!TIP] > To learn more about how to load LoRA weights, see the [LoRA](../../tutorials/using_peft_for_inference) loading guide. +## RFD4ProteinaLoraLoaderMixin + +[[autodoc]] loaders.lora_pipeline.RFD4ProteinaLoraLoaderMixin + ## LoraBaseMixin [[autodoc]] loaders.lora_base.LoraBaseMixin diff --git a/docs/source/en/api/models/rfd4_proteina_autoencoder.md b/docs/source/en/api/models/rfd4_proteina_autoencoder.md new file mode 100644 index 000000000000..121a121b00f2 --- /dev/null +++ b/docs/source/en/api/models/rfd4_proteina_autoencoder.md @@ -0,0 +1,38 @@ + + +# RFD4ProteinaAutoencoder + +The autoencoder of [RFD4-Proteína](../pipelines/rfd4_proteina). `encode` maps atom37 coordinates to per-token local +latents, and `decode` turns latents and backbone positions back into atom37 structure and residue identities. It is +registered in the pipeline as `vae`. + +```py +from diffusers import RFD4ProteinaAutoencoder + +vae = RFD4ProteinaAutoencoder.from_pretrained("nvidia/RFD4-Proteina-Diffusers", subfolder="vae") +``` + +## RFD4ProteinaAutoencoder + +[[autodoc]] RFD4ProteinaAutoencoder + - encode + - decode + - forward + +## RFD4ProteinaEncoderOutput + +[[autodoc]] models.autoencoders.autoencoder_rfd4_proteina.RFD4ProteinaEncoderOutput + +## RFD4ProteinaDecoderOutput + +[[autodoc]] models.autoencoders.autoencoder_rfd4_proteina.RFD4ProteinaDecoderOutput diff --git a/docs/source/en/api/models/rfd4_proteina_transformer.md b/docs/source/en/api/models/rfd4_proteina_transformer.md new file mode 100644 index 000000000000..d07e5dcd98ee --- /dev/null +++ b/docs/source/en/api/models/rfd4_proteina_transformer.md @@ -0,0 +1,35 @@ + + +# RFD4ProteinaTransformer + +The denoiser of [RFD4-Proteína](../pipelines/rfd4_proteina). It predicts velocities for the two modalities the +sampler integrates — backbone CA positions and per-token local latents — from prepared molecular tensors. Inputs are +named tensors plus a conditioning mapping whose contents the checkpoint configuration selects; structure parsing is +outside this class. + +```py +from diffusers import RFD4ProteinaTransformer + +transformer = RFD4ProteinaTransformer.from_pretrained( + "nvidia/RFD4-Proteina-Diffusers", subfolder="transformer" +) +``` + +## RFD4ProteinaTransformer + +[[autodoc]] RFD4ProteinaTransformer + - forward + +## RFD4ProteinaTransformerOutput + +[[autodoc]] models.transformers.transformer_rfd4_proteina.RFD4ProteinaTransformerOutput diff --git a/docs/source/en/api/pipelines/rfd4_proteina.md b/docs/source/en/api/pipelines/rfd4_proteina.md new file mode 100644 index 000000000000..66a5dd592fce --- /dev/null +++ b/docs/source/en/api/pipelines/rfd4_proteina.md @@ -0,0 +1,94 @@ + + +# 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`](https://huggingface.co/nvidia/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. + +```bash +pip install diffusers git+https://huggingface.co/nvidia/RFD4-Proteina-prep@c9959efdebd75ac7fe9129fbbf0e291a6e0487d5 +``` + +```py +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"] +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").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. +* `num_inference_steps` has to match the length of the `recycle_depths` pipeline config. Leave it unset to use the + released recipe of 60 steps and 180 trunk passes. +* 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`. + + ```py + generators = [torch.Generator("cuda").manual_seed(seed) for seed in range(8)] + output = pipe(**inputs, 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`: + + ```py + 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. + +## RFD4ProteinaPipeline + +[[autodoc]] RFD4ProteinaPipeline + - all + - __call__ + +## RFD4ProteinaPipelineOutput + +[[autodoc]] pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput diff --git a/docs/source/en/api/schedulers/rfd4_proteina.md b/docs/source/en/api/schedulers/rfd4_proteina.md new file mode 100644 index 000000000000..5cc156688ca1 --- /dev/null +++ b/docs/source/en/api/schedulers/rfd4_proteina.md @@ -0,0 +1,29 @@ + + +# RFD4ProteinaScheduler + +`RFD4ProteinaScheduler` is the temporal-score-rescaled flow sampler of RFD4-Proteína for one modality. The pipeline uses +two instances, one for the backbone (`bb_ca`, coordinates in nanometers) and one for the per-token local latents +(`local_latents`), each with its own time schedule, noise schedule and rescaling. See [`RFD4ProteinaPipeline`] for the +sampling loop. + +## RFD4ProteinaScheduler + +[[autodoc]] RFD4ProteinaScheduler + - set_timesteps + - add_noise + - step + +## RFD4ProteinaSchedulerOutput + +[[autodoc]] schedulers.scheduling_rfd4_proteina.RFD4ProteinaSchedulerOutput diff --git a/scripts/convert_rfd4_proteina_to_diffusers.py b/scripts/convert_rfd4_proteina_to_diffusers.py new file mode 100644 index 000000000000..586ea7252a65 --- /dev/null +++ b/scripts/convert_rfd4_proteina_to_diffusers.py @@ -0,0 +1,601 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# Licensed under the Apache License, Version 2.0. +"""Convert a native RFD4-Proteina bundle (`config.json` + `model.safetensors`) to a Diffusers pipeline. + +Example: + +```bash +hf download nvidia/RFD4-Proteina config.json model.safetensors --local-dir rfd4-native +python scripts/convert_rfd4_proteina_to_diffusers.py --source rfd4-native --destination rfd4-proteina-diffusers +``` + +The sampler settings are not stored in the trained weights. They default to the released `thinking_medium` recipe +(60 steps, 180 trunk passes); override them with `--sampling-config` and `--recycle-schedule`. + +The script refuses a checkpoint whose `nn` settings it does not recognise rather than converting without them, so a +model trained with a feature this implementation lacks fails at conversion time instead of sampling from a different +conditioning signal. Two cases to expect: + +- `noise_align_chirality` orients the chirality feature towards the noisy backbone. It needs `chiral_frame_atoms` + and `chiral_slot_owner` in the prepared features; without them the pipeline raises rather than falling back. +- `tri_mult_backend: cuequivariance` selects a fused triangle-multiplication kernel that this implementation does + not have. Conversion prints a note and continues: results are unchanged, generation is slower. +""" + +import argparse +import hashlib +import json +import re +from pathlib import Path +from typing import Any + +import torch +from safetensors.torch import load_file + +import diffusers +from diffusers import RFD4ProteinaAutoencoder, RFD4ProteinaPipeline, RFD4ProteinaScheduler, RFD4ProteinaTransformer + + +# Vocabulary tables, generated from the RFD4 repository's canonical encoding +# (`rfproteina.utils.constants.UNIFIED_RESTYPE_ATOM37_MASK` and +# `rfproteina.partial_autoencoder.decoder.get_allowed_residue_type_mask`). Bit `a` of row `r` +# is set when atom37 slot `a` exists for residue type `r`. +# +# The released vocabulary is the 147-class unified encoding: the protein and nucleic residues plus the gap token, +# glycans and non-canonical amino acids. It includes the terminal oxygen (OXT). +_ATOM37_MASK_BITS_147 = ( + 0, 68719476767, 74633447487, 68719575103, + 68719673407, 68719477791, 68820142143, 68920805439, + 68719476759, 68754104383, 68719481055, 68719489087, + 103079741503, 68720263231, 73017602111, 68719478847, + 68719477023, 68719477407, 94780796991, 75165085759, + 68719476959, 23, 68723671039, 72997670911, + 68734156799, 75145154559, 4095, 68723671031, + 72997670903, 68734156791, 83735089143, 4087, + 2, 0, 69541560735, 86503359135, + 103079216159, 76239138879, 75162715167, 104690370623, + 114874535807, 68719674399, 73049071679, 105294334079, + 68719544383, 120528046143, 92403675199, 103281068159, + 68736451615, 68719739935, 69793219615, 68722623583, + 69323456535, 68753228351, 103348176959, 68720003135, + 68820666431, 74650486847, 74633447743, 69793415487, + 68719476799, 86436744223, 68727867455, 68719509535, + 68721672255, 85911957567, 111942338879, 94523031679, + 68719478047, 68754235455, 113149935903, 70910009919, + 68720035903, 68719493151, 73017602111, 68865229855, + 129500471743, 77511033239, 73025990719, 68720264223, + 68719477055, 68787112991, 132603826175, 74633709631, + 69122130975, 69793742975, 75701989439, 68720328767, + 94915014719, 94847905983, 68736266815, 133140300767, + 73014653151, 94847938623, 74643933247, 68724198943, + 76314325055, 68754366527, 71009599511, 107914211519, + 115691865695, 73017634879, 68828547127, 68720003135, + 68720013535, 94780797503, 111686535551, 196671, + 71186254015, 94814351423, 70078611999, 82012242207, + 75299303487, 68987912415, 68719477407, 69256348191, + 122541443391, 23, 68719481055, 103214222879, + 90197602879, 124084963391, 75299303551, 69524980767, + 94623799135, 68719547103, 69994545439, 108993185855, + 68720459839, 75278332415, 88051368599, 68766818495, + 80719380767, 77316797151, 65407, 4095, + 4095, 2047, 4095, 4095, + 4095, 33536703, 65407, 2015, + 4095, 65407, 65407, +) # fmt: skip +_ATOM37_MASK_SHA256 = { + 147: "31670c6a6b5e89b97fcd0c09a24f074faa5460d9bc43deee6d59a9e1f22ef34e", +} + +# Residue classes each token type may decode to; bit `c` of row `t` means token type `t` may +# emit class `c`. Row `unknown` is the union of the others. +_ALLOWED_RESIDUE_BITS_147 = ( + 4194302, + 130023424, + 4160749568, + 4294967296, + 178405961588244985132285746181186883457908734, + 178384183516762045070624090206311258882310144, + 21778071482940061661655974875615985664000, +) # fmt: skip + +# n_residue_types -> (atom37 bitmask rows, allowed-residue bitmask rows) +_VOCABULARIES = { + 147: (_ATOM37_MASK_BITS_147, _ALLOWED_RESIDUE_BITS_147), +} + +# Released sampling recipe, transcribed from `configs/generation/inference_base.yaml` (`model:`) +# of the RFD4 repository. The recipe is versioned separately from the weights, so a checkpoint +# alone does not pin it and this copy can fall behind. +# +# This copy is the fallback for when the RFD4 repository is not to hand. +DEFAULT_MODEL_ARGS = { + "bb_ca": { + "schedule": {"mode": "log", "p": 2.0, "tail_steps": 10, "tail_step_multiplier": 2}, + "gt": {"mode": "1/t", "p": 1.0, "clamp_val": None}, + "simulation_step_params": { + "sampling_mode": "vf_tsr", + "center_every_step": False, + "sc_scale_noise": 1.0, + "t_lim_ode": 0.95, + "t_lim_ode_below": 0.02, + "tsr_k": 16.0, + "tsr_sigma": 3.0, + }, + }, + "local_latents": { + "schedule": {"mode": "power", "p": 2.0, "tail_steps": 10, "tail_step_multiplier": 2}, + "gt": {"mode": "tan", "p": 1.0, "clamp_val": None}, + "simulation_step_params": { + "sampling_mode": "vf_tsr", + "center_every_step": False, + "sc_scale_noise": 1.0, + "t_lim_ode": 0.95, + "t_lim_ode_below": 0.02, + "tsr_k": 4.0, + "tsr_sigma": 12.0, + }, + }, +} +# `[number_of_steps, recycle_depth]` bands of the released `thinking_medium` recycle strategy, which is the +# front-loaded medium schedule: 60 steps and 180 trunk passes. +DEFAULT_RECYCLE_SCHEDULE = [[5, 6], [10, 5], [5, 4], [10, 3], [20, 2], [10, 1]] + +_RMS_NORM_TYPES = ("rmsnorm_torch", "rmsnorm_torch_fused") + + +def load_released_recipe(path: Path) -> dict[str, Any]: + """Read the per-modality sampler recipe from the RFD4 repository's own configuration. + + Accepts `configs/generation/inference_base.yaml` (or a Hydra config composed from it) and the + JSON equivalent. Reading the file beats transcribing it: the recipe moves independently of the + weights, and a stale transcription changes the sampled trajectory without changing anything + visible. + """ + text = path.read_text() + if path.suffix in {".yaml", ".yml"}: + try: + import yaml + except ImportError as error: + raise ImportError("Reading a YAML recipe needs PyYAML; or pass the JSON equivalent.") from error + document = yaml.safe_load(text) + else: + document = json.loads(text) + # `inference_base.yaml` nests the recipe under `model:`; a dumped recipe may already be bare. + recipe = document.get("model", document) if isinstance(document, dict) else None + if not isinstance(recipe, dict) or not {"bb_ca", "local_latents"} <= recipe.keys(): + raise ValueError(f"{path} has no per-modality sampler recipe under `model:`.") + sections = ("schedule", "gt", "simulation_step_params") + return { + modality: {section: dict(recipe[modality][section]) for section in sections} + for modality in ("bb_ca", "local_latents") + } + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _split_qkv(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Native checkpoints fuse the query, key and value projections; Diffusers keeps them separate.""" + converted = {} + for key, value in state.items(): + if ".to_qkv." in key: + prefix, suffix = key.split(".to_qkv.") + for name, part in zip(("to_q", "to_k", "to_v"), value.chunk(3, dim=0)): + converted[f"{prefix}.{name}.{suffix}"] = part.contiguous() + else: + converted[key] = value + return converted + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise ValueError(message) + + +# Settings copied straight into the Diffusers transformer config. +_COPIED_NN_KEYS = ( + "nlayers token_dim nheads dim_cond pair_repr_dim token_type_emb_dim cath_emb_dim t_emb_dim seq_sep_dim " + "xt_pair_dist_dim xt_pair_dist_min xt_pair_dist_max feats_seq feats_cond_seq feats_pair_repr feats_pair_cond " + "update_pair_repr update_pair_repr_every_n use_tri_mult tri_mult_c tri_mult_bias" +).split() + +# Settings checked against what this implementation supports, then dropped because they select +# the one behaviour it implements. +_VALIDATED_NN_KEYS = ("output_parameterization", "use_qkln", "parallel_mha_transition", "norm_type", "pair_write") + +# Settings turned into Diffusers config entries by hand below. +_DERIVED_NN_KEYS = ( + "recycle", + "sequence_readout", + "noise_align_chirality", + "noise_align_chirality_translate", + "noise_align_internal_coords", + "noise_align_internal_coords_translate", +) + +# Settings deliberately not carried, with the reason each one is safe to drop. +_IGNORED_NN_KEYS = { + "name": "selects the native module; this script only converts the local-latents transformer", + "checkpointing": "training-time activation checkpointing; Diffusers exposes its own", + "strict_feats": "native feature-factory strictness; this implementation is always strict", + "use_torch_compile": "runtime option; use `compile_repeated_blocks()` after loading", + "use_atom37_factory": "selects the atom37 feature factory, the only one implemented here", + "idx_emb_dim": "width of a positional embedding no released feature list uses", + "attention_type": "backend selection; Diffusers routes attention through `set_attention_backend()`", + "tri_mult_backend": ( + "kernel selection for triangle multiplication. This implementation has only the PyTorch path, so a " + "checkpoint asking for the fused cuEquivariance kernel still runs, but slower" + ), +} + + +def _check_known_keys(config: dict[str, Any], known: set[str], what: str) -> None: + """Fail on settings this script does not recognise, rather than dropping them silently.""" + unknown = sorted(set(config) - known) + _require( + not unknown, + f"unrecognised {what} settings {unknown}. They were added to the RFD4 repository after this script was " + f"written; port them (or add them to the ignored list with a reason) instead of converting without them", + ) + + +def build_transformer_config(nn_config: dict[str, Any], latent_dim: int, num_token_types: int) -> dict[str, Any]: + known = {*_COPIED_NN_KEYS, *_VALIDATED_NN_KEYS, *_DERIVED_NN_KEYS, *_IGNORED_NN_KEYS} + _check_known_keys(nn_config, known, "nn") + _require(nn_config.get("output_parameterization") == {"bb_ca": "v", "local_latents": "v"}, "velocity checkpoint") + _require(nn_config.get("use_qkln", True), "released denoiser uses query/key normalization") + _require(not nn_config.get("parallel_mha_transition", False), "parallel_mha_transition is unsupported") + _require(nn_config.get("norm_type") in _RMS_NORM_TYPES, f"unsupported norm_type {nn_config.get('norm_type')!r}") + _require(nn_config.get("pair_write") == "hadamard", f"unsupported pair_write {nn_config.get('pair_write')!r}") + _require( + not nn_config.get("noise_align_internal_coords", False) + and not nn_config.get("noise_align_internal_coords_translate", False) + and not nn_config.get("noise_align_chirality_translate", False), + "noise_align_internal_coords and the translate variants are not implemented: the former needs " + "internal-coordinate group tensors that the preprocessing package does not produce, and no released " + "checkpoint uses the latter", + ) + if nn_config.get("tri_mult_backend") == "cuequivariance": + print( + "note: the checkpoint selects the fused cuEquivariance triangle kernel. This implementation has only " + "the PyTorch path, so results are unchanged but generation is slower on NVIDIA hardware.", + ) + config = {key: nn_config[key] for key in _COPIED_NN_KEYS if key in nn_config} + readout = nn_config.get("sequence_readout") or {} + _require(readout.get("enabled", False), "the denoiser has the sequence readout the autoencoder is conditioned on") + recycle = nn_config.get("recycle") or {} + config.update( + latent_dim=latent_dim, + num_token_types=num_token_types, + recycle=bool(recycle.get("enabled", False)), + num_sequence_classes=readout.get("num_classes", 33), + noise_align_chirality=bool(nn_config.get("noise_align_chirality", False)), + ) + if "recycle_remaining" in config.get("feats_cond_seq", ()): + # The width of the one-hot is trained in. It is `nn.recycle.conditioning_dim` when set, and otherwise the depth + # the gradient flows through (one pass for `last`, two for `last2`, `training_max_recycles` for `all`). + grad_depths = {"last": 1, "last2": 2, "all": recycle.get("training_max_recycles")} + width = recycle.get("conditioning_dim", grad_depths.get(recycle.get("grad", "last"))) + _require( + type(width) is int and width >= 1, + "recycle_remaining conditioning needs a positive integer width from nn.recycle.conditioning_dim or " + "nn.recycle.grad", + ) + config["recycle_conditioning_dim"] = width + return config + + +# Native feature creators are a list indexed by position, each holding its own embeddings; Diffusers keeps the +# modules of the features that have weights in one `feature_modules` dict keyed by feature. +_FEATURE_EMBEDDINGS = { + ("token_type", "embedding"): "token_type", + ("secondary_structure", "ss_embedding"): "secondary_structure", + ("atom_charge", "embedding"): "atom_charge", + ("atomic_element", "embedding"): "atomic_element", + ("cath_code", "embedding_C"): "cath_code.c", + ("cath_code", "embedding_A"): "cath_code.a", + ("cath_code", "embedding_T"): "cath_code.t", + ("atom37_element_charge_embedding", "element_embedding"): "atom37_element_charge_embedding.element", + ("atom37_element_charge_embedding", "charge_embedding"): "atom37_element_charge_embedding.charge", +} +_RENAMED_MODULES = ( + (".q_layer_norm.", ".norm_q."), + (".k_layer_norm.", ".norm_k."), + (".to_out_node.", ".to_out."), + (".scale_output.to_adaln_zero_gamma.", ".scale_output."), +) + + +def _rename_keys(state: dict[str, torch.Tensor], factories: dict[str, list[str]]) -> dict[str, torch.Tensor]: + """Map native parameter names onto the Diffusers layout. + + `factories` gives the feature names of each feature factory by its module path, which is what turns a native + `feat_creators.` into a feature name. The decoder's one-hot `cond_residue_type` embedding is a frozen + identity matrix, which Diffusers computes with `F.one_hot` instead of storing. + """ + converted = {} + for key, value in state.items(): + for old, new in _RENAMED_MODULES: + key = key.replace(old, new) + if ".feat_creators." in key: + prefix, rest = key.split(".feat_creators.") + index, attribute, tail = rest.split(".", 2) + name = factories[prefix][int(index)] + if (name, attribute) == ("cond_residue_type", "embedding"): + _require(torch.equal(value, torch.eye(value.shape[0], dtype=value.dtype)), "identity embedding") + continue + key = f"{prefix}.feature_modules.{_FEATURE_EMBEDDINGS[(name, attribute)]}.{tail}" + converted[key] = value + return converted + + +def build_autoencoder_config(autoencoder_config: dict[str, Any], num_token_types: int) -> dict[str, Any]: + ae = autoencoder_config.get("nn_ae", autoencoder_config) + encoder, decoder = ae["encoder"], ae["decoder"] + n_residue_types = ae["n_residue_types"] + _require( + n_residue_types in _VOCABULARIES, + f"unsupported residue vocabulary of {n_residue_types} classes; the released encodings have " + f"{sorted(_VOCABULARIES)} classes", + ) + atom37_bits, allowed_bits = _VOCABULARIES[n_residue_types] + _require( + num_token_types == len(allowed_bits), + f"checkpoint has {num_token_types} token types but the {n_residue_types}-class vocabulary defines " + f"{len(allowed_bits)}", + ) + for side in (encoder, decoder): + _require(not side.get("use_attn_pair_bias", True) and not side.get("update_pair_repr", False), "no pair bias") + _require(not decoder.get("abs_coors", True), "the decoder predicts CA-relative coordinates") + shared = ("nlayers", "token_dim", "nheads", "dim_cond", "pair_repr_dim", "token_type_emb_dim") + _require(all(encoder[key] == decoder[key] for key in shared), "encoder and decoder must share their dimensions") + atom37_mask = [[bool(bits & (1 << atom)) for atom in range(37)] for bits in atom37_bits] + allowed_residue_types = [[bool(bits & (1 << cls)) for cls in range(n_residue_types)] for bits in allowed_bits] + return { + **{key: encoder[key] for key in shared}, + "latent_dim": ae["latent_z_dim"], + "n_residue_types": n_residue_types, + "num_token_types": num_token_types, + "encoder_feats_seq": encoder["feats_seq"], + "decoder_feats_seq": decoder["feats_seq"], + "atom37_mask": atom37_mask, + "allowed_residue_types": allowed_residue_types, + } + + +def _infer_num_token_types(state: dict[str, torch.Tensor]) -> int: + """Read the token-type vocabulary size off the trained embedding. + + The native configuration does not record it: it follows the repository's `NUM_TOKEN_TYPES`, + which grew from 5 to 7 when glycans and non-canonical amino acids were added. The embedding + shape is the only record of what a given checkpoint was trained with. + """ + widths = { + tuple(value.shape)[0] + for key, value in state.items() + if key.endswith("embedding.weight") and value.ndim == 2 and "init_repr_factory" in key + } + candidates = sorted(width for width in widths if width in {len(bits) for _, bits in _VOCABULARIES.values()}) + _require( + len(candidates) == 1, + f"could not read the token-type vocabulary size from the checkpoint (saw embedding widths {sorted(widths)})", + ) + return candidates[0] + + +def _build_scheduler(arguments: dict[str, Any]) -> RFD4ProteinaScheduler: + """Map one modality's native sampler settings to scheduler arguments, rejecting what the scheduler does not do.""" + schedule, noise, step = arguments["schedule"], arguments["gt"], arguments["simulation_step_params"] + _require(schedule["p"] == 2.0 and schedule["mode"] in {"log", "power"}, "the released time schedules") + _require(noise["p"] == 1.0 and noise["clamp_val"] is None, "the released noise schedules") + _require(step.get("sampling_mode", "vf_tsr") == "vf_tsr" and not step["center_every_step"], "the TSR sampler") + _require(step["sc_scale_noise"] == 1.0, "unit noise scale") + return RFD4ProteinaScheduler( + schedule_mode=schedule["mode"], + noise_mode=noise["mode"], + tsr_k=step["tsr_k"], + tsr_sigma=step["tsr_sigma"], + tail_steps=schedule["tail_steps"], + tail_step_multiplier=schedule["tail_step_multiplier"], + t_lim_ode=step["t_lim_ode"], + t_lim_ode_below=step["t_lim_ode_below"], + ) + + +def convert_rfd4_proteina_bundle( + source: str | Path, + destination: str | Path, + *, + model_args: dict[str, Any] | None = None, + recycle_schedule: list[list[int]] | None = None, +) -> Path: + """Convert a local format-1 RFD4 bundle into a Diffusers `RFD4ProteinaPipeline` directory. + + Args: + source: Directory with `config.json` and `model.safetensors`. A recorded `weights_sha256` is checked + when the bundle carries one; older exports omit the field. + destination: New directory; an existing destination is rejected. + model_args: Per-modality sampler settings (`schedule`, `gt`, `simulation_step_params`). + recycle_schedule: `[number_of_steps, recycle_depth]` bands that fix the number of inference steps. + + Every tensor key and shape is checked (`strict=True`), so missing or unexpected weights fail loudly. + """ + source, destination = Path(source), Path(destination) + if destination.exists(): + raise FileExistsError(f"Destination already exists: {destination}") + config_path, weights_path = source / "config.json", source / "model.safetensors" + payload = json.loads(config_path.read_text()) + _require(payload.get("model_type") == "rfd4_proteina" and payload.get("format_version") == 1, "format-1 bundle") + weights_sha256 = _sha256(weights_path) + recorded = payload.get("weights_sha256") + # Bundles exported before the digest was added carry no `weights_sha256`; only mismatches are fatal. + _require(recorded in (None, weights_sha256), "weights_sha256 does not match model.safetensors") + config = payload["config"] + + state = load_file(weights_path) + unexpected = sorted(key for key in state if not key.startswith(("nn.", "autoencoder."))) + _require(not unexpected, f"unexpected tensors: {unexpected}") + num_token_types = _infer_num_token_types(state) + + autoencoder_kwargs = build_autoencoder_config(config["autoencoder_config"], num_token_types) + _require( + any(key.startswith("autoencoder.decoder.atom_presence_linear.") for key in state), + "the decoder has the atom-presence head of the unified autoencoder", + ) + flow_config = config["product_flowmatcher"] + _require(flow_config["local_latents"]["dim"] == autoencoder_kwargs["latent_dim"], "latent dimensions must match") + # The pipeline hardcodes the released noise: zero-center-of-mass backbone noise offset per molecule by 0.2. + backbone_flow, latent_flow = flow_config["bb_ca"], flow_config["local_latents"] + _require( + backbone_flow["zero_com_noise"] + and backbone_flow["stochastic_centering_mode"] == "x_0" + and backbone_flow["stochastic_centering_scope"] == "per_group" + and backbone_flow["stochastic_centering_group_keys"] == ["molecule_iid"] + and backbone_flow["stochastic_centering_scale"] == 0.2 + and not latent_flow["zero_com_noise"] + and latent_flow["stochastic_centering_mode"] == "none", + "the released stochastic centering", + ) + transformer_kwargs = build_transformer_config(config["nn"], autoencoder_kwargs["latent_dim"], num_token_types) + _require( + transformer_kwargs["num_sequence_classes"] == autoencoder_kwargs["n_residue_types"], + "the transformer's sequence readout and the autoencoder's residue vocabulary must have the same size", + ) + transformer = RFD4ProteinaTransformer(**transformer_kwargs) + vae = RFD4ProteinaAutoencoder(**autoencoder_kwargs) + transformer_config = transformer.config + transformer_state = _rename_keys( + { + key.removeprefix("nn."): value + for key, value in _split_qkv({k: v for k, v in state.items() if k.startswith("nn.")}).items() + }, + { + "init_repr_factory": transformer_config.feats_seq, + "cond_factory": transformer_config.feats_cond_seq, + "pair_repr_builder.init_repr_factory": transformer_config.feats_pair_repr, + "pair_repr_builder.cond_factory": transformer_config.feats_pair_cond, + }, + ) + autoencoder_state = _rename_keys( + { + key.removeprefix("autoencoder.").removeprefix("model."): value + for key, value in _split_qkv({k: v for k, v in state.items() if k.startswith("autoencoder.")}).items() + }, + { + "encoder.init_repr_factory": vae.config.encoder_feats_seq, + "decoder.init_repr_factory": vae.config.decoder_feats_seq, + "encoder.cond_factory": [], + "decoder.cond_factory": [], + }, + ) + transformer.load_state_dict(transformer_state, strict=True) + vae.load_state_dict(autoencoder_state, strict=True) + + model_args = model_args or DEFAULT_MODEL_ARGS + depths = [depth for steps, depth in (recycle_schedule or DEFAULT_RECYCLE_SCHEDULE) for _ in range(steps)] + schedulers = {name: _build_scheduler(model_args[name]) for name in ("bb_ca", "local_latents")} + RFD4ProteinaPipeline( + transformer=transformer, + vae=vae, + bb_ca_scheduler=schedulers["bb_ca"], + local_latents_scheduler=schedulers["local_latents"], + recycle_depths=depths, + ).save_pretrained(destination, safe_serialization=True) + release_path = source / "release-checkpoints.json" + provenance = { + "source_model": "nvidia/RFD4-Proteina", + "source_format_version": payload["format_version"], + "source_bundle_revision": json.loads(release_path.read_text()).get("bundle_revision") + if release_path.is_file() + else None, + "source_weights_sha256": weights_sha256, + "source_config_sha256": _sha256(config_path), + "weights": "EMA, float32; the fused query/key/value projections are split and a few modules renamed, every other tensor is unchanged", + "sampler_recipe": "thinking_medium" if recycle_schedule is None else "custom", + "num_inference_steps": len(depths), + "residue_vocabulary_size": autoencoder_kwargs["n_residue_types"], + "token_type_vocabulary_size": num_token_types, + "atom37_mask_sha256": _ATOM37_MASK_SHA256[autoencoder_kwargs["n_residue_types"]], + "path_length_search_cutoff": (config.get("path_length_search_cutoffs") or [64])[0], + "converted_with": f"diffusers {diffusers.__version__}, scripts/convert_rfd4_proteina_to_diffusers.py", + } + (destination / "provenance.json").write_text(json.dumps(provenance, indent=2) + "\n") + return destination + + +def convert_rfd4_proteina_adapter(source: str | Path, destination: str | Path) -> Path: + """Convert a native PEFT adapter (`adapter_config.json` + `adapter_model.safetensors`) to Diffusers LoRA weights. + + The native adapter targets the fused `to_qkv` projection, which Diffusers keeps as separate `to_q`, `to_k` and + `to_v`. The fused LoRA is split exactly: the down projection is shared by the three, and the up projection is cut + into the three row blocks. `to_out_node` is named `to_out`. The result loads with `pipe.load_lora_weights`. + """ + source = Path(source) + config = json.loads((source / "adapter_config.json").read_text()) + _require(set(config["target_modules"]) == {"to_qkv", "to_out_node"}, "an adapter on to_qkv and to_out_node") + _require(not config.get("use_dora") and config.get("bias", "none") == "none", "a plain LoRA adapter") + converted = {} + for key, value in load_file(source / "adapter_model.safetensors").items(): + prefix, module, kind = re.fullmatch(r"(.*)\.(to_qkv|to_out_node)\.(lora_[AB])\.weight", key).groups() + if module == "to_out_node": + converted[f"{prefix}.to_out.{kind}.weight"] = value + elif kind == "lora_A": + converted.update({f"{prefix}.{name}.lora_A.weight": value.clone() for name in ("to_q", "to_k", "to_v")}) + else: + for name, part in zip(("to_q", "to_k", "to_v"), value.chunk(3, dim=0)): + converted[f"{prefix}.{name}.lora_B.weight"] = part.contiguous() + metadata = { + "r": config["r"], + "lora_alpha": config["lora_alpha"], + "lora_dropout": config["lora_dropout"], + "bias": "none", + "target_modules": ["to_q", "to_k", "to_v", "to_out"], + } + RFD4ProteinaPipeline.save_lora_weights( + destination, transformer_lora_layers=converted, transformer_lora_adapter_metadata=metadata + ) + return Path(destination) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--source", + type=Path, + required=True, + help="Directory with config.json and model.safetensors, or with a native adapter when --adapter is set.", + ) + parser.add_argument("--destination", type=Path, required=True, help="New output directory.") + parser.add_argument( + "--sampling-config", + type=Path, + help="The RFD4 repository's configs/generation/inference_base.yaml, or the JSON equivalent. Preferred over " + "the transcribed default.", + ) + parser.add_argument( + "--adapter", + action="store_true", + help="Convert a native PEFT adapter directory into Diffusers LoRA weights instead of a model bundle.", + ) + parser.add_argument( + "--recycle-schedule", type=json.loads, help='JSON list of [steps, depth] bands, e.g. "[[60, 1]]".' + ) + args = parser.parse_args() + if args.adapter: + print(convert_rfd4_proteina_adapter(args.source, args.destination)) + return + model_args = load_released_recipe(args.sampling_config) if args.sampling_config else None + print( + convert_rfd4_proteina_bundle( + args.source, args.destination, model_args=model_args, recycle_schedule=args.recycle_schedule + ) + ) + + +if __name__ == "__main__": + main() diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 6db7a0c90e98..65efb81750ba 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -359,6 +359,8 @@ "QwenImageControlNetModel", "QwenImageMultiControlNetModel", "QwenImageTransformer2DModel", + "RFD4ProteinaAutoencoder", + "RFD4ProteinaTransformer", "SanaControlNetModel", "SanaTransformer2DModel", "SanaVideoTransformer3DModel", @@ -438,6 +440,7 @@ "LDMSuperResolutionPipeline", "PNDMPipeline", "RePaintPipeline", + "RFD4ProteinaPipeline", "ScoreSdeVePipeline", "StableDiffusionMixin", ] @@ -486,6 +489,7 @@ "PiflowScheduler", "PNDMScheduler", "RePaintScheduler", + "RFD4ProteinaScheduler", "SASolverScheduler", "SchedulerMixin", "SCMScheduler", @@ -1275,6 +1279,8 @@ QwenImageControlNetModel, QwenImageMultiControlNetModel, QwenImageTransformer2DModel, + RFD4ProteinaAutoencoder, + RFD4ProteinaTransformer, SanaControlNetModel, SanaTransformer2DModel, SanaVideoTransformer3DModel, @@ -1352,6 +1358,7 @@ LDMSuperResolutionPipeline, PNDMPipeline, RePaintPipeline, + RFD4ProteinaPipeline, ScoreSdeVePipeline, StableDiffusionMixin, ) @@ -1398,6 +1405,7 @@ PiflowScheduler, PNDMScheduler, RePaintScheduler, + RFD4ProteinaScheduler, SASolverScheduler, SchedulerMixin, SCMScheduler, diff --git a/src/diffusers/loaders/__init__.py b/src/diffusers/loaders/__init__.py index 72a24ba3803f..868e70e48736 100644 --- a/src/diffusers/loaders/__init__.py +++ b/src/diffusers/loaders/__init__.py @@ -12,9 +12,10 @@ _import_structure["transformer_sd3"] = ["SD3Transformer2DLoadersMixin"] _import_structure["unet"] = ["UNet2DConditionLoadersMixin"] _import_structure["utils"] = ["AttnProcsLayers"] + _import_structure["lora_pipeline"] = ["RFD4ProteinaLoraLoaderMixin"] if is_transformers_available(): _import_structure["single_file"] = ["FromSingleFileMixin"] - _import_structure["lora_pipeline"] = [ + _import_structure["lora_pipeline"] += [ "AceStepLoraLoaderMixin", "AmusedLoraLoaderMixin", "AnimaLoraLoaderMixin", @@ -60,6 +61,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: if is_torch_available(): + from .lora_pipeline import RFD4ProteinaLoraLoaderMixin from .single_file_model import FromOriginalModelMixin from .transformer_flux import FluxTransformer2DLoadersMixin from .transformer_sd3 import SD3Transformer2DLoadersMixin diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 22d596c6c682..52a16ce33aa1 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -7323,6 +7323,195 @@ def unfuse_lora(self, components: list[str] = ["transformer", "transformer_ref"] super().unfuse_lora(components=components, **kwargs) +class RFD4ProteinaLoraLoaderMixin(LoraBaseMixin): + r""" + Load LoRA layers into [`RFD4ProteinaTransformer`]. + + The autoencoder is paired with the denoiser checkpoint and is not adapted, so the transformer is the only loadable + module. + """ + + _lora_loadable_modules = ["transformer"] + transformer_name = TRANSFORMER_NAME + + @classmethod + @validate_hf_hub_args + def lora_state_dict( + cls, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + **kwargs, + ): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.lora_state_dict`] for more details. + """ + cache_dir = kwargs.pop("cache_dir", None) + force_download = kwargs.pop("force_download", False) + proxies = kwargs.pop("proxies", None) + local_files_only = kwargs.pop("local_files_only", None) + token = kwargs.pop("token", None) + revision = kwargs.pop("revision", None) + subfolder = kwargs.pop("subfolder", None) + weight_name = kwargs.pop("weight_name", None) + use_safetensors = kwargs.pop("use_safetensors", None) + return_lora_metadata = kwargs.pop("return_lora_metadata", False) + + allow_pickle = False + if use_safetensors is None: + use_safetensors = True + allow_pickle = True + + user_agent = {"file_type": "attn_procs_weights", "framework": "pytorch"} + + state_dict, metadata = _fetch_state_dict( + pretrained_model_name_or_path_or_dict=pretrained_model_name_or_path_or_dict, + weight_name=weight_name, + use_safetensors=use_safetensors, + local_files_only=local_files_only, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + token=token, + revision=revision, + subfolder=subfolder, + user_agent=user_agent, + allow_pickle=allow_pickle, + ) + + is_dora_scale_present = any("dora_scale" in k for k in state_dict) + if is_dora_scale_present: + warn_msg = "It seems like you are using a DoRA checkpoint that is not compatible in Diffusers at the moment. So, we are going to filter out the keys associated to 'dora_scale` from the state dict. If you think this is a mistake please open an issue https://github.com/huggingface/diffusers/issues/new." + logger.warning(warn_msg) + state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k} + + out = (state_dict, metadata) if return_lora_metadata else state_dict + return out + + @require_peft_backend + # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights + def load_lora_weights( + self, + pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], + adapter_name: str | None = None, + hotswap: bool = False, + **kwargs, + ): + """ + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for more details. + """ + low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT_LORA) + + # if a dict is passed, copy it instead of modifying it inplace + if isinstance(pretrained_model_name_or_path_or_dict, dict): + pretrained_model_name_or_path_or_dict = pretrained_model_name_or_path_or_dict.copy() + + # First, ensure that the checkpoint is a compatible one and can be successfully loaded. + kwargs["return_lora_metadata"] = True + state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) + + is_correct_format = all("lora" in key for key in state_dict.keys()) + if not is_correct_format: + raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + + self.load_lora_into_transformer( + state_dict, + transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=self, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + ) + + @classmethod + # Copied from diffusers.loaders.lora_pipeline.SD3LoraLoaderMixin.load_lora_into_transformer with SD3Transformer2DModel->RFD4ProteinaTransformer + def load_lora_into_transformer( + cls, + state_dict, + transformer, + adapter_name=None, + _pipeline=None, + low_cpu_mem_usage=False, + hotswap: bool = False, + metadata=None, + ): + """ + See [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_into_unet`] for more details. + """ + # Load the layers corresponding to transformer. + logger.info(f"Loading {cls.transformer_name}.") + transformer.load_lora_adapter( + state_dict, + network_alphas=None, + adapter_name=adapter_name, + metadata=metadata, + _pipeline=_pipeline, + low_cpu_mem_usage=low_cpu_mem_usage, + hotswap=hotswap, + ) + + @classmethod + # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.save_lora_weights + def save_lora_weights( + cls, + save_directory: str | os.PathLike, + transformer_lora_layers: dict[str, torch.nn.Module | torch.Tensor] = None, + is_main_process: bool = True, + weight_name: str = None, + save_function: Callable = None, + safe_serialization: bool = True, + transformer_lora_adapter_metadata: dict | None = None, + ): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for more information. + """ + lora_layers = {} + lora_metadata = {} + + if transformer_lora_layers: + lora_layers[cls.transformer_name] = transformer_lora_layers + lora_metadata[cls.transformer_name] = transformer_lora_adapter_metadata + + if not lora_layers: + raise ValueError("You must pass at least one of `transformer_lora_layers` or `text_encoder_lora_layers`.") + + cls._save_lora_weights( + save_directory=save_directory, + lora_layers=lora_layers, + lora_metadata=lora_metadata, + is_main_process=is_main_process, + weight_name=weight_name, + save_function=save_function, + safe_serialization=safe_serialization, + ) + + # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.fuse_lora + def fuse_lora( + self, + components: list[str] = ["transformer"], + lora_scale: float = 1.0, + safe_fusing: bool = False, + adapter_names: list[str] | None = None, + **kwargs, + ): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.fuse_lora`] for more details. + """ + super().fuse_lora( + components=components, + lora_scale=lora_scale, + safe_fusing=safe_fusing, + adapter_names=adapter_names, + **kwargs, + ) + + # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.unfuse_lora + def unfuse_lora(self, components: list[str] = ["transformer"], **kwargs): + r""" + See [`~loaders.StableDiffusionLoraLoaderMixin.unfuse_lora`] for more details. + """ + super().unfuse_lora(components=components, **kwargs) + + class LoraLoaderMixin(StableDiffusionLoraLoaderMixin): def __init__(self, *args, **kwargs): deprecation_message = "LoraLoaderMixin is deprecated and this will be removed in a future version. Please use `StableDiffusionLoraLoaderMixin`, instead." diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index c49bc35426d7..fb25d3fc7acc 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -59,6 +59,7 @@ _import_structure["autoencoders.autoencoder_mmaudio"] = ["MMAudioVAE"] _import_structure["autoencoders.autoencoder_oobleck"] = ["AutoencoderOobleck"] _import_structure["autoencoders.autoencoder_rae"] = ["AutoencoderRAE"] + _import_structure["autoencoders.autoencoder_rfd4_proteina"] = ["RFD4ProteinaAutoencoder"] _import_structure["autoencoders.autoencoder_same"] = ["AutoencoderSAME"] _import_structure["autoencoders.autoencoder_tiny"] = ["AutoencoderTiny"] _import_structure["autoencoders.autoencoder_vidtok"] = ["AutoencoderVidTok"] @@ -154,6 +155,7 @@ _import_structure["transformers.transformer_prx"] = ["PRXTransformer2DModel"] _import_structure["transformers.transformer_qwenimage"] = ["QwenImageTransformer2DModel"] _import_structure["transformers.transformer_qwenimage21"] = ["QwenImage21Transformer2DModel"] + _import_structure["transformers.transformer_rfd4_proteina"] = ["RFD4ProteinaTransformer"] _import_structure["transformers.transformer_sana_video"] = ["SanaVideoTransformer3DModel"] _import_structure["transformers.transformer_sana_wm"] = ["SanaWMTransformer3DModel"] _import_structure["transformers.transformer_sana_wm_refiner"] = ["SanaWMLTX2RefinerTransformer3DModel"] @@ -224,6 +226,7 @@ MiniMaxMusic3Vocoder, MMAudioVAE, MMAudioVocoder, + RFD4ProteinaAutoencoder, VQModel, ) from .cache_utils import CacheMixin @@ -310,6 +313,7 @@ PRXTransformer2DModel, QwenImage21Transformer2DModel, QwenImageTransformer2DModel, + RFD4ProteinaTransformer, SanaTransformer2DModel, SanaVideoTransformer3DModel, SanaWMLTX2RefinerTransformer3DModel, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 5e5206b3e76a..1cbe189979c9 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -29,6 +29,7 @@ from .autoencoder_mmaudio import MMAudioVAE from .autoencoder_oobleck import AutoencoderOobleck from .autoencoder_rae import AutoencoderRAE +from .autoencoder_rfd4_proteina import RFD4ProteinaAutoencoder from .autoencoder_same import AutoencoderSAME from .autoencoder_tiny import AutoencoderTiny from .autoencoder_vidtok import AutoencoderVidTok diff --git a/src/diffusers/models/autoencoders/autoencoder_rfd4_proteina.py b/src/diffusers/models/autoencoders/autoencoder_rfd4_proteina.py new file mode 100644 index 000000000000..19161aa18a98 --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_rfd4_proteina.py @@ -0,0 +1,690 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# Diffusers integration additions are licensed under the Apache License, Version 2.0. +# +# Adapted from RFD4-Proteina (BSD-3-Clause), with MIT and Apache-2.0 components. +# Copyright (c) 2026, Institute for Protein Design, University of Washington +# Copyright (c) 2022 MattMcPartlon +# Copyright (c) 2024 Phil Wang +# Copyright 2021 AlQuraishi Laboratory +# Copyright 2021 DeepMind Technologies Limited +# Full third-party copyright notices, license conditions and disclaimers for the RFD4-Proteina +# integration are in diffusers/models/transformers/transformer_rfd4_proteina.py. +"""RFD4-Proteina autoencoder: per-token local latents <-> atom37 coordinates and residue identities.""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import partial +from typing import Any + +import torch +from torch import nn +from torch.nn import functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import BaseOutput +from ...utils.accelerate_utils import apply_forward_hook +from ...utils.torch_utils import randn_tensor +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_utils import ModelMixin, get_parameter_dtype + + +NORMS = {"layer_norm": nn.LayerNorm, "rms_norm": nn.RMSNorm} + + +# Copied from diffusers.models.transformers.transformer_rfd4_proteina.RFD4ProteinaTransition +class RFD4ProteinaTransition(nn.Module): + def __init__(self, width: int, expansion: int = 4) -> None: + super().__init__() + self.swish_linear = nn.Sequential(nn.Linear(width, 2 * width * expansion, bias=False)) + self.linear_out = nn.Linear(width * expansion, width, bias=False) + + def forward(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + left, gate = self.swish_linear(values).chunk(2, dim=-1) + return self.linear_out(F.silu(gate) * left) * mask[..., None] + + +# Copied from diffusers.models.transformers.transformer_rfd4_proteina.RFD4ProteinaAdaptiveNorm +class RFD4ProteinaAdaptiveNorm(nn.Module): + def __init__(self, width: int, condition_width: int, norm_type: str) -> None: + super().__init__() + self.norm = NORMS[norm_type](width, eps=1e-5, elementwise_affine=False) + self.norm_cond = NORMS[norm_type](condition_width, eps=1e-5) + self.to_gamma = nn.Sequential(nn.Linear(condition_width, width), nn.Sigmoid()) + self.to_beta = nn.Linear(condition_width, width, bias=False) + + def forward(self, values: torch.Tensor, condition: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + normalized = self.norm(values) + normalized_condition = self.norm_cond(condition) + gamma, beta = self.to_gamma(normalized_condition), self.to_beta(normalized_condition) + if condition.ndim == 2 and values.ndim > 2: + shape = (condition.shape[0], *((1,) * (values.ndim - 2)), values.shape[-1]) + gamma, beta = gamma.reshape(shape), beta.reshape(shape) + return (normalized * gamma + beta) * mask[..., None] + + +def feature_dimensions( + token_type_emb_dim: int, latent_dim: int, n_residue_types: int, charge_emb_dim: int +) -> dict[str, int]: + """Width of every decoder input feature. Unlike `NORMS` this cannot be a constant: four of the + widths are set by the checkpoint configuration, so the table is built per model.""" + return { + "token_type": token_type_emb_dim, + "chain_idx_seq": 1, + "x1_a37coors_nm": 148, + "x1_a37coors_nm_rel": 148, + "ca_coors_nm": 3, + "z_latent_seq": latent_dim, + "cond_residue_type": n_residue_types, + "incident_bond_type_counts": 37 * 10, + "atom37_element_charge_embedding": 37 * (8 + charge_emb_dim), + } + + +# Copied from diffusers.models.transformers.transformer_rfd4_proteina.RFD4ProteinaPairAttnProcessor +class RFD4ProteinaPairAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "RFD4ProteinaPairAttention", + hidden_states: torch.Tensor, + pairs: torch.Tensor | None, + mask: torch.Tensor, + ) -> torch.Tensor: + batch, tokens, width = hidden_states.shape + hidden_states = attn.node_norm(hidden_states) + query, key = attn.to_q(hidden_states), attn.to_k(hidden_states) + query, key = attn.norm_q(query), attn.norm_k(key) + value, gate = attn.to_v(hidden_states), attn.to_g(hidden_states) + query, key, value, gate = ( + tensor.view(batch, tokens, attn.heads, attn.head_width) for tensor in (query, key, value, gate) + ) + pair_mask = mask[:, None, :, None] & mask[:, None, None, :] + if attn.use_pair: + bias = attn.to_bias(attn.pair_norm(pairs)).permute(0, 3, 1, 2).masked_fill(~pair_mask, -1e4) + else: + bias = torch.zeros(batch, attn.heads, tokens, tokens, device=query.device, dtype=query.dtype) + bias = bias.masked_fill(~pair_mask, -1e4) + output = dispatch_attention_fn( + query, key, value, attn_mask=bias, backend=self._attention_backend, parallel_config=self._parallel_config + ) + output = (output * gate.sigmoid()).reshape(batch, tokens, width) + return attn.to_out(output) + + +# Copied from diffusers.models.transformers.transformer_rfd4_proteina.RFD4ProteinaPairAttention +class RFD4ProteinaPairAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = RFD4ProteinaPairAttnProcessor + _available_processors = [RFD4ProteinaPairAttnProcessor] + _supports_qkv_fusion = False + + def __init__(self, width: int, pair_width: int, heads: int, norm_type: str, use_pair: bool) -> None: + super().__init__() + self.heads, self.head_width, self.use_pair = heads, width // heads, use_pair + self.node_norm = NORMS[norm_type](width, eps=1e-5) + self.to_q = nn.Linear(width, width) + self.to_k = nn.Linear(width, width) + self.to_v = nn.Linear(width, width) + self.to_g = nn.Linear(width, width) + self.norm_q = NORMS[norm_type](width, eps=1e-5) + self.norm_k = NORMS[norm_type](width, eps=1e-5) + self.to_out = nn.Linear(width, width) + if use_pair: + self.to_bias = nn.Linear(pair_width, heads, bias=False) + self.pair_norm = NORMS[norm_type](pair_width, eps=1e-5) + self.set_processor(RFD4ProteinaPairAttnProcessor()) + + def forward(self, values: torch.Tensor, pairs: torch.Tensor | None, mask: torch.Tensor) -> torch.Tensor: + return self.processor(self, values, pairs, mask) + + +# Copied from diffusers.models.transformers.transformer_rfd4_proteina.RFD4ProteinaAdaptiveAttention +class RFD4ProteinaAdaptiveAttention(nn.Module): + def __init__( + self, width: int, pair_width: int, heads: int, condition_width: int, norm_type: str, use_pair: bool + ) -> None: + super().__init__() + self.adaln = RFD4ProteinaAdaptiveNorm(width, condition_width, norm_type) + self.mha = RFD4ProteinaPairAttention(width, pair_width, heads, norm_type, use_pair) + self.scale_output = nn.Sequential(nn.Linear(condition_width, width), nn.Sigmoid()) + + def forward( + self, values: torch.Tensor, pairs: torch.Tensor | None, condition: torch.Tensor, mask: torch.Tensor + ) -> torch.Tensor: + output = self.mha(self.adaln(values, condition, mask), pairs, mask) + return output * self.scale_output(condition) * mask[..., None] + + +# Copied from diffusers.models.transformers.transformer_rfd4_proteina.RFD4ProteinaAdaptiveTransition +class RFD4ProteinaAdaptiveTransition(nn.Module): + def __init__(self, width: int, condition_width: int, norm_type: str) -> None: + super().__init__() + self.adaln = RFD4ProteinaAdaptiveNorm(width, condition_width, norm_type) + self.transition = RFD4ProteinaTransition(width) + self.scale_output = nn.Sequential(nn.Linear(condition_width, width), nn.Sigmoid()) + + def forward(self, values: torch.Tensor, condition: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + output = self.transition(self.adaln(values, condition, mask), mask) + return output * self.scale_output(condition) * mask[..., None] + + +# Copied from diffusers.models.transformers.transformer_rfd4_proteina.RFD4ProteinaTrunkLayer +class RFD4ProteinaTrunkLayer(nn.Module): + """One attention + transition block, conditioned on `condition` through adaptive normalization.""" + + def __init__( + self, token_dim: int, pair_repr_dim: int, nheads: int, dim_cond: int, norm_type: str, use_pair: bool + ) -> None: + super().__init__() + self.mhba = RFD4ProteinaAdaptiveAttention(token_dim, pair_repr_dim, nheads, dim_cond, norm_type, use_pair) + self.transition = RFD4ProteinaAdaptiveTransition(token_dim, dim_cond, norm_type) + + def forward( + self, values: torch.Tensor, pairs: torch.Tensor | None, condition: torch.Tensor, mask: torch.Tensor + ) -> torch.Tensor: + values = values * mask[..., None] + values = (values + self.mhba(values, pairs, condition, mask)) * mask[..., None] + return (values + self.transition(values, condition, mask)) * mask[..., None] + + +def _autoencoder_zeros(features: Mapping[str, Any], dim: int) -> torch.Tensor: + mask = features["padding_mask"] + return torch.zeros((*mask.shape, dim), device=mask.device) + + +def _chain_idx_seq(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return features["chain_iid"].float().unsqueeze(-1) + + +def _ca_coors_nm(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return features["ca_coors_nm"] + + +def _z_latent_seq(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return features["z_latent"] + + +def _incident_bond_type_counts(features: Mapping[str, Any], dim: int) -> torch.Tensor: + counts = features["incident_bond_type_counts"].float() + return (counts * features["atom37_mask"].unsqueeze(-1)).flatten(start_dim=-2) + + +def _cond_residue_type(features: Mapping[str, Any], dim: int) -> torch.Tensor: + values = features.get("condition_sequence_1_residue") + if values is None: + return _autoencoder_zeros(features, dim) + valid = values != 0 + return F.one_hot((values * valid).long(), dim).float() * valid[..., None] + + +def _x1_a37coors_nm(features: Mapping[str, Any], dim: int) -> torch.Tensor: + coords, mask = features["coords_nm"], features["atom37_mask"] + coords = coords * mask[..., None] + return torch.cat((coords.flatten(2), mask.float()), dim=-1) + + +def _x1_a37coors_nm_rel(features: Mapping[str, Any], dim: int) -> torch.Tensor: + coords, mask = features["coords_nm"], features["atom37_mask"] + coords = coords * mask[..., None] + coords = (coords - coords[:, :, 1, None]) * mask[..., None] + return torch.cat((coords.flatten(2), mask.float()), dim=-1) + + +# Features without learned weights: plain functions of the prepared tensors and the feature width. +AUTOENCODER_FEATURE_FUNCTIONS = { + "chain_idx_seq": _chain_idx_seq, + "ca_coors_nm": _ca_coors_nm, + "z_latent_seq": _z_latent_seq, + "incident_bond_type_counts": _incident_bond_type_counts, + "cond_residue_type": _cond_residue_type, + "x1_a37coors_nm": _x1_a37coors_nm, + "x1_a37coors_nm_rel": _x1_a37coors_nm_rel, +} + + +class _TokenTypeEmbedding(nn.Embedding): + def __init__(self, dim: int, num_token_types: int, **kwargs) -> None: + super().__init__(num_token_types, dim) + + def forward(self, features: Mapping[str, Any]) -> torch.Tensor: + return super().forward(features["token_type"].long()) + + +class _ElementChargeEmbedding(nn.Module): + def __init__(self, dim: int, charge_emb_dim: int, **kwargs) -> None: + super().__init__() + self.element = nn.Embedding(119, 8) + self.charge = nn.Embedding(11, charge_emb_dim) + + def forward(self, features: Mapping[str, Any]) -> torch.Tensor: + mask = features["atom37_mask"].bool() + elements = features["ground_truth_element"].long().masked_fill(~mask, 0) + # Formal charges are stored signed; the embedding covers -4..6 saturated. + charges = (features["charge"].long().masked_fill(~mask, 0) + 4).clamp(0, 10) + encoded = torch.cat((self.element(elements), self.charge(charges)), dim=-1) + return (encoded * mask.unsqueeze(-1)).flatten(start_dim=-2) + + +# The features with learned weights. +AUTOENCODER_FEATURE_MODULES = { + "token_type": _TokenTypeEmbedding, + "atom37_element_charge_embedding": _ElementChargeEmbedding, +} + + +class RFD4ProteinaAutoencoderFeatureFactory(nn.Module): + """Concatenates named features and projects them to `width`. Each feature is resolved once, to a function + (`AUTOENCODER_FEATURE_FUNCTIONS`) or to a module with weights (`AUTOENCODER_FEATURE_MODULES`, kept in + `self.feature_modules`).""" + + def __init__( + self, + names: Sequence[str], + width: int, + token_type_emb_dim: int, + latent_dim: int, + n_residue_types: int, + num_token_types: int, + charge_emb_dim: int = 16, + ) -> None: + super().__init__() + dimensions = feature_dimensions(token_type_emb_dim, latent_dim, n_residue_types, charge_emb_dim) + self.feature_modules = nn.ModuleDict() + self.feature_fns = [] + for name in names: + if name not in dimensions: + raise ValueError(f"unsupported RFD4 autoencoder feature: {name}") + if name in AUTOENCODER_FEATURE_FUNCTIONS: + self.feature_fns.append(partial(AUTOENCODER_FEATURE_FUNCTIONS[name], dim=dimensions[name])) + else: + module = AUTOENCODER_FEATURE_MODULES[name]( + dimensions[name], num_token_types=num_token_types, charge_emb_dim=charge_emb_dim + ) + self.feature_modules[name] = module + self.feature_fns.append(module) + self.linear_out = nn.Linear(sum(dimensions[name] for name in names), width, bias=False) + + def forward(self, features: Mapping[str, Any]) -> torch.Tensor: + mask = features["padding_mask"].bool() + stacked = torch.cat([feature(features) for feature in self.feature_fns], dim=-1) * mask[..., None] + compute_dtype = features["compute_dtype"] + # The projection may be kept in float32 while the model runs in lower precision. A quantized weight is an integer + # tensor, so its input follows the compute dtype instead. + weight_dtype = get_parameter_dtype(self.linear_out) + projected = self.linear_out(stacked.to(weight_dtype if weight_dtype.is_floating_point else compute_dtype)) + return (projected * mask[..., None]).to(compute_dtype) + + +class RFD4ProteinaEncoder(nn.Module): + """Maps atom37 coordinates to the mean and log-scale of the per-token local latent distribution.""" + + def __init__( + self, + nlayers: int, + token_dim: int, + nheads: int, + dim_cond: int, + pair_repr_dim: int, + token_type_emb_dim: int, + latent_dim: int, + feats_seq: Sequence[str], + n_residue_types: int, + num_token_types: int, + ) -> None: + super().__init__() + vocabulary = (n_residue_types, num_token_types) + self.init_repr_factory = RFD4ProteinaAutoencoderFeatureFactory( + feats_seq, token_dim, token_type_emb_dim, latent_dim, *vocabulary + ) + self.dim_cond = dim_cond + self.transition_c_1 = RFD4ProteinaTransition(dim_cond, 2) + self.transition_c_2 = RFD4ProteinaTransition(dim_cond, 2) + self.transformer_layers = nn.ModuleList( + RFD4ProteinaTrunkLayer(token_dim, pair_repr_dim, nheads, dim_cond, "layer_norm", False) + for _ in range(nlayers) + ) + self.latent_decoder_mean_n_log_scale = nn.Sequential( + nn.LayerNorm(token_dim), nn.Linear(token_dim, 2 * latent_dim, bias=False) + ) + + def forward( + self, features: Mapping[str, Any], generator: torch.Generator | list[torch.Generator] | None = None + ) -> dict[str, torch.Tensor]: + batch_size = features["padding_mask"].shape[0] + generators = ( + None if generator is None else ([generator] if isinstance(generator, torch.Generator) else generator) + ) + if generators is not None and len(generators) != batch_size: + raise ValueError(f"Provide one autoencoder generator per batch row ({batch_size}).") + padding = features["padding_mask"].bool() + attention = padding & ~features["is_in_context_conditioning"].bool() + emit = attention & features["supports_local_latent"].bool() + # The autoencoder is not conditioned: the conditioning vector is zero. + condition = torch.zeros( + (*padding.shape, self.dim_cond), device=padding.device, dtype=features["compute_dtype"] + ) + condition = self.transition_c_2(self.transition_c_1(condition, padding), padding) + values = self.init_repr_factory(features) * padding[..., None] + for layer in self.transformer_layers: + values = layer(values, None, condition, attention) + # The head that decides the sampled latent is kept in float32; give it matching input. + head_dtype = get_parameter_dtype(self.latent_decoder_mean_n_log_scale[1]) + mean, log_scale = self.latent_decoder_mean_n_log_scale(values.to(head_dtype)).chunk(2, dim=-1) + mean, log_scale = mean * emit[..., None], log_scale * emit[..., None] + noise = ( + torch.randn_like(log_scale) + if generators is None + else randn_tensor(log_scale.shape, generator=generators, device=log_scale.device, dtype=log_scale.dtype) + ) + latent = (mean + noise * log_scale.exp()) * emit[..., None] + return {"latent": latent, "mean": mean, "log_scale": log_scale} + + +class RFD4ProteinaDecoder(nn.Module): + """Maps backbone coordinates and local latents to atom37 coordinates and residue identities.""" + + def __init__( + self, + nlayers: int, + token_dim: int, + nheads: int, + dim_cond: int, + pair_repr_dim: int, + token_type_emb_dim: int, + latent_dim: int, + n_residue_types: int, + num_token_types: int, + feats_seq: Sequence[str], + atom37_mask: Sequence[Sequence[bool]], + allowed_residue_types: Sequence[Sequence[bool]], + ) -> None: + super().__init__() + self.init_repr_factory = RFD4ProteinaAutoencoderFeatureFactory( + feats_seq, token_dim, token_type_emb_dim, latent_dim, n_residue_types, num_token_types + ) + self.dim_cond = dim_cond + self.transition_c_1 = RFD4ProteinaTransition(dim_cond, 2) + self.transition_c_2 = RFD4ProteinaTransition(dim_cond, 2) + self.transformer_layers = nn.ModuleList( + RFD4ProteinaTrunkLayer(token_dim, pair_repr_dim, nheads, dim_cond, "layer_norm", False) + for _ in range(nlayers) + ) + self.logit_linear = nn.Sequential(nn.LayerNorm(token_dim), nn.Linear(token_dim, n_residue_types, bias=False)) + self.struct_linear = nn.Sequential(nn.LayerNorm(token_dim), nn.Linear(token_dim, 37 * 3, bias=False)) + self.atom_presence_linear = nn.Sequential(nn.LayerNorm(token_dim), nn.Linear(token_dim, 37, bias=False)) + self.register_buffer("_atom37_mask", torch.tensor(atom37_mask, dtype=torch.bool), persistent=False) + self.register_buffer( + "_allowed_residue_types", torch.tensor(allowed_residue_types, dtype=torch.bool), persistent=False + ) + + def forward(self, features: Mapping[str, Any]) -> dict[str, torch.Tensor]: + padding = features["padding_mask"].bool() + attention = padding & ~features["is_in_context_conditioning"].bool() + emit = attention & features["supports_local_latent"].bool() + # The autoencoder is not conditioned: the conditioning vector is zero. + condition = torch.zeros( + (*padding.shape, self.dim_cond), device=padding.device, dtype=features["compute_dtype"] + ) + condition = self.transition_c_2(self.transition_c_1(condition, padding), padding) + values = self.init_repr_factory(features) * padding[..., None] + for layer in self.transformer_layers: + values = layer(values, None, condition, attention) + logits = self.logit_linear(values) * emit[..., None] + if "token_type" in features: + allowed = self._allowed_residue_types[features["token_type"].long()] + logits = logits.masked_fill(~(allowed | ~emit[..., None]), torch.finfo(logits.dtype).min) + coordinates = self.struct_linear(values).reshape(*values.shape[:2], 37, 3).float() * emit[..., None, None] + ca = features["ca_coors_nm"] + coordinates[..., 1, :] = 0 + coordinates = coordinates + ca[:, :, None] + coordinates = coordinates * emit[..., None, None] + residue_type = logits.argmax(dim=-1) * emit + # A residue the sequence condition names is kept as given; index 0 is the mask token, which leaves the + # decision to the decoder. + sequence = features.get("condition_sequence_1_residue") + if sequence is not None: + residue_type = torch.where((sequence != 0) & emit, sequence, residue_type) + atom_mask = self._atom37_mask.to(residue_type.device)[residue_type] & emit[..., None] + # The decoder predicts which of the residue's slots are actually occupied, so a modified residue can drop atoms + # its canonical form would have. + present = self.atom_presence_linear(values) > 0 + present[..., 1] = True # every emitted token keeps its reference atom + # Atoms that belong to a conditioned coordinate group are required, whatever the presence head says. + groups = features.get("condition_CoordinateGroup_1_atom") + if groups is not None: + present = present | (groups >= 0) + atom_mask = atom_mask & present + return { + "coordinates": coordinates * atom_mask[..., None], + "residue_type": residue_type * emit, + "emit_mask": emit, + "atom_mask": atom_mask, + } + + +@dataclass +class RFD4ProteinaEncoderOutput(BaseOutput): + """Sampled latent and Gaussian parameters, each shaped `(batch, tokens, latent_dim)`. + + Args: + latent (`torch.Tensor`): Sample `mean + noise * exp(log_scale)`, zero outside emitted latent tokens. + mean (`torch.Tensor`): Gaussian mean, zero outside emitted latent tokens. + log_scale (`torch.Tensor`): Log standard deviation, zero outside emitted latent tokens. + """ + + latent: torch.Tensor + mean: torch.Tensor | None = None + log_scale: torch.Tensor | None = None + + +@dataclass +class RFD4ProteinaDecoderOutput(BaseOutput): + """Decoded molecular coordinates and discrete atom37 topology. + + Args: + coordinates (`torch.Tensor`): Coordinates `(batch, tokens, 37, 3)` in nanometers; non-emitted atoms are zero. + residue_type (`torch.Tensor`): Residue vocabulary indices shaped `(batch, tokens)`. + emit_mask (`torch.Tensor`): Boolean emitted-token mask shaped `(batch, tokens)`. + atom_mask (`torch.Tensor`): Boolean valid-atom mask shaped `(batch, tokens, 37)`. + """ + + coordinates: torch.Tensor + residue_type: torch.Tensor | None = None + emit_mask: torch.Tensor | None = None + atom_mask: torch.Tensor | None = None + + +class RFD4ProteinaAutoencoder(ModelMixin, ConfigMixin, AttentionMixin): + """Paired molecular encoder and atom37 decoder. + + Coordinates supplied to and returned by this model are in nanometers. + + Args: + latent_dim (`int`, defaults to `32`): Width of the per-token local latent. + n_residue_types (`int`, defaults to `147`): Residue vocabulary size. + num_token_types (`int`, defaults to `7`): Size of the token-type vocabulary: protein, RNA, DNA, atomized, + unknown, glycan and non-canonical amino acid. + allowed_residue_types (`list[list[bool]]`): `(num_token_types, n_residue_types)` table of the residue classes + each token type may decode to. + nlayers (`int`, defaults to `12`): Trunk layers in each of the encoder and decoder. + token_dim (`int`, defaults to `768`): Token representation width. + nheads (`int`, defaults to `12`): Attention heads per layer. + dim_cond (`int`, defaults to `128`): Width of the (unused, zero) conditioning vector. + pair_repr_dim (`int`, defaults to `256`): Pair width the trunk layers are constructed with. The released + autoencoder uses no pair bias, so this holds no parameters. + token_type_emb_dim (`int`, defaults to `32`): Token-type embedding width. + encoder_feats_seq, decoder_feats_seq (`list[str]`): Features projected into the encoder and decoder inputs. + atom37_mask (`list[list[bool]]`): Valid atom37 slots for each residue type, shaped `(n_residue_types, 37)`. + """ + + _no_split_modules = ["RFD4ProteinaTrunkLayer"] + _repeated_blocks = ["RFD4ProteinaTrunkLayer"] + _skip_layerwise_casting_patterns = ["norm", "embedding"] + # The head producing the latent mean and log-scale decides the sampled latent, so it stays fp32. + _keep_in_fp32_modules = ["latent_decoder_mean_n_log_scale", "norm"] + # Callers reach this model through `encode`/`decode`, never `forward`, so a block-level group that + # hooked the root would never onload. Naming the two submodules as blocks puts the hooks on modules + # whose own forward does run, which is how `AutoencoderKL` supports offloading through the same path. + _group_offload_block_modules = ["encoder", "decoder"] + + @register_to_config + def __init__( + self, + atom37_mask: Sequence[Sequence[bool]], + allowed_residue_types: Sequence[Sequence[bool]], + latent_dim: int = 32, + n_residue_types: int = 147, + num_token_types: int = 7, + nlayers: int = 12, + token_dim: int = 768, + nheads: int = 12, + dim_cond: int = 128, + pair_repr_dim: int = 256, + token_type_emb_dim: int = 32, + encoder_feats_seq: tuple[str, ...] = ("token_type", "x1_a37coors_nm", "x1_a37coors_nm_rel", "chain_idx_seq"), + decoder_feats_seq: tuple[str, ...] = ("ca_coors_nm", "z_latent_seq", "token_type", "cond_residue_type"), + ) -> None: + super().__init__() + shared = { + "nlayers": nlayers, + "token_dim": token_dim, + "nheads": nheads, + "dim_cond": dim_cond, + "pair_repr_dim": pair_repr_dim, + "token_type_emb_dim": token_type_emb_dim, + "latent_dim": latent_dim, + } + self.latent_dim = latent_dim + vocabulary = {"n_residue_types": n_residue_types, "num_token_types": num_token_types} + self.encoder = RFD4ProteinaEncoder(**shared, **vocabulary, feats_seq=encoder_feats_seq) + self.decoder = RFD4ProteinaDecoder( + **shared, + **vocabulary, + feats_seq=decoder_feats_seq, + atom37_mask=atom37_mask, + allowed_residue_types=allowed_residue_types, + ) + + @property + def allowed_residue_types(self) -> torch.Tensor: + """`(num_token_types, n_residue_types)` table of the classes each token type may decode to.""" + return self.decoder._allowed_residue_types + + @apply_forward_hook + def encode( + self, + coords_nm: torch.Tensor, + padding_mask: torch.Tensor, + token_type: torch.Tensor | None = None, + atom37_mask: torch.Tensor | None = None, + conditioning: dict[str, Any] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + return_dict: bool = True, + ) -> RFD4ProteinaEncoderOutput | tuple: + """Encode prepared coordinates and conditioning into sampled per-token local latents. + + Args: + coords_nm (`torch.Tensor`): `(batch, tokens, 37, 3)` atom37 coordinates in nanometers. + padding_mask (`torch.Tensor`): `(batch, tokens)`, true where a token is real rather than padding. + 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. + conditioning (`dict[str, Any]`, *optional*): The remaining topology and conditioning tensors, + under the keys the preprocessor emits. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): One generator per batch row + for independent stochastic latent samples. Without one, sampling draws from the global stream. + return_dict (`bool`, defaults to `True`): Return a typed output, or a tuple of its fields. + + Returns: + [`RFD4ProteinaEncoderOutput`] or `tuple`: Latent, mean and log-scale tensors shaped `(batch, tokens, + latent_dim)`. + """ + features = { + **(conditioning or {}), + "compute_dtype": self.dtype, + "coords_nm": coords_nm, + "padding_mask": padding_mask, + } + if token_type is not None: + features["token_type"] = token_type + if atom37_mask is not None: + features["atom37_mask"] = atom37_mask + output = RFD4ProteinaEncoderOutput(**self.encoder(features, generator=generator)) + return output if return_dict else output.to_tuple() + + @apply_forward_hook + def decode( + self, + z_latent: torch.Tensor, + ca_coors_nm: torch.Tensor, + padding_mask: torch.Tensor, + token_type: torch.Tensor | None = None, + atom37_mask: torch.Tensor | None = None, + conditioning: dict[str, Any] | None = None, + return_dict: bool = True, + ) -> RFD4ProteinaDecoderOutput | tuple: + """Decode per-token latents and backbone positions into atom37 structure. + + Args: + z_latent (`torch.Tensor`): `(batch, tokens, latent_dim)` per-token local latents. + ca_coors_nm (`torch.Tensor`): `(batch, tokens, 3)` backbone CA positions in nanometers. + padding_mask (`torch.Tensor`): `(batch, tokens)`, true where a token is real rather than padding. + 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. + conditioning (`dict[str, Any]`, *optional*): The remaining topology and conditioning tensors, + under the keys the preprocessor emits. The sequence condition is read from here. + return_dict (`bool`, defaults to `True`): Return a typed output, or a tuple of its fields. + + Returns: + [`RFD4ProteinaDecoderOutput`] or `tuple`: Coordinates `(batch, tokens, 37, 3)` in nanometers, residue types + `(batch, tokens)`, emission mask `(batch, tokens)` and atom mask `(batch, tokens, 37)`. Non-emitted atom + coordinates are zero. + """ + features = { + **(conditioning or {}), + "compute_dtype": self.dtype, + "z_latent": z_latent, + "ca_coors_nm": ca_coors_nm, + "padding_mask": padding_mask, + } + if token_type is not None: + features["token_type"] = token_type + if atom37_mask is not None: + features["atom37_mask"] = atom37_mask + output = RFD4ProteinaDecoderOutput(**self.decoder(features)) + return output if return_dict else output.to_tuple() + + def forward( + self, + z_latent: torch.Tensor, + ca_coors_nm: torch.Tensor, + padding_mask: torch.Tensor, + token_type: torch.Tensor | None = None, + atom37_mask: torch.Tensor | None = None, + conditioning: dict[str, Any] | None = None, + return_dict: bool = True, + ) -> RFD4ProteinaDecoderOutput | tuple: + """Decode per-token latents and backbone positions; equivalent to calling [`~RFD4ProteinaAutoencoder.decode`]. + + Args: + z_latent (`torch.Tensor`): `(batch, tokens, latent_dim)` per-token local latents. + ca_coors_nm (`torch.Tensor`): `(batch, tokens, 3)` backbone CA positions in nanometers. + padding_mask (`torch.Tensor`): `(batch, tokens)`, true where a token is real rather than padding. + 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. + conditioning (`dict[str, Any]`, *optional*): The remaining topology and conditioning tensors. + return_dict (`bool`, defaults to `True`): Return a typed output, or a tuple of its fields. + + Returns: + [`RFD4ProteinaDecoderOutput`] or `tuple`: Decoded coordinates, residue types and atom/emit masks. + """ + return self.decode( + z_latent=z_latent, + ca_coors_nm=ca_coors_nm, + padding_mask=padding_mask, + token_type=token_type, + atom37_mask=atom37_mask, + conditioning=conditioning, + return_dict=return_dict, + ) diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 1076b2b33fae..eedf93c14828 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -64,6 +64,7 @@ from .transformer_prx import PRXTransformer2DModel from .transformer_qwenimage import QwenImageTransformer2DModel from .transformer_qwenimage21 import QwenImage21Transformer2DModel + from .transformer_rfd4_proteina import RFD4ProteinaTransformer from .transformer_sana_video import SanaVideoTransformer3DModel from .transformer_sana_wm import SanaWMTransformer3DModel from .transformer_sana_wm_refiner import SanaWMLTX2RefinerTransformer3DModel diff --git a/src/diffusers/models/transformers/transformer_rfd4_proteina.py b/src/diffusers/models/transformers/transformer_rfd4_proteina.py new file mode 100644 index 000000000000..fcde248cdb73 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_rfd4_proteina.py @@ -0,0 +1,1189 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# Diffusers integration additions are licensed under the Apache License, Version 2.0. +# +# This file carries the third-party notices for the whole RFD4-Proteina integration; the +# autoencoder, scheduler and pipeline point here rather than repeating them. +# +# Adapted from RFD4-Proteina, revision 0e05aece83eafc4bf06dad7f9a7e5d50922ec83f. The model and +# sampling code is BSD-3-Clause. Through it, these components are also adapted: +# * pair-biased attention (`RFD4ProteinaPairAttention`): MIT, Copyright (c) 2022 MattMcPartlon +# * adaptive normalization and SwiGLU (`RFD4ProteinaAdaptiveNorm`, the swish-gated linear in +# `RFD4ProteinaTransition`): MIT, Copyright (c) 2024 Phil Wang, from alphafold3-pytorch, +# https://github.com/lucidrains/alphafold3-pytorch +# * pair transition and triangular multiplicative update: Apache-2.0, +# Copyright 2021 AlQuraishi Laboratory, Copyright 2021 DeepMind Technologies Limited. +# Apache-2.0 is the license this repository ships in LICENSE. +# These were adapted to Diffusers model interfaces, attention processors and checkpoint layouts; +# the original implementations are not imported. AtomWorks remains a separate dependency under its +# own license, and model weights have their own distribution terms. The notices below preserve the +# upstream terms and do not relicense the upstream material. +# +# --------------------------------------------------------------------------------------------- +# RFD4-Proteina +# BSD 3-Clause License +# +# Copyright (c) 2026, Institute for Protein Design, University of Washington +# +# Redistribution and use in source and binary forms, with or without modification, are permitted +# provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, this list of conditions +# and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright notice, this list of +# conditions and the following disclaimer in the documentation and/or other materials provided +# with the distribution. +# +# * Neither the name of the copyright holder nor the names of its contributors may be used to +# endorse or promote products derived from this software without specific prior written +# permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR +# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND +# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY +# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# +# --------------------------------------------------------------------------------------------- +# Pair-biased attention +# MIT License +# +# Copyright (c) 2022 MattMcPartlon +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and +# associated documentation files (the "Software"), to deal in the Software without restriction, +# including without limitation the rights to use, copy, modify, merge, publish, distribute, +# sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all copies or +# substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT +# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +# +# --------------------------------------------------------------------------------------------- +# Adaptive normalization and SwiGLU +# MIT License +# +# Copyright (c) 2024 Phil Wang +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and +# associated documentation files (the "Software"), to deal in the Software without restriction, +# including without limitation the rights to use, copy, modify, merge, publish, distribute, +# sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all copies or +# substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT +# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +"""RFD4-Proteina denoiser operating on prepared molecular tensors.""" + +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import partial +from typing import Any, Literal + +import torch +from torch import nn +from torch.nn import functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin +from ...utils import BaseOutput, apply_lora_scale +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_utils import ModelMixin, get_parameter_dtype + + +# Annotation keys of the prepared conditioning tensors, as written by the RFD4 input preparation. +CONDITION_KEYS = { + "sequence": "condition_sequence_1_residue", + "coordinates": "condition_coordinate_1_atom", + "internal_coordinates": "condition_InternalCoordinate_1_atom", + "coordinate_group": "condition_CoordinateGroup_1_atom", + "rasa": "condition_rasa-binned_1_atom", + "secondary_structure": "condition_secondary-structure_1_residue", + "hbond": "condition_hbond_1_atom", + "pistack": "condition_pistack_1_atom", + "pication": "condition_pication_1_atom", + "saltbridge": "condition_saltbridge_1_atom", + "hydrophobic": "condition_hydrophobic_1_atom", + "metal_coord": "condition_metal-coord_1_atom", + "hotspot": "condition_hotspot_1_atom", + "index_group": "condition_index-group_1_token", + "cath": "condition_cath-label_1_residue", + "centroid": "condition_ChainCentroid_1_atom", +} + +SEQUENCE_PROBABILITY_TOKENS = ( + "ALA ARG ASN ASP CYS GLN GLU GLY HIS ILE LEU LYS MET PHE PRO SER THR TRP TYR VAL A C G U DA DC DG DT".split() +) + +NORMS = {"layer_norm": nn.LayerNorm, "rms_norm": nn.RMSNorm} + + +def prepare_coord_group_seq( + features: Mapping[str, Any], generators: Sequence[torch.Generator] | None = None +) -> torch.Tensor | None: + """Randomly relabel the coordinate groups of each design. + + Draws from one generator per batch row when given, and from the global stream otherwise. + """ + values = features.get(CONDITION_KEYS["coordinate_group"]) + if values is None: + return None + values = values.long().clone() + rows = list(generators) if generators is not None else [None] * values.shape[0] + for row, generator in enumerate(rows): + present = values[row] != -1 + unique, inverse = torch.unique(values[row][present], return_inverse=True) + device = generator.device if generator is not None else values.device + permutation = torch.randperm(len(unique), device=device, generator=generator) + values[row][present] = permutation.to(values.device)[inverse] + return values.float() + + +def get_rfd4_time_embedding(values: torch.Tensor, width: int, max_positions: int = 2000) -> torch.Tensor: + """Sinusoidal embedding of times in [0, 1], in the original model's operation order. + + This is the same function as [`~models.embeddings.get_timestep_embedding`] called with `flip_sin_to_cos=False`, + `downscale_freq_shift=1` and `max_period=max_positions`, and differs from it only in float32 rounding: this scales + the time by `max_positions` before multiplying by the frequencies, where the shared helper scales afterwards. + Because the sine arguments reach `max_positions` radians, that reordering moves the embedding by up to 1.2e-4, + about a thousand times float32 eps. + + Over the released 60-step recipe that compounds to 0.17 angstrom of maximum coordinate deviation (0.004 angstrom + mean) with identical sequence output, so it is small but not noise, and it would cost the exact parity this port + holds against the original runtime. Kept deliberately. + """ + values = values.float() * max_positions + half = width // 2 + frequencies = torch.exp( + torch.arange(half, dtype=torch.float32, device=values.device) * (-math.log(max_positions) / (half - 1)) + ) + result = torch.cat((torch.sin(values[:, None] * frequencies), torch.cos(values[:, None] * frequencies)), dim=-1) + return F.pad(result, (0, width % 2)) + + +def align_chirality( + chirality_vectors: torch.Tensor, + features: Mapping[str, Any], +) -> torch.Tensor: + """Orient packed stereochemical vectors towards the current noisy backbone. + + Checkpoints trained with `noise_align_chirality` see the chirality feature rotated onto the noisy neighbour + directions rather than in its canonical frame, so the two are different inputs to the same weights. Needs + `chiral_frame_atoms` `(batch, tokens, slots, 2 + K)` and `chiral_slot_owner` `(batch, tokens, slots)` alongside the + `chirality` feature. + + Args: + chirality_vectors (`torch.Tensor`): Packed unit vectors `(batch, tokens, 3 * slots)`. + features (`Mapping`): Prepared tensors, including the noisy state under `x_t`. + + Returns: + `torch.Tensor`: Aligned vectors; slots without a frame keep their input value. + """ + missing = sorted({"chiral_frame_atoms", "chiral_slot_owner"} - features.keys()) + if missing: + raise KeyError( + f"noise_align_chirality needs {missing} in the prepared features. The released checkpoint was trained " + "with chirality vectors oriented towards the noisy backbone; without these tensors the model would " + "silently receive a different conditioning signal." + ) + batch_size, num_tokens, width = chirality_vectors.shape + num_slots = width // 3 + anchors = features["x_t"]["bb_ca"].float() + valid = features["padding_mask"].bool() + frames = features["chiral_frame_atoms"] + rows, tokens, slots = torch.where((frames[..., 0] >= 0) & valid[..., None]) + if rows.numel() == 0: + return chirality_vectors * valid[..., None] + origins, neighbors = frames[rows, tokens, slots, :2], frames[rows, tokens, slots, 2:] + present = neighbors >= 0 + neighbors = neighbors.clamp_min(0) + packed = chirality_vectors.reshape(batch_size, num_tokens, num_slots, 3) + supplied = packed[rows[:, None], neighbors, slots[:, None]].float() + origin = anchors[rows[:, None], origins].mean(-2, keepdim=True) + noisy = anchors[rows[:, None], neighbors] - origin + with torch.autocast(device_type=chirality_vectors.device.type, enabled=False): + # Owner ids map every slot, including copies, onto the frame that defines it; zero is "no frame". + frame_index = torch.zeros(batch_size, num_tokens * num_slots, device=packed.device, dtype=torch.long) + frame_index[rows, tokens * num_slots + slots] = torch.arange(1, len(rows) + 1, device=packed.device) + owners = features["chiral_slot_owner"] + slot_frames = frame_index[torch.arange(batch_size, device=packed.device)[:, None, None], owners.clamp_min(0)] + supplied = torch.where(present[..., None], supplied, 0.0) + noisy = torch.where(present[..., None], noisy, 0.0) + noisy = noisy / torch.linalg.vector_norm(noisy, dim=-1, keepdim=True).clamp_min(1e-8) + # Fit proper rotations by SVD of the covariance, applied as `source @ rotation`. Where the fit is not unique the + # arbitrary optimum SVD returns is used, matching the reference implementation. + u, singular, vh = torch.linalg.svd((supplied.transpose(-1, -2) @ noisy).float()) + signs = torch.ones_like(singular) + signs[..., -1] = torch.linalg.det(u @ vh).sign() + rotation = (u * signs[..., None, :]) @ vh + rotations = torch.cat([torch.eye(3, device=packed.device)[None], rotation]) + aligned = torch.einsum("bnsx,bnsxy->bnsy", packed.float(), rotations[slot_frames]) + aligned = torch.where((owners >= 0)[..., None], aligned, packed) + return aligned.reshape_as(chirality_vectors) * valid[..., None] + + +def feature_dimensions( + t_emb_dim: int, + latent_dim: int, + token_type_emb_dim: int, + recycle_conditioning_dim: int, + cath_emb_dim: int, + seq_sep_dim: int, + xt_pair_dist_dim: int, +) -> dict[str, int]: + """Width of every denoiser input feature. Several are set by the checkpoint configuration, so the table is + built per model.""" + return { + "time_emb_bb_ca": t_emb_dim, + "time_emb_local_latents": t_emb_dim, + "recycle_remaining": recycle_conditioning_dim, + "token_type": token_type_emb_dim, + "xt_bb_ca": 3, + "xt_local_latents": latent_dim, + "rasa_binned": 37, + "chain_centroid_xyz": 3, + "secondary_structure": 16, + "hbond_modernized": 74, + "pistack": 37, + "pication": 74, + "saltbridge": 37, + "hydrophobic": 37, + "metal_coord": 37, + "hotspot": 37, + "sequence_probabilities": 56, + "na_ss_state": 2, + "coord_group_seq": 37, + "atom_charge": 16, + "atomic_element": 16, + "chirality": 18, + "cond_a37coors_nm": 148, + "cond_a37coors_internal_nm": 148, + "token_bond_path_length": seq_sep_dim, + "xt_bb_ca_pair_dists": xt_pair_dist_dim, + "molecule_idx_pair": 1, + "molecule_entity_pair": 1, + "index_group_pair": 1, + "bond_mask": 1, + "bond_order": 5, + "is_aromatic": 1, + "na_ss_pair": 2, + "canonical_polymer_ordering": 1, + "cath_code": cath_emb_dim * 3, + } + + +def _zeros(features: Mapping[str, Any], dim: int, pair: bool = False) -> torch.Tensor: + mask = features["padding_mask"] + shape = (*mask.shape, mask.shape[-1], dim) if pair else (*mask.shape, dim) + return torch.zeros(shape, device=mask.device) + + +def _time_embedding_bb_ca(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return get_rfd4_time_embedding(features["t"]["bb_ca"], dim) + + +def _time_embedding_local_latents(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return get_rfd4_time_embedding(features["t"]["local_latents"], dim) + + +def _recycle_remaining(features: Mapping[str, Any], dim: int) -> torch.Tensor: + # One-hot over the passes still to come, saturating at the trained width. All-zero means the conditioning is + # absent, which is distinct from "no passes remain". + mask = features["padding_mask"] + remaining = features.get("recycle_remaining") + if remaining is None: + return torch.zeros((mask.shape[0], dim), device=mask.device) + index = torch.full((mask.shape[0],), min(int(remaining), dim - 1), device=mask.device) + return F.one_hot(index.long(), dim).float() + + +def _xt_bb_ca(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return features["x_t"]["bb_ca"] + + +def _xt_local_latents(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return features["x_t"]["local_latents"] + + +def _cond_a37coors_nm(features: Mapping[str, Any], dim: int) -> torch.Tensor: + values = features.get(CONDITION_KEYS["coordinates"]) + if values is None: + return _zeros(features, dim) + mask = torch.isfinite(values).any(dim=-1) + coords = torch.nan_to_num(values) / 10.0 * mask[..., None] + return torch.cat((coords.flatten(2), mask.float()), dim=-1) + + +def _cond_a37coors_internal_nm(features: Mapping[str, Any], dim: int) -> torch.Tensor: + values = features.get(CONDITION_KEYS["internal_coordinates"]) + groups = features.get(CONDITION_KEYS["coordinate_group"]) + if values is None or groups is None: + return _zeros(features, dim) + mask = (groups != -1) & (groups != 0) + coords = torch.nan_to_num(values) / 10.0 * mask[..., None] + return torch.cat((coords.flatten(2), mask.float()), dim=-1) + + +def _rasa_binned(features: Mapping[str, Any], dim: int) -> torch.Tensor: + values = features.get(CONDITION_KEYS["rasa"]) + return _zeros(features, dim) if values is None else values.float() * features["atom37_mask"].float() + + +def _chain_centroid_xyz(features: Mapping[str, Any], dim: int) -> torch.Tensor: + values = features.get(CONDITION_KEYS["centroid"]) + if values is None: + return _zeros(features, dim) + finite = torch.isfinite(values).all(dim=-1, keepdim=True) + return torch.where(finite, values.float() / 10.0, torch.zeros_like(values, dtype=torch.float32)) + + +def _atom_condition( + features: Mapping[str, Any], dim: int, name: str, categorical: bool = False, required: bool = True +) -> torch.Tensor: + value = features.get(CONDITION_KEYS[name]) + if value is None: + if required: + raise KeyError(f"required conditioning annotation for {name} is missing") + return _zeros(features, dim) + mask = features["atom37_mask"].float() + if name in {"hbond", "pication"}: + first = ((value == 1) | (value == 3)).float() * mask + second = ((value == 2) | (value == 3)).float() * mask + return torch.cat((first, second), dim=-1) + return ((value != 0).float() if categorical else value.float()) * mask + + +def _hbond_modernized(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _atom_condition(features, dim, "hbond") + + +def _pistack(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _atom_condition(features, dim, "pistack") + + +def _pication(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _atom_condition(features, dim, "pication") + + +def _saltbridge(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _atom_condition(features, dim, "saltbridge", categorical=True) + + +def _hydrophobic(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _atom_condition(features, dim, "hydrophobic") + + +def _metal_coord(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _atom_condition(features, dim, "metal_coord", categorical=True) + + +def _hotspot(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _atom_condition(features, dim, "hotspot", required=False) + + +def _sequence_probabilities(features: Mapping[str, Any], dim: int) -> torch.Tensor: + by_token = {} + for key, value in features.items(): + if "sequence-probability" in key and key.endswith("_1_residue"): + by_token.setdefault(key.rsplit("-", 1)[-1].removesuffix("_1_residue"), value) + missing = torch.full(features["padding_mask"].shape, torch.nan, device=features["padding_mask"].device) + probabilities = torch.stack( + [by_token.get(token, missing) for token in SEQUENCE_PROBABILITY_TOKENS], dim=-1 + ).float() + present = ~probabilities.isnan() + return torch.cat((probabilities.masked_fill(~present, 0), present.float()), dim=-1) + + +def _na_ss_state(features: Mapping[str, Any], dim: int) -> torch.Tensor: + if "na_ss_state" not in features: + return _zeros(features, dim) + return torch.stack((features["na_ss_state"].float(), features["na_ss_state_mask"].float()), dim=-1) + + +def _coord_group_seq(features: Mapping[str, Any], dim: int) -> torch.Tensor: + values = features.get("_prepared_coord_group_seq") + return _zeros(features, dim) if values is None else values + + +def _chirality(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return features["chirality"].float() if "chirality" in features else _zeros(features, dim) + + +def _token_bond_path_length(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return F.one_hot(features["token_bond_path_length"].long().clamp(0, dim - 1), dim).float() + + +def _pair_match(values: torch.Tensor) -> torch.Tensor: + return (values[:, :, None] == values[:, None, :]).float().unsqueeze(-1) + + +def _molecule_idx_pair(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _pair_match(features["molecule_iid"]) + + +def _molecule_entity_pair(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _pair_match(features["molecule_entity_id"]) + + +def _index_group_pair(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _pair_match(features.get(CONDITION_KEYS["index_group"])) + + +def _pair_values(features: Mapping[str, Any], name: str, dim: int) -> torch.Tensor: + return features[name] if name in features else _zeros(features, dim, pair=True)[..., 0] + + +def _bond_mask(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _pair_values(features, "bond_mask", dim).float().unsqueeze(-1) + + +def _is_aromatic(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _pair_values(features, "is_aromatic", dim).float().unsqueeze(-1) + + +def _canonical_polymer_ordering(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return _pair_values(features, "canonical_polymer_ordering", dim).float().unsqueeze(-1) + + +def _bond_order(features: Mapping[str, Any], dim: int) -> torch.Tensor: + return F.one_hot(_pair_values(features, "bond_order", dim).long().clamp(0, 4), 5).float() + + +def _na_ss_pair(features: Mapping[str, Any], dim: int) -> torch.Tensor: + if "na_ss_pair" not in features: + return _zeros(features, dim, pair=True) + return torch.stack((features["na_ss_pair"].float(), features["na_ss_pair_mask"].float()), dim=-1) + + +# Features without learned weights: plain functions of the prepared tensors and the feature width. +FEATURE_FUNCTIONS = { + "time_emb_bb_ca": _time_embedding_bb_ca, + "time_emb_local_latents": _time_embedding_local_latents, + "recycle_remaining": _recycle_remaining, + "xt_bb_ca": _xt_bb_ca, + "xt_local_latents": _xt_local_latents, + "cond_a37coors_nm": _cond_a37coors_nm, + "cond_a37coors_internal_nm": _cond_a37coors_internal_nm, + "rasa_binned": _rasa_binned, + "chain_centroid_xyz": _chain_centroid_xyz, + "hbond_modernized": _hbond_modernized, + "pistack": _pistack, + "pication": _pication, + "saltbridge": _saltbridge, + "hydrophobic": _hydrophobic, + "metal_coord": _metal_coord, + "hotspot": _hotspot, + "sequence_probabilities": _sequence_probabilities, + "na_ss_state": _na_ss_state, + "coord_group_seq": _coord_group_seq, + "chirality": _chirality, + "token_bond_path_length": _token_bond_path_length, + "molecule_idx_pair": _molecule_idx_pair, + "molecule_entity_pair": _molecule_entity_pair, + "index_group_pair": _index_group_pair, + "bond_mask": _bond_mask, + "is_aromatic": _is_aromatic, + "canonical_polymer_ordering": _canonical_polymer_ordering, + "bond_order": _bond_order, + "na_ss_pair": _na_ss_pair, +} + + +class _TokenTypeEmbedding(nn.Embedding): + def __init__(self, dim: int, num_token_types: int, **kwargs) -> None: + super().__init__(num_token_types, dim) + + def forward(self, features: Mapping[str, Any]) -> torch.Tensor: + return super().forward(features["token_type"].long()) + + +class _SecondaryStructureEmbedding(nn.Embedding): + def __init__(self, dim: int, **kwargs) -> None: + super().__init__(4, dim, padding_idx=0) + + def forward(self, features: Mapping[str, Any]) -> torch.Tensor: + values = features.get(CONDITION_KEYS["secondary_structure"]) + return _zeros(features, self.embedding_dim) if values is None else super().forward(values.long()) + + +class _AtomChargeEmbedding(nn.Embedding): + def __init__(self, dim: int, **kwargs) -> None: + super().__init__(11, dim) + + def forward(self, features: Mapping[str, Any]) -> torch.Tensor: + if "charge" not in features: + return _zeros(features, self.embedding_dim) + values = (features["charge"][:, :, 1].long() + 4).clamp(0, 10) + return super().forward(values) * (features["token_type"] == 3)[..., None] + + +class _AtomicElementEmbedding(nn.Embedding): + def __init__(self, dim: int, **kwargs) -> None: + super().__init__(128, dim) + + def forward(self, features: Mapping[str, Any]) -> torch.Tensor: + if "element" not in features: + return _zeros(features, self.embedding_dim) + values = features["element"][:, :, 1].long().clamp(0, self.num_embeddings - 1) + return super().forward(values) * (features["token_type"] == 3)[..., None] + + +class _CathEmbedding(nn.Module): + def __init__(self, dim: int, cath_emb_dim: int, **kwargs) -> None: + super().__init__() + self.dim = dim + self.c = nn.Embedding(6, cath_emb_dim) + self.a = nn.Embedding(44, cath_emb_dim) + self.t = nn.Embedding(1473, cath_emb_dim) + + def forward(self, features: Mapping[str, Any]) -> torch.Tensor: + values = features.get(CONDITION_KEYS["cath"]) + if values is None: + return _zeros(features, self.dim) + values = values.long() + 1 + return torch.cat((self.c(values[..., 0]), self.a(values[..., 1]), self.t(values[..., 2])), dim=-1) + + +class _PairDistances(nn.Module): + """One-hot bins of the pairwise distances of the noisy backbone. It has no weights; it needs the bin range.""" + + def __init__(self, dim: int, xt_pair_dist_min: float, xt_pair_dist_max: float, **kwargs) -> None: + super().__init__() + self.dim, self.limits = dim, (xt_pair_dist_min, xt_pair_dist_max) + + def forward(self, features: Mapping[str, Any]) -> torch.Tensor: + coords = features["x_t"]["bb_ca"] + distances = torch.linalg.norm(coords[:, :, None] - coords[:, None, :], dim=-1) + limits = torch.linspace(*self.limits, self.dim - 1, device=coords.device) + return F.one_hot(torch.bucketize(distances, limits), self.dim).float() + + +# The features that need a module: the ones with learned weights, and the one with a setting of its own. +FEATURE_MODULES = { + "token_type": _TokenTypeEmbedding, + "secondary_structure": _SecondaryStructureEmbedding, + "atom_charge": _AtomChargeEmbedding, + "atomic_element": _AtomicElementEmbedding, + "cath_code": _CathEmbedding, + "xt_bb_ca_pair_dists": _PairDistances, +} + + +class RFD4ProteinaFeatureFactory(nn.Module): + """Concatenates named features and projects them to `width`, computing the projection in the model dtype. + + Each configured feature is resolved once, to a function (`FEATURE_FUNCTIONS`) or to a module (`FEATURE_MODULES`, + kept in `self.feature_modules`). + """ + + def __init__( + self, + names: Sequence[str] | None, + width: int, + mode: Literal["seq", "pair"], + t_emb_dim: int, + latent_dim: int, + token_type_emb_dim: int, + num_token_types: int, + recycle_conditioning_dim: int, + cath_emb_dim: int, + seq_sep_dim: int, + xt_pair_dist_dim: int, + xt_pair_dist_min: float, + xt_pair_dist_max: float, + normalize: bool = False, + compact: bool = False, + ) -> None: + super().__init__() + self.mode, self.width = mode, width + self.compact = compact + self.names = list(names or ()) + dimensions = feature_dimensions( + t_emb_dim, + latent_dim, + token_type_emb_dim, + recycle_conditioning_dim, + cath_emb_dim, + seq_sep_dim, + xt_pair_dist_dim, + ) + self.feature_modules = nn.ModuleDict() + self.feature_fns = [] + for name in self.names: + if name not in dimensions: + raise ValueError(f"unsupported RFD4 feature: {name}") + if name in FEATURE_FUNCTIONS: + self.feature_fns.append(partial(FEATURE_FUNCTIONS[name], dim=dimensions[name])) + else: + module = FEATURE_MODULES[name]( + dimensions[name], + num_token_types=num_token_types, + cath_emb_dim=cath_emb_dim, + xt_pair_dist_min=xt_pair_dist_min, + xt_pair_dist_max=xt_pair_dist_max, + ) + self.feature_modules[name] = module + self.feature_fns.append(module) + if self.names: + self.ln_out = nn.RMSNorm(width, eps=1e-5) if normalize else nn.Identity() + self.linear_out = nn.Linear(sum(dimensions[name] for name in self.names), width, bias=False) + + def forward(self, features: Mapping[str, Any]) -> torch.Tensor: + mask = features["padding_mask"].bool() + if not self.names: + shape = (*mask.shape, mask.shape[-1], self.width) if self.mode == "pair" else (*mask.shape, self.width) + return torch.zeros(shape, device=mask.device, dtype=features["compute_dtype"]) + spatial = (mask.shape[-1],) if self.mode == "seq" else (mask.shape[-1], mask.shape[-1]) + tensors = [feature(features) for feature in self.feature_fns] + if self.compact and all(value.ndim == 2 for value in tensors): + active = mask.any(dim=-1) + else: + tensors = [ + value.reshape(value.shape[0], *((1,) * len(spatial)), -1).expand(-1, *spatial, -1) + if value.ndim == 2 + else value + for value in tensors + ] + active = mask if self.mode == "seq" else mask[:, :, None] & mask[:, None, :] + stacked = torch.cat(tensors, dim=-1) * active[..., None] + compute_dtype = features["compute_dtype"] + # The projection may be kept in float32 while the model runs in lower precision. A quantized weight is an integer + # tensor, so its input follows the compute dtype instead. + weight_dtype = get_parameter_dtype(self.linear_out) + projected = self.linear_out(stacked.to(weight_dtype if weight_dtype.is_floating_point else compute_dtype)) + return (self.ln_out(projected) * active[..., None]).to(compute_dtype) + + +class RFD4ProteinaTransition(nn.Module): + def __init__(self, width: int, expansion: int = 4) -> None: + super().__init__() + self.swish_linear = nn.Sequential(nn.Linear(width, 2 * width * expansion, bias=False)) + self.linear_out = nn.Linear(width * expansion, width, bias=False) + + def forward(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + left, gate = self.swish_linear(values).chunk(2, dim=-1) + return self.linear_out(F.silu(gate) * left) * mask[..., None] + + +class RFD4ProteinaAdaptiveNorm(nn.Module): + def __init__(self, width: int, condition_width: int, norm_type: str) -> None: + super().__init__() + self.norm = NORMS[norm_type](width, eps=1e-5, elementwise_affine=False) + self.norm_cond = NORMS[norm_type](condition_width, eps=1e-5) + self.to_gamma = nn.Sequential(nn.Linear(condition_width, width), nn.Sigmoid()) + self.to_beta = nn.Linear(condition_width, width, bias=False) + + def forward(self, values: torch.Tensor, condition: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + normalized = self.norm(values) + normalized_condition = self.norm_cond(condition) + gamma, beta = self.to_gamma(normalized_condition), self.to_beta(normalized_condition) + if condition.ndim == 2 and values.ndim > 2: + shape = (condition.shape[0], *((1,) * (values.ndim - 2)), values.shape[-1]) + gamma, beta = gamma.reshape(shape), beta.reshape(shape) + return (normalized * gamma + beta) * mask[..., None] + + +class RFD4ProteinaPairAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "RFD4ProteinaPairAttention", + hidden_states: torch.Tensor, + pairs: torch.Tensor | None, + mask: torch.Tensor, + ) -> torch.Tensor: + batch, tokens, width = hidden_states.shape + hidden_states = attn.node_norm(hidden_states) + query, key = attn.to_q(hidden_states), attn.to_k(hidden_states) + query, key = attn.norm_q(query), attn.norm_k(key) + value, gate = attn.to_v(hidden_states), attn.to_g(hidden_states) + query, key, value, gate = ( + tensor.view(batch, tokens, attn.heads, attn.head_width) for tensor in (query, key, value, gate) + ) + pair_mask = mask[:, None, :, None] & mask[:, None, None, :] + if attn.use_pair: + bias = attn.to_bias(attn.pair_norm(pairs)).permute(0, 3, 1, 2).masked_fill(~pair_mask, -1e4) + else: + bias = torch.zeros(batch, attn.heads, tokens, tokens, device=query.device, dtype=query.dtype) + bias = bias.masked_fill(~pair_mask, -1e4) + output = dispatch_attention_fn( + query, key, value, attn_mask=bias, backend=self._attention_backend, parallel_config=self._parallel_config + ) + output = (output * gate.sigmoid()).reshape(batch, tokens, width) + return attn.to_out(output) + + +class RFD4ProteinaPairAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = RFD4ProteinaPairAttnProcessor + _available_processors = [RFD4ProteinaPairAttnProcessor] + _supports_qkv_fusion = False + + def __init__(self, width: int, pair_width: int, heads: int, norm_type: str, use_pair: bool) -> None: + super().__init__() + self.heads, self.head_width, self.use_pair = heads, width // heads, use_pair + self.node_norm = NORMS[norm_type](width, eps=1e-5) + self.to_q = nn.Linear(width, width) + self.to_k = nn.Linear(width, width) + self.to_v = nn.Linear(width, width) + self.to_g = nn.Linear(width, width) + self.norm_q = NORMS[norm_type](width, eps=1e-5) + self.norm_k = NORMS[norm_type](width, eps=1e-5) + self.to_out = nn.Linear(width, width) + if use_pair: + self.to_bias = nn.Linear(pair_width, heads, bias=False) + self.pair_norm = NORMS[norm_type](pair_width, eps=1e-5) + self.set_processor(RFD4ProteinaPairAttnProcessor()) + + def forward(self, values: torch.Tensor, pairs: torch.Tensor | None, mask: torch.Tensor) -> torch.Tensor: + return self.processor(self, values, pairs, mask) + + +class RFD4ProteinaAdaptiveAttention(nn.Module): + def __init__( + self, width: int, pair_width: int, heads: int, condition_width: int, norm_type: str, use_pair: bool + ) -> None: + super().__init__() + self.adaln = RFD4ProteinaAdaptiveNorm(width, condition_width, norm_type) + self.mha = RFD4ProteinaPairAttention(width, pair_width, heads, norm_type, use_pair) + self.scale_output = nn.Sequential(nn.Linear(condition_width, width), nn.Sigmoid()) + + def forward( + self, values: torch.Tensor, pairs: torch.Tensor | None, condition: torch.Tensor, mask: torch.Tensor + ) -> torch.Tensor: + output = self.mha(self.adaln(values, condition, mask), pairs, mask) + return output * self.scale_output(condition) * mask[..., None] + + +class RFD4ProteinaAdaptiveTransition(nn.Module): + def __init__(self, width: int, condition_width: int, norm_type: str) -> None: + super().__init__() + self.adaln = RFD4ProteinaAdaptiveNorm(width, condition_width, norm_type) + self.transition = RFD4ProteinaTransition(width) + self.scale_output = nn.Sequential(nn.Linear(condition_width, width), nn.Sigmoid()) + + def forward(self, values: torch.Tensor, condition: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + output = self.transition(self.adaln(values, condition, mask), mask) + return output * self.scale_output(condition) * mask[..., None] + + +class RFD4ProteinaTrunkLayer(nn.Module): + """One attention + transition block, conditioned on `condition` through adaptive normalization.""" + + def __init__( + self, token_dim: int, pair_repr_dim: int, nheads: int, dim_cond: int, norm_type: str, use_pair: bool + ) -> None: + super().__init__() + self.mhba = RFD4ProteinaAdaptiveAttention(token_dim, pair_repr_dim, nheads, dim_cond, norm_type, use_pair) + self.transition = RFD4ProteinaAdaptiveTransition(token_dim, dim_cond, norm_type) + + def forward( + self, values: torch.Tensor, pairs: torch.Tensor | None, condition: torch.Tensor, mask: torch.Tensor + ) -> torch.Tensor: + values = values * mask[..., None] + values = (values + self.mhba(values, pairs, condition, mask)) * mask[..., None] + return (values + self.transition(values, condition, mask)) * mask[..., None] + + +class RFD4ProteinaTriangleMultiplication(nn.Module): + def __init__(self, width: int, hidden: int, outgoing: bool, bias: bool) -> None: + super().__init__() + self.outgoing = outgoing + self.linear_a_p = nn.Linear(width, hidden, bias=bias) + self.linear_a_g = nn.Linear(width, hidden, bias=bias) + self.linear_b_p = nn.Linear(width, hidden, bias=bias) + self.linear_b_g = nn.Linear(width, hidden, bias=bias) + self.linear_g = nn.Linear(width, width, bias=bias) + self.linear_z = nn.Linear(hidden, width, bias=bias) + self.layer_norm_in = nn.LayerNorm(width, eps=1e-5, bias=bias) + self.layer_norm_out = nn.LayerNorm(hidden, eps=1e-5, bias=bias) + + def forward(self, pairs: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor: + normalized = self.layer_norm_in(pairs) + left = self.linear_a_p(normalized) * self.linear_a_g(normalized).sigmoid() * pair_mask[..., None] + right = self.linear_b_p(normalized) * self.linear_b_g(normalized).sigmoid() * pair_mask[..., None] + if self.outgoing: + output = torch.matmul(left.permute(0, 3, 1, 2), right.permute(0, 3, 2, 1)) + else: + output = torch.matmul(left.permute(0, 3, 2, 1), right.permute(0, 3, 1, 2)) + output = self.linear_z(self.layer_norm_out(output.permute(0, 2, 3, 1))) + return output * self.linear_g(normalized).sigmoid() + + +class RFD4ProteinaPairTransition(nn.Module): + def __init__(self, width: int) -> None: + super().__init__() + self.layer_norm = nn.LayerNorm(width, eps=1e-5) + self.linear_1 = nn.Linear(width, 2 * width) + self.linear_2 = nn.Linear(2 * width, width) + + def forward(self, pairs: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor: + return self.linear_2(F.relu(self.linear_1(self.layer_norm(pairs)))) * pair_mask[..., None] + + +class RFD4ProteinaPairUpdate(nn.Module): + def __init__( + self, + token_dim: int, + pair_repr_dim: int, + use_tri_mult: bool, + tri_mult_c: int, + tri_mult_bias: bool, + ) -> None: + super().__init__() + self.layer_norm_in = nn.LayerNorm(token_dim) + self.linear_x = nn.Linear(token_dim, 2 * pair_repr_dim, bias=False) + self.use_tri_mult = use_tri_mult + if use_tri_mult: + self.tri_mult_out = RFD4ProteinaTriangleMultiplication(pair_repr_dim, tri_mult_c, True, tri_mult_bias) + self.tri_mult_in = RFD4ProteinaTriangleMultiplication(pair_repr_dim, tri_mult_c, False, tri_mult_bias) + self.transition_out = RFD4ProteinaPairTransition(pair_repr_dim) + + def forward(self, values: torch.Tensor, pairs: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + pair_mask = mask[:, :, None] & mask[:, None, :] + left, right = self.linear_x(self.layer_norm_in(values * mask[..., None])).chunk(2, dim=-1) + pairs = (pairs + left[:, None] * right[:, :, None]) * pair_mask[..., None] + if self.use_tri_mult: + pairs = (pairs + self.tri_mult_out(pairs, pair_mask)) * pair_mask[..., None] + pairs = (pairs + self.tri_mult_in(pairs, pair_mask)) * pair_mask[..., None] + return (pairs + self.transition_out(pairs, pair_mask)) * pair_mask[..., None] + + +class RFD4ProteinaPairBuilder(nn.Module): + def __init__( + self, + feats_pair_repr: Sequence[str], + feats_pair_cond: Sequence[str], + pair_repr_dim: int, + dim_cond: int, + feature_kwargs: Mapping[str, Any], + ) -> None: + super().__init__() + self.init_repr_factory = RFD4ProteinaFeatureFactory( + feats_pair_repr, pair_repr_dim, "pair", **feature_kwargs, normalize=True + ) + self.cond_factory = None + self.adaln = None + if feats_pair_cond: + self.cond_factory = RFD4ProteinaFeatureFactory( + feats_pair_cond, dim_cond, "pair", **feature_kwargs, normalize=True, compact=True + ) + self.adaln = RFD4ProteinaAdaptiveNorm(pair_repr_dim, dim_cond, "rms_norm") + + def forward(self, features: Mapping[str, Any]) -> torch.Tensor: + pairs = self.init_repr_factory(features) + if self.cond_factory is not None: + mask = features["padding_mask"].bool() + pairs = self.adaln(pairs, self.cond_factory(features), mask[:, :, None] & mask[:, None, :]) + return pairs + + +@dataclass +class RFD4ProteinaTransformerOutput(BaseOutput): + """Denoiser predictions before scheduler integration. + + Args: + backbone (`torch.Tensor`): Backbone velocity shaped `(batch, tokens, 3)`, in nanometers per unit time. + local_latents (`torch.Tensor`): Latent velocity shaped `(batch, tokens, latent_dim)`. + sequence_logits (`torch.Tensor`, *optional*): Residue logits shaped `(batch, tokens, 147)`. + insertion_logits (`torch.Tensor`, *optional*): Insertion assignment logits shaped `(batch, tokens, tokens)`. + """ + + backbone: torch.Tensor + local_latents: torch.Tensor | None = None + sequence_logits: torch.Tensor | None = None + insertion_logits: torch.Tensor | None = None + + +class RFD4ProteinaTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin): + """RFD4 velocity denoiser with paired backbone and local-latent modalities. + + The model runs in its own dtype: load it with `dtype=torch.bfloat16` for reduced-precision inference. The sampling + state passed to `forward` stays in float32; feature embeddings are computed in float32 and cast at the input + projection. + + Args: + nlayers (`int`, defaults to `14`): Number of trunk layers. + token_dim (`int`, defaults to `768`): Token representation width. + nheads (`int`, defaults to `12`): Attention heads per layer. + dim_cond (`int`, defaults to `256`): Width of the conditioning vector (time, recycle index, CATH label). + pair_repr_dim (`int`, defaults to `384`): Pair representation width. + latent_dim (`int`, defaults to `32`): Width of the per-token local latent. + token_type_emb_dim (`int`, defaults to `32`): Token-type embedding width. + cath_emb_dim (`int`, defaults to `32`): Width of each CATH class/architecture/topology embedding. + t_emb_dim (`int`, defaults to `256`): Sinusoidal time embedding width. + seq_sep_dim (`int`, defaults to `64`): Number of bond-path-length bins. + xt_pair_dist_dim (`int`, defaults to `30`): Number of backbone pair-distance bins. + xt_pair_dist_min (`float`, defaults to `0.1`), xt_pair_dist_max (`float`, defaults to `3.0`): + Pair-distance range in nanometers. + feats_seq, feats_cond_seq, feats_pair_repr, feats_pair_cond (`list[str]`): + Names of the features projected into the token, conditioning, pair and pair-conditioning inputs. + update_pair_repr (`bool`, defaults to `True`), update_pair_repr_every_n (`int`, defaults to `2`): + Whether and how often the pair representation is updated from the tokens between layers. + use_tri_mult (`bool`, defaults to `True`), tri_mult_c (`int`, defaults to `256`), tri_mult_bias + (`bool`, defaults to `False`): Triangle multiplication settings. + recycle (`bool`, defaults to `True`): Whether the trunk is re-run `recycle_depth` times per call. + num_token_types (`int`, defaults to `7`): Size of the token-type vocabulary: protein, RNA, DNA, atomized, + unknown, glycan and non-canonical amino acid. + recycle_conditioning_dim (`int`, defaults to `1`): Width of the `recycle_remaining` one-hot, when that + feature is in `feats_cond_seq`. + noise_align_chirality (`bool`, defaults to `False`): Orient the chirality feature towards the noisy backbone. + Checkpoints trained this way need `chiral_frame_atoms` and `chiral_slot_owner` features. + num_sequence_classes (`int`, defaults to `147`): Size of the residue vocabulary the sequence readout predicts. + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["RFD4ProteinaTrunkLayer"] + _repeated_blocks = ["RFD4ProteinaTrunkLayer", "RFD4ProteinaPairUpdate"] + _skip_layerwise_casting_patterns = ["norm", "ln_out", "embedding"] + # Conditioning and normalization stay in full precision: the sampler integrates sixty steps and + # amplifies rounding carried in the time and recycle embeddings. + _keep_in_fp32_modules = ["cond_factory", "norm", "ln_out"] + + @register_to_config + def __init__( + self, + nlayers: int = 14, + token_dim: int = 768, + nheads: int = 12, + dim_cond: int = 256, + pair_repr_dim: int = 384, + latent_dim: int = 32, + token_type_emb_dim: int = 32, + cath_emb_dim: int = 32, + t_emb_dim: int = 256, + seq_sep_dim: int = 64, + xt_pair_dist_dim: int = 30, + xt_pair_dist_min: float = 0.1, + xt_pair_dist_max: float = 3.0, + feats_seq: tuple[str, ...] = ( + "xt_bb_ca", + "xt_local_latents", + "token_type", + "sequence_probabilities", + "cond_a37coors_nm", + "cond_a37coors_internal_nm", + "rasa_binned", + "secondary_structure", + "hbond_modernized", + "pistack", + "pication", + "saltbridge", + "hydrophobic", + "metal_coord", + "na_ss_state", + "hotspot", + "atom_charge", + "atomic_element", + "coord_group_seq", + "chirality", + "chain_centroid_xyz", + ), + feats_cond_seq: tuple[str, ...] = ( + "time_emb_bb_ca", + "time_emb_local_latents", + "cath_code", + "recycle_remaining", + ), + feats_pair_repr: tuple[str, ...] = ( + "xt_bb_ca_pair_dists", + "bond_mask", + "bond_order", + "is_aromatic", + "index_group_pair", + "molecule_idx_pair", + "molecule_entity_pair", + "token_bond_path_length", + "canonical_polymer_ordering", + "na_ss_pair", + ), + feats_pair_cond: tuple[str, ...] = ("time_emb_bb_ca", "time_emb_local_latents"), + update_pair_repr: bool = True, + update_pair_repr_every_n: int = 2, + use_tri_mult: bool = True, + tri_mult_c: int = 256, + tri_mult_bias: bool = False, + recycle: bool = True, + num_sequence_classes: int = 147, + num_token_types: int = 7, + recycle_conditioning_dim: int = 1, + noise_align_chirality: bool = False, + ) -> None: + super().__init__() + feature_kwargs = { + "t_emb_dim": t_emb_dim, + "latent_dim": latent_dim, + "token_type_emb_dim": token_type_emb_dim, + "num_token_types": num_token_types, + "recycle_conditioning_dim": recycle_conditioning_dim, + "cath_emb_dim": cath_emb_dim, + "seq_sep_dim": seq_sep_dim, + "xt_pair_dist_dim": xt_pair_dist_dim, + "xt_pair_dist_min": xt_pair_dist_min, + "xt_pair_dist_max": xt_pair_dist_max, + } + self.nlayers = nlayers + self.update_pair_repr = update_pair_repr + self.gradient_checkpointing = False + self.init_repr_factory = RFD4ProteinaFeatureFactory(feats_seq, token_dim, "seq", **feature_kwargs) + self.cond_factory = RFD4ProteinaFeatureFactory(feats_cond_seq, dim_cond, "seq", **feature_kwargs) + self.transition_c_1 = RFD4ProteinaTransition(dim_cond, 2) + self.transition_c_2 = RFD4ProteinaTransition(dim_cond, 2) + self.pair_repr_builder = RFD4ProteinaPairBuilder( + feats_pair_repr, feats_pair_cond, pair_repr_dim, dim_cond, feature_kwargs + ) + self.transformer_layers = nn.ModuleList( + RFD4ProteinaTrunkLayer(token_dim, pair_repr_dim, nheads, dim_cond, "rms_norm", True) + for _ in range(nlayers) + ) + if update_pair_repr: + # Keyed by the index of the trunk layer each update follows, so only those layers have an entry. + self.pair_update_layers = nn.ModuleDict( + { + str(index): RFD4ProteinaPairUpdate( + token_dim, pair_repr_dim, use_tri_mult, tri_mult_c, tri_mult_bias + ) + for index in range(0, nlayers - 1, update_pair_repr_every_n) + } + ) + self.recycle_enabled = recycle + # The released checkpoint conditions the trunk on the passes still to come, so the conditioning vector is rebuilt + # for every pass instead of once. + self.condition_on_recycle_remaining = "recycle_remaining" in feats_cond_seq + self.noise_align_chirality = noise_align_chirality + if recycle: + self.s_cycle_norm = nn.RMSNorm(token_dim, eps=1e-5) + self.s_cycle_proj = nn.Linear(token_dim, token_dim, bias=False) + self.z_cycle_norm = nn.RMSNorm(pair_repr_dim, eps=1e-5) + self.z_cycle_proj = nn.Linear(pair_repr_dim, pair_repr_dim, bias=False) + self.local_latents_linear = nn.Sequential( + nn.RMSNorm(token_dim, eps=1e-5), nn.Linear(token_dim, latent_dim, bias=False) + ) + self.ca_linear = nn.Sequential(nn.RMSNorm(token_dim, eps=1e-5), nn.Linear(token_dim, 3, bias=False)) + self.sequence_readout = nn.Sequential( + nn.RMSNorm(token_dim, eps=1e-5), nn.Linear(token_dim, num_sequence_classes, bias=False) + ) + self.insertion_q_proj = nn.Linear(token_dim, token_dim, bias=False) + self.insertion_k_proj = nn.Linear(token_dim, token_dim, bias=False) + self._insertion_head_scale = token_dim**-0.5 + + @apply_lora_scale("attention_kwargs") + def forward( + self, + bb_ca: torch.Tensor, + local_latents: torch.Tensor, + bb_ca_time: torch.Tensor, + local_latents_time: torch.Tensor, + padding_mask: torch.Tensor, + token_type: torch.Tensor | None = None, + atom37_mask: torch.Tensor | None = None, + conditioning: dict[str, Any] | None = None, + recycle_depth: int = 0, + attention_kwargs: dict[str, Any] | None = None, + return_dict: bool = True, + ) -> RFD4ProteinaTransformerOutput | tuple: + """Predict velocities from prepared molecular tensors without advancing the sampler. + + Args: + bb_ca (`torch.Tensor`): Noisy backbone coordinates `(batch, tokens, 3)` in nanometers. + local_latents (`torch.Tensor`): Noisy local latents `(batch, tokens, latent_dim)`. + bb_ca_time (`torch.Tensor`): `(batch,)` backbone time in [0, 1]. + local_latents_time (`torch.Tensor`): `(batch,)` local-latent time in [0, 1]. + padding_mask (`torch.Tensor`): `(batch, tokens)`, true where a token is real rather than padding. + 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. + conditioning (`dict[str, Any]`, *optional*): The remaining topology and conditioning tensors, + under the keys the preprocessor emits. Which of them are read is fixed by the checkpoint's `feats_seq`, + `feats_cond_seq`, `feats_pair_repr` and `feats_pair_cond`; a key a configured feature asks for but that + is absent leaves that feature unconditioned rather than raising. + recycle_depth (`int`, defaults to `0`): Number of trunk passes; zero still executes one pass. + attention_kwargs (`dict`, *optional*): Only `scale` is read, as the LoRA scale for this call. + return_dict (`bool`, defaults to `True`): Return a typed output, or a tuple of non-None fields. + + Returns: + [`RFD4ProteinaTransformerOutput`] or `tuple`: Backbone/local-latent velocities and optional + sequence/insertion logits, before scheduler integration. + """ + features = { + **(conditioning or {}), + "padding_mask": padding_mask, + "x_t": {"bb_ca": bb_ca, "local_latents": local_latents}, + "t": {"bb_ca": bb_ca_time, "local_latents": local_latents_time}, + "compute_dtype": self.dtype, + } + for name, value in (("token_type", token_type), ("atom37_mask", atom37_mask)): + if value is not None: + features[name] = value + mask = padding_mask.bool() + emitted = mask & ~features["is_in_context_conditioning"].bool() + # Checkpoints trained with `noise_align_chirality` see the chirality feature rotated onto the noisy neighbour + # frame, so the feature the trunk reads is replaced before any layer runs. + if self.noise_align_chirality and "chirality" in features: + features["chirality"] = align_chirality(features["chirality"].float(), features) + + cycles = max(1, recycle_depth) if self.recycle_enabled else 1 + values = self.init_repr_factory(features) * mask[..., None] + pairs = self.pair_repr_builder(features) + if self.recycle_enabled: + previous_values, previous_pairs = torch.zeros_like(values), torch.zeros_like(pairs) + for cycle in range(cycles): + if self.recycle_enabled: + values_in = values + self.s_cycle_proj(self.s_cycle_norm(previous_values)) + pairs_in = (pairs + self.z_cycle_proj(self.z_cycle_norm(previous_pairs))).contiguous() + else: + values_in, pairs_in = values, pairs + # Before the first pass the remaining-passes feature is absent (all-zero), which is distinct from "no passes + # remain". Later passes carry the passes still to come. + if cycle == 0 or self.condition_on_recycle_remaining: + condition = self.cond_factory( + { + **features, + "recycle_remaining": None if cycle == 0 else cycles - cycle - 1, + } + ) + condition = self.transition_c_2(self.transition_c_1(condition, mask), mask) + hidden, hidden_pairs = values_in, pairs_in + for index, layer in enumerate(self.transformer_layers): + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden = self._gradient_checkpointing_func(layer, hidden, hidden_pairs, condition, mask) + else: + hidden = layer(hidden, hidden_pairs, condition, mask) + if self.update_pair_repr and str(index) in self.pair_update_layers: + hidden_pairs = self.pair_update_layers[str(index)](hidden, hidden_pairs, mask) + previous_values, previous_pairs = hidden, hidden_pairs + values = previous_values + + backbone = self.ca_linear(values) * emitted[..., None] + local_latents = ( + self.local_latents_linear(values) * (emitted & features["supports_local_latent"].bool())[..., None] + ) + sequence_logits = self.sequence_readout(values.detach()) * mask[..., None] + detached = values.detach() + query, key = self.insertion_q_proj(detached), self.insertion_k_proj(detached) + insertion_logits = torch.bmm(query, key.transpose(1, 2)) * self._insertion_head_scale + output = RFD4ProteinaTransformerOutput( + backbone=backbone, + local_latents=local_latents, + sequence_logits=sequence_logits, + insertion_logits=insertion_logits, + ) + return output if return_dict else output.to_tuple() diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 4f1f600e4faf..a9d1f805c1e3 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -49,6 +49,7 @@ _import_structure["consistency_models"] = ["ConsistencyModelPipeline"] _import_structure["ddim"] = ["DDIMPipeline"] _import_structure["ddpm"] = ["DDPMPipeline"] + _import_structure["rfd4_proteina"] = ["RFD4ProteinaPipeline"] _import_structure["dit"] = ["DiTPipeline"] _import_structure["latent_diffusion"].extend(["LDMSuperResolutionPipeline"]) _import_structure["pipeline_utils"] = [ @@ -580,6 +581,7 @@ ImagePipelineOutput, StableDiffusionMixin, ) + from .rfd4_proteina import RFD4ProteinaPipeline try: if not (is_torch_available() and is_librosa_available()): diff --git a/src/diffusers/pipelines/rfd4_proteina/__init__.py b/src/diffusers/pipelines/rfd4_proteina/__init__.py new file mode 100644 index 000000000000..166ba1d00739 --- /dev/null +++ b/src/diffusers/pipelines/rfd4_proteina/__init__.py @@ -0,0 +1,35 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + is_torch_available, +) + + +_import_structure = {} + +try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["pipeline_rfd4_proteina"] = ["RFD4ProteinaPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .pipeline_rfd4_proteina import RFD4ProteinaPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/src/diffusers/pipelines/rfd4_proteina/pipeline_output.py b/src/diffusers/pipelines/rfd4_proteina/pipeline_output.py new file mode 100644 index 000000000000..a41890b40aa2 --- /dev/null +++ b/src/diffusers/pipelines/rfd4_proteina/pipeline_output.py @@ -0,0 +1,41 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +from dataclasses import dataclass + +import torch + +from ...utils import BaseOutput + + +@dataclass +class RFD4ProteinaPipelineOutput(BaseOutput): + """Molecular inference results. Coordinates are in nanometers. + + Args: + 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. + state (`dict[str, torch.Tensor]`): + Final product-space state: `bb_ca` coordinates `(B, N, 3)` and `local_latents` `(B, N, D)`. + clean (`dict[str, torch.Tensor]`): + Predicted clean product-space state from 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. + """ + + coords_nm: torch.Tensor + residue_type: torch.Tensor + emit_mask: torch.Tensor + atom_mask: torch.Tensor + state: dict[str, torch.Tensor] + clean: dict[str, torch.Tensor] + sequence_logits: torch.Tensor + insertion_logits: torch.Tensor diff --git a/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py b/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py new file mode 100644 index 000000000000..240df1359d1d --- /dev/null +++ b/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py @@ -0,0 +1,564 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# Diffusers integration additions are licensed under the Apache License, Version 2.0. +# +# RFD4 sampling is adapted from RFD4-Proteina (BSD-3-Clause). +# Copyright (c) 2026, Institute for Protein Design, University of Washington +# Full third-party copyright notices, license conditions and disclaimers for the RFD4-Proteina +# integration are in diffusers/models/transformers/transformer_rfd4_proteina.py. +from collections.abc import Callable +from typing import Any + +import torch + +from ...loaders import RFD4ProteinaLoraLoaderMixin +from ...models.autoencoders.autoencoder_rfd4_proteina import RFD4ProteinaAutoencoder +from ...models.transformers.transformer_rfd4_proteina import ( + CONDITION_KEYS, + RFD4ProteinaTransformer, + prepare_coord_group_seq, +) +from ...schedulers.scheduling_rfd4_proteina import RFD4ProteinaScheduler +from ...utils import is_torch_xla_available, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import RFD4ProteinaPipelineOutput + + +# Trunk recycling passes of the released front-loaded schedule: 60 steps and 180 passes. +RELEASED_RECYCLE_DEPTHS = (6,) * 5 + (5,) * 10 + (4,) * 5 + (3,) * 10 + (2,) * 20 + (1,) * 10 + +# The released recipe offsets the backbone noise of each molecule by one shared Gaussian vector of this scale. +STOCHASTIC_CENTERING_SCALE = 0.2 + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> 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( + ... **{name: features[name] for name in required}, + ... 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") + ``` +""" + + +class RFD4ProteinaPipeline(DiffusionPipeline, RFD4ProteinaLoraLoaderMixin): + """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. + + Args: + 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. + """ + + model_cpu_offload_seq = "transformer->vae" + _callback_tensor_inputs = ["sample", "last_clean", "sequence_logits", "insertion_logits"] + + @property + def num_timesteps(self) -> int: + """Steps the current call will take, for progress reporting from a callback.""" + return self._num_timesteps + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def interrupt(self) -> bool: + """Set from a callback to stop the denoising loop at the next step boundary.""" + return self._interrupt + + def __init__( + self, + transformer: RFD4ProteinaTransformer, + vae: RFD4ProteinaAutoencoder, + bb_ca_scheduler: RFD4ProteinaScheduler, + local_latents_scheduler: RFD4ProteinaScheduler, + recycle_depths: list[int] | None = None, + ): + super().__init__() + if transformer.config.latent_dim != vae.config.latent_dim: + raise ValueError("Transformer and autoencoder latent dimensions must match.") + self.register_modules( + transformer=transformer, + vae=vae, + bb_ca_scheduler=bb_ca_scheduler, + local_latents_scheduler=local_latents_scheduler, + ) + self.register_to_config( + recycle_depths=list(RELEASED_RECYCLE_DEPTHS if recycle_depths is None else recycle_depths) + ) + + def check_inputs( + self, + padding_mask, + is_in_context_conditioning, + supports_local_latent, + atom37_mask, + coords_nm, + conditioning_features, + num_designs_per_prompt, + generator, + partial_t, + inverse_folding, + latents, + callback_on_step_end_tensor_inputs, + output_type, + ): + """Validate molecular masks, per-design generators, partial diffusion and callbacks.""" + if type(num_designs_per_prompt) is not int or num_designs_per_prompt < 1: + raise ValueError(f"num_designs_per_prompt must be a positive integer, got {num_designs_per_prompt!r}.") + masks = { + "padding_mask": padding_mask, + "is_in_context_conditioning": is_in_context_conditioning, + "supports_local_latent": supports_local_latent, + } + if any(not isinstance(mask, torch.Tensor) for mask in masks.values()): + raise TypeError("Prepared feature masks must be torch tensors.") + if padding_mask.ndim != 2 or min(padding_mask.shape) <= 0: + raise ValueError("padding_mask must have nonempty (batch, tokens) shape.") + if any( + mask.dtype != torch.bool or mask.shape != padding_mask.shape or mask.device != padding_mask.device + for mask in masks.values() + ): + raise ValueError("Prepared sampling masks must be boolean tensors with matching shape and device.") + for key in (CONDITION_KEYS["sequence"], "molecule_iid"): + if key not in (conditioning_features or {}): + raise KeyError(key) + + if generator is not None: + generators = [generator] if isinstance(generator, torch.Generator) else list(generator) + if any(not isinstance(item, torch.Generator) for item in generators): + raise TypeError("generator lists must contain only torch.Generator objects.") + batch_size = padding_mask.shape[0] * num_designs_per_prompt + if len(generators) != batch_size or len({id(item) for item in generators}) != batch_size: + raise ValueError("Provide one distinct generator per batch row.") + + if any(value is not None for value in partial_t.values()): + if any(value is None for value in partial_t.values()): + raise ValueError("Partial diffusion needs a start time for both modalities.") + if inverse_folding and any(float(value) != 0 for value in partial_t.values()): + raise ValueError("inverse folding requires zero start time for both modalities.") + if coords_nm is None: + raise KeyError("partial and inverse generation require clean coords_nm.") + if atom37_mask is None: + raise ValueError("partial and inverse generation require atom37_mask.") + # These paths encode the input, and a non-finite coordinate propagates through the encoder + # until the whole latent is NaN. `atom37_mask` says which atom37 slots a residue uses, not + # whether it has coordinates, so an unresolved residue is masked in and still NaN. Reject it + # here rather than return a structure of NaNs. + unusable = (~torch.isfinite(coords_nm).all(dim=-1)) & atom37_mask.bool() + if unusable.any(): + raise ValueError( + f"coords_nm is not finite at {int(unusable.sum())} atom37 slots that atom37_mask marks " + "as used, which would make the encoded latent entirely NaN. These are usually " + "unresolved residues; mask them out with is_resolved_mask, or supply coordinates for " + "them, before using partial_t or inverse_folding." + ) + elif inverse_folding: + raise ValueError("inverse_folding requires partial_t with zero start times.") + + if any(value is not None for value in latents.values()) and any(value is None for value in latents.values()): + raise ValueError("Starting latents must be given for both modalities.") + if not all(name in self._callback_tensor_inputs for name in callback_on_step_end_tensor_inputs): + raise ValueError(f"callback_on_step_end_tensor_inputs has to be in {self._callback_tensor_inputs}.") + if output_type not in {"pt", "np"}: + raise ValueError(f"output_type must be 'pt' or 'np', got {output_type!r}.") + + def _sample_noise( + self, + shape: tuple[int, ...], + mask: torch.Tensor, + generators: list[torch.Generator] | None, + zero_com: bool = False, + ) -> torch.Tensor: + """Draw Gaussian noise on the masked tokens, optionally with zero center of mass.""" + noise = randn_tensor(shape, generator=generators, device=mask.device, dtype=torch.float32) * mask[..., None] + if zero_com: + count = mask.sum(dim=-1, keepdim=True).clamp(min=1) + noise -= noise.sum(dim=-2, keepdim=True) / count[..., None] + return noise * mask[..., None] + + def _center_noise( + self, + noise: torch.Tensor, + mask: torch.Tensor, + generators: list[torch.Generator] | None, + conditioning: dict[str, torch.Tensor], + ) -> torch.Tensor: + """Offset the backbone noise of each molecule by one shared random vector (stochastic centering).""" + offset_mask = mask + if "coord_cond_mask" in conditioning: + offset_mask = offset_mask & ~conditioning["coord_cond_mask"].bool() + _, molecules = torch.unique(conditioning["molecule_iid"].long(), return_inverse=True) + rows = [] + for row, generator in enumerate(generators if generators is not None else [None] * noise.shape[0]): + selected = offset_mask[row : row + 1] + updated = noise[row : row + 1] + for molecule in molecules[row][selected[0]].unique(): + in_molecule = selected & (molecules[row : row + 1] == molecule) + offset = randn_tensor((1, noise.shape[-1]), generator=generator, device=noise.device) + updated = updated + STOCHASTIC_CENTERING_SCALE * offset.unsqueeze(-2) * in_molecule[..., None].float() + rows.append(updated * mask[row : row + 1, :, None]) + return torch.cat(rows, dim=0) + + def prepare_latents( + self, + bb_mask: torch.Tensor, + latent_mask: torch.Tensor, + conditioning: dict[str, torch.Tensor], + generators: list[torch.Generator] | None, + clean_bb: torch.Tensor | None = None, + clean_latents: torch.Tensor | None = None, + start_times: tuple[float, float] | None = None, + ) -> dict[str, torch.Tensor]: + """Draw the starting state: pure noise, or the clean input noised to `start_times` (backbone, latents). + + 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. + """ + bb_shape = (*bb_mask.shape, 3) + latent_shape = (*latent_mask.shape, self.transformer.config.latent_dim) + bb_noise = self._sample_noise(bb_shape, bb_mask, generators, zero_com=True) + if start_times is None: + bb_noise = self._center_noise(bb_noise, bb_mask, generators, conditioning) + return {"bb_ca": bb_noise, "local_latents": self._sample_noise(latent_shape, latent_mask, generators)} + latent_noise = self._sample_noise(latent_shape, latent_mask, generators) + bb_noise = self._center_noise(bb_noise, bb_mask, generators, conditioning) + bb = self.bb_ca_scheduler.add_noise(clean_bb * bb_mask[..., None], bb_noise, start_times[0]) + latents = self.local_latents_scheduler.add_noise( + clean_latents * latent_mask[..., None], latent_noise, start_times[1] + ) + return {"bb_ca": bb * bb_mask[..., None], "local_latents": latents * latent_mask[..., None]} + + def _noise_backbone(self, target_backbone, bb_mask, conditioning, generators, index): + """Noise the preserved backbone to the scheduled time of step `index`, for inverse folding.""" + noise = self._sample_noise(tuple(target_backbone.shape), bb_mask, generators, zero_com=True) + noise = self._center_noise(noise, bb_mask, generators, conditioning) + noised = self.bb_ca_scheduler.add_noise( + target_backbone * bb_mask[..., None], noise, self.bb_ca_scheduler.timesteps[index] + ) + return noised * bb_mask[..., None] + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + padding_mask: torch.Tensor, + is_in_context_conditioning: torch.Tensor, + supports_local_latent: torch.Tensor, + token_type: torch.Tensor | None = None, + atom37_mask: torch.Tensor | None = None, + coords_nm: torch.Tensor | None = None, + conditioning_features: dict[str, Any] | None = None, + num_designs_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + num_inference_steps: int | None = None, + partial_t_bb_ca: float | None = None, + partial_t_local_latents: float | None = None, + inverse_folding: bool = False, + latents_bb_ca: torch.Tensor | None = None, + latents_local_latents: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[DiffusionPipeline, int, torch.Tensor, dict], dict | None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["sample"], + output_type: str = "pt", + return_dict: bool = True, + ) -> RFD4ProteinaPipelineOutput | tuple: + """Generate structures and sequences from prepared molecular tensors. + + Args: + 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. + 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. + num_inference_steps (`int`, *optional*): + Number of denoising steps. Defaults to, and must equal, the length of the `recycle_depths` config. + 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 `"sample"`. + 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 `["sample"]`): + Tensors passed to the callback: `"sample"`, `"last_clean"`, `"sequence_logits"`, `"insertion_logits"`. + return_dict (`bool`, defaults to `True`): Return a typed output, or a tuple of its non-None fields. + + Examples: + + 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. + """ + partial_t = {"bb_ca": partial_t_bb_ca, "local_latents": partial_t_local_latents} + latents = {"bb_ca": latents_bb_ca, "local_latents": latents_local_latents} + self.check_inputs( + padding_mask, + is_in_context_conditioning, + supports_local_latent, + atom37_mask, + coords_nm, + conditioning_features, + num_designs_per_prompt, + generator, + partial_t, + inverse_folding, + latents, + callback_on_step_end_tensor_inputs, + output_type, + ) + # One stream per design when generators are supplied; the global stream otherwise. + generators = None + if generator is not None: + generators = [generator] if isinstance(generator, torch.Generator) else list(generator) + device = self._execution_device + padding_mask, is_in_context_conditioning, supports_local_latent = ( + tensor.to(device) for tensor in (padding_mask, is_in_context_conditioning, supports_local_latent) + ) + token_type = None if token_type is None else token_type.to(device) + atom37_mask = None if atom37_mask is None else atom37_mask.to(device) + coords_nm = None if coords_nm is None else coords_nm.to(device) + conditioning = {name: value.to(device) for name, value in conditioning_features.items()} + if num_designs_per_prompt > 1: + # Every input row is repeated for its designs, one after the other. + padding_mask, is_in_context_conditioning, supports_local_latent = ( + tensor.repeat_interleave(num_designs_per_prompt, dim=0) + for tensor in (padding_mask, is_in_context_conditioning, supports_local_latent) + ) + token_type, atom37_mask, coords_nm = ( + None if tensor is None else tensor.repeat_interleave(num_designs_per_prompt, dim=0) + for tensor in (token_type, atom37_mask, coords_nm) + ) + conditioning = { + name: value.repeat_interleave(num_designs_per_prompt, dim=0) for name, value in conditioning.items() + } + conditioning = { + **conditioning, + "is_in_context_conditioning": is_in_context_conditioning, + "supports_local_latent": supports_local_latent, + } + batch_size = padding_mask.shape[0] + emitted = padding_mask & ~is_in_context_conditioning + emit_masks = {"bb_ca": emitted, "local_latents": emitted & supports_local_latent} + recycle_depths = self.config.recycle_depths + steps = len(recycle_depths) if num_inference_steps is None else num_inference_steps + if steps != len(recycle_depths): + raise ValueError(f"num_inference_steps must equal the {len(recycle_depths)} entries of recycle_depths.") + is_partial = partial_t_bb_ca is not None + start_times = (float(partial_t_bb_ca), float(partial_t_local_latents)) if is_partial else (None, None) + self.bb_ca_scheduler.set_timesteps(steps, device=device, start_time=start_times[0]) + self.local_latents_scheduler.set_timesteps(steps, device=device, start_time=start_times[1]) + schedulers = {"bb_ca": self.bb_ca_scheduler, "local_latents": self.local_latents_scheduler} + + use_coord_group_seq = "coord_group_seq" in self.transformer.config.feats_seq + + target_backbone = None + clean_state = None + if is_partial: + encoded = self.vae.encode( + coords_nm=coords_nm, + padding_mask=padding_mask, + token_type=token_type, + atom37_mask=atom37_mask, + conditioning=conditioning, + generator=generators, + ) + clean_state = {"bb_ca": coords_nm[..., 1, :].float(), "local_latents": encoded.latent.float()} + target_backbone = clean_state["bb_ca"].clone() if inverse_folding else None + sample = self.prepare_latents( + emit_masks["bb_ca"], + emit_masks["local_latents"], + conditioning, + generators, + *((clean_state["bb_ca"], clean_state["local_latents"], start_times) if is_partial else ()), + ) + + if inverse_folding: + sample["bb_ca"] = self._noise_backbone(target_backbone, emit_masks["bb_ca"], conditioning, generators, 0) + + if latents_bb_ca is not None: + # A supplied starting point replaces the sampled one; the schedule is unchanged. + sample = {"bb_ca": latents_bb_ca.to(device), "local_latents": latents_local_latents.to(device)} + + sequence_logits = insertion_logits = last_clean = None + self._num_timesteps = steps + self._interrupt = False + self._attention_kwargs = attention_kwargs + with self.progress_bar(total=steps) as progress_bar: + for index in range(steps): + if self.interrupt: + continue + timesteps = {name: scheduler.timesteps[index] for name, scheduler in schedulers.items()} + times = {name: timestep.expand(batch_size) for name, timestep in timesteps.items()} + step_conditioning = conditioning + if use_coord_group_seq: + # Draws from the row generators each step, so it must stay inside the loop. + coord_group_seq = prepare_coord_group_seq(conditioning, generators) + if coord_group_seq is not None: + step_conditioning = {**conditioning, "_prepared_coord_group_seq": coord_group_seq} + prediction = self.transformer( + bb_ca=sample["bb_ca"], + local_latents=sample["local_latents"], + bb_ca_time=times["bb_ca"], + local_latents_time=times["local_latents"], + padding_mask=padding_mask, + token_type=token_type, + atom37_mask=atom37_mask, + conditioning=step_conditioning, + recycle_depth=recycle_depths[index], + attention_kwargs=self.attention_kwargs, + ) + velocity = {"bb_ca": prediction.backbone.float(), "local_latents": prediction.local_latents.float()} + model_input = sample + sample = { + name: scheduler.step( + velocity[name], timesteps[name], sample[name], emit_masks[name], generator=generators + ).prev_sample + for name, scheduler in schedulers.items() + } + sequence_logits = prediction.sequence_logits + insertion_logits = prediction.insertion_logits + last_clean = { + name: (model_input[name] + (1 - times[name][..., None, None]) * velocity[name]) + * emit_masks[name][..., None] + for name in schedulers + } + if target_backbone is not None: + if index == steps - 1: + sample["bb_ca"] = target_backbone.clone() + else: + sample["bb_ca"] = self._noise_backbone( + target_backbone, emit_masks["bb_ca"], conditioning, generators, index + 1 + ) + if callback_on_step_end is not None: + callback_kwargs = {} + for name in callback_on_step_end_tensor_inputs: + callback_kwargs[name] = locals()[name] + callback_outputs = callback_on_step_end(self, index, timesteps["bb_ca"], callback_kwargs) + if callback_outputs is not None: + sample = callback_outputs.pop("sample", sample) + + progress_bar.update() + if XLA_AVAILABLE: + xm.mark_step() + + # Decode the terminal state, keeping CA-only tokens that have no latent. + decode_conditioning = dict(conditioning) + key = CONDITION_KEYS["sequence"] + latent_mask = emit_masks["local_latents"] + if token_type is not None: + # Restrict the residue classes each token type may decode to. + allowed = self.vae.allowed_residue_types.to(sequence_logits.device)[token_type.long()] + sequence_logits = sequence_logits.masked_fill( + ~(allowed | ~latent_mask[..., None]), torch.finfo(sequence_logits.dtype).min + ) + # Index 0 is the mask token: only unspecified sequence positions are filled from the prediction. + decode_conditioning[key] = torch.where( + latent_mask & (decode_conditioning[key] == 0), sequence_logits.argmax(dim=-1), decode_conditioning[key] + ) + decoded_output = self.vae.decode( + z_latent=sample["local_latents"], + ca_coors_nm=sample["bb_ca"], + padding_mask=padding_mask, + token_type=token_type, + atom37_mask=atom37_mask, + conditioning=decode_conditioning, + ) + emit_mask = decoded_output.emit_mask + no_latent = padding_mask & ~emit_mask + coords_nm = decoded_output.coordinates.clone() + coords_nm[..., 1, :] = torch.where(no_latent[..., None], sample["bb_ca"], coords_nm[..., 1, :]) + residue_type = decoded_output.residue_type.long().masked_fill(~emit_mask, -1) + atom_mask = decoded_output.atom_mask.clone() + atom_mask[..., 1] |= no_latent + self.maybe_free_model_hooks() + if output_type == "np": + + def to_numpy(tensor): + # numpy has no bfloat16, which is the dtype the released recipe runs in. + tensor = tensor.float() if tensor.dtype in (torch.bfloat16, torch.float16) else tensor + return tensor.cpu().numpy() + + coords_nm, residue_type, emit_mask, atom_mask = ( + to_numpy(tensor) for tensor in (coords_nm, residue_type, emit_mask, atom_mask) + ) + sample = {name: to_numpy(value) for name, value in sample.items()} + last_clean = None if last_clean is None else {k: to_numpy(v) for k, v in last_clean.items()} + sequence_logits = None if sequence_logits is None else to_numpy(sequence_logits) + insertion_logits = None if insertion_logits is None else to_numpy(insertion_logits) + output = RFD4ProteinaPipelineOutput( + coords_nm=coords_nm, + residue_type=residue_type, + emit_mask=emit_mask, + atom_mask=atom_mask, + state=sample, + clean=last_clean, + sequence_logits=sequence_logits, + insertion_logits=insertion_logits, + ) + return output if return_dict else output.to_tuple() diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index 47f4a0857090..3909159729d6 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -76,6 +76,7 @@ _import_structure["scheduling_piflow"] = ["PiflowScheduler"] _import_structure["scheduling_pndm"] = ["PNDMScheduler"] _import_structure["scheduling_repaint"] = ["RePaintScheduler"] + _import_structure["scheduling_rfd4_proteina"] = ["RFD4ProteinaScheduler"] _import_structure["scheduling_sasolver"] = ["SASolverScheduler"] _import_structure["scheduling_scm"] = ["SCMScheduler"] _import_structure["scheduling_sde_ve"] = ["ScoreSdeVeScheduler"] @@ -161,6 +162,7 @@ from .scheduling_piflow import PiflowScheduler from .scheduling_pndm import PNDMScheduler from .scheduling_repaint import RePaintScheduler + from .scheduling_rfd4_proteina import RFD4ProteinaScheduler from .scheduling_sasolver import SASolverScheduler from .scheduling_scm import SCMScheduler from .scheduling_sde_ve import ScoreSdeVeScheduler diff --git a/src/diffusers/schedulers/scheduling_rfd4_proteina.py b/src/diffusers/schedulers/scheduling_rfd4_proteina.py new file mode 100644 index 000000000000..44c6025f3db9 --- /dev/null +++ b/src/diffusers/schedulers/scheduling_rfd4_proteina.py @@ -0,0 +1,276 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# Diffusers integration additions are licensed under the Apache License, Version 2.0. +# +# RFD4 sampling is adapted from RFD4-Proteina (BSD-3-Clause). +# Copyright (c) 2026, Institute for Protein Design, University of Washington +# Full third-party copyright notices, license conditions and disclaimers for the RFD4-Proteina +# integration are in diffusers/models/transformers/transformer_rfd4_proteina.py. +from dataclasses import dataclass + +import torch + +from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import BaseOutput +from ..utils.torch_utils import randn_tensor +from .scheduling_utils import SchedulerMixin + + +@dataclass +class RFD4ProteinaSchedulerOutput(BaseOutput): + """Updated state of one modality. + + Args: + prev_sample (`torch.Tensor`): + The state after the step, shaped `(B, N, D)`: `D` is 3 for the backbone (coordinates in nanometers) and the + latent width for the local latents. + """ + + prev_sample: torch.Tensor + + +def _schedule( + mode: str, steps: int, *, start: float = 0.0, tail_steps: int = 0, tail_multiplier: int = 1 +) -> torch.Tensor: + """Times from `start` to 1 with `steps` intervals; the last `tail_steps` intervals are `tail_multiplier` times finer.""" + if tail_steps: + if type(tail_multiplier) is not int or tail_multiplier < 2: + raise ValueError("tail_step_multiplier must be an integer >= 2.") + base_steps = steps - (tail_multiplier - 1) * tail_steps + if base_steps < tail_steps: + raise ValueError("steps must be at least tail_steps * tail_step_multiplier.") + fine = _schedule(mode, base_steps * tail_multiplier, start=start) + tail_start = (base_steps - tail_steps) * tail_multiplier + indices = [*range(0, tail_start, tail_multiplier), *range(tail_start, len(fine))] + return fine[torch.tensor(indices)] + if mode == "power": + return torch.linspace(start**0.5, 1, steps + 1) ** 2 + times = 1 - torch.logspace(-2, start, steps + 1).flip(0) + times -= times.min() - start + return times / times.max() + + +def _noise_schedule(times: torch.Tensor, mode: str) -> torch.Tensor: + times = times.clamp(0, 1 - 1e-5) + if mode == "tan": + angle = (1 - times) * torch.pi / 2 + noise = (torch.pi / 2) * torch.sin(angle) / (torch.cos(angle) + 1e-2) + else: + noise = 1 / (times + 1e-2) + return noise.clamp(min=0) + + +def _score(sample: torch.Tensor, velocity: torch.Tensor, times: torch.Tensor) -> torch.Tensor: + return (times[:, None, None] * velocity - sample) / (1 - times)[:, None, None] + + +def _velocity(sample: torch.Tensor, score: torch.Tensor, times: torch.Tensor) -> torch.Tensor: + times = times[:, None, None] + return (sample + (1 - times) * score) / times + + +def _tsr_ratio(times: torch.Tensor, k: float, sigma: float) -> torch.Tensor: + snr = times.square() / (1 - times).square() + return (snr * sigma**2 + 1) / (snr * sigma**2 / k + 1) + + +class RFD4ProteinaScheduler(SchedulerMixin, ConfigMixin): + """Temporal-score-rescaled flow sampler for one modality of RFD4-Proteína. + + The pipeline integrates two modalities, the backbone and the per-token local latents, with one scheduler each. They + differ in their time schedule, noise schedule and rescaling. + + Args: + schedule_mode (`str`): + Spacing of the times between noise and data: `"log"` or `"power"`. + noise_mode (`str`): + Shape of the noise injected in the stochastic window: `"1/t"` or `"tan"`. + tsr_k (`float`): + How low the sampling temperature goes; higher is lower. + tsr_sigma (`float`): + How fast the low temperature is reached; higher is faster. + tail_steps (`int`, defaults to `10`): + Number of final intervals that are finer than the rest by `tail_step_multiplier`. + tail_step_multiplier (`int`, defaults to `2`): + How much finer the final intervals are. + t_lim_ode (`float`, defaults to `0.95`): + Above this time the step is deterministic. + t_lim_ode_below (`float`, defaults to `0.02`): + Below this time the step is deterministic. + + Call `set_timesteps` before sampling. Times increase from noise toward clean data. A scheduler tracks the current + interval and should not be shared by concurrent generations. + """ + + order = 1 + _compatibles = [] + + @register_to_config + def __init__( + self, + schedule_mode: str, + noise_mode: str, + tsr_k: float, + tsr_sigma: float, + tail_steps: int = 10, + tail_step_multiplier: int = 2, + t_lim_ode: float = 0.95, + t_lim_ode_below: float = 0.02, + ): + if schedule_mode not in {"log", "power"}: + raise ValueError(f"schedule_mode must be 'log' or 'power', got {schedule_mode!r}.") + if noise_mode not in {"1/t", "tan"}: + raise ValueError(f"noise_mode must be '1/t' or 'tan', got {noise_mode!r}.") + self.timesteps = None + self._step_index = None + self._times = None + self._noise = None + + @property + def step_index(self) -> int | None: + """Index of the next integration interval.""" + return self._step_index + + def set_timesteps( + self, + num_inference_steps: int, + device: str | torch.device | None = None, + start_time: float | None = None, + ) -> None: + """Resolve the schedule and reset the integration cursor. + + Args: + num_inference_steps (`int`): + Number of integration intervals. + device (`str` or `torch.device`, *optional*): + Device for `timesteps`. The interval lengths and noise coefficients stay on the CPU. + start_time (`float`, *optional*): + Partial-diffusion start time in `[0, 1)`. Omit for full generation. + """ + if type(num_inference_steps) is not int or num_inference_steps < 1: + raise ValueError("num_inference_steps must be a positive integer.") + if start_time is not None and not 0 <= float(start_time) < 1: + raise ValueError("start_time must be in [0, 1).") + + def schedule_for(steps: int) -> torch.Tensor: + return _schedule( + self.config.schedule_mode, + steps, + start=1e-3, + tail_steps=self.config.tail_steps, + tail_multiplier=self.config.tail_step_multiplier, + ) + + if start_time is None: + full_times = times = schedule_for(num_inference_steps) + else: + dense_steps = max(num_inference_steps * 100, 1000) + above = schedule_for(dense_steps) >= float(start_time) - 1e-6 + first = above.nonzero(as_tuple=True)[0][0].item() if above.any() else dense_steps + points_above = max(dense_steps - first, 1) + full_times = schedule_for(round(num_inference_steps * dense_steps / points_above)) + times = full_times[-(num_inference_steps + 1) :] + # The noise schedule carries the temporal score rescaling twice, once for the drift and once for the noise. + noise = _noise_schedule(full_times[:-1], self.config.noise_mode) + noise /= _tsr_ratio(full_times[:-1], self.config.tsr_k, self.config.tsr_sigma) ** 2 + # The interval lengths and noise coefficients stay CPU scalars; only the model times move to the device. + self._times, self._noise = times, noise[-num_inference_steps:] + self.timesteps = times[:-1].to(device) + self._step_index = 0 + + @staticmethod + def _generators( + generator: torch.Generator | list[torch.Generator] | None, batch_size: int + ) -> tuple[torch.Generator, ...] | None: + """Normalise per-row generators; `None` means draw from the global stream.""" + if generator is None: + return None + generators = (generator,) if isinstance(generator, torch.Generator) else tuple(generator) + if len(generators) != batch_size or any(not isinstance(item, torch.Generator) for item in generators): + raise ValueError(f"Provide one torch.Generator per batch row ({batch_size}).") + if len({id(item) for item in generators}) != len(generators): + raise ValueError("Each batch row requires a distinct torch.Generator.") + return generators + + def add_noise( + self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: float | torch.Tensor + ) -> torch.Tensor: + """Sample the forward marginal `(1 - t) * noise + t * clean`. + + Args: + original_samples (`torch.Tensor`): Clean state `(batch, tokens, D)`. + noise (`torch.Tensor`): Noise of the same shape, drawn by the caller. + timesteps (`float` or `torch.Tensor`): Time in `[0, 1)`, a scalar or `(batch,)`. + """ + batch_size = original_samples.shape[0] + time = torch.as_tensor(timesteps, device=original_samples.device, dtype=original_samples.dtype) + if time.ndim == 0: + time = time.expand(batch_size) + if time.shape != (batch_size,) or torch.any((time < 0) | (time >= 1)): + raise ValueError(f"timesteps must have shape ({batch_size},) with values in [0, 1).") + return (1 - time[:, None, None]) * noise + time[:, None, None] * original_samples + + def step( + self, + model_output: torch.Tensor, + timestep: torch.Tensor, + sample: torch.Tensor, + emit_mask: torch.Tensor, + generator: torch.Generator | list[torch.Generator] | None = None, + return_dict: bool = True, + ) -> RFD4ProteinaSchedulerOutput | tuple[torch.Tensor]: + """Advance one integration interval and update `step_index`. + + Args: + model_output (`torch.Tensor`): + Predicted velocity, shaped like `sample`. + timestep (`torch.Tensor`): + The next element of `timesteps`. Out-of-order steps are rejected. + sample (`torch.Tensor`): + Current state `(B, N, D)`. + emit_mask (`torch.Tensor`): + Boolean `(B, N)`, true for the tokens that are updated (not context or padding). + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + A generator for a single row or distinct generators for each batch row. Omitting them uses the global + stream. + return_dict (`bool`, defaults to `True`): + Return `RFD4ProteinaSchedulerOutput`; otherwise a one-item tuple holding the state. + + Returns: + `RFD4ProteinaSchedulerOutput` or `tuple`: The updated state. + """ + if self._step_index is None: + raise ValueError("Call set_timesteps before step.") + if self._step_index >= len(self.timesteps): + raise ValueError("The trajectory is complete; call set_timesteps to start another.") + if not torch.equal(timestep, self.timesteps[self._step_index]): + raise ValueError("timestep must be the next element of timesteps.") + + generators = self._generators(generator, sample.shape[0]) + index = self._step_index + t = timestep.expand(sample.shape[0]) + dt = self._times[index + 1] - self._times[index] + gt = self._noise[index] + ratio = _tsr_ratio(t[0], self.config.tsr_k, self.config.tsr_sigma) + if t[0] == 0: + # The score is undefined at zero time; the released sampler holds the state. + prev_sample = sample * emit_mask[..., None] + else: + score = _score(sample, model_output, t) + rescaled = _velocity(sample, score * ratio, t) + if t[0] > self.config.t_lim_ode or t[0] < self.config.t_lim_ode_below: + change = rescaled * dt + else: + random = randn_tensor( + sample.shape, + generator=list(generators) if generators is not None else None, + device=sample.device, + dtype=sample.dtype, + ) + # The drift and the noise schedule both carry the temporal score rescaling, which keeps the injected + # noise consistent with the lower temperature. + change = (rescaled + gt * score * ratio) * dt + torch.sqrt(2 * gt * dt) * random + prev_sample = (sample + change) * emit_mask[..., None] + self._step_index += 1 + if not return_dict: + return (prev_sample,) + return RFD4ProteinaSchedulerOutput(prev_sample=prev_sample) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 2896212c1edb..422314fe935e 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -2239,6 +2239,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class RFD4ProteinaAutoencoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + +class RFD4ProteinaTransformer(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class SanaControlNetModel(metaclass=DummyObject): _backends = ["torch"] @@ -3216,6 +3246,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class RFD4ProteinaPipeline(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class ScoreSdeVePipeline(metaclass=DummyObject): _backends = ["torch"] @@ -3876,6 +3921,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class RFD4ProteinaScheduler(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class SASolverScheduler(metaclass=DummyObject): _backends = ["torch"] diff --git a/tests/models/autoencoders/test_models_autoencoder_rfd4_proteina.py b/tests/models/autoencoders/test_models_autoencoder_rfd4_proteina.py new file mode 100644 index 000000000000..3e873b62c57a --- /dev/null +++ b/tests/models/autoencoders/test_models_autoencoder_rfd4_proteina.py @@ -0,0 +1,101 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers import RFD4ProteinaAutoencoder +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin + + +enable_full_determinism() + + +class RFD4ProteinaAutoencoderTesterConfig(BaseModelTesterConfig): + main_input_name = "z_latent" + + @property + def model_class(self): + return RFD4ProteinaAutoencoder + + def get_init_dict(self): + return { + "atom37_mask": [[True] * 37 for _ in range(33)], + "allowed_residue_types": [[True] * 33 for _ in range(5)], + "latent_dim": 4, + "n_residue_types": 33, + "num_token_types": 5, + "nlayers": 1, + "token_dim": 8, + "nheads": 2, + "dim_cond": 8, + "pair_repr_dim": 4, + "encoder_feats_seq": [], + "decoder_feats_seq": ["z_latent_seq"], + } + + def get_dummy_inputs(self): + generator = torch.Generator("cpu").manual_seed(0) + padding_mask = torch.tensor([[True, True, False]], device=torch_device) + return { + "z_latent": randn_tensor((1, 3, 4), generator=generator, device=torch_device), + "ca_coors_nm": randn_tensor((1, 3, 3), generator=generator, device=torch_device), + "padding_mask": padding_mask, + "atom37_mask": torch.ones(1, 3, 37, dtype=torch.bool, device=torch_device), + "conditioning": { + "is_in_context_conditioning": torch.tensor([[False, True, False]], device=torch_device), + "supports_local_latent": padding_mask, + "condition_sequence_1_residue": torch.zeros(1, 3, dtype=torch.long, device=torch_device), + }, + } + + @property + def input_shape(self) -> tuple[int, ...]: + return (1, 3, 4) + + @property + def output_shape(self) -> tuple[int, ...]: + return (3, 37, 3) + + +class TestRFD4ProteinaAutoencoderModel(RFD4ProteinaAutoencoderTesterConfig, ModelTesterMixin): + def test_decode_keeps_the_conditioned_residue_and_its_coordinate_group_atoms(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + conditioning = inputs["conditioning"] + conditioning["is_in_context_conditioning"] = torch.zeros(1, 3, dtype=torch.bool, device=torch_device) + conditioning["condition_sequence_1_residue"] = torch.tensor([[5, 0, 0]], device=torch_device) + groups = torch.full((1, 3, 37), -1, device=torch_device) + groups[0, 0, 10] = 0 + conditioning["condition_CoordinateGroup_1_atom"] = groups + with torch.no_grad(): + output = model.decode(**inputs) + assert output.residue_type[0, 0] == 5 + assert output.atom_mask[0, 0, 10] + + +class TestRFD4ProteinaAutoencoderMemory(RFD4ProteinaAutoencoderTesterConfig, MemoryTesterMixin): + def get_init_dict(self): + # The layerwise-casting memory test needs layers large enough for the compute dtype to change peak memory. + return { + **super().get_init_dict(), + "nlayers": 2, + "token_dim": 512, + "nheads": 8, + "pair_repr_dim": 64, + "dim_cond": 64, + } diff --git a/tests/models/transformers/test_models_transformer_rfd4_proteina.py b/tests/models/transformers/test_models_transformer_rfd4_proteina.py new file mode 100644 index 000000000000..baf60c726734 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_rfd4_proteina.py @@ -0,0 +1,143 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers import RFD4ProteinaTransformer +from diffusers.models.transformers.transformer_rfd4_proteina import CONDITION_KEYS, FEATURE_FUNCTIONS +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + BaseModelTesterConfig, + LoraTesterMixin, + MemoryTesterMixin, + ModelTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class RFD4ProteinaTransformerTesterConfig(BaseModelTesterConfig): + # The shared tests read the batch size from this input, so it has to be a tensor, not a mapping. + main_input_name = "padding_mask" + + @property + def model_class(self): + return RFD4ProteinaTransformer + + def get_init_dict(self): + return { + "nlayers": 1, + "token_dim": 8, + "pair_repr_dim": 4, + "nheads": 2, + "feats_seq": ["xt_bb_ca", "xt_local_latents"], + "feats_cond_seq": [], + "feats_pair_repr": [], + "feats_pair_cond": [], + "dim_cond": 8, + "latent_dim": 4, + "update_pair_repr": False, + "use_tri_mult": False, + "recycle": False, + } + + def get_dummy_inputs(self): + generator = torch.Generator("cpu").manual_seed(0) + padding_mask = torch.tensor([[True, True, False]], device=torch_device) + return { + "bb_ca": randn_tensor((1, 3, 3), generator=generator, device=torch_device), + "local_latents": randn_tensor((1, 3, 4), generator=generator, device=torch_device), + "bb_ca_time": torch.tensor([0.5], device=torch_device), + "local_latents_time": torch.tensor([0.5], device=torch_device), + "padding_mask": padding_mask, + "conditioning": { + "is_in_context_conditioning": torch.tensor([[False, True, False]], device=torch_device), + "supports_local_latent": padding_mask, + }, + "recycle_depth": 1, + } + + @property + def input_shape(self) -> tuple[int, ...]: + return (1, 3, 3) + + @property + def output_shape(self) -> tuple[int, ...]: + return (3, 3) + + +class TestRFD4ProteinaTransformerModel(RFD4ProteinaTransformerTesterConfig, ModelTesterMixin): + def test_transformer_executes_requested_recycle_depth(self): + init_dict = {**self.get_init_dict(), "recycle": True, "feats_cond_seq": ["recycle_remaining"], "t_emb_dim": 8} + transformer = RFD4ProteinaTransformer(**init_dict).to(torch_device).eval() + inputs = self.get_dummy_inputs() + cycles = [] + handle = transformer.transformer_layers[0].register_forward_hook(lambda *_: cycles.append(None)) + for depth, expected_cycles in ((0, 1), (1, 1), (3, 3)): + cycles.clear() + transformer(**{**inputs, "recycle_depth": depth}) + assert len(cycles) == expected_cycles + handle.remove() + + def test_condition_features_follow_the_annotation_semantics(self): + atom_mask = torch.ones(1, 3, 37, dtype=torch.bool) + atom_mask[:, :, 3] = False + features = { + "padding_mask": torch.tensor([[True, True, False]]), + "atom37_mask": atom_mask, + CONDITION_KEYS["internal_coordinates"]: torch.full((1, 3, 37, 3), 10.0), + CONDITION_KEYS["coordinate_group"]: torch.tensor( + [[[-1, 0, 7] + [-1] * 34, [-1, 0, 9] + [-1] * 34, [-1] * 37]] + ), + CONDITION_KEYS["rasa"]: torch.full((1, 3, 37), 3), + CONDITION_KEYS["centroid"]: torch.tensor([[[10.0, 20.0, 30.0], [float("nan")] * 3, [float("inf")] * 3]]), + } + + internal = FEATURE_FUNCTIONS["cond_a37coors_internal_nm"](features, 148) + assert internal[..., 111:].sum().item() == 2 + assert internal[..., :111].sum().item() == 6 + rasa = FEATURE_FUNCTIONS["rasa_binned"](features, 37) + assert torch.equal(rasa[~atom_mask], torch.zeros_like(rasa[~atom_mask])) + centroid = FEATURE_FUNCTIONS["chain_centroid_xyz"](features, 3) + assert torch.equal(centroid, torch.tensor([[[1.0, 2.0, 3.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]])) + + +class TestRFD4ProteinaTransformerMemory(RFD4ProteinaTransformerTesterConfig, MemoryTesterMixin): + def get_init_dict(self): + # The layerwise-casting memory test needs layers large enough for the compute dtype to change peak memory. Three + # layers with the pair update after every second one keep a layer without an update, as the released model has. + return { + **super().get_init_dict(), + "nlayers": 3, + "update_pair_repr": True, + "token_dim": 512, + "nheads": 8, + "pair_repr_dim": 64, + "dim_cond": 64, + } + + +class TestRFD4ProteinaTransformerTraining(RFD4ProteinaTransformerTesterConfig, TrainingTesterMixin): + def test_gradient_checkpointing_is_applied(self): + # The shared default expects a `*Model` class name; this transformer is named `RFD4ProteinaTransformer`. + super().test_gradient_checkpointing_is_applied(expected_set={"RFD4ProteinaTransformer"}) + + +class TestRFD4ProteinaTransformerLoRA(RFD4ProteinaTransformerTesterConfig, LoraTesterMixin): + pass diff --git a/tests/pipelines/rfd4_proteina/__init__.py b/tests/pipelines/rfd4_proteina/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py b/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py new file mode 100644 index 000000000000..49a375656070 --- /dev/null +++ b/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py @@ -0,0 +1,407 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 + +import numpy as np +import pytest +import torch + +from diffusers import ( + RFD4ProteinaAutoencoder, + RFD4ProteinaPipeline, + RFD4ProteinaScheduler, + RFD4ProteinaTransformer, +) +from diffusers.models.transformers.transformer_rfd4_proteina import ( + CONDITION_KEYS, + RFD4ProteinaTransformerOutput, + prepare_coord_group_seq, +) + +from ...testing_utils import torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin + + +SEQUENCE_KEY = CONDITION_KEYS["sequence"] + + +def prepared_inputs(device="cpu"): + padding = torch.tensor([[True, True, False]], device=device) + return { + "padding_mask": padding, + "is_in_context_conditioning": torch.tensor([[False, True, False]], device=device), + "supports_local_latent": padding, + "atom37_mask": torch.ones(1, 3, 37, dtype=torch.bool, device=device), + "conditioning_features": { + SEQUENCE_KEY: torch.zeros(1, 3, dtype=torch.long, device=device), + "molecule_iid": torch.zeros(1, 3, dtype=torch.long, device=device), + "chain_iid": torch.zeros(1, 3, dtype=torch.long, device=device), + }, + } + + +def allowed_residue_types(): + """The 33-class protein and nucleic encoding: each token type may decode to its own classes, never the mask token.""" + allowed = [[False] * 33 for _ in range(5)] + for index in range(1, 22): # protein: twenty amino acids and the unknown residue + allowed[0][index] = True + for index in range(22, 27): # RNA, including the unknown base + allowed[1][index] = True + for index in range(27, 32): # DNA, including the unknown base + allowed[2][index] = True + allowed[3][32] = True # atomized tokens decode to the atomized class only + for index in range(1, 33): # unknown token type: any class except the mask token + allowed[4][index] = True + return allowed + + +class RFD4ProteinaPipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = RFD4ProteinaPipeline + required_input_params_in_call_signature = frozenset(["padding_mask", "generator", "num_inference_steps"]) + # `generator` is required: every design needs its own random stream. + optional_input_params = frozenset(["num_inference_steps", "return_dict"]) + batch_input_params = frozenset( + ["padding_mask", "is_in_context_conditioning", "supports_local_latent", "atom37_mask", "conditioning_features"] + ) + callback_cfg_params = frozenset() + output_shape = (3, 37, 3) + + def get_dummy_components(self, depths=(1, 2), feats_seq=("xt_bb_ca", "xt_local_latents")): + torch.manual_seed(0) + transformer = RFD4ProteinaTransformer( + nlayers=1, + token_dim=8, + pair_repr_dim=4, + nheads=2, + feats_seq=list(feats_seq), + feats_cond_seq=[], + feats_pair_repr=[], + feats_pair_cond=[], + dim_cond=8, + latent_dim=4, + num_sequence_classes=33, + num_token_types=5, + update_pair_repr=False, + use_tri_mult=False, + recycle=True, + ) + torch.manual_seed(0) + vae = RFD4ProteinaAutoencoder( + atom37_mask=[[True] * 37 for _ in range(33)], + allowed_residue_types=allowed_residue_types(), + latent_dim=4, + n_residue_types=33, + num_token_types=5, + nlayers=1, + token_dim=8, + nheads=2, + dim_cond=8, + pair_repr_dim=4, + encoder_feats_seq=["chain_idx_seq"], + decoder_feats_seq=["z_latent_seq"], + ) + # `t_lim_ode_below=1` keeps every interval on the noise-free branch, so integration draws no random numbers. + scheduler_kwargs = {"tail_steps": 0, "t_lim_ode": 0.98, "t_lim_ode_below": 1.0} + return { + "transformer": transformer, + "vae": vae, + "bb_ca_scheduler": RFD4ProteinaScheduler("log", "1/t", 16.0, 3.0, **scheduler_kwargs), + "local_latents_scheduler": RFD4ProteinaScheduler("power", "tan", 4.0, 8.0, **scheduler_kwargs), + "recycle_depths": list(depths), + } + + def get_dummy_inputs(self): + return { + **prepared_inputs(torch_device), + "generator": self.get_generator(0), + "num_inference_steps": 2, + } + + def batch_input(self, value, batch_size): + # Inputs are `(batch, ...)` tensors, or a dictionary of them: batching repeats along the batch axis. + if isinstance(value, dict): + return {name: self.batch_input(tensor, batch_size) for name, tensor in value.items()} + return value.repeat(batch_size, *([1] * (value.ndim - 1))) + + +class TestRFD4ProteinaPipeline(RFD4ProteinaPipelineTesterConfig, PipelineTesterMixin): + def test_components_function(self): + # `recycle_depths` is a config value, not a component. + init_components = self.get_dummy_components() + init_components = {k: v for k, v in init_components.items() if k != "recycle_depths"} + pipe = self.get_pipeline(**init_components) + assert set(pipe.components.keys()) == set(init_components.keys()) + + @pytest.mark.skip(reason="Callbacks receive dictionaries of tensors; covered by `test_callback_on_step_end`.") + def test_callback_inputs(self): + pass + + def get_pipeline_with(self, **kwargs): + pipe = self.get_pipeline(**self.get_dummy_components(**kwargs)).to(torch_device) + return pipe + + def assert_rejected_before_rng(self, pipe, error, match, seed=7, **overrides): + inputs = self.get_dummy_inputs() + inputs["generator"] = self.get_generator(seed) + inputs.update(overrides) + before = inputs["generator"].get_state().clone() + with pytest.raises(error, match=match): + pipe(**inputs) + assert torch.equal(inputs["generator"].get_state(), before) + + def test_pipeline_calls_model_once_per_step_and_recycles_inside_the_call(self): + pipe = self.get_pipeline_with() + calls, cycles = [], [] + handles = [ + pipe.transformer.register_forward_hook(lambda *_: calls.append(None)), + pipe.transformer.transformer_layers[0].register_forward_hook(lambda *_: cycles.append(None)), + ] + pipe(**self.get_dummy_inputs()) + for handle in handles: + handle.remove() + assert len(calls) == 2 + assert len(cycles) == 3 # recycle depths (1, 2) execute 1 + 2 trunk passes + + def test_callback_on_step_end(self): + pipe = self.get_pipeline_with() + seen = [] + + def callback(pipe, step, timestep, callback_kwargs): + assert set(callback_kwargs) == {"sample", "last_clean"} + seen.append(step) + if step == 0: + return {"sample": {name: torch.zeros_like(value) for name, value in callback_kwargs["sample"].items()}} + return None + + output = pipe( + **self.get_dummy_inputs(), + callback_on_step_end=callback, + callback_on_step_end_tensor_inputs=["sample", "last_clean"], + ) + assert seen == [0, 1] + assert torch.isfinite(output.state["bb_ca"]).all() + with pytest.raises(ValueError, match="callback_on_step_end_tensor_inputs"): + pipe( + **self.get_dummy_inputs(), + callback_on_step_end=callback, + callback_on_step_end_tensor_inputs=["latents"], + ) + + def test_missing_sequence_condition_fails_before_sampling(self): + self.assert_rejected_before_rng(self.get_pipeline_with(), KeyError, SEQUENCE_KEY, conditioning_features={}) + + def coord_group_inputs(self): + inputs = prepared_inputs(torch_device) + groups = torch.zeros(1, 3, 37, dtype=torch.long) + groups[:, 1] = 1 + inputs["conditioning_features"][CONDITION_KEYS["coordinate_group"]] = groups.to(torch_device) + return inputs, groups + + def test_pipeline_prepares_random_coord_groups_outside_transformer(self): + pipe = self.get_pipeline_with(depths=(0, 1), feats_seq=["xt_bb_ca", "xt_local_latents", "coord_group_seq"]) + inputs, groups = self.coord_group_inputs() + expected_generator = self.get_generator(7) + torch.randn((1, 3, 3), generator=expected_generator) + torch.randn((1, 3), generator=expected_generator) # the stochastic centering offset of the one molecule + torch.randn((1, 3, 4), generator=expected_generator) + expected = [torch.randperm(2, generator=expected_generator)[groups] for _ in range(2)] + captured = [] + handle = pipe.transformer.register_forward_pre_hook( + lambda module, args, kwargs: captured.append(kwargs["conditioning"]["_prepared_coord_group_seq"].cpu()), + with_kwargs=True, + ) + actual_generator = self.get_generator(7) + pipe(**inputs, generator=actual_generator, num_inference_steps=2) + handle.remove() + assert len(captured) == 2 + assert all(torch.equal(actual, wanted) for actual, wanted in zip(captured, expected)) + assert torch.equal(actual_generator.get_state(), expected_generator.get_state()) + + def test_coordinate_group_relabeling_uses_row_generators_without_global_rng_mutation(self): + conditioning = { + CONDITION_KEYS["coordinate_group"]: torch.tensor( + [[[-1, 0, 7] + [-1] * 34, [-1, 0, 9] + [-1] * 34, [-1] * 37]] + ) + } + before = torch.get_rng_state().clone() + first = prepare_coord_group_seq(conditioning, [self.get_generator(23)]) + replay = prepare_coord_group_seq(conditioning, [self.get_generator(23)]) + assert torch.equal(first, replay) + assert torch.equal(before, torch.get_rng_state()) + + def test_partial_and_inverse_initialization_use_clean_tensor_inputs(self): + pipe = self.get_pipeline_with() + coords_nm = torch.zeros(1, 3, 37, 3, device=torch_device) + coords_nm[:, 0, 1] = torch.tensor([1.0, 2.0, 3.0], device=torch_device) + inputs = {**prepared_inputs(torch_device), "coords_nm": coords_nm, "num_inference_steps": 2} + partial = pipe(**inputs, generator=self.get_generator(13), partial_t_bb_ca=0.5, partial_t_local_latents=0.5) + inverse = pipe( + **inputs, + generator=self.get_generator(13), + partial_t_bb_ca=0.0, + partial_t_local_latents=0.0, + inverse_folding=True, + ) + assert all(torch.isfinite(value).all() for value in partial.state.values()) + assert torch.equal(inverse.state["bb_ca"][:, 0], coords_nm[:, 0, 1]) + self.assert_rejected_before_rng( + pipe, ValueError, "inverse_folding requires partial_t", coords_nm=coords_nm, inverse_folding=True + ) + self.assert_rejected_before_rng( + pipe, + ValueError, + "zero start time", + coords_nm=coords_nm, + partial_t_bb_ca=0.5, + partial_t_local_latents=0.0, + inverse_folding=True, + ) + self.assert_rejected_before_rng(pipe, ValueError, "both modalities", coords_nm=coords_nm, partial_t_bb_ca=0.5) + + def test_decoder_sequence_prediction_obeys_token_type_and_user_condition(self): + pipe = self.get_pipeline_with() + inputs = prepared_inputs(torch_device) + inputs["conditioning_features"][SEQUENCE_KEY] = torch.tensor([[0, 8, 0]], device=torch_device) + inputs["token_type"] = torch.zeros(1, 3, dtype=torch.long, device=torch_device) + + def fixed_logits(module, args, kwargs, output): + logits = torch.zeros_like(output.sequence_logits) + logits[..., 0], logits[..., 22], logits[..., 1] = 100, 99, 3 + return RFD4ProteinaTransformerOutput( + backbone=output.backbone, + local_latents=output.local_latents, + sequence_logits=logits, + insertion_logits=output.insertion_logits, + ) + + decoded_features = [] + handles = [ + pipe.transformer.register_forward_hook(fixed_logits, with_kwargs=True), + pipe.vae.decoder.register_forward_pre_hook( + lambda module, args: decoded_features.append(args[0][SEQUENCE_KEY].cpu()) + ), + ] + pipe(**inputs, generator=self.get_generator(7), num_inference_steps=2) + for handle in handles: + handle.remove() + # Position 0 is filled with the best protein residue (not the mask token or an RNA type); position 1 is user-fixed. + assert torch.equal(decoded_features[0], torch.tensor([[1, 8, 0]])) + + def test_pipeline_preserves_inputs_and_independent_batch_rows(self): + pipe = self.get_pipeline_with() + inputs = self.get_dummy_inputs() + original = {name: value.clone() for name, value in inputs["conditioning_features"].items()} + single_generator = self.get_generator(7) + single = pipe(**{**inputs, "generator": single_generator}) + for name, value in original.items(): + assert torch.equal(inputs["conditioning_features"][name], value) + generators = [self.get_generator(7), self.get_generator(8)] + batched = {name: self.batch_input(inputs[name], 2) for name in self.batch_input_params} + output = pipe(**{**inputs, **batched, "generator": generators}) + for name, value in single.state.items(): + torch.testing.assert_close(output.state[name][:1], value, rtol=0, atol=1e-6) + assert torch.equal(generators[0].get_state(), single_generator.get_state()) + + def test_num_designs_per_prompt_repeats_every_input_row(self): + pipe = self.get_pipeline_with() + inputs = self.get_dummy_inputs() + inputs.pop("generator") + two_rows = {name: self.batch_input(inputs[name], 2) for name in self.batch_input_params} + # The second input row designs one more token, which shows up in the emit mask of its designs. + two_rows["is_in_context_conditioning"][1, 1] = False + generators = [self.get_generator(seed) for seed in range(4)] + output = pipe(**{**inputs, **two_rows}, num_designs_per_prompt=2, generator=generators) + assert output.coords_nm.shape[0] == 4 + # Rows 0 and 1 are the two designs of input row 0; rows 2 and 3 those of input row 1. + assert output.emit_mask[0].equal(output.emit_mask[1]) + assert output.emit_mask[2].equal(output.emit_mask[3]) + assert not output.emit_mask[0].equal(output.emit_mask[2]) + # Repeating by hand with the same generators gives the same designs. + repeated = { + name: value.repeat_interleave(2, dim=0) + if isinstance(value, torch.Tensor) + else {key: tensor.repeat_interleave(2, dim=0) for key, tensor in value.items()} + for name, value in two_rows.items() + } + by_hand = pipe(**{**inputs, **repeated}, generator=[self.get_generator(seed) for seed in range(4)]) + torch.testing.assert_close(output.state["bb_ca"], by_hand.state["bb_ca"], rtol=0, atol=0) + with pytest.raises(ValueError, match="one distinct generator per batch row"): + pipe(**{**inputs, **two_rows}, num_designs_per_prompt=2, generator=generators[:2]) + with pytest.raises(ValueError, match="num_designs_per_prompt"): + pipe(**inputs, num_designs_per_prompt=0) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + def test_numpy_output_from_reduced_precision(self, dtype): + """numpy has no bfloat16 or float16 path, and the released recipe runs in bfloat16.""" + pipe = self.get_pipeline_with().to(dtype) + output = pipe(**self.get_dummy_inputs(), output_type="np") + assert isinstance(output.coords_nm, np.ndarray) + assert output.coords_nm.dtype == np.float32 + assert np.isfinite(output.coords_nm).all() + + def test_non_finite_coordinates_rejected_before_rng(self): + """Unresolved residues are masked in by `atom37_mask` but have no coordinates, and encoding them + turns the whole latent NaN. The encode paths must say so rather than return a NaN structure.""" + coords_nm = torch.zeros(1, 3, 37, 3, device=torch_device) + coords_nm[0, 0, 1] = float("nan") + self.assert_rejected_before_rng( + self.get_pipeline_with(), + ValueError, + "not finite", + coords_nm=coords_nm, + partial_t_bb_ca=0.3, + partial_t_local_latents=0.6, + ) + + @pytest.mark.parametrize("invalid", ["rank", "dtype", "shape", "device", "empty"]) + def test_invalid_sampling_masks_rejected_before_rng(self, invalid): + inputs = prepared_inputs(torch_device) + if invalid == "rank": + inputs["padding_mask"] = inputs["padding_mask"][0] + elif invalid == "dtype": + inputs["supports_local_latent"] = inputs["supports_local_latent"].float() + elif invalid == "shape": + inputs["is_in_context_conditioning"] = inputs["is_in_context_conditioning"][:, :1] + elif invalid == "device": + inputs["supports_local_latent"] = inputs["supports_local_latent"].to("meta") + else: + inputs["padding_mask"] = inputs["padding_mask"][:, :0] + self.assert_rejected_before_rng(self.get_pipeline_with(), ValueError, "sampling masks|padding_mask", **inputs) + + def test_generators_must_be_distinct_per_row(self): + inputs = self.get_dummy_inputs() + shared = self.get_generator(1) + batched = {name: self.batch_input(inputs[name], 2) for name in self.batch_input_params} + with pytest.raises(ValueError, match="one distinct generator per batch row"): + self.get_pipeline_with()(**{**inputs, **batched, "generator": [shared, shared]}) + + def test_stochastic_centering_offsets_each_molecule_by_one_shared_vector(self): + pipe = self.get_pipeline_with() + mask = torch.ones(1, 4, dtype=torch.bool, device=torch_device) + conditioning = {"molecule_iid": torch.tensor([[0, 0, 1, 1]], device=torch_device)} + centered = pipe.prepare_latents(mask, mask, conditioning, [self.get_generator(3)])["bb_ca"] + base = pipe._sample_noise((1, 4, 3), mask, [self.get_generator(3)], zero_com=True) + offset = centered - base + # Tokens of one molecule move together, and the two molecules move differently. + torch.testing.assert_close(offset[:, 0], offset[:, 1]) + torch.testing.assert_close(offset[:, 2], offset[:, 3]) + assert not torch.allclose(offset[:, 0], offset[:, 2]) + + def test_prepare_latents_replays_distinct_row_generators(self): + pipe = self.get_pipeline_with() + mask = torch.tensor([[True, False], [True, False]], device=torch_device) + conditioning = {"molecule_iid": torch.zeros(2, 2, dtype=torch.long, device=torch_device)} + left = [torch.Generator().manual_seed(7), torch.Generator().manual_seed(11)] + right = [torch.Generator().manual_seed(7), torch.Generator().manual_seed(11)] + first = pipe.prepare_latents(mask, mask, conditioning, left) + second = pipe.prepare_latents(mask, mask, conditioning, right) + for name in first: + torch.testing.assert_close(first[name], second[name], rtol=0, atol=0) + assert first[name][:, 1].count_nonzero() == 0 + assert all(a.get_state().equal(b.get_state()) for a, b in zip(left, right)) + + +class TestRFD4ProteinaPipelineMemory(RFD4ProteinaPipelineTesterConfig, MemoryTesterMixin): + pass diff --git a/tests/schedulers/test_scheduler_rfd4_proteina.py b/tests/schedulers/test_scheduler_rfd4_proteina.py new file mode 100644 index 000000000000..1c2f1a8da700 --- /dev/null +++ b/tests/schedulers/test_scheduler_rfd4_proteina.py @@ -0,0 +1,157 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import RFD4ProteinaScheduler + + +def scheduler(stochastic=False, **kwargs): + """A tiny backbone-style scheduler. With `stochastic=False`, `t_lim_ode_below=1` keeps every interval on the + noise-free branch; `stochastic=True` uses the released thresholds.""" + config = { + "schedule_mode": "log", + "noise_mode": "1/t", + "tsr_k": 16.0, + "tsr_sigma": 3.0, + "tail_steps": 0, + "t_lim_ode": 0.98, + "t_lim_ode_below": 0.02 if stochastic else 1.0, + } + return RFD4ProteinaScheduler(**{**config, **kwargs}) + + +def inputs(): + sample = torch.tensor([[[0.2, -0.4, 0.6], [2.0, 3.0, 4.0]]]) + mask = torch.tensor([[True, False]]) + return sample, mask, torch.full_like(sample, 0.25) + + +def test_schedule_and_config_round_trip(tmp_path): + original = scheduler(stochastic=True) + original.save_pretrained(tmp_path) + loaded = RFD4ProteinaScheduler.from_pretrained(tmp_path, local_files_only=True) + original.set_timesteps(2) + loaded.set_timesteps(2) + torch.testing.assert_close(loaded.timesteps, original.timesteps, rtol=0, atol=0) + torch.testing.assert_close(loaded._noise, original._noise, rtol=0, atol=0) + # Released schedules start near zero (below `t_lim_ode_below`) and end at 1. + assert loaded.timesteps[0] < 0.02 + assert float(loaded._times[-1]) == 1.0 + assert loaded.step_index == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA scalar-placement coverage") +def test_cuda_schedule_keeps_coefficients_on_cpu(): + wrapper = scheduler(stochastic=True) + wrapper.set_timesteps(2, device="cuda") + + assert wrapper.timesteps.device.type == "cuda" + assert wrapper._times.device.type == "cpu" + assert wrapper._noise.device.type == "cpu" + + +@pytest.mark.parametrize("schedule_mode, noise_mode", [("log", "1/t"), ("power", "tan")]) +def test_partial_schedule_is_exact_nonlinear_full_grid_tail(schedule_mode, noise_mode): + kwargs = {"schedule_mode": schedule_mode, "noise_mode": noise_mode} + full = scheduler(stochastic=True, **kwargs) + full.set_timesteps(4) + partial = scheduler(stochastic=True, **kwargs) + partial.set_timesteps(2, start_time=float(full._times[2])) + torch.testing.assert_close(partial._times, full._times[2:], rtol=0, atol=0) + torch.testing.assert_close(partial._noise, full._noise[2:], rtol=0, atol=0) + + +def test_first_step_matches_temporal_score_rescaling(): + wrapper = scheduler(stochastic=True) + wrapper.set_timesteps(2) + sample, mask, velocity = inputs() + actual = wrapper.step(velocity, wrapper.timesteps[0], sample, mask, return_dict=False)[0] + # The first time (1e-3) lies below `t_lim_ode_below`, so the interval is noise free: the velocity is converted to a + # score, rescaled by the temporal score ratio, converted back and integrated with one Euler step. + t = wrapper.timesteps[0] + snr = t**2 / (1 - t) ** 2 + k, sigma = wrapper.config.tsr_k, wrapper.config.tsr_sigma + ratio = (snr * sigma**2 + 1) / (snr * sigma**2 / k + 1) + score = (t * velocity - sample) / (1 - t) + rescaled = (sample + (1 - t) * score * ratio) / t + expected = (sample + rescaled * (wrapper._times[1] - wrapper._times[0])) * mask[..., None] + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5) + assert torch.equal(actual[:, 1], torch.zeros_like(actual[:, 1])) + + +def test_add_noise_interpolates_between_noise_and_data(): + wrapper = scheduler() + clean = torch.ones(2, 3, 3) + noise = torch.zeros_like(clean) + torch.testing.assert_close(wrapper.add_noise(clean, noise, 0.25), torch.full_like(clean, 0.25)) + torch.testing.assert_close( + wrapper.add_noise(clean, noise, torch.tensor([0.0, 0.5]))[:, 0, 0], torch.tensor([0.0, 0.5]) + ) + with pytest.raises(ValueError, match=r"\[0, 1\)"): + wrapper.add_noise(clean, noise, 1.0) + + +def test_step_order_and_generator_validation(): + wrapper = scheduler() + sample, mask, velocity = inputs() + with pytest.raises(ValueError, match="set_timesteps"): + wrapper.step(velocity, torch.zeros(()), sample, mask) + wrapper.set_timesteps(2) + with pytest.raises(ValueError, match="next element"): + wrapper.step(velocity, wrapper.timesteps[1], sample, mask) + assert wrapper.step_index == 0 + repeated = torch.Generator() + with pytest.raises(ValueError, match="distinct"): + wrapper.step( + velocity.expand(2, -1, -1), + wrapper.timesteps[0], + sample.expand(2, -1, -1), + mask.expand(2, -1), + generator=[repeated, repeated], + ) + + +@pytest.mark.parametrize( + "device", + ["cpu", pytest.param("cuda:0", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"))], +) +def test_stochastic_step_default_generator_matches_explicit_stream(device): + explicit, implicit = scheduler(stochastic=True), scheduler(stochastic=True) + explicit.set_timesteps(2, device=device) + implicit.set_timesteps(2, device=device) + sample, mask, velocity = (value.to(device) for value in inputs()) + generator = torch.Generator(device=device).manual_seed(71) + initial_rng = generator.get_state().clone() + devices = [torch.device(device)] if device.startswith("cuda") else [] + with torch.random.fork_rng(devices=devices): + if devices: + torch.cuda.set_rng_state(initial_rng, device) + else: + torch.set_rng_state(initial_rng) + left, right = sample, sample + for timestep in explicit.timesteps: + left = explicit.step(velocity, timestep, left, mask, generator=generator).prev_sample + right = implicit.step(velocity, timestep, right, mask).prev_sample + torch.testing.assert_close(left, right, rtol=0, atol=0) + endpoint = torch.cuda.get_rng_state(device) if devices else torch.get_rng_state() + assert torch.equal(endpoint, generator.get_state()) + assert not torch.equal(endpoint, initial_rng), "The test must exercise a stochastic interval." + + +@pytest.mark.parametrize("kwargs", [{"schedule_mode": "uniform"}, {"noise_mode": "ode"}]) +def test_options_outside_the_released_recipe_are_rejected(kwargs): + with pytest.raises(ValueError, match="must be"): + scheduler(**kwargs) From 5df5db4ee80aa53a201389ba358a1974fa6f5c2c Mon Sep 17 00:00:00 2001 From: Kieran Didi Date: Tue, 20 Oct 2026 15:13:41 +0000 Subject: [PATCH 2/8] Address review: pass tensors instead of dicts, name the pipeline example inputs prepare_latents returns a tuple and the sampling loop works on named tensors. The pipeline output has bb_ca, local_latents and clean_* fields instead of dicts, and the callback tensors are named. The pipeline properties move below __init__, the transformer return annotation names its tuple, LoRA moves from the loaders page to the pipeline docs, and the pipeline tests spell out their helpers. --- docs/source/en/api/loaders/lora.md | 5 - docs/source/en/api/pipelines/rfd4_proteina.md | 18 ++ .../transformers/transformer_rfd4_proteina.py | 2 +- .../rfd4_proteina/pipeline_output.py | 18 +- .../rfd4_proteina/pipeline_rfd4_proteina.py | 156 ++++++++++-------- .../test_pipeline_rfd4_proteina.py | 101 ++++++------ 6 files changed, 167 insertions(+), 133 deletions(-) diff --git a/docs/source/en/api/loaders/lora.md b/docs/source/en/api/loaders/lora.md index 1a4fded1fa10..4c8eb20b299a 100644 --- a/docs/source/en/api/loaders/lora.md +++ b/docs/source/en/api/loaders/lora.md @@ -39,16 +39,11 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi - [`ErnieImageLoraLoaderMixin`] provides similar functions for [Ernie-Image](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ernie_image). - [`LTX2LoraLoaderMixin`] provides similar functions for [Flux2](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ltx2). - [`MiniMaxH3LoraLoaderMixin`] provides similar functions for [MiniMax-H3](https://huggingface.co/docs/diffusers/main/en/api/pipelines/minimax_h3). -- [`RFD4ProteinaLoraLoaderMixin`] provides similar functions for [RFD4-Proteína](https://huggingface.co/docs/diffusers/main/en/api/pipelines/rfd4_proteina). - [`LoraBaseMixin`] provides a base class with several utility methods to fuse, unfuse, unload, LoRAs and more. > [!TIP] > To learn more about how to load LoRA weights, see the [LoRA](../../tutorials/using_peft_for_inference) loading guide. -## RFD4ProteinaLoraLoaderMixin - -[[autodoc]] loaders.lora_pipeline.RFD4ProteinaLoraLoaderMixin - ## LoraBaseMixin [[autodoc]] loaders.lora_base.LoraBaseMixin diff --git a/docs/source/en/api/pipelines/rfd4_proteina.md b/docs/source/en/api/pipelines/rfd4_proteina.md index 66a5dd592fce..6054644c723a 100644 --- a/docs/source/en/api/pipelines/rfd4_proteina.md +++ b/docs/source/en/api/pipelines/rfd4_proteina.md @@ -83,6 +83,24 @@ write_cif(output, features, atoms, "design.cif") 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`: + +```py +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 [[autodoc]] RFD4ProteinaPipeline diff --git a/src/diffusers/models/transformers/transformer_rfd4_proteina.py b/src/diffusers/models/transformers/transformer_rfd4_proteina.py index fcde248cdb73..f17b0fcdbea5 100644 --- a/src/diffusers/models/transformers/transformer_rfd4_proteina.py +++ b/src/diffusers/models/transformers/transformer_rfd4_proteina.py @@ -1100,7 +1100,7 @@ def forward( recycle_depth: int = 0, attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, - ) -> RFD4ProteinaTransformerOutput | tuple: + ) -> RFD4ProteinaTransformerOutput | tuple[torch.Tensor]: """Predict velocities from prepared molecular tensors without advancing the sampler. Args: diff --git a/src/diffusers/pipelines/rfd4_proteina/pipeline_output.py b/src/diffusers/pipelines/rfd4_proteina/pipeline_output.py index a41890b40aa2..b87dce24d617 100644 --- a/src/diffusers/pipelines/rfd4_proteina/pipeline_output.py +++ b/src/diffusers/pipelines/rfd4_proteina/pipeline_output.py @@ -21,10 +21,14 @@ class RFD4ProteinaPipelineOutput(BaseOutput): Boolean mask `(B, N)` of the tokens that were designed. atom_mask (`torch.Tensor`): Boolean mask `(B, N, 37)` of the valid atoms. - state (`dict[str, torch.Tensor]`): - Final product-space state: `bb_ca` coordinates `(B, N, 3)` and `local_latents` `(B, N, D)`. - clean (`dict[str, torch.Tensor]`): - Predicted clean product-space state from the last denoising step. + 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`): @@ -35,7 +39,9 @@ class RFD4ProteinaPipelineOutput(BaseOutput): residue_type: torch.Tensor emit_mask: torch.Tensor atom_mask: torch.Tensor - state: dict[str, torch.Tensor] - clean: dict[str, torch.Tensor] + bb_ca: torch.Tensor + local_latents: torch.Tensor + clean_bb_ca: torch.Tensor + clean_local_latents: torch.Tensor sequence_logits: torch.Tensor insertion_logits: torch.Tensor diff --git a/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py b/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py index 240df1359d1d..30b0e0e071c5 100644 --- a/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py +++ b/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py @@ -52,7 +52,11 @@ >>> required = ["padding_mask", "is_in_context_conditioning", "supports_local_latent", "token_type"] >>> required += ["atom37_mask"] >>> output = pipe( - ... **{name: features[name] for name in required}, + ... 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), ... ) @@ -79,21 +83,14 @@ class RFD4ProteinaPipeline(DiffusionPipeline, RFD4ProteinaLoraLoaderMixin): """ model_cpu_offload_seq = "transformer->vae" - _callback_tensor_inputs = ["sample", "last_clean", "sequence_logits", "insertion_logits"] - - @property - def num_timesteps(self) -> int: - """Steps the current call will take, for progress reporting from a callback.""" - return self._num_timesteps - - @property - def attention_kwargs(self): - return self._attention_kwargs - - @property - def interrupt(self) -> bool: - """Set from a callback to stop the denoising loop at the next step boundary.""" - return self._interrupt + _callback_tensor_inputs = [ + "bb_ca", + "local_latents", + "last_clean_bb_ca", + "last_clean_local_latents", + "sequence_logits", + "insertion_logits", + ] def __init__( self, @@ -116,6 +113,20 @@ def __init__( recycle_depths=list(RELEASED_RECYCLE_DEPTHS if recycle_depths is None else recycle_depths) ) + @property + def num_timesteps(self) -> int: + """Steps the current call will take, for progress reporting from a callback.""" + return self._num_timesteps + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def interrupt(self) -> bool: + """Set from a callback to stop the denoising loop at the next step boundary.""" + return self._interrupt + def check_inputs( self, padding_mask, @@ -238,7 +249,7 @@ def prepare_latents( clean_bb: torch.Tensor | None = None, clean_latents: torch.Tensor | None = None, start_times: tuple[float, float] | None = None, - ) -> dict[str, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: """Draw the starting state: pure noise, or the clean input noised to `start_times` (backbone, latents). The draws follow the native sampler's order on every row generator: for a full generation the backbone noise @@ -250,14 +261,14 @@ def prepare_latents( bb_noise = self._sample_noise(bb_shape, bb_mask, generators, zero_com=True) if start_times is None: bb_noise = self._center_noise(bb_noise, bb_mask, generators, conditioning) - return {"bb_ca": bb_noise, "local_latents": self._sample_noise(latent_shape, latent_mask, generators)} + return bb_noise, self._sample_noise(latent_shape, latent_mask, generators) latent_noise = self._sample_noise(latent_shape, latent_mask, generators) bb_noise = self._center_noise(bb_noise, bb_mask, generators, conditioning) bb = self.bb_ca_scheduler.add_noise(clean_bb * bb_mask[..., None], bb_noise, start_times[0]) latents = self.local_latents_scheduler.add_noise( clean_latents * latent_mask[..., None], latent_noise, start_times[1] ) - return {"bb_ca": bb * bb_mask[..., None], "local_latents": latents * latent_mask[..., None]} + return bb * bb_mask[..., None], latents * latent_mask[..., None] def _noise_backbone(self, target_backbone, bb_mask, conditioning, generators, index): """Noise the preserved backbone to the scheduled time of step `index`, for inverse folding.""" @@ -289,7 +300,7 @@ def __call__( latents_local_latents: torch.Tensor | None = None, attention_kwargs: dict[str, Any] | None = None, callback_on_step_end: Callable[[DiffusionPipeline, int, torch.Tensor, dict], dict | None] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["sample"], + callback_on_step_end_tensor_inputs: list[str] = ["bb_ca", "local_latents"], output_type: str = "pt", return_dict: bool = True, ) -> RFD4ProteinaPipelineOutput | tuple: @@ -313,7 +324,9 @@ def __call__( 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. + 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 @@ -341,11 +354,12 @@ def __call__( 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 `"sample"`. + 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 `["sample"]`): - Tensors passed to the callback: `"sample"`, `"last_clean"`, `"sequence_logits"`, `"insertion_logits"`. + 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. Examples: @@ -403,8 +417,8 @@ def __call__( "supports_local_latent": supports_local_latent, } batch_size = padding_mask.shape[0] - emitted = padding_mask & ~is_in_context_conditioning - emit_masks = {"bb_ca": emitted, "local_latents": emitted & supports_local_latent} + bb_mask = padding_mask & ~is_in_context_conditioning + latent_mask = bb_mask & supports_local_latent recycle_depths = self.config.recycle_depths steps = len(recycle_depths) if num_inference_steps is None else num_inference_steps if steps != len(recycle_depths): @@ -413,12 +427,10 @@ def __call__( start_times = (float(partial_t_bb_ca), float(partial_t_local_latents)) if is_partial else (None, None) self.bb_ca_scheduler.set_timesteps(steps, device=device, start_time=start_times[0]) self.local_latents_scheduler.set_timesteps(steps, device=device, start_time=start_times[1]) - schedulers = {"bb_ca": self.bb_ca_scheduler, "local_latents": self.local_latents_scheduler} use_coord_group_seq = "coord_group_seq" in self.transformer.config.feats_seq - target_backbone = None - clean_state = None + target_backbone = clean_bb = clean_latents = None if is_partial: encoded = self.vae.encode( coords_nm=coords_nm, @@ -428,24 +440,26 @@ def __call__( conditioning=conditioning, generator=generators, ) - clean_state = {"bb_ca": coords_nm[..., 1, :].float(), "local_latents": encoded.latent.float()} - target_backbone = clean_state["bb_ca"].clone() if inverse_folding else None - sample = self.prepare_latents( - emit_masks["bb_ca"], - emit_masks["local_latents"], + clean_bb, clean_latents = coords_nm[..., 1, :].float(), encoded.latent.float() + target_backbone = clean_bb.clone() if inverse_folding else None + bb_ca, local_latents = self.prepare_latents( + bb_mask, + latent_mask, conditioning, generators, - *((clean_state["bb_ca"], clean_state["local_latents"], start_times) if is_partial else ()), + clean_bb, + clean_latents, + start_times if is_partial else None, ) if inverse_folding: - sample["bb_ca"] = self._noise_backbone(target_backbone, emit_masks["bb_ca"], conditioning, generators, 0) + bb_ca = self._noise_backbone(target_backbone, bb_mask, conditioning, generators, 0) if latents_bb_ca is not None: # A supplied starting point replaces the sampled one; the schedule is unchanged. - sample = {"bb_ca": latents_bb_ca.to(device), "local_latents": latents_local_latents.to(device)} + bb_ca, local_latents = latents_bb_ca.to(device), latents_local_latents.to(device) - sequence_logits = insertion_logits = last_clean = None + sequence_logits = insertion_logits = last_clean_bb_ca = last_clean_local_latents = None self._num_timesteps = steps self._interrupt = False self._attention_kwargs = attention_kwargs @@ -453,8 +467,8 @@ def __call__( for index in range(steps): if self.interrupt: continue - timesteps = {name: scheduler.timesteps[index] for name, scheduler in schedulers.items()} - times = {name: timestep.expand(batch_size) for name, timestep in timesteps.items()} + bb_time = self.bb_ca_scheduler.timesteps[index] + latents_time = self.local_latents_scheduler.timesteps[index] step_conditioning = conditioning if use_coord_group_seq: # Draws from the row generators each step, so it must stay inside the loop. @@ -462,10 +476,10 @@ def __call__( if coord_group_seq is not None: step_conditioning = {**conditioning, "_prepared_coord_group_seq": coord_group_seq} prediction = self.transformer( - bb_ca=sample["bb_ca"], - local_latents=sample["local_latents"], - bb_ca_time=times["bb_ca"], - local_latents_time=times["local_latents"], + bb_ca=bb_ca, + local_latents=local_latents, + bb_ca_time=bb_time.expand(batch_size), + local_latents_time=latents_time.expand(batch_size), padding_mask=padding_mask, token_type=token_type, atom37_mask=atom37_mask, @@ -473,35 +487,32 @@ def __call__( recycle_depth=recycle_depths[index], attention_kwargs=self.attention_kwargs, ) - velocity = {"bb_ca": prediction.backbone.float(), "local_latents": prediction.local_latents.float()} - model_input = sample - sample = { - name: scheduler.step( - velocity[name], timesteps[name], sample[name], emit_masks[name], generator=generators - ).prev_sample - for name, scheduler in schedulers.items() - } + bb_velocity, latents_velocity = prediction.backbone.float(), prediction.local_latents.float() + last_clean_bb_ca = (bb_ca + (1 - bb_time) * bb_velocity) * bb_mask[..., None] + last_clean_local_latents = (local_latents + (1 - latents_time) * latents_velocity) * latent_mask[ + ..., None + ] + bb_ca = self.bb_ca_scheduler.step( + bb_velocity, bb_time, bb_ca, bb_mask, generator=generators + ).prev_sample + local_latents = self.local_latents_scheduler.step( + latents_velocity, latents_time, local_latents, latent_mask, generator=generators + ).prev_sample sequence_logits = prediction.sequence_logits insertion_logits = prediction.insertion_logits - last_clean = { - name: (model_input[name] + (1 - times[name][..., None, None]) * velocity[name]) - * emit_masks[name][..., None] - for name in schedulers - } if target_backbone is not None: if index == steps - 1: - sample["bb_ca"] = target_backbone.clone() + bb_ca = target_backbone.clone() else: - sample["bb_ca"] = self._noise_backbone( - target_backbone, emit_masks["bb_ca"], conditioning, generators, index + 1 - ) + bb_ca = self._noise_backbone(target_backbone, bb_mask, conditioning, generators, index + 1) if callback_on_step_end is not None: callback_kwargs = {} for name in callback_on_step_end_tensor_inputs: callback_kwargs[name] = locals()[name] - callback_outputs = callback_on_step_end(self, index, timesteps["bb_ca"], callback_kwargs) + callback_outputs = callback_on_step_end(self, index, bb_time, callback_kwargs) if callback_outputs is not None: - sample = callback_outputs.pop("sample", sample) + bb_ca = callback_outputs.pop("bb_ca", bb_ca) + local_latents = callback_outputs.pop("local_latents", local_latents) progress_bar.update() if XLA_AVAILABLE: @@ -510,7 +521,6 @@ def __call__( # Decode the terminal state, keeping CA-only tokens that have no latent. decode_conditioning = dict(conditioning) key = CONDITION_KEYS["sequence"] - latent_mask = emit_masks["local_latents"] if token_type is not None: # Restrict the residue classes each token type may decode to. allowed = self.vae.allowed_residue_types.to(sequence_logits.device)[token_type.long()] @@ -522,8 +532,8 @@ def __call__( latent_mask & (decode_conditioning[key] == 0), sequence_logits.argmax(dim=-1), decode_conditioning[key] ) decoded_output = self.vae.decode( - z_latent=sample["local_latents"], - ca_coors_nm=sample["bb_ca"], + z_latent=local_latents, + ca_coors_nm=bb_ca, padding_mask=padding_mask, token_type=token_type, atom37_mask=atom37_mask, @@ -532,7 +542,7 @@ def __call__( emit_mask = decoded_output.emit_mask no_latent = padding_mask & ~emit_mask coords_nm = decoded_output.coordinates.clone() - coords_nm[..., 1, :] = torch.where(no_latent[..., None], sample["bb_ca"], coords_nm[..., 1, :]) + coords_nm[..., 1, :] = torch.where(no_latent[..., None], bb_ca, coords_nm[..., 1, :]) residue_type = decoded_output.residue_type.long().masked_fill(~emit_mask, -1) atom_mask = decoded_output.atom_mask.clone() atom_mask[..., 1] |= no_latent @@ -547,8 +557,10 @@ def to_numpy(tensor): coords_nm, residue_type, emit_mask, atom_mask = ( to_numpy(tensor) for tensor in (coords_nm, residue_type, emit_mask, atom_mask) ) - sample = {name: to_numpy(value) for name, value in sample.items()} - last_clean = None if last_clean is None else {k: to_numpy(v) for k, v in last_clean.items()} + bb_ca, local_latents, last_clean_bb_ca, last_clean_local_latents = ( + None if tensor is None else to_numpy(tensor) + for tensor in (bb_ca, local_latents, last_clean_bb_ca, last_clean_local_latents) + ) sequence_logits = None if sequence_logits is None else to_numpy(sequence_logits) insertion_logits = None if insertion_logits is None else to_numpy(insertion_logits) output = RFD4ProteinaPipelineOutput( @@ -556,8 +568,10 @@ def to_numpy(tensor): residue_type=residue_type, emit_mask=emit_mask, atom_mask=atom_mask, - state=sample, - clean=last_clean, + bb_ca=bb_ca, + local_latents=local_latents, + clean_bb_ca=last_clean_bb_ca, + clean_local_latents=last_clean_local_latents, sequence_logits=sequence_logits, insertion_logits=insertion_logits, ) diff --git a/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py b/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py index 49a375656070..b2e8c0aeb6e0 100644 --- a/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py +++ b/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py @@ -128,31 +128,28 @@ def batch_input(self, value, batch_size): class TestRFD4ProteinaPipeline(RFD4ProteinaPipelineTesterConfig, PipelineTesterMixin): def test_components_function(self): - # `recycle_depths` is a config value, not a component. + # The shared test expects every dummy component to be a module, but `recycle_depths` is a list of ints that the + # pipeline stores in its config, so it is left out of the comparison. init_components = self.get_dummy_components() init_components = {k: v for k, v in init_components.items() if k != "recycle_depths"} pipe = self.get_pipeline(**init_components) assert set(pipe.components.keys()) == set(init_components.keys()) - @pytest.mark.skip(reason="Callbacks receive dictionaries of tensors; covered by `test_callback_on_step_end`.") + @pytest.mark.skip( + reason="The shared test passes `latents` and `output_type='latent'`, which this pipeline does not have; " + "`test_callback_on_step_end` covers callbacks." + ) def test_callback_inputs(self): pass - def get_pipeline_with(self, **kwargs): - pipe = self.get_pipeline(**self.get_dummy_components(**kwargs)).to(torch_device) - return pipe - - def assert_rejected_before_rng(self, pipe, error, match, seed=7, **overrides): - inputs = self.get_dummy_inputs() - inputs["generator"] = self.get_generator(seed) - inputs.update(overrides) + def assert_rejected_before_rng(self, pipe, inputs, error, match): before = inputs["generator"].get_state().clone() with pytest.raises(error, match=match): pipe(**inputs) assert torch.equal(inputs["generator"].get_state(), before) def test_pipeline_calls_model_once_per_step_and_recycles_inside_the_call(self): - pipe = self.get_pipeline_with() + pipe = self.get_pipeline().to(torch_device) calls, cycles = [], [] handles = [ pipe.transformer.register_forward_hook(lambda *_: calls.append(None)), @@ -165,23 +162,23 @@ def test_pipeline_calls_model_once_per_step_and_recycles_inside_the_call(self): assert len(cycles) == 3 # recycle depths (1, 2) execute 1 + 2 trunk passes def test_callback_on_step_end(self): - pipe = self.get_pipeline_with() + pipe = self.get_pipeline().to(torch_device) seen = [] def callback(pipe, step, timestep, callback_kwargs): - assert set(callback_kwargs) == {"sample", "last_clean"} + assert set(callback_kwargs) == {"bb_ca", "last_clean_bb_ca"} seen.append(step) if step == 0: - return {"sample": {name: torch.zeros_like(value) for name, value in callback_kwargs["sample"].items()}} + return {"bb_ca": torch.zeros_like(callback_kwargs["bb_ca"])} return None output = pipe( **self.get_dummy_inputs(), callback_on_step_end=callback, - callback_on_step_end_tensor_inputs=["sample", "last_clean"], + callback_on_step_end_tensor_inputs=["bb_ca", "last_clean_bb_ca"], ) assert seen == [0, 1] - assert torch.isfinite(output.state["bb_ca"]).all() + assert torch.isfinite(output.bb_ca).all() with pytest.raises(ValueError, match="callback_on_step_end_tensor_inputs"): pipe( **self.get_dummy_inputs(), @@ -190,7 +187,8 @@ def callback(pipe, step, timestep, callback_kwargs): ) def test_missing_sequence_condition_fails_before_sampling(self): - self.assert_rejected_before_rng(self.get_pipeline_with(), KeyError, SEQUENCE_KEY, conditioning_features={}) + inputs = {**self.get_dummy_inputs(), "generator": self.get_generator(7), "conditioning_features": {}} + self.assert_rejected_before_rng(self.get_pipeline().to(torch_device), inputs, KeyError, SEQUENCE_KEY) def coord_group_inputs(self): inputs = prepared_inputs(torch_device) @@ -200,7 +198,9 @@ def coord_group_inputs(self): return inputs, groups def test_pipeline_prepares_random_coord_groups_outside_transformer(self): - pipe = self.get_pipeline_with(depths=(0, 1), feats_seq=["xt_bb_ca", "xt_local_latents", "coord_group_seq"]) + pipe = self.get_pipeline( + **self.get_dummy_components(depths=(0, 1), feats_seq=["xt_bb_ca", "xt_local_latents", "coord_group_seq"]) + ).to(torch_device) inputs, groups = self.coord_group_inputs() expected_generator = self.get_generator(7) torch.randn((1, 3, 3), generator=expected_generator) @@ -232,7 +232,7 @@ def test_coordinate_group_relabeling_uses_row_generators_without_global_rng_muta assert torch.equal(before, torch.get_rng_state()) def test_partial_and_inverse_initialization_use_clean_tensor_inputs(self): - pipe = self.get_pipeline_with() + pipe = self.get_pipeline().to(torch_device) coords_nm = torch.zeros(1, 3, 37, 3, device=torch_device) coords_nm[:, 0, 1] = torch.tensor([1.0, 2.0, 3.0], device=torch_device) inputs = {**prepared_inputs(torch_device), "coords_nm": coords_nm, "num_inference_steps": 2} @@ -244,24 +244,22 @@ def test_partial_and_inverse_initialization_use_clean_tensor_inputs(self): partial_t_local_latents=0.0, inverse_folding=True, ) - assert all(torch.isfinite(value).all() for value in partial.state.values()) - assert torch.equal(inverse.state["bb_ca"][:, 0], coords_nm[:, 0, 1]) + assert torch.isfinite(partial.bb_ca).all() and torch.isfinite(partial.local_latents).all() + assert torch.equal(inverse.bb_ca[:, 0], coords_nm[:, 0, 1]) + rejected = {**self.get_dummy_inputs(), "generator": self.get_generator(7), "coords_nm": coords_nm} self.assert_rejected_before_rng( - pipe, ValueError, "inverse_folding requires partial_t", coords_nm=coords_nm, inverse_folding=True + pipe, {**rejected, "inverse_folding": True}, ValueError, "inverse_folding requires partial_t" ) self.assert_rejected_before_rng( pipe, + {**rejected, "partial_t_bb_ca": 0.5, "partial_t_local_latents": 0.0, "inverse_folding": True}, ValueError, "zero start time", - coords_nm=coords_nm, - partial_t_bb_ca=0.5, - partial_t_local_latents=0.0, - inverse_folding=True, ) - self.assert_rejected_before_rng(pipe, ValueError, "both modalities", coords_nm=coords_nm, partial_t_bb_ca=0.5) + self.assert_rejected_before_rng(pipe, {**rejected, "partial_t_bb_ca": 0.5}, ValueError, "both modalities") def test_decoder_sequence_prediction_obeys_token_type_and_user_condition(self): - pipe = self.get_pipeline_with() + pipe = self.get_pipeline().to(torch_device) inputs = prepared_inputs(torch_device) inputs["conditioning_features"][SEQUENCE_KEY] = torch.tensor([[0, 8, 0]], device=torch_device) inputs["token_type"] = torch.zeros(1, 3, dtype=torch.long, device=torch_device) @@ -290,7 +288,7 @@ def fixed_logits(module, args, kwargs, output): assert torch.equal(decoded_features[0], torch.tensor([[1, 8, 0]])) def test_pipeline_preserves_inputs_and_independent_batch_rows(self): - pipe = self.get_pipeline_with() + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() original = {name: value.clone() for name, value in inputs["conditioning_features"].items()} single_generator = self.get_generator(7) @@ -300,12 +298,12 @@ def test_pipeline_preserves_inputs_and_independent_batch_rows(self): generators = [self.get_generator(7), self.get_generator(8)] batched = {name: self.batch_input(inputs[name], 2) for name in self.batch_input_params} output = pipe(**{**inputs, **batched, "generator": generators}) - for name, value in single.state.items(): - torch.testing.assert_close(output.state[name][:1], value, rtol=0, atol=1e-6) + torch.testing.assert_close(output.bb_ca[:1], single.bb_ca, rtol=0, atol=1e-6) + torch.testing.assert_close(output.local_latents[:1], single.local_latents, rtol=0, atol=1e-6) assert torch.equal(generators[0].get_state(), single_generator.get_state()) def test_num_designs_per_prompt_repeats_every_input_row(self): - pipe = self.get_pipeline_with() + pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() inputs.pop("generator") two_rows = {name: self.batch_input(inputs[name], 2) for name in self.batch_input_params} @@ -326,7 +324,7 @@ def test_num_designs_per_prompt_repeats_every_input_row(self): for name, value in two_rows.items() } by_hand = pipe(**{**inputs, **repeated}, generator=[self.get_generator(seed) for seed in range(4)]) - torch.testing.assert_close(output.state["bb_ca"], by_hand.state["bb_ca"], rtol=0, atol=0) + torch.testing.assert_close(output.bb_ca, by_hand.bb_ca, rtol=0, atol=0) with pytest.raises(ValueError, match="one distinct generator per batch row"): pipe(**{**inputs, **two_rows}, num_designs_per_prompt=2, generator=generators[:2]) with pytest.raises(ValueError, match="num_designs_per_prompt"): @@ -335,7 +333,7 @@ def test_num_designs_per_prompt_repeats_every_input_row(self): @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) def test_numpy_output_from_reduced_precision(self, dtype): """numpy has no bfloat16 or float16 path, and the released recipe runs in bfloat16.""" - pipe = self.get_pipeline_with().to(dtype) + pipe = self.get_pipeline().to(torch_device).to(dtype) output = pipe(**self.get_dummy_inputs(), output_type="np") assert isinstance(output.coords_nm, np.ndarray) assert output.coords_nm.dtype == np.float32 @@ -346,14 +344,14 @@ def test_non_finite_coordinates_rejected_before_rng(self): turns the whole latent NaN. The encode paths must say so rather than return a NaN structure.""" coords_nm = torch.zeros(1, 3, 37, 3, device=torch_device) coords_nm[0, 0, 1] = float("nan") - self.assert_rejected_before_rng( - self.get_pipeline_with(), - ValueError, - "not finite", - coords_nm=coords_nm, - partial_t_bb_ca=0.3, - partial_t_local_latents=0.6, - ) + inputs = { + **self.get_dummy_inputs(), + "generator": self.get_generator(7), + "coords_nm": coords_nm, + "partial_t_bb_ca": 0.3, + "partial_t_local_latents": 0.6, + } + self.assert_rejected_before_rng(self.get_pipeline().to(torch_device), inputs, ValueError, "not finite") @pytest.mark.parametrize("invalid", ["rank", "dtype", "shape", "device", "empty"]) def test_invalid_sampling_masks_rejected_before_rng(self, invalid): @@ -368,20 +366,23 @@ def test_invalid_sampling_masks_rejected_before_rng(self, invalid): inputs["supports_local_latent"] = inputs["supports_local_latent"].to("meta") else: inputs["padding_mask"] = inputs["padding_mask"][:, :0] - self.assert_rejected_before_rng(self.get_pipeline_with(), ValueError, "sampling masks|padding_mask", **inputs) + inputs = {**self.get_dummy_inputs(), "generator": self.get_generator(7), **inputs} + self.assert_rejected_before_rng( + self.get_pipeline().to(torch_device), inputs, ValueError, "sampling masks|padding_mask" + ) def test_generators_must_be_distinct_per_row(self): inputs = self.get_dummy_inputs() shared = self.get_generator(1) batched = {name: self.batch_input(inputs[name], 2) for name in self.batch_input_params} with pytest.raises(ValueError, match="one distinct generator per batch row"): - self.get_pipeline_with()(**{**inputs, **batched, "generator": [shared, shared]}) + self.get_pipeline().to(torch_device)(**{**inputs, **batched, "generator": [shared, shared]}) def test_stochastic_centering_offsets_each_molecule_by_one_shared_vector(self): - pipe = self.get_pipeline_with() + pipe = self.get_pipeline().to(torch_device) mask = torch.ones(1, 4, dtype=torch.bool, device=torch_device) conditioning = {"molecule_iid": torch.tensor([[0, 0, 1, 1]], device=torch_device)} - centered = pipe.prepare_latents(mask, mask, conditioning, [self.get_generator(3)])["bb_ca"] + centered, _ = pipe.prepare_latents(mask, mask, conditioning, [self.get_generator(3)]) base = pipe._sample_noise((1, 4, 3), mask, [self.get_generator(3)], zero_com=True) offset = centered - base # Tokens of one molecule move together, and the two molecules move differently. @@ -390,16 +391,16 @@ def test_stochastic_centering_offsets_each_molecule_by_one_shared_vector(self): assert not torch.allclose(offset[:, 0], offset[:, 2]) def test_prepare_latents_replays_distinct_row_generators(self): - pipe = self.get_pipeline_with() + pipe = self.get_pipeline().to(torch_device) mask = torch.tensor([[True, False], [True, False]], device=torch_device) conditioning = {"molecule_iid": torch.zeros(2, 2, dtype=torch.long, device=torch_device)} left = [torch.Generator().manual_seed(7), torch.Generator().manual_seed(11)] right = [torch.Generator().manual_seed(7), torch.Generator().manual_seed(11)] first = pipe.prepare_latents(mask, mask, conditioning, left) second = pipe.prepare_latents(mask, mask, conditioning, right) - for name in first: - torch.testing.assert_close(first[name], second[name], rtol=0, atol=0) - assert first[name][:, 1].count_nonzero() == 0 + for first_value, second_value in zip(first, second): + torch.testing.assert_close(first_value, second_value, rtol=0, atol=0) + assert first_value[:, 1].count_nonzero() == 0 assert all(a.get_state().equal(b.get_state()) for a, b in zip(left, right)) From 6db68a5f8735aa89a2f57afca4d2d7052612657e Mon Sep 17 00:00:00 2001 From: Kieran Didi Date: Tue, 20 Oct 2026 15:30:48 +0000 Subject: [PATCH 3/8] Remove the constructor options that are constant for the released checkpoint The transformer always recycles the trunk, updates the pair representation with triangle multiplication (without bias) and aligns the chirality feature, as the released checkpoint does. The converter now checks that a checkpoint has exactly these settings instead of carrying them as options. --- scripts/convert_rfd4_proteina_to_diffusers.py | 29 ++++- .../transformers/transformer_rfd4_proteina.py | 106 +++++++----------- .../test_models_transformer_rfd4_proteina.py | 6 +- .../test_pipeline_rfd4_proteina.py | 3 - 4 files changed, 62 insertions(+), 82 deletions(-) diff --git a/scripts/convert_rfd4_proteina_to_diffusers.py b/scripts/convert_rfd4_proteina_to_diffusers.py index 586ea7252a65..5347b00c4608 100644 --- a/scripts/convert_rfd4_proteina_to_diffusers.py +++ b/scripts/convert_rfd4_proteina_to_diffusers.py @@ -16,8 +16,9 @@ model trained with a feature this implementation lacks fails at conversion time instead of sampling from a different conditioning signal. Two cases to expect: -- `noise_align_chirality` orients the chirality feature towards the noisy backbone. It needs `chiral_frame_atoms` - and `chiral_slot_owner` in the prepared features; without them the pipeline raises rather than falling back. +- The chirality feature is oriented towards the noisy backbone, as in the released checkpoint. It needs + `chiral_frame_atoms` and `chiral_slot_owner` in the prepared features; without them the pipeline raises rather than + falling back. - `tri_mult_backend: cuequivariance` selects a fused triangle-multiplication kernel that this implementation does not have. Conversion prints a note and continues: results are unchanged, generation is slower. """ @@ -201,12 +202,21 @@ def _require(condition: bool, message: str) -> None: _COPIED_NN_KEYS = ( "nlayers token_dim nheads dim_cond pair_repr_dim token_type_emb_dim cath_emb_dim t_emb_dim seq_sep_dim " "xt_pair_dist_dim xt_pair_dist_min xt_pair_dist_max feats_seq feats_cond_seq feats_pair_repr feats_pair_cond " - "update_pair_repr update_pair_repr_every_n use_tri_mult tri_mult_c tri_mult_bias" + "update_pair_repr_every_n tri_mult_c" ).split() # Settings checked against what this implementation supports, then dropped because they select # the one behaviour it implements. -_VALIDATED_NN_KEYS = ("output_parameterization", "use_qkln", "parallel_mha_transition", "norm_type", "pair_write") +_VALIDATED_NN_KEYS = ( + "output_parameterization", + "use_qkln", + "parallel_mha_transition", + "norm_type", + "pair_write", + "update_pair_repr", + "use_tri_mult", + "tri_mult_bias", +) # Settings turned into Diffusers config entries by hand below. _DERIVED_NN_KEYS = ( @@ -252,6 +262,15 @@ def build_transformer_config(nn_config: dict[str, Any], latent_dim: int, num_tok _require(not nn_config.get("parallel_mha_transition", False), "parallel_mha_transition is unsupported") _require(nn_config.get("norm_type") in _RMS_NORM_TYPES, f"unsupported norm_type {nn_config.get('norm_type')!r}") _require(nn_config.get("pair_write") == "hadamard", f"unsupported pair_write {nn_config.get('pair_write')!r}") + # The released checkpoint has these on (and no triangle-multiplication bias); the model implements only that. + _require( + nn_config.get("update_pair_repr", False) + and nn_config.get("use_tri_mult", False) + and not nn_config.get("tri_mult_bias", False) + and nn_config.get("noise_align_chirality", False) + and (nn_config.get("recycle") or {}).get("enabled", False), + "pair updates with triangle multiplication, chirality alignment and trunk recycling", + ) _require( not nn_config.get("noise_align_internal_coords", False) and not nn_config.get("noise_align_internal_coords_translate", False) @@ -272,9 +291,7 @@ def build_transformer_config(nn_config: dict[str, Any], latent_dim: int, num_tok config.update( latent_dim=latent_dim, num_token_types=num_token_types, - recycle=bool(recycle.get("enabled", False)), num_sequence_classes=readout.get("num_classes", 33), - noise_align_chirality=bool(nn_config.get("noise_align_chirality", False)), ) if "recycle_remaining" in config.get("feats_cond_seq", ()): # The width of the one-hot is trained in. It is `nn.recycle.conditioning_dim` when set, and otherwise the depth diff --git a/src/diffusers/models/transformers/transformer_rfd4_proteina.py b/src/diffusers/models/transformers/transformer_rfd4_proteina.py index f17b0fcdbea5..bb680dfde696 100644 --- a/src/diffusers/models/transformers/transformer_rfd4_proteina.py +++ b/src/diffusers/models/transformers/transformer_rfd4_proteina.py @@ -185,7 +185,7 @@ def align_chirality( ) -> torch.Tensor: """Orient packed stereochemical vectors towards the current noisy backbone. - Checkpoints trained with `noise_align_chirality` see the chirality feature rotated onto the noisy neighbour + The released checkpoint sees the chirality feature rotated onto the noisy neighbour directions rather than in its canonical frame, so the two are different inputs to the same weights. Needs `chiral_frame_atoms` `(batch, tokens, slots, 2 + K)` and `chiral_slot_owner` `(batch, tokens, slots)` alongside the `chirality` feature. @@ -200,7 +200,7 @@ def align_chirality( missing = sorted({"chiral_frame_atoms", "chiral_slot_owner"} - features.keys()) if missing: raise KeyError( - f"noise_align_chirality needs {missing} in the prepared features. The released checkpoint was trained " + f"The chirality feature needs {missing} in the prepared features. The released checkpoint was trained " "with chirality vectors oriented towards the noisy backbone; without these tensors the model would " "silently receive a different conditioning signal." ) @@ -805,17 +805,17 @@ def forward( class RFD4ProteinaTriangleMultiplication(nn.Module): - def __init__(self, width: int, hidden: int, outgoing: bool, bias: bool) -> None: + def __init__(self, width: int, hidden: int, outgoing: bool) -> None: super().__init__() self.outgoing = outgoing - self.linear_a_p = nn.Linear(width, hidden, bias=bias) - self.linear_a_g = nn.Linear(width, hidden, bias=bias) - self.linear_b_p = nn.Linear(width, hidden, bias=bias) - self.linear_b_g = nn.Linear(width, hidden, bias=bias) - self.linear_g = nn.Linear(width, width, bias=bias) - self.linear_z = nn.Linear(hidden, width, bias=bias) - self.layer_norm_in = nn.LayerNorm(width, eps=1e-5, bias=bias) - self.layer_norm_out = nn.LayerNorm(hidden, eps=1e-5, bias=bias) + self.linear_a_p = nn.Linear(width, hidden, bias=False) + self.linear_a_g = nn.Linear(width, hidden, bias=False) + self.linear_b_p = nn.Linear(width, hidden, bias=False) + self.linear_b_g = nn.Linear(width, hidden, bias=False) + self.linear_g = nn.Linear(width, width, bias=False) + self.linear_z = nn.Linear(hidden, width, bias=False) + self.layer_norm_in = nn.LayerNorm(width, eps=1e-5, bias=False) + self.layer_norm_out = nn.LayerNorm(hidden, eps=1e-5, bias=False) def forward(self, pairs: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor: normalized = self.layer_norm_in(pairs) @@ -841,30 +841,20 @@ def forward(self, pairs: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor: class RFD4ProteinaPairUpdate(nn.Module): - def __init__( - self, - token_dim: int, - pair_repr_dim: int, - use_tri_mult: bool, - tri_mult_c: int, - tri_mult_bias: bool, - ) -> None: + def __init__(self, token_dim: int, pair_repr_dim: int, tri_mult_c: int) -> None: super().__init__() self.layer_norm_in = nn.LayerNorm(token_dim) self.linear_x = nn.Linear(token_dim, 2 * pair_repr_dim, bias=False) - self.use_tri_mult = use_tri_mult - if use_tri_mult: - self.tri_mult_out = RFD4ProteinaTriangleMultiplication(pair_repr_dim, tri_mult_c, True, tri_mult_bias) - self.tri_mult_in = RFD4ProteinaTriangleMultiplication(pair_repr_dim, tri_mult_c, False, tri_mult_bias) + self.tri_mult_out = RFD4ProteinaTriangleMultiplication(pair_repr_dim, tri_mult_c, True) + self.tri_mult_in = RFD4ProteinaTriangleMultiplication(pair_repr_dim, tri_mult_c, False) self.transition_out = RFD4ProteinaPairTransition(pair_repr_dim) def forward(self, values: torch.Tensor, pairs: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: pair_mask = mask[:, :, None] & mask[:, None, :] left, right = self.linear_x(self.layer_norm_in(values * mask[..., None])).chunk(2, dim=-1) pairs = (pairs + left[:, None] * right[:, :, None]) * pair_mask[..., None] - if self.use_tri_mult: - pairs = (pairs + self.tri_mult_out(pairs, pair_mask)) * pair_mask[..., None] - pairs = (pairs + self.tri_mult_in(pairs, pair_mask)) * pair_mask[..., None] + pairs = (pairs + self.tri_mult_out(pairs, pair_mask)) * pair_mask[..., None] + pairs = (pairs + self.tri_mult_in(pairs, pair_mask)) * pair_mask[..., None] return (pairs + self.transition_out(pairs, pair_mask)) * pair_mask[..., None] @@ -937,17 +927,13 @@ class RFD4ProteinaTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, Attenti Pair-distance range in nanometers. feats_seq, feats_cond_seq, feats_pair_repr, feats_pair_cond (`list[str]`): Names of the features projected into the token, conditioning, pair and pair-conditioning inputs. - update_pair_repr (`bool`, defaults to `True`), update_pair_repr_every_n (`int`, defaults to `2`): - Whether and how often the pair representation is updated from the tokens between layers. - use_tri_mult (`bool`, defaults to `True`), tri_mult_c (`int`, defaults to `256`), tri_mult_bias - (`bool`, defaults to `False`): Triangle multiplication settings. - recycle (`bool`, defaults to `True`): Whether the trunk is re-run `recycle_depth` times per call. + update_pair_repr_every_n (`int`, defaults to `2`): How often the pair representation is updated from the tokens + between layers. + tri_mult_c (`int`, defaults to `256`): Hidden width of the triangle multiplication in those updates. num_token_types (`int`, defaults to `7`): Size of the token-type vocabulary: protein, RNA, DNA, atomized, unknown, glycan and non-canonical amino acid. recycle_conditioning_dim (`int`, defaults to `1`): Width of the `recycle_remaining` one-hot, when that feature is in `feats_cond_seq`. - noise_align_chirality (`bool`, defaults to `False`): Orient the chirality feature towards the noisy backbone. - Checkpoints trained this way need `chiral_frame_atoms` and `chiral_slot_owner` features. num_sequence_classes (`int`, defaults to `147`): Size of the residue vocabulary the sequence readout predicts. """ @@ -1017,16 +1003,11 @@ def __init__( "na_ss_pair", ), feats_pair_cond: tuple[str, ...] = ("time_emb_bb_ca", "time_emb_local_latents"), - update_pair_repr: bool = True, update_pair_repr_every_n: int = 2, - use_tri_mult: bool = True, tri_mult_c: int = 256, - tri_mult_bias: bool = False, - recycle: bool = True, num_sequence_classes: int = 147, num_token_types: int = 7, recycle_conditioning_dim: int = 1, - noise_align_chirality: bool = False, ) -> None: super().__init__() feature_kwargs = { @@ -1042,7 +1023,6 @@ def __init__( "xt_pair_dist_max": xt_pair_dist_max, } self.nlayers = nlayers - self.update_pair_repr = update_pair_repr self.gradient_checkpointing = False self.init_repr_factory = RFD4ProteinaFeatureFactory(feats_seq, token_dim, "seq", **feature_kwargs) self.cond_factory = RFD4ProteinaFeatureFactory(feats_cond_seq, dim_cond, "seq", **feature_kwargs) @@ -1055,26 +1035,20 @@ def __init__( RFD4ProteinaTrunkLayer(token_dim, pair_repr_dim, nheads, dim_cond, "rms_norm", True) for _ in range(nlayers) ) - if update_pair_repr: - # Keyed by the index of the trunk layer each update follows, so only those layers have an entry. - self.pair_update_layers = nn.ModuleDict( - { - str(index): RFD4ProteinaPairUpdate( - token_dim, pair_repr_dim, use_tri_mult, tri_mult_c, tri_mult_bias - ) - for index in range(0, nlayers - 1, update_pair_repr_every_n) - } - ) - self.recycle_enabled = recycle + # Keyed by the index of the trunk layer each update follows, so only those layers have an entry. + self.pair_update_layers = nn.ModuleDict( + { + str(index): RFD4ProteinaPairUpdate(token_dim, pair_repr_dim, tri_mult_c) + for index in range(0, nlayers - 1, update_pair_repr_every_n) + } + ) # The released checkpoint conditions the trunk on the passes still to come, so the conditioning vector is rebuilt # for every pass instead of once. self.condition_on_recycle_remaining = "recycle_remaining" in feats_cond_seq - self.noise_align_chirality = noise_align_chirality - if recycle: - self.s_cycle_norm = nn.RMSNorm(token_dim, eps=1e-5) - self.s_cycle_proj = nn.Linear(token_dim, token_dim, bias=False) - self.z_cycle_norm = nn.RMSNorm(pair_repr_dim, eps=1e-5) - self.z_cycle_proj = nn.Linear(pair_repr_dim, pair_repr_dim, bias=False) + self.s_cycle_norm = nn.RMSNorm(token_dim, eps=1e-5) + self.s_cycle_proj = nn.Linear(token_dim, token_dim, bias=False) + self.z_cycle_norm = nn.RMSNorm(pair_repr_dim, eps=1e-5) + self.z_cycle_proj = nn.Linear(pair_repr_dim, pair_repr_dim, bias=False) self.local_latents_linear = nn.Sequential( nn.RMSNorm(token_dim, eps=1e-5), nn.Linear(token_dim, latent_dim, bias=False) ) @@ -1135,22 +1109,18 @@ def forward( features[name] = value mask = padding_mask.bool() emitted = mask & ~features["is_in_context_conditioning"].bool() - # Checkpoints trained with `noise_align_chirality` see the chirality feature rotated onto the noisy neighbour - # frame, so the feature the trunk reads is replaced before any layer runs. - if self.noise_align_chirality and "chirality" in features: + # The released checkpoint sees the chirality feature rotated onto the noisy neighbour frame, so the feature the + # trunk reads is replaced before any layer runs. + if "chirality" in features: features["chirality"] = align_chirality(features["chirality"].float(), features) - cycles = max(1, recycle_depth) if self.recycle_enabled else 1 + cycles = max(1, recycle_depth) values = self.init_repr_factory(features) * mask[..., None] pairs = self.pair_repr_builder(features) - if self.recycle_enabled: - previous_values, previous_pairs = torch.zeros_like(values), torch.zeros_like(pairs) + previous_values, previous_pairs = torch.zeros_like(values), torch.zeros_like(pairs) for cycle in range(cycles): - if self.recycle_enabled: - values_in = values + self.s_cycle_proj(self.s_cycle_norm(previous_values)) - pairs_in = (pairs + self.z_cycle_proj(self.z_cycle_norm(previous_pairs))).contiguous() - else: - values_in, pairs_in = values, pairs + values_in = values + self.s_cycle_proj(self.s_cycle_norm(previous_values)) + pairs_in = (pairs + self.z_cycle_proj(self.z_cycle_norm(previous_pairs))).contiguous() # Before the first pass the remaining-passes feature is absent (all-zero), which is distinct from "no passes # remain". Later passes carry the passes still to come. if cycle == 0 or self.condition_on_recycle_remaining: @@ -1167,7 +1137,7 @@ def forward( hidden = self._gradient_checkpointing_func(layer, hidden, hidden_pairs, condition, mask) else: hidden = layer(hidden, hidden_pairs, condition, mask) - if self.update_pair_repr and str(index) in self.pair_update_layers: + if str(index) in self.pair_update_layers: hidden_pairs = self.pair_update_layers[str(index)](hidden, hidden_pairs, mask) previous_values, previous_pairs = hidden, hidden_pairs values = previous_values diff --git a/tests/models/transformers/test_models_transformer_rfd4_proteina.py b/tests/models/transformers/test_models_transformer_rfd4_proteina.py index baf60c726734..afd464f7a504 100644 --- a/tests/models/transformers/test_models_transformer_rfd4_proteina.py +++ b/tests/models/transformers/test_models_transformer_rfd4_proteina.py @@ -52,9 +52,6 @@ def get_init_dict(self): "feats_pair_cond": [], "dim_cond": 8, "latent_dim": 4, - "update_pair_repr": False, - "use_tri_mult": False, - "recycle": False, } def get_dummy_inputs(self): @@ -84,7 +81,7 @@ def output_shape(self) -> tuple[int, ...]: class TestRFD4ProteinaTransformerModel(RFD4ProteinaTransformerTesterConfig, ModelTesterMixin): def test_transformer_executes_requested_recycle_depth(self): - init_dict = {**self.get_init_dict(), "recycle": True, "feats_cond_seq": ["recycle_remaining"], "t_emb_dim": 8} + init_dict = {**self.get_init_dict(), "feats_cond_seq": ["recycle_remaining"], "t_emb_dim": 8} transformer = RFD4ProteinaTransformer(**init_dict).to(torch_device).eval() inputs = self.get_dummy_inputs() cycles = [] @@ -125,7 +122,6 @@ def get_init_dict(self): return { **super().get_init_dict(), "nlayers": 3, - "update_pair_repr": True, "token_dim": 512, "nheads": 8, "pair_repr_dim": 64, diff --git a/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py b/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py index b2e8c0aeb6e0..70ab29b3ef8f 100644 --- a/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py +++ b/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py @@ -83,9 +83,6 @@ def get_dummy_components(self, depths=(1, 2), feats_seq=("xt_bb_ca", "xt_local_l latent_dim=4, num_sequence_classes=33, num_token_types=5, - update_pair_repr=False, - use_tri_mult=False, - recycle=True, ) torch.manual_seed(0) vae = RFD4ProteinaAutoencoder( From 603bd926d7ff9290d883066d2b72cd4c0688540d Mon Sep 17 00:00:00 2001 From: Kieran Didi Date: Tue, 20 Oct 2026 16:18:32 +0000 Subject: [PATCH 4/8] Spell out the scheduler test arguments and the docs examples, default the autoencoder to its released features The scheduler test helper takes its options by name, and the usage and batching examples pass the five required inputs by name. The autoencoder's default encoder features now match the released checkpoint, and the decoder moves its allowed-residue table to the token types' device. --- docs/source/en/api/pipelines/rfd4_proteina.md | 22 ++++++++++++++++--- .../autoencoders/autoencoder_rfd4_proteina.py | 12 ++++++++-- .../transformers/transformer_rfd4_proteina.py | 11 +++++----- .../test_scheduler_rfd4_proteina.py | 19 ++++++++-------- 4 files changed, 43 insertions(+), 21 deletions(-) diff --git a/docs/source/en/api/pipelines/rfd4_proteina.md b/docs/source/en/api/pipelines/rfd4_proteina.md index 6054644c723a..09ae90ae99b6 100644 --- a/docs/source/en/api/pipelines/rfd4_proteina.md +++ b/docs/source/en/api/pipelines/rfd4_proteina.md @@ -44,10 +44,17 @@ 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"] -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").manual_seed(42)) +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") ``` @@ -67,7 +74,16 @@ write_cif(output, features, atoms, "design.cif") ```py generators = [torch.Generator("cuda").manual_seed(seed) for seed in range(8)] - output = pipe(**inputs, conditioning_features=conditioning, num_designs_per_prompt=8, generator=generators) + 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 diff --git a/src/diffusers/models/autoencoders/autoencoder_rfd4_proteina.py b/src/diffusers/models/autoencoders/autoencoder_rfd4_proteina.py index 19161aa18a98..14584c89cdfd 100644 --- a/src/diffusers/models/autoencoders/autoencoder_rfd4_proteina.py +++ b/src/diffusers/models/autoencoders/autoencoder_rfd4_proteina.py @@ -433,7 +433,8 @@ def forward(self, features: Mapping[str, Any]) -> dict[str, torch.Tensor]: values = layer(values, None, condition, attention) logits = self.logit_linear(values) * emit[..., None] if "token_type" in features: - allowed = self._allowed_residue_types[features["token_type"].long()] + token_type = features["token_type"].long() + allowed = self._allowed_residue_types.to(token_type.device)[token_type] logits = logits.masked_fill(~(allowed | ~emit[..., None]), torch.finfo(logits.dtype).min) coordinates = self.struct_linear(values).reshape(*values.shape[:2], 37, 3).float() * emit[..., None, None] ca = features["ca_coors_nm"] @@ -543,7 +544,14 @@ def __init__( dim_cond: int = 128, pair_repr_dim: int = 256, token_type_emb_dim: int = 32, - encoder_feats_seq: tuple[str, ...] = ("token_type", "x1_a37coors_nm", "x1_a37coors_nm_rel", "chain_idx_seq"), + encoder_feats_seq: tuple[str, ...] = ( + "token_type", + "x1_a37coors_nm", + "x1_a37coors_nm_rel", + "chain_idx_seq", + "incident_bond_type_counts", + "atom37_element_charge_embedding", + ), decoder_feats_seq: tuple[str, ...] = ("ca_coors_nm", "z_latent_seq", "token_type", "cond_residue_type"), ) -> None: super().__init__() diff --git a/src/diffusers/models/transformers/transformer_rfd4_proteina.py b/src/diffusers/models/transformers/transformer_rfd4_proteina.py index bb680dfde696..d84627b6ef22 100644 --- a/src/diffusers/models/transformers/transformer_rfd4_proteina.py +++ b/src/diffusers/models/transformers/transformer_rfd4_proteina.py @@ -185,10 +185,9 @@ def align_chirality( ) -> torch.Tensor: """Orient packed stereochemical vectors towards the current noisy backbone. - The released checkpoint sees the chirality feature rotated onto the noisy neighbour - directions rather than in its canonical frame, so the two are different inputs to the same weights. Needs - `chiral_frame_atoms` `(batch, tokens, slots, 2 + K)` and `chiral_slot_owner` `(batch, tokens, slots)` alongside the - `chirality` feature. + The released checkpoint sees the chirality feature rotated onto the noisy neighbour directions rather than in its + canonical frame, so the two are different inputs to the same weights. Needs `chiral_frame_atoms` `(batch, tokens, + slots, 2 + K)` and `chiral_slot_owner` `(batch, tokens, slots)` alongside the `chirality` feature. Args: chirality_vectors (`torch.Tensor`): Packed unit vectors `(batch, tokens, 3 * slots)`. @@ -927,8 +926,8 @@ class RFD4ProteinaTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, Attenti Pair-distance range in nanometers. feats_seq, feats_cond_seq, feats_pair_repr, feats_pair_cond (`list[str]`): Names of the features projected into the token, conditioning, pair and pair-conditioning inputs. - update_pair_repr_every_n (`int`, defaults to `2`): How often the pair representation is updated from the tokens - between layers. + update_pair_repr_every_n (`int`, defaults to `2`): + How often the pair representation is updated from the tokens between layers. tri_mult_c (`int`, defaults to `256`): Hidden width of the triangle multiplication in those updates. num_token_types (`int`, defaults to `7`): Size of the token-type vocabulary: protein, RNA, DNA, atomized, unknown, glycan and non-canonical amino acid. diff --git a/tests/schedulers/test_scheduler_rfd4_proteina.py b/tests/schedulers/test_scheduler_rfd4_proteina.py index 1c2f1a8da700..002ead7eda17 100644 --- a/tests/schedulers/test_scheduler_rfd4_proteina.py +++ b/tests/schedulers/test_scheduler_rfd4_proteina.py @@ -18,19 +18,19 @@ from diffusers import RFD4ProteinaScheduler -def scheduler(stochastic=False, **kwargs): +def scheduler(schedule_mode="log", noise_mode="1/t", stochastic=False): """A tiny backbone-style scheduler. With `stochastic=False`, `t_lim_ode_below=1` keeps every interval on the noise-free branch; `stochastic=True` uses the released thresholds.""" config = { - "schedule_mode": "log", - "noise_mode": "1/t", + "schedule_mode": schedule_mode, + "noise_mode": noise_mode, "tsr_k": 16.0, "tsr_sigma": 3.0, "tail_steps": 0, "t_lim_ode": 0.98, "t_lim_ode_below": 0.02 if stochastic else 1.0, } - return RFD4ProteinaScheduler(**{**config, **kwargs}) + return RFD4ProteinaScheduler(**config) def inputs(): @@ -65,10 +65,9 @@ def test_cuda_schedule_keeps_coefficients_on_cpu(): @pytest.mark.parametrize("schedule_mode, noise_mode", [("log", "1/t"), ("power", "tan")]) def test_partial_schedule_is_exact_nonlinear_full_grid_tail(schedule_mode, noise_mode): - kwargs = {"schedule_mode": schedule_mode, "noise_mode": noise_mode} - full = scheduler(stochastic=True, **kwargs) + full = scheduler(schedule_mode, noise_mode, stochastic=True) full.set_timesteps(4) - partial = scheduler(stochastic=True, **kwargs) + partial = scheduler(schedule_mode, noise_mode, stochastic=True) partial.set_timesteps(2, start_time=float(full._times[2])) torch.testing.assert_close(partial._times, full._times[2:], rtol=0, atol=0) torch.testing.assert_close(partial._noise, full._noise[2:], rtol=0, atol=0) @@ -151,7 +150,7 @@ def test_stochastic_step_default_generator_matches_explicit_stream(device): assert not torch.equal(endpoint, initial_rng), "The test must exercise a stochastic interval." -@pytest.mark.parametrize("kwargs", [{"schedule_mode": "uniform"}, {"noise_mode": "ode"}]) -def test_options_outside_the_released_recipe_are_rejected(kwargs): +@pytest.mark.parametrize("schedule_mode, noise_mode", [("uniform", "1/t"), ("log", "ode")]) +def test_options_outside_the_released_recipe_are_rejected(schedule_mode, noise_mode): with pytest.raises(ValueError, match="must be"): - scheduler(**kwargs) + scheduler(schedule_mode, noise_mode) From 5ecd352b87652a0f587152057bce266669a38994 Mon Sep 17 00:00:00 2001 From: Kieran Didi Date: Tue, 20 Oct 2026 16:41:12 +0000 Subject: [PATCH 5/8] Follow the repository's test and docs checks Restore the LoRA mixin on the loaders page, which utils/check_support_list.py requires for every LoRA loader mixin, and keep the mention in the pipeline docs. Compose the attention and torch.compile testers on the models, drop the model-level LoRA tests that the first pull request should not carry, and make the pipeline's backbone-noising step public as a lifecycle method. --- docs/source/en/api/loaders/lora.md | 5 +++++ .../rfd4_proteina/pipeline_rfd4_proteina.py | 6 +++--- .../test_models_autoencoder_rfd4_proteina.py | 16 +++++++++++++++- .../test_models_transformer_rfd4_proteina.py | 11 +++++++++-- 4 files changed, 32 insertions(+), 6 deletions(-) diff --git a/docs/source/en/api/loaders/lora.md b/docs/source/en/api/loaders/lora.md index 4c8eb20b299a..1a4fded1fa10 100644 --- a/docs/source/en/api/loaders/lora.md +++ b/docs/source/en/api/loaders/lora.md @@ -39,11 +39,16 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi - [`ErnieImageLoraLoaderMixin`] provides similar functions for [Ernie-Image](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ernie_image). - [`LTX2LoraLoaderMixin`] provides similar functions for [Flux2](https://huggingface.co/docs/diffusers/main/en/api/pipelines/ltx2). - [`MiniMaxH3LoraLoaderMixin`] provides similar functions for [MiniMax-H3](https://huggingface.co/docs/diffusers/main/en/api/pipelines/minimax_h3). +- [`RFD4ProteinaLoraLoaderMixin`] provides similar functions for [RFD4-Proteína](https://huggingface.co/docs/diffusers/main/en/api/pipelines/rfd4_proteina). - [`LoraBaseMixin`] provides a base class with several utility methods to fuse, unfuse, unload, LoRAs and more. > [!TIP] > To learn more about how to load LoRA weights, see the [LoRA](../../tutorials/using_peft_for_inference) loading guide. +## RFD4ProteinaLoraLoaderMixin + +[[autodoc]] loaders.lora_pipeline.RFD4ProteinaLoraLoaderMixin + ## LoraBaseMixin [[autodoc]] loaders.lora_base.LoraBaseMixin diff --git a/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py b/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py index 30b0e0e071c5..a5a14db427a5 100644 --- a/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py +++ b/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py @@ -270,7 +270,7 @@ def prepare_latents( ) return bb * bb_mask[..., None], latents * latent_mask[..., None] - def _noise_backbone(self, target_backbone, bb_mask, conditioning, generators, index): + def noise_backbone(self, target_backbone, bb_mask, conditioning, generators, index): """Noise the preserved backbone to the scheduled time of step `index`, for inverse folding.""" noise = self._sample_noise(tuple(target_backbone.shape), bb_mask, generators, zero_com=True) noise = self._center_noise(noise, bb_mask, generators, conditioning) @@ -453,7 +453,7 @@ def __call__( ) if inverse_folding: - bb_ca = self._noise_backbone(target_backbone, bb_mask, conditioning, generators, 0) + bb_ca = self.noise_backbone(target_backbone, bb_mask, conditioning, generators, 0) if latents_bb_ca is not None: # A supplied starting point replaces the sampled one; the schedule is unchanged. @@ -504,7 +504,7 @@ def __call__( if index == steps - 1: bb_ca = target_backbone.clone() else: - bb_ca = self._noise_backbone(target_backbone, bb_mask, conditioning, generators, index + 1) + bb_ca = self.noise_backbone(target_backbone, bb_mask, conditioning, generators, index + 1) if callback_on_step_end is not None: callback_kwargs = {} for name in callback_on_step_end_tensor_inputs: diff --git a/tests/models/autoencoders/test_models_autoencoder_rfd4_proteina.py b/tests/models/autoencoders/test_models_autoencoder_rfd4_proteina.py index 3e873b62c57a..38d9388de142 100644 --- a/tests/models/autoencoders/test_models_autoencoder_rfd4_proteina.py +++ b/tests/models/autoencoders/test_models_autoencoder_rfd4_proteina.py @@ -19,7 +19,13 @@ from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, torch_device -from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, +) enable_full_determinism() @@ -99,3 +105,11 @@ def get_init_dict(self): "pair_repr_dim": 64, "dim_cond": 64, } + + +class TestRFD4ProteinaAutoencoderAttention(RFD4ProteinaAutoencoderTesterConfig, AttentionTesterMixin): + pass + + +class TestRFD4ProteinaAutoencoderCompile(RFD4ProteinaAutoencoderTesterConfig, TorchCompileTesterMixin): + pass diff --git a/tests/models/transformers/test_models_transformer_rfd4_proteina.py b/tests/models/transformers/test_models_transformer_rfd4_proteina.py index afd464f7a504..53941a1cad58 100644 --- a/tests/models/transformers/test_models_transformer_rfd4_proteina.py +++ b/tests/models/transformers/test_models_transformer_rfd4_proteina.py @@ -21,10 +21,11 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import ( + AttentionTesterMixin, BaseModelTesterConfig, - LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, + TorchCompileTesterMixin, TrainingTesterMixin, ) @@ -135,5 +136,11 @@ def test_gradient_checkpointing_is_applied(self): super().test_gradient_checkpointing_is_applied(expected_set={"RFD4ProteinaTransformer"}) -class TestRFD4ProteinaTransformerLoRA(RFD4ProteinaTransformerTesterConfig, LoraTesterMixin): +class TestRFD4ProteinaTransformerAttention(RFD4ProteinaTransformerTesterConfig, AttentionTesterMixin): + def get_init_dict(self): + # The processor-count test needs more than one attention layer to build a dict of the wrong size. + return {**super().get_init_dict(), "nlayers": 2} + + +class TestRFD4ProteinaTransformerCompile(RFD4ProteinaTransformerTesterConfig, TorchCompileTesterMixin): pass From f9164067a37eea67280ca06b998846d7906e41b6 Mon Sep 17 00:00:00 2001 From: Kieran Didi Date: Tue, 20 Oct 2026 17:09:55 +0000 Subject: [PATCH 6/8] Keep the pipeline test helpers at the top of the test class --- .../test_pipeline_rfd4_proteina.py | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py b/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py index 70ab29b3ef8f..afd9da2adec9 100644 --- a/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py +++ b/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py @@ -124,6 +124,19 @@ def batch_input(self, value, batch_size): class TestRFD4ProteinaPipeline(RFD4ProteinaPipelineTesterConfig, PipelineTesterMixin): + def assert_rejected_before_rng(self, pipe, inputs, error, match): + before = inputs["generator"].get_state().clone() + with pytest.raises(error, match=match): + pipe(**inputs) + assert torch.equal(inputs["generator"].get_state(), before) + + def coord_group_inputs(self): + inputs = prepared_inputs(torch_device) + groups = torch.zeros(1, 3, 37, dtype=torch.long) + groups[:, 1] = 1 + inputs["conditioning_features"][CONDITION_KEYS["coordinate_group"]] = groups.to(torch_device) + return inputs, groups + def test_components_function(self): # The shared test expects every dummy component to be a module, but `recycle_depths` is a list of ints that the # pipeline stores in its config, so it is left out of the comparison. @@ -139,12 +152,6 @@ def test_components_function(self): def test_callback_inputs(self): pass - def assert_rejected_before_rng(self, pipe, inputs, error, match): - before = inputs["generator"].get_state().clone() - with pytest.raises(error, match=match): - pipe(**inputs) - assert torch.equal(inputs["generator"].get_state(), before) - def test_pipeline_calls_model_once_per_step_and_recycles_inside_the_call(self): pipe = self.get_pipeline().to(torch_device) calls, cycles = [], [] @@ -187,13 +194,6 @@ def test_missing_sequence_condition_fails_before_sampling(self): inputs = {**self.get_dummy_inputs(), "generator": self.get_generator(7), "conditioning_features": {}} self.assert_rejected_before_rng(self.get_pipeline().to(torch_device), inputs, KeyError, SEQUENCE_KEY) - def coord_group_inputs(self): - inputs = prepared_inputs(torch_device) - groups = torch.zeros(1, 3, 37, dtype=torch.long) - groups[:, 1] = 1 - inputs["conditioning_features"][CONDITION_KEYS["coordinate_group"]] = groups.to(torch_device) - return inputs, groups - def test_pipeline_prepares_random_coord_groups_outside_transformer(self): pipe = self.get_pipeline( **self.get_dummy_components(depths=(0, 1), feats_seq=["xt_bb_ca", "xt_local_latents", "coord_group_seq"]) From da3df15ee83f39b2a1b6c4890c92330221e96f82 Mon Sep 17 00:00:00 2001 From: Kieran Didi Date: Tue, 20 Oct 2026 19:22:26 +0000 Subject: [PATCH 7/8] Drop the training-only detach calls and make the noise helpers static methods --- .../models/transformers/transformer_rfd4_proteina.py | 5 ++--- .../pipelines/rfd4_proteina/pipeline_rfd4_proteina.py | 4 ++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_rfd4_proteina.py b/src/diffusers/models/transformers/transformer_rfd4_proteina.py index d84627b6ef22..526a770af5cc 100644 --- a/src/diffusers/models/transformers/transformer_rfd4_proteina.py +++ b/src/diffusers/models/transformers/transformer_rfd4_proteina.py @@ -1145,9 +1145,8 @@ def forward( local_latents = ( self.local_latents_linear(values) * (emitted & features["supports_local_latent"].bool())[..., None] ) - sequence_logits = self.sequence_readout(values.detach()) * mask[..., None] - detached = values.detach() - query, key = self.insertion_q_proj(detached), self.insertion_k_proj(detached) + sequence_logits = self.sequence_readout(values) * mask[..., None] + query, key = self.insertion_q_proj(values), self.insertion_k_proj(values) insertion_logits = torch.bmm(query, key.transpose(1, 2)) * self._insertion_head_scale output = RFD4ProteinaTransformerOutput( backbone=backbone, diff --git a/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py b/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py index a5a14db427a5..babc946b58f5 100644 --- a/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py +++ b/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py @@ -203,8 +203,8 @@ def check_inputs( if output_type not in {"pt", "np"}: raise ValueError(f"output_type must be 'pt' or 'np', got {output_type!r}.") + @staticmethod def _sample_noise( - self, shape: tuple[int, ...], mask: torch.Tensor, generators: list[torch.Generator] | None, @@ -217,8 +217,8 @@ def _sample_noise( noise -= noise.sum(dim=-2, keepdim=True) / count[..., None] return noise * mask[..., None] + @staticmethod def _center_noise( - self, noise: torch.Tensor, mask: torch.Tensor, generators: list[torch.Generator] | None, From 07084aca8cfd122e1fb5cd21f37497f1bf5c8ac8 Mon Sep 17 00:00:00 2001 From: Kieran Didi Date: Tue, 20 Oct 2026 19:56:13 +0000 Subject: [PATCH 8/8] Remove num_inference_steps and skip the unused encode for inverse folding The number of steps is fixed by the released recipe and is the length of recycle_depths. Inverse folding starts both modalities at time zero, so its latents are noise and the encoded input was only consuming random draws. --- docs/source/en/api/pipelines/rfd4_proteina.md | 4 +- .../rfd4_proteina/pipeline_rfd4_proteina.py | 42 ++++++++++--------- .../test_pipeline_rfd4_proteina.py | 15 ++++--- 3 files changed, 34 insertions(+), 27 deletions(-) diff --git a/docs/source/en/api/pipelines/rfd4_proteina.md b/docs/source/en/api/pipelines/rfd4_proteina.md index 09ae90ae99b6..daba2cfbcdae 100644 --- a/docs/source/en/api/pipelines/rfd4_proteina.md +++ b/docs/source/en/api/pipelines/rfd4_proteina.md @@ -66,8 +66,8 @@ write_cif(output, features, atoms, "design.cif") 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. -* `num_inference_steps` has to match the length of the `recycle_depths` pipeline config. Leave it unset to use the - released recipe of 60 steps and 180 trunk passes. +* 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`. diff --git a/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py b/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py index babc946b58f5..016a6c51f869 100644 --- a/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py +++ b/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py @@ -251,6 +251,7 @@ def prepare_latents( start_times: tuple[float, float] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """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 @@ -265,9 +266,13 @@ def prepare_latents( latent_noise = self._sample_noise(latent_shape, latent_mask, generators) bb_noise = self._center_noise(bb_noise, bb_mask, generators, conditioning) bb = self.bb_ca_scheduler.add_noise(clean_bb * bb_mask[..., None], bb_noise, start_times[0]) - latents = self.local_latents_scheduler.add_noise( - clean_latents * latent_mask[..., None], latent_noise, start_times[1] - ) + if clean_latents is None: + # Inverse folding starts both modalities at time zero, where the noised latents are the noise itself. + latents = latent_noise + else: + latents = self.local_latents_scheduler.add_noise( + clean_latents * latent_mask[..., None], latent_noise, start_times[1] + ) return bb * bb_mask[..., None], latents * latent_mask[..., None] def noise_backbone(self, target_backbone, bb_mask, conditioning, generators, index): @@ -292,7 +297,6 @@ def __call__( conditioning_features: dict[str, Any] | None = None, num_designs_per_prompt: int = 1, generator: torch.Generator | list[torch.Generator] | None = None, - num_inference_steps: int | None = None, partial_t_bb_ca: float | None = None, partial_t_local_latents: float | None = None, inverse_folding: bool = False, @@ -334,8 +338,6 @@ def __call__( 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. - num_inference_steps (`int`, *optional*): - Number of denoising steps. Defaults to, and must equal, the length of the `recycle_depths` config. 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`. @@ -420,9 +422,7 @@ def __call__( bb_mask = padding_mask & ~is_in_context_conditioning latent_mask = bb_mask & supports_local_latent recycle_depths = self.config.recycle_depths - steps = len(recycle_depths) if num_inference_steps is None else num_inference_steps - if steps != len(recycle_depths): - raise ValueError(f"num_inference_steps must equal the {len(recycle_depths)} entries of recycle_depths.") + steps = len(recycle_depths) is_partial = partial_t_bb_ca is not None start_times = (float(partial_t_bb_ca), float(partial_t_local_latents)) if is_partial else (None, None) self.bb_ca_scheduler.set_timesteps(steps, device=device, start_time=start_times[0]) @@ -432,16 +432,20 @@ def __call__( target_backbone = clean_bb = clean_latents = None if is_partial: - encoded = self.vae.encode( - coords_nm=coords_nm, - padding_mask=padding_mask, - token_type=token_type, - atom37_mask=atom37_mask, - conditioning=conditioning, - generator=generators, - ) - clean_bb, clean_latents = coords_nm[..., 1, :].float(), encoded.latent.float() - target_backbone = clean_bb.clone() if inverse_folding else None + clean_bb = coords_nm[..., 1, :].float() + if inverse_folding: + # Both start times are zero, so the starting latents are pure noise and need no encoded input. + target_backbone = clean_bb.clone() + else: + encoded = self.vae.encode( + coords_nm=coords_nm, + padding_mask=padding_mask, + token_type=token_type, + atom37_mask=atom37_mask, + conditioning=conditioning, + generator=generators, + ) + clean_latents = encoded.latent.float() bb_ca, local_latents = self.prepare_latents( bb_mask, latent_mask, diff --git a/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py b/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py index afd9da2adec9..03979395cde2 100644 --- a/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py +++ b/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py @@ -59,9 +59,9 @@ def allowed_residue_types(): class RFD4ProteinaPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = RFD4ProteinaPipeline - required_input_params_in_call_signature = frozenset(["padding_mask", "generator", "num_inference_steps"]) + required_input_params_in_call_signature = frozenset(["padding_mask", "generator"]) # `generator` is required: every design needs its own random stream. - optional_input_params = frozenset(["num_inference_steps", "return_dict"]) + optional_input_params = frozenset(["return_dict"]) batch_input_params = frozenset( ["padding_mask", "is_in_context_conditioning", "supports_local_latent", "atom37_mask", "conditioning_features"] ) @@ -113,7 +113,6 @@ def get_dummy_inputs(self): return { **prepared_inputs(torch_device), "generator": self.get_generator(0), - "num_inference_steps": 2, } def batch_input(self, value, batch_size): @@ -137,6 +136,10 @@ def coord_group_inputs(self): inputs["conditioning_features"][CONDITION_KEYS["coordinate_group"]] = groups.to(torch_device) return inputs, groups + def test_inference_batch_single_identical(self): + # The shared test copies `num_inference_steps` into the batched inputs, and this pipeline has no such argument. + super().test_inference_batch_single_identical(additional_params_copy_to_batched_inputs=[]) + def test_components_function(self): # The shared test expects every dummy component to be a module, but `recycle_depths` is a list of ints that the # pipeline stores in its config, so it is left out of the comparison. @@ -210,7 +213,7 @@ def test_pipeline_prepares_random_coord_groups_outside_transformer(self): with_kwargs=True, ) actual_generator = self.get_generator(7) - pipe(**inputs, generator=actual_generator, num_inference_steps=2) + pipe(**inputs, generator=actual_generator) handle.remove() assert len(captured) == 2 assert all(torch.equal(actual, wanted) for actual, wanted in zip(captured, expected)) @@ -232,7 +235,7 @@ def test_partial_and_inverse_initialization_use_clean_tensor_inputs(self): pipe = self.get_pipeline().to(torch_device) coords_nm = torch.zeros(1, 3, 37, 3, device=torch_device) coords_nm[:, 0, 1] = torch.tensor([1.0, 2.0, 3.0], device=torch_device) - inputs = {**prepared_inputs(torch_device), "coords_nm": coords_nm, "num_inference_steps": 2} + inputs = {**prepared_inputs(torch_device), "coords_nm": coords_nm} partial = pipe(**inputs, generator=self.get_generator(13), partial_t_bb_ca=0.5, partial_t_local_latents=0.5) inverse = pipe( **inputs, @@ -278,7 +281,7 @@ def fixed_logits(module, args, kwargs, output): lambda module, args: decoded_features.append(args[0][SEQUENCE_KEY].cpu()) ), ] - pipe(**inputs, generator=self.get_generator(7), num_inference_steps=2) + pipe(**inputs, generator=self.get_generator(7)) for handle in handles: handle.remove() # Position 0 is filled with the best protein residue (not the mask token or an RNA type); position 1 is user-fixed.