diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 39fe7ccc6a61..4384bc183da2 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -415,6 +415,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 @@ -714,6 +718,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 @@ -844,6 +852,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..daba2cfbcdae --- /dev/null +++ b/docs/source/en/api/pipelines/rfd4_proteina.md @@ -0,0 +1,128 @@ + + +# 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"] +conditioning = {name: value for name, value in features.items() if name not in required} + +output = pipe( + padding_mask=features["padding_mask"], + is_in_context_conditioning=features["is_in_context_conditioning"], + supports_local_latent=features["supports_local_latent"], + token_type=features["token_type"], + atom37_mask=features["atom37_mask"], + conditioning_features=conditioning, + generator=torch.Generator("cuda").manual_seed(42), +) +write_cif(output, features, atoms, "design.cif") +``` + +## Tips + +* The model repository ships an annotated input and a generated design for each task under `examples/`: protein binder, + antibody, VHH, enzyme, motif scaffolding, small-molecule binder, unconditional, inverse folding and partial diffusion. +* Pass `partial_t_bb_ca` and `partial_t_local_latents` with start times in `[0, 1)` to redesign an input from a noised + start, or `inverse_folding=True` with both start times at `0.0` to keep the backbone and design the sequence. Both + encode the input, so they need clean `coords_nm` and reject coordinates that are not finite, which are usually + unresolved residues. +* The number of steps is the length of the `recycle_depths` pipeline config, 60 for the released recipe (180 trunk + passes). Change the schedule by passing a different `recycle_depths` list when constructing the pipeline. +* To generate several designs of each input, set `num_designs_per_prompt` and pass one distinct generator per + design. The output has `batch * num_designs_per_prompt` rows, and row `r` is design `r % num_designs_per_prompt` of + input row `r // num_designs_per_prompt`. + + ```py + generators = [torch.Generator("cuda").manual_seed(seed) for seed in range(8)] + output = pipe( + padding_mask=features["padding_mask"], + is_in_context_conditioning=features["is_in_context_conditioning"], + supports_local_latent=features["supports_local_latent"], + token_type=features["token_type"], + atom37_mask=features["atom37_mask"], + conditioning_features=conditioning, + num_designs_per_prompt=8, + generator=generators, + ) + ``` + + `write_cif` reads the features of the row it writes, so repeat them the same way, with `repeat_interleave` rather + than `repeat`: + + ```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. + +## 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 + - 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..5347b00c4608 --- /dev/null +++ b/scripts/convert_rfd4_proteina_to_diffusers.py @@ -0,0 +1,618 @@ +# 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: + +- 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. +""" + +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_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", + "update_pair_repr", + "use_tri_mult", + "tri_mult_bias", +) + +# 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}") + # 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) + 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, + num_sequence_classes=readout.get("num_classes", 33), + ) + 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 d20c85f5344d..ac677917242b 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -360,6 +360,8 @@ "QwenImageControlNetModel", "QwenImageMultiControlNetModel", "QwenImageTransformer2DModel", + "RFD4ProteinaAutoencoder", + "RFD4ProteinaTransformer", "SanaControlNetModel", "SanaTransformer2DModel", "SanaVideoTransformer3DModel", @@ -439,6 +441,7 @@ "LDMSuperResolutionPipeline", "PNDMPipeline", "RePaintPipeline", + "RFD4ProteinaPipeline", "ScoreSdeVePipeline", "StableDiffusionMixin", ] @@ -487,6 +490,7 @@ "PiflowScheduler", "PNDMScheduler", "RePaintScheduler", + "RFD4ProteinaScheduler", "SASolverScheduler", "SchedulerMixin", "SCMScheduler", @@ -1278,6 +1282,8 @@ QwenImageControlNetModel, QwenImageMultiControlNetModel, QwenImageTransformer2DModel, + RFD4ProteinaAutoencoder, + RFD4ProteinaTransformer, SanaControlNetModel, SanaTransformer2DModel, SanaVideoTransformer3DModel, @@ -1355,6 +1361,7 @@ LDMSuperResolutionPipeline, PNDMPipeline, RePaintPipeline, + RFD4ProteinaPipeline, ScoreSdeVePipeline, StableDiffusionMixin, ) @@ -1401,6 +1408,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 78d0c98df0ee..3713cdfec22c 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"] @@ -155,6 +156,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"] @@ -225,6 +227,7 @@ MiniMaxMusic3Vocoder, MMAudioVAE, MMAudioVocoder, + RFD4ProteinaAutoencoder, VQModel, ) from .cache_utils import CacheMixin @@ -312,6 +315,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..14584c89cdfd --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_rfd4_proteina.py @@ -0,0 +1,698 @@ +# 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: + 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"] + 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", + "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__() + 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 3697c953c243..dd7014096e54 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -65,6 +65,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..526a770af5cc --- /dev/null +++ b/src/diffusers/models/transformers/transformer_rfd4_proteina.py @@ -0,0 +1,1157 @@ +# 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. + + 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)`. + 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"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." + ) + 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) -> None: + super().__init__() + self.outgoing = outgoing + 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) + 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, 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.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] + 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_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`. + 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_every_n: int = 2, + tri_mult_c: int = 256, + num_sequence_classes: int = 147, + num_token_types: int = 7, + recycle_conditioning_dim: int = 1, + ) -> 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.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) + ) + # 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.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[torch.Tensor]: + """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() + # 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) + values = self.init_repr_factory(features) * mask[..., None] + pairs = self.pair_repr_builder(features) + previous_values, previous_pairs = torch.zeros_like(values), torch.zeros_like(pairs) + for cycle in range(cycles): + 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: + 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 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) * 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, + 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 ff8772b73fc6..5a2c9fb8ef57 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"] = [ @@ -581,6 +582,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..b87dce24d617 --- /dev/null +++ b/src/diffusers/pipelines/rfd4_proteina/pipeline_output.py @@ -0,0 +1,47 @@ +# 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. + bb_ca (`torch.Tensor`): + Final backbone state `(B, N, 3)` in nanometers. + local_latents (`torch.Tensor`): + Final local-latent state `(B, N, D)`. + clean_bb_ca (`torch.Tensor`): + Backbone prediction of the clean state at the last denoising step. + clean_local_latents (`torch.Tensor`): + Local-latent prediction of the clean state at the last denoising step. + sequence_logits (`torch.Tensor`): + Final per-token residue logits `(B, N, 147)`. + insertion_logits (`torch.Tensor`): + Final insertion predictions `(B, N, N)` for unindexed motifs. + """ + + coords_nm: torch.Tensor + residue_type: torch.Tensor + emit_mask: torch.Tensor + atom_mask: 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 new file mode 100644 index 000000000000..016a6c51f869 --- /dev/null +++ b/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py @@ -0,0 +1,582 @@ +# 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( + ... padding_mask=features["padding_mask"], + ... is_in_context_conditioning=features["is_in_context_conditioning"], + ... supports_local_latent=features["supports_local_latent"], + ... token_type=features["token_type"], + ... atom37_mask=features["atom37_mask"], + ... conditioning_features={name: value for name, value in features.items() if name not in required}, + ... generator=torch.Generator("cuda").manual_seed(42), + ... ) + >>> write_cif(output, features, atoms, "design.cif") + ``` +""" + + +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 = [ + "bb_ca", + "local_latents", + "last_clean_bb_ca", + "last_clean_local_latents", + "sequence_logits", + "insertion_logits", + ] + + 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) + ) + + @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, + 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}.") + + @staticmethod + def _sample_noise( + 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] + + @staticmethod + def _center_noise( + 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, + ) -> 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 + 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_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]) + 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): + """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, + 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] = ["bb_ca", "local_latents"], + 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. It is one mapping because the preprocessor emits about + 30 optional tensors under fixed names, and a signature with that many arguments would tie the pipeline + to the preprocessor's naming. + num_designs_per_prompt (`int`, defaults to `1`): + Number of designs to generate for each input row. The inputs are repeated along the batch axis, so the + output has `batch * num_designs_per_prompt` rows and row `r` is design `r % num_designs_per_prompt` of + input row `r // num_designs_per_prompt`. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One distinct generator per output row, for reproducible designs. The generators advance and the global + PyTorch random stream is left untouched. Omit it to draw from the global stream. + partial_t_bb_ca (`float`, *optional*): + Start time in [0, 1) for the backbone, to start from a noised version of the input instead of pure + noise. Give it together with `partial_t_local_latents`; requires `coords_nm` and `atom37_mask`. + partial_t_local_latents (`float`, *optional*): + Start time in [0, 1) for the local latents; see `partial_t_bb_ca`. + inverse_folding (`bool`, defaults to `False`): + Preserve the input backbone and design the sequence. Requires zero for both partial start times. + latents_bb_ca (`torch.Tensor`, *optional*): + Pre-generated starting backbone state `(batch * num_designs_per_prompt, tokens, 3)`, to replace the one + this call would otherwise sample. The schedule is unchanged. Give it together with + `latents_local_latents`. + latents_local_latents (`torch.Tensor`, *optional*): + Pre-generated starting local latents `(batch * num_designs_per_prompt, tokens, latent_dim)`. + attention_kwargs (`dict`, *optional*): + Passed to the transformer; `scale` sets the LoRA scale. + callback_on_step_end (`Callable`, *optional*): + Called as `callback_on_step_end(pipeline, step, timestep, callback_kwargs)` after each step. + `callback_kwargs` holds the tensors named in `callback_on_step_end_tensor_inputs`; a returned dict may + replace `"bb_ca"` and `"local_latents"`. + output_type (`str`, defaults to `"pt"`): + `"pt"` returns torch tensors and `"np"` returns numpy arrays. + callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["bb_ca", "local_latents"]`): + Tensors passed to the callback: `"bb_ca"`, `"local_latents"`, `"last_clean_bb_ca"`, + `"last_clean_local_latents"`, `"sequence_logits"`, `"insertion_logits"`. + return_dict (`bool`, defaults to `True`): Return a typed output, or a tuple of its non-None fields. + + 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] + 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) + 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]) + + use_coord_group_seq = "coord_group_seq" in self.transformer.config.feats_seq + + target_backbone = clean_bb = clean_latents = None + if is_partial: + 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, + conditioning, + generators, + clean_bb, + clean_latents, + start_times if is_partial else None, + ) + + if inverse_folding: + 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. + bb_ca, local_latents = latents_bb_ca.to(device), latents_local_latents.to(device) + + 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 + with self.progress_bar(total=steps) as progress_bar: + for index in range(steps): + if self.interrupt: + continue + 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. + 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=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, + conditioning=step_conditioning, + recycle_depth=recycle_depths[index], + attention_kwargs=self.attention_kwargs, + ) + 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 + if target_backbone is not None: + if index == steps - 1: + bb_ca = target_backbone.clone() + else: + 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, bb_time, callback_kwargs) + if callback_outputs is not None: + bb_ca = callback_outputs.pop("bb_ca", bb_ca) + local_latents = callback_outputs.pop("local_latents", local_latents) + + 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"] + 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=local_latents, + ca_coors_nm=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], 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) + ) + 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( + coords_nm=coords_nm, + residue_type=residue_type, + emit_mask=emit_mask, + atom_mask=atom_mask, + 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, + ) + 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 a3f37dd90b48..22e653a403cc 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -2254,6 +2254,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"] @@ -3231,6 +3261,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"] @@ -3891,6 +3936,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..38d9388de142 --- /dev/null +++ b/tests/models/autoencoders/test_models_autoencoder_rfd4_proteina.py @@ -0,0 +1,115 @@ +# 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 ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, +) + + +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, + } + + +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 new file mode 100644 index 000000000000..53941a1cad58 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_rfd4_proteina.py @@ -0,0 +1,146 @@ +# 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 ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + 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, + } + + 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(), "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, + "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 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 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..03979395cde2 --- /dev/null +++ b/tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py @@ -0,0 +1,408 @@ +# 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"]) + # `generator` is required: every design needs its own random stream. + optional_input_params = frozenset(["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, + ) + 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), + } + + 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 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_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. + 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="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 test_pipeline_calls_model_once_per_step_and_recycles_inside_the_call(self): + pipe = self.get_pipeline().to(torch_device) + 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().to(torch_device) + seen = [] + + def callback(pipe, step, timestep, callback_kwargs): + assert set(callback_kwargs) == {"bb_ca", "last_clean_bb_ca"} + seen.append(step) + if step == 0: + 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=["bb_ca", "last_clean_bb_ca"], + ) + assert seen == [0, 1] + assert torch.isfinite(output.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): + 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 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"]) + ).to(torch_device) + 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) + 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().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} + 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 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, {**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", + ) + 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().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) + + 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)) + 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().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) + 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}) + 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().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} + # 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.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"): + 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().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 + 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") + 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): + 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] + 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().to(torch_device)(**{**inputs, **batched, "generator": [shared, shared]}) + + def test_stochastic_centering_offsets_each_molecule_by_one_shared_vector(self): + 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)]) + 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().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 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)) + + +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..002ead7eda17 --- /dev/null +++ b/tests/schedulers/test_scheduler_rfd4_proteina.py @@ -0,0 +1,156 @@ +# 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(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": 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) + + +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): + full = scheduler(schedule_mode, noise_mode, stochastic=True) + full.set_timesteps(4) + 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) + + +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("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(schedule_mode, noise_mode)