mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-23 09:09:01 +03:00
Merge branch 'main' into return-type-fix-models
This commit is contained in:
Sayak Paul
GitHub
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
38
docs/source/en/api/models/rfd4_proteina_autoencoder.md
Normal file
38
docs/source/en/api/models/rfd4_proteina_autoencoder.md
Normal file
@@ -0,0 +1,38 @@
|
||||
<!--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.
|
||||
-->
|
||||
|
||||
# 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
|
||||
35
docs/source/en/api/models/rfd4_proteina_transformer.md
Normal file
35
docs/source/en/api/models/rfd4_proteina_transformer.md
Normal file
@@ -0,0 +1,35 @@
|
||||
<!--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.
|
||||
-->
|
||||
|
||||
# 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
|
||||
128
docs/source/en/api/pipelines/rfd4_proteina.md
Normal file
128
docs/source/en/api/pipelines/rfd4_proteina.md
Normal file
@@ -0,0 +1,128 @@
|
||||
<!--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.
|
||||
-->
|
||||
|
||||
# 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
|
||||
29
docs/source/en/api/schedulers/rfd4_proteina.md
Normal file
29
docs/source/en/api/schedulers/rfd4_proteina.md
Normal file
@@ -0,0 +1,29 @@
|
||||
<!--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.
|
||||
-->
|
||||
|
||||
# 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
|
||||
618
scripts/convert_rfd4_proteina_to_diffusers.py
Normal file
618
scripts/convert_rfd4_proteina_to_diffusers.py
Normal file
@@ -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.<index>` 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()
|
||||
@@ -360,6 +360,8 @@ else:
|
||||
"QwenImageControlNetModel",
|
||||
"QwenImageMultiControlNetModel",
|
||||
"QwenImageTransformer2DModel",
|
||||
"RFD4ProteinaAutoencoder",
|
||||
"RFD4ProteinaTransformer",
|
||||
"SanaControlNetModel",
|
||||
"SanaTransformer2DModel",
|
||||
"SanaVideoTransformer3DModel",
|
||||
@@ -439,6 +441,7 @@ else:
|
||||
"LDMSuperResolutionPipeline",
|
||||
"PNDMPipeline",
|
||||
"RePaintPipeline",
|
||||
"RFD4ProteinaPipeline",
|
||||
"ScoreSdeVePipeline",
|
||||
"StableDiffusionMixin",
|
||||
]
|
||||
@@ -487,6 +490,7 @@ else:
|
||||
"PiflowScheduler",
|
||||
"PNDMScheduler",
|
||||
"RePaintScheduler",
|
||||
"RFD4ProteinaScheduler",
|
||||
"SASolverScheduler",
|
||||
"SchedulerMixin",
|
||||
"SCMScheduler",
|
||||
@@ -1278,6 +1282,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
||||
QwenImageControlNetModel,
|
||||
QwenImageMultiControlNetModel,
|
||||
QwenImageTransformer2DModel,
|
||||
RFD4ProteinaAutoencoder,
|
||||
RFD4ProteinaTransformer,
|
||||
SanaControlNetModel,
|
||||
SanaTransformer2DModel,
|
||||
SanaVideoTransformer3DModel,
|
||||
@@ -1355,6 +1361,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
||||
LDMSuperResolutionPipeline,
|
||||
PNDMPipeline,
|
||||
RePaintPipeline,
|
||||
RFD4ProteinaPipeline,
|
||||
ScoreSdeVePipeline,
|
||||
StableDiffusionMixin,
|
||||
)
|
||||
@@ -1401,6 +1408,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
||||
PiflowScheduler,
|
||||
PNDMScheduler,
|
||||
RePaintScheduler,
|
||||
RFD4ProteinaScheduler,
|
||||
SASolverScheduler,
|
||||
SchedulerMixin,
|
||||
SCMScheduler,
|
||||
|
||||
@@ -12,9 +12,10 @@ if is_torch_available():
|
||||
_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 @@ _import_structure["peft"] = ["PeftAdapterMixin"]
|
||||
|
||||
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
|
||||
|
||||
@@ -7323,6 +7323,195 @@ class MiniMaxH3LoraLoaderMixin(LoraBaseMixin):
|
||||
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."
|
||||
|
||||
@@ -59,6 +59,7 @@ if is_torch_available():
|
||||
_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 @@ if is_torch_available():
|
||||
_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 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
||||
MiniMaxMusic3Vocoder,
|
||||
MMAudioVAE,
|
||||
MMAudioVocoder,
|
||||
RFD4ProteinaAutoencoder,
|
||||
VQModel,
|
||||
)
|
||||
from .cache_utils import CacheMixin
|
||||
@@ -312,6 +315,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
||||
PRXTransformer2DModel,
|
||||
QwenImage21Transformer2DModel,
|
||||
QwenImageTransformer2DModel,
|
||||
RFD4ProteinaTransformer,
|
||||
SanaTransformer2DModel,
|
||||
SanaVideoTransformer3DModel,
|
||||
SanaWMLTX2RefinerTransformer3DModel,
|
||||
|
||||
@@ -29,6 +29,7 @@ from .autoencoder_longcat_audio_dit import LongCatAudioDiTVae
|
||||
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
|
||||
|
||||
698
src/diffusers/models/autoencoders/autoencoder_rfd4_proteina.py
Normal file
698
src/diffusers/models/autoencoders/autoencoder_rfd4_proteina.py
Normal file
@@ -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,
|
||||
)
|
||||
@@ -65,6 +65,7 @@ if is_torch_available():
|
||||
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
|
||||
|
||||
1157
src/diffusers/models/transformers/transformer_rfd4_proteina.py
Normal file
1157
src/diffusers/models/transformers/transformer_rfd4_proteina.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -49,6 +49,7 @@ else:
|
||||
_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 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
||||
ImagePipelineOutput,
|
||||
StableDiffusionMixin,
|
||||
)
|
||||
from .rfd4_proteina import RFD4ProteinaPipeline
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_librosa_available()):
|
||||
|
||||
35
src/diffusers/pipelines/rfd4_proteina/__init__.py
Normal file
35
src/diffusers/pipelines/rfd4_proteina/__init__.py
Normal file
@@ -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__)
|
||||
47
src/diffusers/pipelines/rfd4_proteina/pipeline_output.py
Normal file
47
src/diffusers/pipelines/rfd4_proteina/pipeline_output.py
Normal file
@@ -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
|
||||
582
src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py
Normal file
582
src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py
Normal file
@@ -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()
|
||||
@@ -76,6 +76,7 @@ else:
|
||||
_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 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
||||
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
|
||||
|
||||
276
src/diffusers/schedulers/scheduling_rfd4_proteina.py
Normal file
276
src/diffusers/schedulers/scheduling_rfd4_proteina.py
Normal file
@@ -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)
|
||||
@@ -2254,6 +2254,36 @@ class QwenImageTransformer2DModel(metaclass=DummyObject):
|
||||
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 @@ class RePaintPipeline(metaclass=DummyObject):
|
||||
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 @@ class RePaintScheduler(metaclass=DummyObject):
|
||||
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"]
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from diffusers import Kandinsky6SRVAE
|
||||
@@ -30,6 +31,24 @@ from ..testing_utils import (
|
||||
enable_full_determinism()
|
||||
|
||||
|
||||
UNSPLITTABLE_ENCODER_DECODER = pytest.mark.xfail(
|
||||
reason=(
|
||||
"`_no_split_modules` keeps the whole encoder and decoder together, "
|
||||
"preventing the test's required GPU/CPU split."
|
||||
),
|
||||
raises=AssertionError,
|
||||
strict=True,
|
||||
)
|
||||
DISK_OFFLOAD_OUTPUT_DEVICE = pytest.mark.xfail(
|
||||
reason=(
|
||||
"`_no_split_modules` keeps each encoder/decoder intact, forcing all-disk dispatch under the test's budgets; "
|
||||
"the encode/decode hooks leave the output on CPU, but the reference is on GPU."
|
||||
),
|
||||
raises=RuntimeError,
|
||||
strict=True,
|
||||
)
|
||||
|
||||
|
||||
class Kandinsky6SRVAETesterConfig(BaseModelTesterConfig):
|
||||
@property
|
||||
def model_class(self):
|
||||
@@ -98,7 +117,17 @@ class TestKandinsky6SRVAEModel(Kandinsky6SRVAETesterConfig, ModelTesterMixin):
|
||||
|
||||
|
||||
class TestKandinsky6SRVAEMemory(Kandinsky6SRVAETesterConfig, MemoryTesterMixin):
|
||||
pass
|
||||
@UNSPLITTABLE_ENCODER_DECODER
|
||||
def test_cpu_offload(self, base_model_output, tmp_path):
|
||||
super().test_cpu_offload(base_model_output, tmp_path)
|
||||
|
||||
@DISK_OFFLOAD_OUTPUT_DEVICE
|
||||
def test_disk_offload_without_safetensors(self, base_model_output, tmp_path):
|
||||
super().test_disk_offload_without_safetensors(base_model_output, tmp_path)
|
||||
|
||||
@DISK_OFFLOAD_OUTPUT_DEVICE
|
||||
def test_disk_offload_with_safetensors(self, base_model_output, tmp_path):
|
||||
super().test_disk_offload_with_safetensors(base_model_output, tmp_path)
|
||||
|
||||
|
||||
class TestKandinsky6SRVAETorchCompile(Kandinsky6SRVAETesterConfig, TorchCompileTesterMixin):
|
||||
|
||||
@@ -13,12 +13,13 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from diffusers import MMAudioVAE
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
from ...testing_utils import enable_full_determinism, torch_device
|
||||
from ...testing_utils import enable_full_determinism, require_accelerator, torch_device
|
||||
from ..testing_utils import (
|
||||
BaseModelTesterConfig,
|
||||
MemoryTesterMixin,
|
||||
@@ -30,6 +31,28 @@ from ..testing_utils import (
|
||||
enable_full_determinism()
|
||||
|
||||
|
||||
MEL_DTYPE = pytest.mark.xfail(
|
||||
reason="MMAudio's STFT and mel-filter multiplication require float32.",
|
||||
raises=RuntimeError,
|
||||
strict=True,
|
||||
)
|
||||
LAYERWISE_CASTING_BACKWARD = pytest.mark.xfail(
|
||||
reason="MMAudio's gain multiplication saves weights that are cast to float8 before backward.",
|
||||
raises=RuntimeError,
|
||||
strict=True,
|
||||
)
|
||||
NORMALIZATION_BUFFER_OFFLOAD = pytest.mark.xfail(
|
||||
reason="MMAudio encode/decode bypass the offloading hooks for normalization buffers.",
|
||||
raises=RuntimeError,
|
||||
strict=True,
|
||||
)
|
||||
INCOMPLETE_DEVICE_MAP = pytest.mark.xfail(
|
||||
reason="Automatic device maps omit MMAudio's normalization buffers.",
|
||||
raises=ValueError,
|
||||
strict=True,
|
||||
)
|
||||
|
||||
|
||||
class MMAudioVAETesterConfig(BaseModelTesterConfig):
|
||||
@property
|
||||
def model_class(self):
|
||||
@@ -79,6 +102,16 @@ class MMAudioVAETesterConfig(BaseModelTesterConfig):
|
||||
|
||||
|
||||
class TestMMAudioVAEModel(MMAudioVAETesterConfig, ModelTesterMixin):
|
||||
@MEL_DTYPE
|
||||
@require_accelerator
|
||||
@pytest.mark.skipif(
|
||||
torch_device not in ["cuda", "xpu"],
|
||||
reason="float16 and bfloat16 can only be use for inference with an accelerator",
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
|
||||
def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype):
|
||||
super().test_from_save_pretrained_dtype_inference(tmp_path, dtype)
|
||||
|
||||
def test_latent_shape(self):
|
||||
model = self.model_class(**self.get_init_dict()).to(torch_device).eval()
|
||||
with torch.no_grad():
|
||||
@@ -89,7 +122,44 @@ class TestMMAudioVAEModel(MMAudioVAETesterConfig, ModelTesterMixin):
|
||||
|
||||
|
||||
class TestMMAudioVAEMemory(MMAudioVAETesterConfig, MemoryTesterMixin):
|
||||
pass
|
||||
@MEL_DTYPE
|
||||
def test_layerwise_casting_memory(self):
|
||||
super().test_layerwise_casting_memory()
|
||||
|
||||
@LAYERWISE_CASTING_BACKWARD
|
||||
def test_layerwise_casting_training(self):
|
||||
super().test_layerwise_casting_training()
|
||||
|
||||
@NORMALIZATION_BUFFER_OFFLOAD
|
||||
@pytest.mark.parametrize("record_stream", [False, True])
|
||||
def test_group_offloading(self, base_model_output, record_stream):
|
||||
super().test_group_offloading(base_model_output, record_stream)
|
||||
|
||||
@pytest.mark.parametrize("record_stream", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"offload_type", ["block_level", pytest.param("leaf_level", marks=NORMALIZATION_BUFFER_OFFLOAD)]
|
||||
)
|
||||
def test_group_offloading_with_layerwise_casting(self, record_stream, offload_type):
|
||||
super().test_group_offloading_with_layerwise_casting(record_stream, offload_type)
|
||||
|
||||
@pytest.mark.parametrize("record_stream", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"offload_type", ["block_level", pytest.param("leaf_level", marks=NORMALIZATION_BUFFER_OFFLOAD)]
|
||||
)
|
||||
def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type):
|
||||
super().test_group_offloading_with_disk(tmp_path, record_stream, offload_type)
|
||||
|
||||
@INCOMPLETE_DEVICE_MAP
|
||||
def test_cpu_offload(self, base_model_output, tmp_path):
|
||||
super().test_cpu_offload(base_model_output, tmp_path)
|
||||
|
||||
@INCOMPLETE_DEVICE_MAP
|
||||
def test_disk_offload_without_safetensors(self, base_model_output, tmp_path):
|
||||
super().test_disk_offload_without_safetensors(base_model_output, tmp_path)
|
||||
|
||||
@INCOMPLETE_DEVICE_MAP
|
||||
def test_disk_offload_with_safetensors(self, base_model_output, tmp_path):
|
||||
super().test_disk_offload_with_safetensors(base_model_output, tmp_path)
|
||||
|
||||
|
||||
class TestMMAudioVAETorchCompile(MMAudioVAETesterConfig, TorchCompileTesterMixin):
|
||||
|
||||
@@ -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
|
||||
@@ -13,6 +13,7 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from diffusers import MMAudioVocoder
|
||||
@@ -23,13 +24,19 @@ from ..testing_utils import (
|
||||
BaseModelTesterConfig,
|
||||
MemoryTesterMixin,
|
||||
ModelTesterMixin,
|
||||
TorchCompileTesterMixin,
|
||||
)
|
||||
|
||||
|
||||
enable_full_determinism()
|
||||
|
||||
|
||||
NESTED_UPSAMPLER_OFFLOAD = pytest.mark.xfail(
|
||||
reason="Block offloading hooks on MMAudio's nested upsampler ModuleLists are never called.",
|
||||
raises=RuntimeError,
|
||||
strict=True,
|
||||
)
|
||||
|
||||
|
||||
class MMAudioVocoderTesterConfig(BaseModelTesterConfig):
|
||||
@property
|
||||
def model_class(self):
|
||||
@@ -79,8 +86,21 @@ class TestMMAudioVocoderModel(MMAudioVocoderTesterConfig, ModelTesterMixin):
|
||||
|
||||
|
||||
class TestMMAudioVocoderMemory(MMAudioVocoderTesterConfig, MemoryTesterMixin):
|
||||
pass
|
||||
@NESTED_UPSAMPLER_OFFLOAD
|
||||
@pytest.mark.parametrize("record_stream", [False, True])
|
||||
def test_group_offloading(self, base_model_output, record_stream):
|
||||
super().test_group_offloading(base_model_output, record_stream)
|
||||
|
||||
@pytest.mark.parametrize("record_stream", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"offload_type", [pytest.param("block_level", marks=NESTED_UPSAMPLER_OFFLOAD), "leaf_level"]
|
||||
)
|
||||
def test_group_offloading_with_layerwise_casting(self, record_stream, offload_type):
|
||||
super().test_group_offloading_with_layerwise_casting(record_stream, offload_type)
|
||||
|
||||
class TestMMAudioVocoderTorchCompile(MMAudioVocoderTesterConfig, TorchCompileTesterMixin):
|
||||
pass
|
||||
@pytest.mark.parametrize("record_stream", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"offload_type", [pytest.param("block_level", marks=NESTED_UPSAMPLER_OFFLOAD), "leaf_level"]
|
||||
)
|
||||
def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type):
|
||||
super().test_group_offloading_with_disk(tmp_path, record_stream, offload_type)
|
||||
@@ -13,6 +13,7 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from diffusers import Kandinsky6SRLatentUpscalerBank
|
||||
@@ -23,13 +24,30 @@ from ..testing_utils import (
|
||||
BaseModelTesterConfig,
|
||||
MemoryTesterMixin,
|
||||
ModelTesterMixin,
|
||||
TorchCompileTesterMixin,
|
||||
)
|
||||
|
||||
|
||||
enable_full_determinism()
|
||||
|
||||
|
||||
UNSPLITTABLE_UPSCALER = pytest.mark.xfail(
|
||||
reason=(
|
||||
"`_no_split_modules` keeps each Kandinsky6SRLatentUpscaler intact, "
|
||||
"preventing the test's required GPU/CPU split."
|
||||
),
|
||||
raises=AssertionError,
|
||||
strict=True,
|
||||
)
|
||||
DISK_OFFLOAD_NUMERICS = pytest.mark.xfail(
|
||||
reason=(
|
||||
"`_no_split_modules` keeps each upscaler intact, forcing all-disk dispatch and CPU execution under the test's "
|
||||
"budgets; numerical differences from the GPU reference can exceed the output tolerance."
|
||||
),
|
||||
raises=AssertionError,
|
||||
strict=False,
|
||||
)
|
||||
|
||||
|
||||
class Kandinsky6SRLatentUpscalerBankTesterConfig(BaseModelTesterConfig):
|
||||
@property
|
||||
def model_class(self):
|
||||
@@ -87,10 +105,14 @@ class TestKandinsky6SRLatentUpscalerBankModel(Kandinsky6SRLatentUpscalerBankTest
|
||||
|
||||
|
||||
class TestKandinsky6SRLatentUpscalerBankMemory(Kandinsky6SRLatentUpscalerBankTesterConfig, MemoryTesterMixin):
|
||||
pass
|
||||
@UNSPLITTABLE_UPSCALER
|
||||
def test_cpu_offload(self, base_model_output, tmp_path):
|
||||
super().test_cpu_offload(base_model_output, tmp_path)
|
||||
|
||||
@DISK_OFFLOAD_NUMERICS
|
||||
def test_disk_offload_without_safetensors(self, base_model_output, tmp_path):
|
||||
super().test_disk_offload_without_safetensors(base_model_output, tmp_path)
|
||||
|
||||
class TestKandinsky6SRLatentUpscalerBankTorchCompile(
|
||||
Kandinsky6SRLatentUpscalerBankTesterConfig, TorchCompileTesterMixin
|
||||
):
|
||||
pass
|
||||
@DISK_OFFLOAD_NUMERICS
|
||||
def test_disk_offload_with_safetensors(self, base_model_output, tmp_path):
|
||||
super().test_disk_offload_with_safetensors(base_model_output, tmp_path)
|
||||
@@ -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
|
||||
0
tests/pipelines/rfd4_proteina/__init__.py
Normal file
0
tests/pipelines/rfd4_proteina/__init__.py
Normal file
408
tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py
Normal file
408
tests/pipelines/rfd4_proteina/test_pipeline_rfd4_proteina.py
Normal file
@@ -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
|
||||
156
tests/schedulers/test_scheduler_rfd4_proteina.py
Normal file
156
tests/schedulers/test_scheduler_rfd4_proteina.py
Normal file
@@ -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)
|
||||
Reference in New Issue
Block a user