Upload folder using huggingface_hub

README.md

@@ -0,0 +1,290 @@
+# Protein Design with Diffusers
+
+A [diffusers](https://github.com/huggingface/diffusers) `ModularPipeline` wrapper for protein design, combining structure generation ([RFdiffusion3](https://www.biorxiv.org/content/10.1101/2025.09.18.676967v2)) and sequence design ([ProteinMPNN](https://www.science.org/doi/10.1126/science.add2187) / [LigandMPNN](https://www.nature.com/articles/s41592-025-02626-1)) into composable, swappable pipeline blocks.
+
+All three models — RFD3, ProteinMPNN, and LigandMPNN — rely on [Foundry](https://github.com/RosettaCommons/foundry) for their underlying implementations and [AtomWorks](https://github.com/RosettaCommons/atomworks) for structure I/O. This package adds only the thin wrappers needed for diffusers integration.
+
+## Getting Started
+
+### Installation
+
+```bash
+# Install foundry (provides model implementations + AtomWorks)
+pip install rc-foundry[all]
+
+# Install diffusers with modular pipeline support
+pip install diffusers
+```
+
+### Running with Diffusers
+
+```python
+import torch
+from diffusers import ModularPipeline
+
+# Load the pipeline
+pipe = ModularPipeline.from_pretrained(
+    "dn6/RFDiffusion-3",
+    trust_remote_code=True,
+)
+pipe.load_components(
+    device_map="cuda",
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
+)
+
+# Generate a 100-residue protein backbone
+state = pipe(contigs="100")
+
+# Access coordinates
+print(state.output.xyz.shape)  # [B, L, 3]
+```
+
+## Why Diffusers?
+
+Wrapping RFdiffusion3 as a diffusers `ModularPipeline` gives you access to the full diffusers ecosystem out of the box:
+
+### CPU Offloading
+
+Run large models on limited VRAM by offloading components to CPU when not in use:
+
+```python
+pipe.enable_model_cpu_offload()
+state = pipe(contigs="100")
+```
+
+### Hub Integration
+
+All models are hosted on the Hugging Face Hub. Load by repo ID, share fine-tuned variants, and version your checkpoints:
+
+```python
+pipe = ModularPipeline.from_pretrained("dn6/RFDiffusion-3", trust_remote_code=True)
+```
+
+### LoRA Fine-Tuning
+
+Fine-tune RFdiffusion3 on custom datasets with LoRA — supported natively by `ModelMixin`:
+
+```python
+from peft import LoraConfig
+
+lora_config = LoraConfig(r=16, lora_alpha=16, target_modules=["to_q", "to_k", "to_v"])
+pipe.transformer.add_adapter(lora_config)
+
+# After training
+pipe.transformer.save_pretrained("my-rfd3-lora")
+```
+
+### Composable Workflows
+
+Inspect, swap, and extend pipeline blocks at runtime:
+
+```python
+# Inspect the pipeline
+print(pipe.blocks)
+
+# Swap ProteinMPNN for LigandMPNN
+mpnn = AutoModel.from_pretrained("dn6/RFDiffusion-3", subfolder="mpnn_ligand", trust_remote_code=True)
+pipe.update_components(mpnn=mpnn)
+
+# Add a custom post-processing block
+from diffusers.modular_pipelines import ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, OutputParam
+
+class ScoreDesignStep(ModularPipelineBlocks):
+    @property
+    def inputs(self):
+        return [InputParam("xyz", required=True)]
+
+    @property
+    def intermediate_outputs(self):
+        return [OutputParam("radius_of_gyration")]
+
+    def __call__(self, components, state):
+        block_state = self.get_block_state(state)
+        xyz = block_state.xyz
+        centroid = xyz.mean(dim=-2, keepdim=True)
+        block_state.radius_of_gyration = ((xyz - centroid) ** 2).sum(-1).mean().sqrt()
+        self.set_block_state(state, block_state)
+        return components, state
+
+# Insert after the decoder
+pipe._blocks.sub_blocks.insert("score", ScoreDesignStep(), index=3)
+```
+
+## Output Types
+
+All output types return the base tensors (`xyz`, `sequence_indices`, `sequence_logits`). The `output_type` parameter controls what additional format is produced:
+
+| `output_type` | Additional output | Writes to disk |
+|---|---|---|
+| `"tensor"` | — | — |
+| `"pdb"` | `pdb_string` | `.pdb` file |
+| `"cif"` | `atom_array`, `atom_array_stack`, `trajectory_stack` | `.cif` via AtomWorks |
+| `"cif.gz"` | Same as `"cif"` | `.cif.gz` compressed |
+
+CIF outputs use [AtomWorks](https://github.com/RosettaCommons/atomworks) `to_cif_file` and return [biotite](https://www.biotite-python.org/) `AtomArray` / `AtomArrayStack` objects, matching the foundry output format.
+
+```python
+# Save as compressed CIF (matches foundry output format)
+state = pipe(contigs="100", output_type="cif.gz", output_path="design_0")
+
+# AtomArray is available directly
+atom_array = state.output.atom_array
+print(atom_array)  # biotite AtomArray with CA coords + residue names
+
+# Denoising trajectory as AtomArrayStack (one model per step)
+trajectory = state.output.trajectory_stack
+
+# PDB string output
+state = pipe(contigs="100", output_type="pdb", output_path="design_0.pdb")
+print(state.output.pdb_string[:200])
+```
+
+## Models
+
+By default, `load_components` loads the RFdiffusion3 transformer and scheduler. MPNN models are optional — load them separately when you need sequence design.
+
+### RFdiffusion3 (RFD3)
+
+[RFdiffusion3](https://www.biorxiv.org/content/10.1101/2025.09.18.676967v2) is an all-atom generative model that designs protein structures via iterative denoising. Uses an EDM noise schedule with 200 steps. Loaded automatically by `load_components`.
+
+| Component | Subfolder | Description |
+|-----------|-----------|-------------|
+| `transformer` | `transformer/` | `RFDiffusionTransformerModel` (168M params) |
+| `scheduler` | `scheduler/` | `RFDiffusionScheduler` (EDM noise schedule + Euler stepping) |
+
+### ProteinMPNN / LigandMPNN
+
+[ProteinMPNN](https://www.science.org/doi/10.1126/science.add2187) and [LigandMPNN](https://www.nature.com/articles/s41592-025-02626-1) are inverse-folding models that design amino acid sequences for a given protein backbone. These are **not** loaded by default — load them with `AutoModel` and register via `update_components`:
+
+```python
+from diffusers import AutoModel
+
+mpnn = AutoModel.from_pretrained("dn6/RFDiffusion-3", subfolder="mpnn", trust_remote_code=True)
+pipe.update_components(mpnn=mpnn)
+```
+
+Three variants are available:
+
+| Subfolder | Variant | Params | Description |
+|-----------|---------|--------|-------------|
+| `mpnn/` | ProteinMPNN | 1.66M | Standard protein sequence design |
+| `mpnn_ligand/` | LigandMPNN | 2.62M | Ligand-aware sequence design |
+| `mpnn_soluble/` | SolubleMPNN | 1.66M | Optimized for soluble proteins |
+
+## Workflows
+
+The active workflow is selected automatically based on which inputs you provide. Passing `temperature` triggers the MPNN sequence design step; passing `input_xyz` enables motif conditioning.
+
+| Workflow | Trigger inputs | What runs |
+|----------|---------------|-----------|
+| `structure_only` | `contigs` | RFdiffusion3 |
+| `structure_and_sequence` | `contigs`, `temperature` | RFdiffusion3 → MPNN |
+| `motif_structure_and_sequence` | `contigs`, `input_xyz`, `temperature` | Motif-conditioned RFdiffusion3 → MPNN |
+
+> Workflows that include MPNN require loading an MPNN variant first (see above).
+
+You can also select a workflow explicitly:
+
+```python
+workflow = pipe.get_workflow("structure_and_sequence")
+```
+
+### Structure Only
+
+```python
+state = pipe(contigs="100")
+print(state.output.xyz.shape)  # [1, 100, 3]
+```
+
+### Structure + Sequence Design
+
+```python
+from diffusers import AutoModel
+
+# Load an MPNN variant and register it
+mpnn = AutoModel.from_pretrained("dn6/RFDiffusion-3", subfolder="mpnn", trust_remote_code=True)
+pipe.update_components(mpnn=mpnn)
+
+# Passing temperature triggers the MPNN step
+state = pipe(contigs="100", temperature=0.1)
+print(state.mpnn_output.designed_sequence)  # e.g. "MKVLSEG..."
+```
+
+### Motif-Conditioned Design
+
+```python
+import torch
+
+motif_coords = torch.randn(16, 3)  # [N_motif, 3]
+state = pipe(
+    contigs="A10-25/50",
+    input_xyz=motif_coords,
+    temperature=0.1,
+)
+```
+
+## Full Design Pipeline
+
+The three pipelines can be composed into a complete protein design workflow:
+
+```
+RFD3 (design backbone) → MPNN (design sequence) → RF3 (validate fold)
+```
+
+Each is a standalone `ModularPipeline` that can run independently. Here's the full end-to-end flow:
+
+```python
+import torch
+from diffusers import AutoModel, ModularPipeline
+
+# 1. Design a backbone with RFdiffusion3
+design_pipe = ModularPipeline.from_pretrained("dn6/RFDiffusion-3", trust_remote_code=True)
+design_pipe.load_components(device_map="cuda", torch_dtype=torch.bfloat16, trust_remote_code=True)
+
+mpnn = AutoModel.from_pretrained("dn6/RFDiffusion-3", subfolder="mpnn", trust_remote_code=True)
+design_pipe.update_components(mpnn=mpnn)
+
+state = design_pipe(contigs="100", temperature=0.1, output_type="cif.gz", output_path="design")
+designed_sequence = state.mpnn_output.designed_sequence
+
+# 2. Validate the design with RF3 (structure prediction)
+fold_pipe = ModularPipeline.from_pretrained("dn6/RosettaFold-3", trust_remote_code=True)
+fold_pipe.load_components(device_map="cuda", torch_dtype=torch.bfloat16, trust_remote_code=True)
+
+state = fold_pipe(sequence=designed_sequence, output_type="cif.gz", output_path="prediction")
+```
+
+> [RF3](https://www.biorxiv.org/content/10.1101/2025.08.14.670328) (RosettaFold3) is available as a separate pipeline at [`dn6/RosettaFold-3`](https://huggingface.co/dn6/RosettaFold-3).
+
+## Citation
+
+If you use this code, please cite the relevant work:
+
+```bibtex
+@article{butcher2025_rfdiffusion3,
+    author = {Butcher, Jasper and Krishna, Rohith and Mitra, Raktim and Brent, Rafael Isaac and Li, Yanjing and Corley, Nathaniel and Kim, Paul T and Funk, Jonathan and Mathis, Simon Valentin and Salike, Saman and Muraishi, Aiko and Eisenach, Helen and Thompson, Tuscan Rock and Chen, Jie and Politanska, Yuliya and Sehgal, Enisha and Coventry, Brian and Zhang, Odin and Qiang, Bo and Didi, Kieran and Kazman, Maxwell and DiMaio, Frank and Baker, David},
+    title = {De novo Design of All-atom Biomolecular Interactions with RFdiffusion3},
+    journal = {bioRxiv},
+    year = {2025},
+    doi = {10.1101/2025.09.18.676967},
+}
+
+@article{dauparas2022robust,
+    author = {Dauparas, Justas and Anishchenko, Ivan and Bennett, Nathaniel and Bai, Hua and Ragotte, Robert J and Milles, Lukas F and Wicky, Basile IM and Courbet, Alexis and de Haas, Rob J and Bethel, Neville and others},
+    title = {Robust deep learning--based protein sequence design using ProteinMPNN},
+    journal = {Science},
+    volume = {378},
+    number = {6615},
+    pages = {49--56},
+    year = {2022},
+}
+
+@article{dauparas2025atomic,
+    author = {Dauparas, Justas and Lee, Gyu Rie and Pecoraro, Robert and An, Linna and Anishchenko, Ivan and Glasscock, Cameron and Baker, David},
+    title = {Atomic context-conditioned protein sequence design using LigandMPNN},
+    journal = {Nature Methods},
+    year = {2025},
+}
+```

__init__.py

@@ -0,0 +1,34 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+from .transformer import RFDiffusionTransformerModel, RFDiffusionTransformerOutput
+from .scheduler import RFDiffusionScheduler
+from .modular_blocks import (
+    ALL_BLOCKS,
+    AUTO_BLOCKS,
+    UNCONDITIONAL_BLOCKS,
+    RFDiffusionAutoBeforeDenoiseStep,
+    RFDiffusionAutoBlocks,
+    RFDiffusionAutoDecodeStep,
+    RFDiffusionAutoDenoiseStep,
+)
+from .before_denoise import (
+    RFDiffusionInputStep,
+    RFDiffusionPrepareLatentsStep,
+    RFDiffusionSetTimestepsStep,
+)
+from .denoise import RFDiffusionDenoiseStep
+from .decoders import RFDiffusionDecodeStep, RFDiffusionPipelineOutput
+from .mpnn import MPNNModel, MPNNModelOutput
+from .modular_blocks import MPNNAutoDesignStep, MPNNPipelineOutput, MPNNSequenceDesignStep

before_denoise.py

@@ -0,0 +1,312 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+from typing import List, Tuple, Union
+
+import torch
+
+from diffusers.utils import logging
+from diffusers.modular_pipelines import ModularPipeline, ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__)
+
+
+def parse_contig_string(contig_str: str) -> Tuple[int, List[Tuple[int, int]]]:
+    """
+    Parse contig specification string.
+
+    Supports formats like:
+    - "100" -> 100 residues to design
+    - "50-100" -> random length between 50-100
+    - "A10-25/50" -> motif from chain A residues 10-25, plus 50 designed
+
+    Returns:
+        total_length: Total protein length
+        motif_ranges: List of (start, end) for motif residues (0-indexed)
+    """
+    parts = contig_str.split("/")
+    total_length = 0
+    motif_ranges = []
+
+    for part in parts:
+        part = part.strip()
+        if not part:
+            continue
+
+        if part[0].isalpha():
+            chain = part[0]
+            residue_spec = part[1:]
+            if "-" in residue_spec:
+                start, end = map(int, residue_spec.split("-"))
+            else:
+                start = end = int(residue_spec)
+            motif_len = end - start + 1
+            motif_ranges.append((total_length, total_length + motif_len))
+            total_length += motif_len
+        else:
+            if "-" in part:
+                min_len, max_len = map(int, part.split("-"))
+                add_len = (min_len + max_len) // 2
+            else:
+                add_len = int(part)
+            total_length += add_len
+
+    return total_length, motif_ranges
+
+
+class RFDiffusionInputStep(ModularPipelineBlocks):
+    """
+    Input processing step for RFDiffusion.
+
+    Parses contigs to prepare features for structure generation.
+    """
+
+    model_name = "rfdiffusion"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Input processing step that:\n"
+            "  1. Parses contig specification to determine protein length and design regions\n"
+            "  2. Generates masks for motif positions\n"
+        )
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam(
+                "contigs",
+                required=True,
+                type_hint=Union[str, List[str]],
+                description="Contig specification defining design regions (e.g., '100' or 'A10-25/50-100')",
+            ),
+            InputParam(
+                "input_xyz",
+                type_hint=torch.Tensor,
+                description="Input coordinates for motif residues [N_motif, 3]",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam(
+                "motif_mask",
+                type_hint=torch.Tensor,
+                description="Boolean mask for motif (fixed) positions",
+            ),
+            OutputParam(
+                "motif_xyz",
+                type_hint=torch.Tensor,
+                description="Coordinates for motif residues",
+            ),
+            OutputParam(
+                "L",
+                type_hint=int,
+                description="Total length of the protein being designed",
+            ),
+            OutputParam(
+                "batch_size",
+                type_hint=int,
+                description="Batch size (typically 1 for RFDiffusion)",
+            ),
+            OutputParam(
+                "dtype",
+                type_hint=torch.dtype,
+                description="Data type for tensors",
+            ),
+        ]
+
+    def check_inputs(self, components, block_state):
+        if block_state.contigs is None:
+            raise ValueError("`contigs` must be provided to specify protein design regions")
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        self.check_inputs(components, block_state)
+
+        contigs = block_state.contigs
+        input_xyz = block_state.input_xyz
+
+        if isinstance(contigs, list):
+            contig_str = "/".join(contigs)
+        else:
+            contig_str = contigs
+
+        L, motif_ranges = parse_contig_string(contig_str)
+
+        motif_mask = torch.zeros(L, dtype=torch.bool)
+        for start, end in motif_ranges:
+            motif_mask[start:end] = True
+
+        if input_xyz is not None:
+            motif_xyz = input_xyz
+        else:
+            motif_xyz = None
+
+        block_state.motif_mask = motif_mask
+        block_state.motif_xyz = motif_xyz
+        block_state.L = L
+        block_state.batch_size = 1
+        block_state.dtype = torch.float32
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class RFDiffusionSetTimestepsStep(ModularPipelineBlocks):
+    """
+    Set up the EDM noise schedule for RFDiffusion3.
+    """
+
+    model_name = "rfdiffusion"
+
+    @property
+    def description(self) -> str:
+        return "Sets up the EDM noise schedule matching the original inference sampler."
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec("scheduler", description="RFDiffusion3 EDM scheduler"),
+        ]
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam(
+                "num_inference_steps",
+                default=None,
+                type_hint=int,
+                description="Number of denoising steps (default: use scheduler config)",
+            ),
+            InputParam("L", required=True, type_hint=int, description="Protein length"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam(
+                "noise_schedule",
+                type_hint=torch.Tensor,
+                description="EDM noise schedule [num_timesteps] from high to low noise",
+            ),
+            OutputParam(
+                "num_inference_steps",
+                type_hint=int,
+                description="Number of inference steps",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        if hasattr(components, "scheduler") and components.scheduler is not None:
+            noise_schedule = components.scheduler.get_noise_schedule()
+        else:
+            # Fallback: simple linear schedule
+            noise_schedule = torch.linspace(160.0 * 16.0, 4e-4 * 16.0, 200)
+
+        block_state.noise_schedule = noise_schedule
+        block_state.num_inference_steps = len(noise_schedule)
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class RFDiffusionPrepareLatentsStep(ModularPipelineBlocks):
+    """
+    Prepare initial noised coordinates for RFDiffusion3.
+
+    Matches the original _get_initial_structure:
+        noise = c0 * randn(D, L, 3)
+        noise[..., is_motif, :] = 0
+        X_L = noise + coord_motif
+    """
+
+    model_name = "rfdiffusion"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Prepares initial coordinates by sampling Gaussian noise scaled by "
+            "the first noise schedule value, matching the original sampler."
+        )
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec("scheduler", description="RFDiffusion3 EDM scheduler"),
+            ComponentSpec("transformer", description="RFDiffusion transformer model"),
+        ]
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam("generator", type_hint=torch.Generator, description="Random generator for reproducibility"),
+            InputParam("diffusion_batch_size", default=1, type_hint=int, description="Number of samples to generate in parallel"),
+            InputParam("L", required=True, type_hint=int, description="Protein length"),
+            InputParam("motif_mask", required=True, type_hint=torch.Tensor),
+            InputParam("motif_xyz", type_hint=torch.Tensor),
+            InputParam("noise_schedule", required=True, type_hint=torch.Tensor),
+            InputParam("dtype", type_hint=torch.dtype),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam("xyz", type_hint=torch.Tensor, description="Initial noised coordinates [D, L, 3]"),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        L = block_state.L
+        motif_mask = block_state.motif_mask
+        motif_xyz = block_state.motif_xyz
+        noise_schedule = block_state.noise_schedule
+        dtype = block_state.dtype or torch.float32
+        generator = block_state.generator
+        D = block_state.diffusion_batch_size or 1
+
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # Initial noise scaled by first noise level (c0), matching original:
+        #   noise = c0 * randn(D, L, 3)
+        c0 = noise_schedule[0]
+        noise = c0 * torch.randn((D, L, 3), dtype=dtype, device=device, generator=generator)
+
+        # Zero out noise for motif atoms
+        if motif_mask is not None:
+            noise[:, motif_mask] = 0.0
+
+        # Build initial coordinates: motif coords + noise
+        coord_motif = torch.zeros((D, L, 3), dtype=dtype, device=device)
+        if motif_xyz is not None and motif_mask is not None:
+            motif_indices = motif_mask.nonzero(as_tuple=True)[0]
+            for i, idx in enumerate(motif_indices):
+                if i < motif_xyz.shape[0]:
+                    coord_motif[:, idx] = motif_xyz[i].to(dtype=dtype, device=device)
+
+        xyz = noise + coord_motif
+
+        block_state.xyz = xyz
+
+        self.set_block_state(state, block_state)
+        return components, state

decoders.py

@@ -0,0 +1,281 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+import numpy as np
+import torch
+from atomworks.io.utils.io_utils import to_cif_file
+from biotite.structure import AtomArray, AtomArrayStack, stack
+
+from diffusers.utils import logging
+from diffusers.modular_pipelines import ModularPipeline, ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__)
+
+
+AA_NAMES = [
+    "ALA", "ARG", "ASN", "ASP", "CYS", "GLN", "GLU", "GLY", "HIS", "ILE",
+    "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL",
+    "UNK",
+]
+
+
+def _build_atom_array(xyz: torch.Tensor, seq: Optional[torch.Tensor] = None) -> AtomArray:
+    """
+    Build a biotite AtomArray from CA coordinates and optional sequence.
+
+    Args:
+        xyz: Coordinates for a single structure [L, 3].
+        seq: Sequence indices [L] (indexes into AA_NAMES).
+    """
+    xyz_np = xyz.detach().cpu().float().numpy()
+    L = xyz_np.shape[0]
+
+    arr = AtomArray(L)
+    arr.coord = xyz_np
+    arr.atom_name = np.full(L, "CA")
+    arr.element = np.full(L, "C")
+    arr.chain_id = np.full(L, "A")
+    arr.res_id = np.arange(1, L + 1)
+
+    if seq is not None:
+        seq_np = seq.detach().cpu().numpy()
+        arr.res_name = np.array([
+            AA_NAMES[int(idx)] if int(idx) < len(AA_NAMES) else "UNK"
+            for idx in seq_np
+        ])
+    else:
+        arr.res_name = np.full(L, "ALA")
+
+    return arr
+
+
+def _build_atom_array_stack(
+    xyz: torch.Tensor,
+    seq: Optional[torch.Tensor] = None,
+) -> AtomArrayStack:
+    """
+    Build an AtomArrayStack from batched coordinates [B, L, 3].
+
+    Matches foundry ``build_stack_from_atom_array_and_batched_coords``.
+    """
+    template = _build_atom_array(xyz[0], seq[0] if seq is not None else None)
+    B = xyz.shape[0]
+    arr_stack = stack([template for _ in range(B)])
+    arr_stack.coord = xyz.detach().cpu().float().numpy()
+    return arr_stack
+
+
+def _build_trajectory_stack(
+    trajectory: List[torch.Tensor],
+    seq: Optional[torch.Tensor] = None,
+) -> AtomArrayStack:
+    """
+    Build an AtomArrayStack from a denoising trajectory.
+
+    Each entry is [B, L, 3]; takes the first batch element per step.
+    """
+    coords = torch.stack([t[0] for t in trajectory])  # [N_steps, L, 3]
+    template = _build_atom_array(coords[0], seq[0] if seq is not None else None)
+    arr_stack = stack([template for _ in range(coords.shape[0])])
+    arr_stack.coord = coords.detach().cpu().float().numpy()
+    return arr_stack
+
+
+@dataclass
+class RFDiffusionPipelineOutput:
+    """Output class for RFDiffusion pipeline."""
+
+    xyz: torch.Tensor
+    atom_array: Optional[AtomArray] = None
+    atom_array_stack: Optional[AtomArrayStack] = None
+    trajectory_stack: Optional[AtomArrayStack] = None
+    sequence_indices: Optional[torch.Tensor] = None
+    sequence_logits: Optional[torch.Tensor] = None
+    single: Optional[torch.Tensor] = None
+    pair: Optional[torch.Tensor] = None
+    pdb_string: Optional[str] = None
+    trajectory: Optional[List[torch.Tensor]] = None
+
+
+class RFDiffusionDecodeStep(ModularPipelineBlocks):
+    """
+    Decode step for RFDiffusion.
+
+    Converts denoised coordinates to final output format.
+
+    Supported ``output_type`` values:
+
+    - ``"tensor"`` — raw tensors only
+    - ``"pdb"`` — tensors + PDB format string
+    - ``"cif"`` — tensors + AtomArray via AtomWorks, writes ``.cif``
+    - ``"cif.gz"`` — same as ``"cif"`` but compressed
+    """
+
+    model_name = "rfdiffusion"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Decode step that converts denoised coordinates to final output, "
+            "supporting tensor, PDB, and CIF (via AtomWorks) formats."
+        )
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam(
+                "output_type",
+                default="tensor",
+                type_hint=str,
+                description="Output format: 'tensor', 'pdb', 'cif', or 'cif.gz'",
+            ),
+            InputParam(
+                "output_path",
+                type_hint=str,
+                description="Path to save output structure",
+            ),
+            InputParam("xyz", required=True, type_hint=torch.Tensor, description="Denoised coordinates [B, L, 3]"),
+            InputParam("sequence_indices", type_hint=torch.Tensor, description="Predicted sequence [B, L]"),
+            InputParam("sequence_logits", type_hint=torch.Tensor, description="Sequence logits [B, L, n_aa]"),
+            InputParam("single", type_hint=torch.Tensor, description="Single representation"),
+            InputParam("pair", type_hint=torch.Tensor, description="Pair representation"),
+            InputParam("trajectory", type_hint=List[torch.Tensor], description="Denoising trajectory"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam("output", type_hint=RFDiffusionPipelineOutput, description="Final pipeline output"),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        xyz = block_state.xyz
+        sequence_indices = block_state.sequence_indices
+        sequence_logits = block_state.sequence_logits
+        single = block_state.single
+        pair = block_state.pair
+        trajectory = block_state.trajectory
+        output_type = block_state.output_type or "tensor"
+        output_path = block_state.output_path
+
+        pdb_string = None
+        atom_array = None
+        atom_array_stack = None
+        trajectory_stack = None
+
+        # Build AtomArray for CIF output types
+        if output_type in ("cif", "cif.gz"):
+            atom_array = _build_atom_array(
+                xyz[0], sequence_indices[0] if sequence_indices is not None else None,
+            )
+            if xyz.shape[0] > 1:
+                atom_array_stack = _build_atom_array_stack(xyz, sequence_indices)
+            if trajectory:
+                trajectory_stack = _build_trajectory_stack(trajectory, sequence_indices)
+
+        # Build PDB string
+        if output_type == "pdb":
+            pdb_string = self._coords_to_pdb(
+                xyz.squeeze(0),
+                sequence_indices.squeeze(0) if sequence_indices is not None else None,
+            )
+
+        # Write to disk
+        if output_path is not None:
+            import os
+            os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+
+            if output_type in ("cif", "cif.gz"):
+                to_write = atom_array_stack if atom_array_stack is not None else atom_array
+                base = output_path.rsplit(".", 1)[0] if "." in output_path else output_path
+                to_cif_file(to_write, base, file_type=output_type, include_entity_poly=False)
+                if trajectory_stack is not None:
+                    to_cif_file(trajectory_stack, base + "_trajectory", file_type=output_type, include_entity_poly=False)
+            elif output_type == "pdb" and pdb_string is not None:
+                with open(output_path, "w") as f:
+                    f.write(pdb_string)
+
+        output = RFDiffusionPipelineOutput(
+            xyz=xyz,
+            atom_array=atom_array,
+            atom_array_stack=atom_array_stack,
+            trajectory_stack=trajectory_stack,
+            sequence_indices=sequence_indices,
+            sequence_logits=sequence_logits,
+            single=single,
+            pair=pair,
+            pdb_string=pdb_string,
+            trajectory=trajectory,
+        )
+
+        block_state.output = output
+        self.set_block_state(state, block_state)
+        return components, state
+
+    def _coords_to_pdb(
+        self,
+        xyz: torch.Tensor,
+        seq: Optional[torch.Tensor] = None,
+    ) -> str:
+        """Convert coordinates to PDB format string."""
+        xyz_np = xyz.cpu().numpy()
+        L = xyz_np.shape[0]
+
+        if seq is not None:
+            seq_np = seq.cpu().numpy()
+        else:
+            seq_np = None
+
+        lines = []
+        atom_idx = 1
+
+        for i in range(L):
+            if seq_np is not None:
+                aa_idx = int(seq_np[i])
+                aa_name = AA_NAMES[aa_idx] if aa_idx < len(AA_NAMES) else "UNK"
+            else:
+                aa_name = "ALA"
+
+            if xyz_np.ndim == 2:
+                x, y, z = xyz_np[i, :]
+                line = (
+                    f"ATOM  {atom_idx:5d}  CA  {aa_name:3s} A"
+                    f"{i+1:4d}    {x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00           C  "
+                )
+                lines.append(line)
+                atom_idx += 1
+            else:
+                for j, atom_name in enumerate(["N", "CA", "C"]):
+                    if j >= xyz_np.shape[1]:
+                        break
+                    x, y, z = xyz_np[i, j, :]
+
+                    line = (
+                        f"ATOM  {atom_idx:5d}  {atom_name:<3s} {aa_name:3s} A"
+                        f"{i+1:4d}    {x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00           "
+                        f"{atom_name[0]:>2s}  "
+                    )
+                    lines.append(line)
+                    atom_idx += 1
+
+        lines.append("END")
+        return "\n".join(lines)

denoise.py

@@ -0,0 +1,193 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+"""
+Denoising loop for RFDiffusion3.
+
+Implements the iterative denoising procedure from the original
+inference_sampler.py (SampleDiffusionWithMotif.sample_diffusion_like_af3).
+
+The loop iterates over consecutive pairs (c_t_minus_1, c_t) in the noise schedule:
+1. Inject stochastic noise: t_hat = c_t_minus_1 * (gamma + 1), epsilon ~ N(0, noise_scale * sqrt(t_hat^2 - c_t_minus_1^2))
+2. Call model: X_denoised = model(X_noisy, t_hat)
+3. Euler update: X = X_noisy + step_scale * (c_t - t_hat) * (X_noisy - X_denoised) / t_hat
+"""
+
+from typing import Callable, List
+
+import torch
+
+from diffusers.utils import logging
+from diffusers.modular_pipelines import ModularPipeline, ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+
+logger = logging.get_logger(__name__)
+
+
+class RFDiffusionDenoiseStep(ModularPipelineBlocks):
+    """
+    Iterative denoising step for RFDiffusion3.
+
+    Implements the EDM stochastic sampling loop matching the original
+    SampleDiffusionWithMotif.sample_diffusion_like_af3.
+    """
+
+    model_name = "rfdiffusion"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Iteratively denoise protein structure through reverse diffusion. "
+            "Uses EDM stochastic sampling with gamma noise injection and step scaling."
+        )
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec("transformer", description="RFDiffusion transformer for structure prediction"),
+            ComponentSpec("scheduler", description="Scheduler for noise injection and stepping"),
+        ]
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam(
+                "n_recycle",
+                default=None,
+                type_hint=int,
+                description="Number of recycling iterations (None uses model default)",
+            ),
+            InputParam(
+                "callback",
+                type_hint=Callable,
+                description="Optional callback function called at each step",
+            ),
+            InputParam(
+                "callback_steps",
+                default=1,
+                type_hint=int,
+                description="Frequency of callback invocation",
+            ),
+            InputParam("xyz", required=True, type_hint=torch.Tensor, description="Initial noised coordinates [D, L, 3]"),
+            InputParam("noise_schedule", required=True, type_hint=torch.Tensor, description="EDM noise schedule"),
+            InputParam("motif_mask", required=True, type_hint=torch.Tensor, description="Mask for fixed motif positions"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam("xyz", type_hint=torch.Tensor, description="Denoised coordinates [D, L, 3]"),
+            OutputParam("single", type_hint=torch.Tensor, description="Single representation"),
+            OutputParam("pair", type_hint=torch.Tensor, description="Pair representation"),
+            OutputParam("sequence_logits", type_hint=torch.Tensor, description="Predicted sequence logits"),
+            OutputParam("sequence_indices", type_hint=torch.Tensor, description="Predicted sequence indices"),
+            OutputParam("trajectory", type_hint=List[torch.Tensor], description="Denoising trajectory"),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        xyz = block_state.xyz
+        noise_schedule = block_state.noise_schedule
+        motif_mask = block_state.motif_mask
+
+        n_recycle = block_state.n_recycle
+        callback = block_state.callback
+        callback_steps = block_state.callback_steps or 1
+
+        X_denoised_L_traj = []
+        X_L = xyz.clone()
+        D = X_L.shape[0]
+        device = X_L.device
+
+        # Ensure all tensors are on the same device as xyz
+        noise_schedule = noise_schedule.to(device)
+        if motif_mask is not None:
+            motif_mask = motif_mask.to(device)
+
+        single = None
+        pair = None
+        sequence_logits = None
+        sequence_indices = None
+
+        has_transformer = hasattr(components, "transformer") and components.transformer is not None
+        has_scheduler = hasattr(components, "scheduler") and components.scheduler is not None
+
+        # Iterate over consecutive pairs (c_t_minus_1, c_t) in the noise schedule
+        # noise_schedule goes from high noise to low noise
+        for step_num in range(len(noise_schedule) - 1):
+            c_t_minus_1 = noise_schedule[step_num]
+            c_t = noise_schedule[step_num + 1]
+
+            # Step 1: Inject stochastic noise (matching original sampler)
+            if has_scheduler:
+                X_noisy_L, t_hat = components.scheduler.add_noise(
+                    X_L, c_t_minus_1, c_t, motif_mask=motif_mask
+                )
+            else:
+                X_noisy_L = X_L
+                t_hat = c_t_minus_1
+
+            # Step 2: Model forward pass
+            if has_transformer:
+                # t_hat is a scalar, tile to batch dimension
+                t_batch = (t_hat.to(device).expand(D) if isinstance(t_hat, torch.Tensor)
+                          else torch.full((D,), t_hat, device=device))
+
+                output = components.transformer(
+                    xyz_noisy=X_noisy_L,
+                    t=t_batch,
+                    motif_mask=motif_mask,
+                    n_recycle=n_recycle,
+                )
+
+                X_denoised_L = output.xyz
+                single = output.single
+                pair = output.pair
+                sequence_logits = output.sequence_logits
+                sequence_indices = output.sequence_indices
+            else:
+                X_denoised_L = X_noisy_L
+
+            # Step 3: Euler update with step_scale (matching original sampler)
+            if has_scheduler:
+                X_L = components.scheduler.step(
+                    xyz_pred=X_denoised_L,
+                    xyz_noisy=X_noisy_L,
+                    c_t_minus_1=c_t_minus_1,
+                    c_t=c_t,
+                    motif_mask=motif_mask,
+                )
+            else:
+                # Fallback simple Euler step
+                delta_L = (X_noisy_L - X_denoised_L) / (t_hat + 1e-8)
+                d_t = c_t - t_hat
+                X_L = X_noisy_L + d_t * delta_L
+
+            X_denoised_L_traj.append(X_denoised_L.clone())
+
+            if callback is not None and step_num % callback_steps == 0:
+                callback(step_num, c_t_minus_1, X_L)
+
+        block_state.xyz = X_L
+        block_state.single = single
+        block_state.pair = pair
+        block_state.sequence_logits = sequence_logits
+        block_state.sequence_indices = sequence_indices
+        block_state.trajectory = X_denoised_L_traj
+
+        self.set_block_state(state, block_state)
+        return components, state

modular_blocks.py

@@ -0,0 +1,413 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+
+from diffusers.utils import logging
+from diffusers.modular_pipelines import (
+    AutoPipelineBlocks,
+    ModularPipeline,
+    ModularPipelineBlocks,
+    PipelineState,
+    SequentialPipelineBlocks,
+)
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, InsertableDict, OutputParam
+
+from .before_denoise import (
+    RFDiffusionInputStep,
+    RFDiffusionPrepareLatentsStep,
+    RFDiffusionSetTimestepsStep,
+)
+from .decoders import RFDiffusionDecodeStep
+from .denoise import RFDiffusionDenoiseStep
+
+
+logger = logging.get_logger(__name__)
+
+
+# ─── Amino acid mappings (used by MPNN blocks) ─────────────────────────
+
+THREE_TO_ONE = {
+    "ALA": "A", "ARG": "R", "ASN": "N", "ASP": "D", "CYS": "C",
+    "GLN": "Q", "GLU": "E", "GLY": "G", "HIS": "H", "ILE": "I",
+    "LEU": "L", "LYS": "K", "MET": "M", "PHE": "F", "PRO": "P",
+    "SER": "S", "THR": "T", "TRP": "W", "TYR": "Y", "VAL": "V",
+    "UNK": "X",
+}
+
+AA_NAMES = list(THREE_TO_ONE.keys())
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+#  RFDiffusion blocks
+# ═══════════════════════════════════════════════════════════════════════════
+
+
+class RFDiffusionBeforeDenoiseStep(SequentialPipelineBlocks):
+    """Sequential block for pre-denoising preparation."""
+
+    block_classes = [
+        RFDiffusionInputStep,
+        RFDiffusionSetTimestepsStep,
+        RFDiffusionPrepareLatentsStep,
+    ]
+    block_names = ["input", "set_timesteps", "prepare_latents"]
+
+    @property
+    def description(self):
+        return (
+            "Before denoise step that prepares the inputs for the denoise step.\n"
+            "This is a sequential pipeline blocks:\n"
+            " - `RFDiffusionInputStep` processes contigs and prepares input features\n"
+            " - `RFDiffusionSetTimestepsStep` sets up the diffusion timesteps\n"
+            " - `RFDiffusionPrepareLatentsStep` initializes noised coordinates\n"
+        )
+
+
+class RFDiffusionAutoBeforeDenoiseStep(AutoPipelineBlocks):
+    """Auto-select before denoise step based on task."""
+
+    block_classes = [RFDiffusionBeforeDenoiseStep]
+    block_names = ["unconditional"]
+    block_trigger_inputs = [None]
+
+    @property
+    def description(self):
+        return (
+            "Before denoise step that prepares the inputs for the denoise step.\n"
+            "This is an auto pipeline block for protein structure generation.\n"
+            " - `RFDiffusionBeforeDenoiseStep` (unconditional) is used.\n"
+        )
+
+
+class RFDiffusionAutoDenoiseStep(AutoPipelineBlocks):
+    """Auto-select denoise step."""
+
+    block_classes = [RFDiffusionDenoiseStep]
+    block_names = ["denoise"]
+    block_trigger_inputs = [None]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Denoise step that iteratively denoises the protein structure. "
+            "This is an auto pipeline block for protein structure generation. "
+            " - `RFDiffusionDenoiseStep` (denoise) for structure generation."
+        )
+
+
+class RFDiffusionAutoDecodeStep(AutoPipelineBlocks):
+    """Auto-select decode step."""
+
+    block_classes = [RFDiffusionDecodeStep]
+    block_names = ["decode"]
+    block_trigger_inputs = [None]
+
+    @property
+    def description(self):
+        return "Decode step that converts denoised coordinates to PDB output.\n - `RFDiffusionDecodeStep`"
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+#  MPNN blocks (defined before RFDiffusionAutoBlocks which references them)
+# ═══════════════════════════════════════════════════════════════════════════
+
+
+@dataclass
+class MPNNPipelineOutput:
+    """Output from ProteinMPNN / LigandMPNN sequence design."""
+
+    designed_sequence: str
+    sequence_indices: torch.Tensor          # [B, L] token indices
+    sequence_logits: torch.Tensor           # [B, L, n_vocab] logits
+    xyz: torch.Tensor                       # [B, L, 3] input structure (passed through)
+    pdb_string: Optional[str] = None        # PDB with designed sequence
+    sequence_recovery: Optional[float] = None
+
+
+class MPNNSequenceDesignStep(ModularPipelineBlocks):
+    """
+    Design sequences for a protein backbone using ProteinMPNN / LigandMPNN.
+
+    Takes ``xyz`` coordinates (typically from an upstream RFDiffusion denoise
+    step) and runs the ``MPNNModel`` to produce amino acid sequences for
+    the designable regions.
+
+    When no ``mpnn`` component is loaded, falls back to using the sequence
+    predictions from upstream RFDiffusion (or glycine everywhere).
+    """
+
+    model_name = "mpnn"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Design amino acid sequences for protein backbones using "
+            "ProteinMPNN or LigandMPNN. Accepts structure coordinates "
+            "from an upstream diffusion step."
+        )
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec("mpnn", description="MPNNModel (ProteinMPNN or LigandMPNN)"),
+        ]
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam(
+                "xyz", required=True, type_hint=torch.Tensor,
+                description="Protein backbone coordinates [B, L, 3] (CA atoms)",
+            ),
+            InputParam(
+                "motif_mask", type_hint=torch.Tensor,
+                description="Mask for fixed/motif positions [L]. True = fixed sequence.",
+            ),
+            InputParam(
+                "sequence_indices", type_hint=torch.Tensor,
+                description="Known sequence indices for motif positions [B, L] (from RFDiffusion)",
+            ),
+            InputParam(
+                "temperature", default=0.1, type_hint=float,
+                description="Sampling temperature (lower = more deterministic)",
+            ),
+            InputParam(
+                "num_designs", default=1, type_hint=int,
+                description="Number of sequence designs to generate per structure",
+            ),
+            InputParam(
+                "output_type", default="tensor", type_hint=str,
+                description="'tensor', 'pdb', or 'cif'",
+            ),
+            InputParam(
+                "output_path", type_hint=str,
+                description="Path to save designed PDB",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam(
+                "mpnn_output", type_hint=MPNNPipelineOutput,
+                description="MPNN sequence design output",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        xyz = block_state.xyz
+        motif_mask = block_state.motif_mask
+        known_seq = block_state.sequence_indices
+        temperature = block_state.temperature or 0.1
+        output_type = block_state.output_type or "tensor"
+        output_path = block_state.output_path
+
+        B, L, _ = xyz.shape
+        device = xyz.device
+
+        has_mpnn = hasattr(components, "mpnn") and components.mpnn is not None
+
+        if has_mpnn:
+            sequence_logits, sequence_indices = self._run_mpnn(
+                components.mpnn, xyz, motif_mask, known_seq, temperature,
+            )
+        else:
+            if known_seq is not None:
+                sequence_indices = known_seq
+            else:
+                sequence_indices = torch.full((B, L), 7, dtype=torch.long, device=device)  # GLY
+            sequence_logits = torch.zeros(B, L, len(AA_NAMES), device=device)
+            sequence_logits.scatter_(2, sequence_indices.unsqueeze(-1), 1.0)
+
+        seq_list = sequence_indices[0].cpu().tolist()
+        designed_sequence = "".join(
+            THREE_TO_ONE.get(AA_NAMES[min(idx, len(AA_NAMES) - 1)], "X")
+            for idx in seq_list
+        )
+
+        pdb_string = None
+        if output_type in ("pdb",):
+            pdb_string = self._coords_to_pdb(xyz[0], sequence_indices[0])
+            if output_path:
+                import os
+                os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+                with open(output_path, "w") as f:
+                    f.write(pdb_string)
+
+        output = MPNNPipelineOutput(
+            designed_sequence=designed_sequence,
+            sequence_indices=sequence_indices,
+            sequence_logits=sequence_logits,
+            xyz=xyz,
+            pdb_string=pdb_string,
+        )
+
+        block_state.mpnn_output = output
+        self.set_block_state(state, block_state)
+        return components, state
+
+    def _run_mpnn(self, mpnn, xyz, motif_mask, known_seq, temperature):
+        """Run the MPNNModel wrapper on backbone coordinates."""
+        B, L, _ = xyz.shape
+        device = xyz.device
+        dtype = xyz.dtype
+
+        ca = xyz
+        n_offset = torch.tensor([-1.458, 0.0, 0.0], device=device, dtype=dtype)
+        c_offset = torch.tensor([0.550, 1.424, 0.0], device=device, dtype=dtype)
+        o_offset = torch.tensor([0.550, 2.500, 0.0], device=device, dtype=dtype)
+
+        X = torch.stack([
+            ca + n_offset, ca, ca + c_offset, ca + o_offset,
+        ], dim=2)
+
+        if motif_mask is not None:
+            designed_mask = ~motif_mask.unsqueeze(0).expand(B, -1)
+        else:
+            designed_mask = None
+
+        output = mpnn(
+            X=X, S=known_seq, designed_residue_mask=designed_mask, temperature=temperature,
+        )
+
+        logits = output.sequence_logits
+        indices = output.sequence_indices
+
+        if motif_mask is not None and known_seq is not None:
+            indices[:, motif_mask] = known_seq[:, motif_mask]
+
+        return logits, indices
+
+    def _coords_to_pdb(self, xyz: torch.Tensor, seq: torch.Tensor) -> str:
+        xyz_np = xyz.cpu().numpy()
+        seq_np = seq.cpu().numpy()
+        L = xyz_np.shape[0]
+        lines = []
+        for i in range(L):
+            aa_idx = int(seq_np[i])
+            aa_name = AA_NAMES[min(aa_idx, len(AA_NAMES) - 1)]
+            x, y, z = xyz_np[i, :]
+            lines.append(
+                f"ATOM  {i+1:5d}  CA  {aa_name:3s} A{i+1:4d}    "
+                f"{x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00           C  "
+            )
+        lines.append("END")
+        return "\n".join(lines)
+
+
+class MPNNAutoDesignStep(AutoPipelineBlocks):
+    """Auto-select MPNN design step."""
+
+    block_classes = [MPNNSequenceDesignStep]
+    block_names = ["sequence_design"]
+    block_trigger_inputs = [None]
+
+    @property
+    def description(self) -> str:
+        return "Sequence design using ProteinMPNN or LigandMPNN."
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+#  Top-level pipeline blocks
+# ═══════════════════════════════════════════════════════════════════════════
+
+
+class RFDiffusionAutoBlocks(SequentialPipelineBlocks):
+    """
+    Full protein design pipeline: RFDiffusion3 + optional ProteinMPNN/LigandMPNN.
+
+    The active workflow is selected by trigger inputs:
+      - ``contigs`` only → structure generation
+      - ``contigs`` + ``temperature`` → structure + sequence design
+      - ``contigs`` + ``input_xyz`` + ``temperature`` → motif-conditioned + sequence design
+
+    The MPNN step is skipped when ``temperature`` is not provided or when
+    no ``mpnn`` component is loaded.
+    """
+
+    block_classes = [
+        RFDiffusionAutoBeforeDenoiseStep,
+        RFDiffusionAutoDenoiseStep,
+        RFDiffusionAutoDecodeStep,
+        MPNNAutoDesignStep,
+    ]
+    block_names = [
+        "before_denoise",
+        "denoise",
+        "decoder",
+        "sequence_design",
+    ]
+
+    _workflow_map = {
+        "structure_only": {
+            "contigs": True,
+        },
+        "structure_and_sequence": {
+            "contigs": True,
+            "temperature": True,
+        },
+        "motif_structure_and_sequence": {
+            "contigs": True,
+            "input_xyz": True,
+            "temperature": True,
+        },
+    }
+
+    @property
+    def description(self):
+        return (
+            "Modular pipeline for protein design using RFDiffusion3.\n"
+            "Workflows:\n"
+            "  - structure_only: backbone generation\n"
+            "  - structure_and_sequence: backbone + MPNN sequence design\n"
+            "  - motif_structure_and_sequence: motif-conditioned + MPNN\n"
+        )
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+#  Block registries
+# ═══════════════════════════════════════════════════════════════════════════
+
+
+UNCONDITIONAL_BLOCKS = InsertableDict(
+    [
+        ("input", RFDiffusionInputStep),
+        ("set_timesteps", RFDiffusionSetTimestepsStep),
+        ("prepare_latents", RFDiffusionPrepareLatentsStep),
+        ("denoise", RFDiffusionDenoiseStep),
+        ("decode", RFDiffusionDecodeStep),
+        ("sequence_design", MPNNSequenceDesignStep),
+    ]
+)
+
+AUTO_BLOCKS = InsertableDict(
+    [
+        ("before_denoise", RFDiffusionAutoBeforeDenoiseStep),
+        ("denoise", RFDiffusionAutoDenoiseStep),
+        ("decode", RFDiffusionAutoDecodeStep),
+        ("sequence_design", MPNNAutoDesignStep),
+    ]
+)
+
+ALL_BLOCKS = {
+    "unconditional": UNCONDITIONAL_BLOCKS,
+    "auto": AUTO_BLOCKS,
+}

modular_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "RFDiffusionAutoBlocks",
+  "_diffusers_version": "0.37.0.dev0",
+  "auto_map": {
+    "ModularPipelineBlocks": "modular_blocks.RFDiffusionAutoBlocks"
+  }
+}

modular_model_index.json

@@ -0,0 +1,37 @@
+{
+    "_blocks_class_name": "RFDiffusionAutoBlocks",
+    "_class_name": "ModularPipeline",
+    "_diffusers_version": "0.37.0.dev0",
+    "default_num_residues": 100,
+    "default_num_timesteps": 200,
+    "sigma_data": 16.0,
+    "transformer": [
+        null,
+        null,
+        {
+            "pretrained_model_name_or_path": "dn6/RFDiffusion-3",
+            "subfolder": "transformer",
+            "type_hint": [
+                "diffusers",
+                "AutoModel"
+            ],
+            "revision": null,
+            "variant": null
+        }
+    ],
+    "scheduler": [
+        null,
+        null,
+        {
+            "pretrained_model_name_or_path": "dn6/RFDiffusion-3",
+            "subfolder": "scheduler",
+            "type_hint": [
+                "diffusers",
+                "AutoModel"
+            ],
+            "revision": null,
+            "variant": null,
+            "default_creation_method": "from_config"
+        }
+    ]
+}

mpnn/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+from .model_mpnn import MPNNModel, MPNNModelOutput

mpnn/config.json

@@ -0,0 +1,17 @@
+{
+  "_class_name": "MPNNModel",
+  "_diffusers_version": "0.37.0.dev0",
+  "auto_map": {
+    "AutoModel": "model_mpnn.MPNNModel"
+  },
+  "model_type": "protein_mpnn",
+  "hidden_dim": 128,
+  "num_encoder_layers": 3,
+  "num_decoder_layers": 3,
+  "num_neighbors": 48,
+  "dropout_rate": 0.1,
+  "num_positional_embeddings": 16,
+  "min_rbf_mean": 2.0,
+  "max_rbf_mean": 22.0,
+  "num_rbf": 16
+}

mpnn/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f94bbbaa904c3554d3b10397dce3c90e01adb725fd80e3275ff48f11cc4745f
+size 6653924

mpnn/model_mpnn.py

@@ -0,0 +1,178 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+"""
+ProteinMPNN / LigandMPNN model wrapper.
+
+A thin diffusers-compatible wrapper around the foundry MPNN model,
+following the same pattern as the transformer and scheduler wrappers.
+Reuses the foundry model implementation directly, adding only the
+ModelMixin/ConfigMixin interface for diffusers integration.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+from mpnn.model.mpnn import LigandMPNN, ProteinMPNN
+
+
+MODEL_CLASSES = {
+    "protein_mpnn": ProteinMPNN,
+    "ligand_mpnn": LigandMPNN,
+}
+
+
+@dataclass
+class MPNNModelOutput:
+    """Output from the MPNN model wrapper."""
+
+    sequence_logits: torch.Tensor           # [B, L, n_vocab]
+    sequence_indices: torch.Tensor          # [B, L]
+    decoder_features: dict                  # full decoder output dict
+
+
+class MPNNModel(ModelMixin, ConfigMixin):
+    """
+    Diffusers-compatible wrapper around the foundry ProteinMPNN / LigandMPNN.
+
+    Wraps `mpnn.model.mpnn.ProteinMPNN` (or `LigandMPNN`) to provide a
+    diffusers ModelMixin/ConfigMixin interface. All model logic is delegated
+    to the foundry implementation.
+
+    State dict keys match the foundry checkpoint format via the `model.*`
+    prefix (stripped on load).
+    """
+
+    config_name = "config.json"
+
+    @register_to_config
+    def __init__(
+        self,
+        model_type: str = "protein_mpnn",
+        hidden_dim: int = 128,
+        num_encoder_layers: int = 3,
+        num_decoder_layers: int = 3,
+        num_neighbors: int = 48,
+        dropout_rate: float = 0.1,
+        num_positional_embeddings: int = 16,
+        min_rbf_mean: float = 2.0,
+        max_rbf_mean: float = 22.0,
+        num_rbf: int = 16,
+        # LigandMPNN-specific
+        num_context_atoms: int = 25,
+        num_context_encoding_layers: int = 2,
+    ):
+        super().__init__()
+
+        model_cls = MODEL_CLASSES.get(model_type)
+        if model_cls is None:
+            raise ValueError(
+                f"Unknown model_type '{model_type}'. "
+                f"Choose from: {list(MODEL_CLASSES.keys())}"
+            )
+
+        common_kwargs = dict(
+            num_node_features=hidden_dim,
+            num_edge_features=hidden_dim,
+            hidden_dim=hidden_dim,
+            num_encoder_layers=num_encoder_layers,
+            num_decoder_layers=num_decoder_layers,
+            num_neighbors=num_neighbors,
+            dropout_rate=dropout_rate,
+            num_positional_embeddings=num_positional_embeddings,
+            min_rbf_mean=min_rbf_mean,
+            max_rbf_mean=max_rbf_mean,
+            num_rbf=num_rbf,
+        )
+
+        if model_type == "ligand_mpnn":
+            common_kwargs["num_context_atoms"] = num_context_atoms
+            common_kwargs["num_context_encoding_layers"] = num_context_encoding_layers
+
+        self.model = model_cls(**common_kwargs)
+
+    def forward(
+        self,
+        X: torch.Tensor,
+        S: Optional[torch.Tensor] = None,
+        residue_mask: Optional[torch.Tensor] = None,
+        designed_residue_mask: Optional[torch.Tensor] = None,
+        chain_labels: Optional[torch.Tensor] = None,
+        R_idx: Optional[torch.Tensor] = None,
+        temperature: float = 0.1,
+        **kwargs,
+    ) -> MPNNModelOutput:
+        """
+        Run ProteinMPNN / LigandMPNN sequence design.
+
+        Args:
+            X: Backbone atom coordinates [B, L, num_atoms, 3].
+                For ProteinMPNN: num_atoms=4 (N, CA, C, O).
+            S: Ground-truth sequence tokens [B, L] (optional, for teacher forcing).
+            residue_mask: Valid residue mask [B, L] (default: all valid).
+            designed_residue_mask: Which residues to design [B, L] (default: all).
+            chain_labels: Chain identifiers [B, L] (default: single chain).
+            R_idx: Residue indices [B, L] (default: 0..L-1).
+            temperature: Sampling temperature (default: 0.1).
+
+        Returns:
+            MPNNModelOutput with sequence logits and sampled indices.
+        """
+        B, L = X.shape[0], X.shape[1]
+        device = X.device
+
+        if S is None:
+            S = torch.zeros(B, L, dtype=torch.long, device=device)
+        if residue_mask is None:
+            residue_mask = torch.ones(B, L, dtype=torch.bool, device=device)
+        if designed_residue_mask is None:
+            designed_residue_mask = torch.ones(B, L, dtype=torch.bool, device=device)
+        if chain_labels is None:
+            chain_labels = torch.zeros(B, L, dtype=torch.long, device=device)
+        if R_idx is None:
+            R_idx = torch.arange(L, device=device).unsqueeze(0).expand(B, -1)
+
+        # Atom mask: mark all atoms as valid based on coordinate presence
+        X_m = (X.abs().sum(dim=-1) > 0).float()  # [B, L, num_atoms]
+
+        network_input = {
+            "X": X,
+            "X_m": X_m,
+            "S": S,
+            "R_idx": R_idx,
+            "chain_labels": chain_labels,
+            "residue_mask": residue_mask,
+            "designed_residue_mask": designed_residue_mask,
+            "temperature": temperature,
+            **kwargs,
+        }
+
+        output = self.model(network_input)
+
+        logits = output["decoder_features"]["logits"]  # [B, L, n_vocab]
+        S_sampled = output["decoder_features"].get(
+            "S_sampled", logits.argmax(dim=-1)
+        )
+
+        return MPNNModelOutput(
+            sequence_logits=logits,
+            sequence_indices=S_sampled,
+            decoder_features=output["decoder_features"],
+        )

mpnn_ligand/config.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": "MPNNModel",
+  "_diffusers_version": "0.37.0.dev0",
+  "auto_map": {
+    "AutoModel": "model_mpnn.MPNNModel"
+  },
+  "model_type": "ligand_mpnn",
+  "hidden_dim": 128,
+  "num_encoder_layers": 3,
+  "num_decoder_layers": 3,
+  "num_neighbors": 32,
+  "dropout_rate": 0.1,
+  "num_positional_embeddings": 16,
+  "min_rbf_mean": 2.0,
+  "max_rbf_mean": 22.0,
+  "num_rbf": 16,
+  "num_context_atoms": 25,
+  "num_context_encoding_layers": 2
+}

mpnn_ligand/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63d94e809089821f2f01c965b8f3a5e736d6b9b53c378fc27d42f47376e18964
+size 10497908

mpnn_ligand/model_mpnn.py

@@ -0,0 +1,178 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+"""
+ProteinMPNN / LigandMPNN model wrapper.
+
+A thin diffusers-compatible wrapper around the foundry MPNN model,
+following the same pattern as the transformer and scheduler wrappers.
+Reuses the foundry model implementation directly, adding only the
+ModelMixin/ConfigMixin interface for diffusers integration.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+from mpnn.model.mpnn import LigandMPNN, ProteinMPNN
+
+
+MODEL_CLASSES = {
+    "protein_mpnn": ProteinMPNN,
+    "ligand_mpnn": LigandMPNN,
+}
+
+
+@dataclass
+class MPNNModelOutput:
+    """Output from the MPNN model wrapper."""
+
+    sequence_logits: torch.Tensor           # [B, L, n_vocab]
+    sequence_indices: torch.Tensor          # [B, L]
+    decoder_features: dict                  # full decoder output dict
+
+
+class MPNNModel(ModelMixin, ConfigMixin):
+    """
+    Diffusers-compatible wrapper around the foundry ProteinMPNN / LigandMPNN.
+
+    Wraps `mpnn.model.mpnn.ProteinMPNN` (or `LigandMPNN`) to provide a
+    diffusers ModelMixin/ConfigMixin interface. All model logic is delegated
+    to the foundry implementation.
+
+    State dict keys match the foundry checkpoint format via the `model.*`
+    prefix (stripped on load).
+    """
+
+    config_name = "config.json"
+
+    @register_to_config
+    def __init__(
+        self,
+        model_type: str = "protein_mpnn",
+        hidden_dim: int = 128,
+        num_encoder_layers: int = 3,
+        num_decoder_layers: int = 3,
+        num_neighbors: int = 48,
+        dropout_rate: float = 0.1,
+        num_positional_embeddings: int = 16,
+        min_rbf_mean: float = 2.0,
+        max_rbf_mean: float = 22.0,
+        num_rbf: int = 16,
+        # LigandMPNN-specific
+        num_context_atoms: int = 25,
+        num_context_encoding_layers: int = 2,
+    ):
+        super().__init__()
+
+        model_cls = MODEL_CLASSES.get(model_type)
+        if model_cls is None:
+            raise ValueError(
+                f"Unknown model_type '{model_type}'. "
+                f"Choose from: {list(MODEL_CLASSES.keys())}"
+            )
+
+        common_kwargs = dict(
+            num_node_features=hidden_dim,
+            num_edge_features=hidden_dim,
+            hidden_dim=hidden_dim,
+            num_encoder_layers=num_encoder_layers,
+            num_decoder_layers=num_decoder_layers,
+            num_neighbors=num_neighbors,
+            dropout_rate=dropout_rate,
+            num_positional_embeddings=num_positional_embeddings,
+            min_rbf_mean=min_rbf_mean,
+            max_rbf_mean=max_rbf_mean,
+            num_rbf=num_rbf,
+        )
+
+        if model_type == "ligand_mpnn":
+            common_kwargs["num_context_atoms"] = num_context_atoms
+            common_kwargs["num_context_encoding_layers"] = num_context_encoding_layers
+
+        self.model = model_cls(**common_kwargs)
+
+    def forward(
+        self,
+        X: torch.Tensor,
+        S: Optional[torch.Tensor] = None,
+        residue_mask: Optional[torch.Tensor] = None,
+        designed_residue_mask: Optional[torch.Tensor] = None,
+        chain_labels: Optional[torch.Tensor] = None,
+        R_idx: Optional[torch.Tensor] = None,
+        temperature: float = 0.1,
+        **kwargs,
+    ) -> MPNNModelOutput:
+        """
+        Run ProteinMPNN / LigandMPNN sequence design.
+
+        Args:
+            X: Backbone atom coordinates [B, L, num_atoms, 3].
+                For ProteinMPNN: num_atoms=4 (N, CA, C, O).
+            S: Ground-truth sequence tokens [B, L] (optional, for teacher forcing).
+            residue_mask: Valid residue mask [B, L] (default: all valid).
+            designed_residue_mask: Which residues to design [B, L] (default: all).
+            chain_labels: Chain identifiers [B, L] (default: single chain).
+            R_idx: Residue indices [B, L] (default: 0..L-1).
+            temperature: Sampling temperature (default: 0.1).
+
+        Returns:
+            MPNNModelOutput with sequence logits and sampled indices.
+        """
+        B, L = X.shape[0], X.shape[1]
+        device = X.device
+
+        if S is None:
+            S = torch.zeros(B, L, dtype=torch.long, device=device)
+        if residue_mask is None:
+            residue_mask = torch.ones(B, L, dtype=torch.bool, device=device)
+        if designed_residue_mask is None:
+            designed_residue_mask = torch.ones(B, L, dtype=torch.bool, device=device)
+        if chain_labels is None:
+            chain_labels = torch.zeros(B, L, dtype=torch.long, device=device)
+        if R_idx is None:
+            R_idx = torch.arange(L, device=device).unsqueeze(0).expand(B, -1)
+
+        # Atom mask: mark all atoms as valid based on coordinate presence
+        X_m = (X.abs().sum(dim=-1) > 0).float()  # [B, L, num_atoms]
+
+        network_input = {
+            "X": X,
+            "X_m": X_m,
+            "S": S,
+            "R_idx": R_idx,
+            "chain_labels": chain_labels,
+            "residue_mask": residue_mask,
+            "designed_residue_mask": designed_residue_mask,
+            "temperature": temperature,
+            **kwargs,
+        }
+
+        output = self.model(network_input)
+
+        logits = output["decoder_features"]["logits"]  # [B, L, n_vocab]
+        S_sampled = output["decoder_features"].get(
+            "S_sampled", logits.argmax(dim=-1)
+        )
+
+        return MPNNModelOutput(
+            sequence_logits=logits,
+            sequence_indices=S_sampled,
+            decoder_features=output["decoder_features"],
+        )

mpnn_soluble/config.json

@@ -0,0 +1,17 @@
+{
+  "_class_name": "MPNNModel",
+  "_diffusers_version": "0.37.0.dev0",
+  "auto_map": {
+    "AutoModel": "model_mpnn.MPNNModel"
+  },
+  "model_type": "protein_mpnn",
+  "hidden_dim": 128,
+  "num_encoder_layers": 3,
+  "num_decoder_layers": 3,
+  "num_neighbors": 48,
+  "dropout_rate": 0.1,
+  "num_positional_embeddings": 16,
+  "min_rbf_mean": 2.0,
+  "max_rbf_mean": 22.0,
+  "num_rbf": 16
+}

mpnn_soluble/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d7ce5491820a26c7cfae3ab46421c40d1af645accb6c10b4818243eea2f0165
+size 6653924

mpnn_soluble/model_mpnn.py

@@ -0,0 +1,178 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+"""
+ProteinMPNN / LigandMPNN model wrapper.
+
+A thin diffusers-compatible wrapper around the foundry MPNN model,
+following the same pattern as the transformer and scheduler wrappers.
+Reuses the foundry model implementation directly, adding only the
+ModelMixin/ConfigMixin interface for diffusers integration.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+from mpnn.model.mpnn import LigandMPNN, ProteinMPNN
+
+
+MODEL_CLASSES = {
+    "protein_mpnn": ProteinMPNN,
+    "ligand_mpnn": LigandMPNN,
+}
+
+
+@dataclass
+class MPNNModelOutput:
+    """Output from the MPNN model wrapper."""
+
+    sequence_logits: torch.Tensor           # [B, L, n_vocab]
+    sequence_indices: torch.Tensor          # [B, L]
+    decoder_features: dict                  # full decoder output dict
+
+
+class MPNNModel(ModelMixin, ConfigMixin):
+    """
+    Diffusers-compatible wrapper around the foundry ProteinMPNN / LigandMPNN.
+
+    Wraps `mpnn.model.mpnn.ProteinMPNN` (or `LigandMPNN`) to provide a
+    diffusers ModelMixin/ConfigMixin interface. All model logic is delegated
+    to the foundry implementation.
+
+    State dict keys match the foundry checkpoint format via the `model.*`
+    prefix (stripped on load).
+    """
+
+    config_name = "config.json"
+
+    @register_to_config
+    def __init__(
+        self,
+        model_type: str = "protein_mpnn",
+        hidden_dim: int = 128,
+        num_encoder_layers: int = 3,
+        num_decoder_layers: int = 3,
+        num_neighbors: int = 48,
+        dropout_rate: float = 0.1,
+        num_positional_embeddings: int = 16,
+        min_rbf_mean: float = 2.0,
+        max_rbf_mean: float = 22.0,
+        num_rbf: int = 16,
+        # LigandMPNN-specific
+        num_context_atoms: int = 25,
+        num_context_encoding_layers: int = 2,
+    ):
+        super().__init__()
+
+        model_cls = MODEL_CLASSES.get(model_type)
+        if model_cls is None:
+            raise ValueError(
+                f"Unknown model_type '{model_type}'. "
+                f"Choose from: {list(MODEL_CLASSES.keys())}"
+            )
+
+        common_kwargs = dict(
+            num_node_features=hidden_dim,
+            num_edge_features=hidden_dim,
+            hidden_dim=hidden_dim,
+            num_encoder_layers=num_encoder_layers,
+            num_decoder_layers=num_decoder_layers,
+            num_neighbors=num_neighbors,
+            dropout_rate=dropout_rate,
+            num_positional_embeddings=num_positional_embeddings,
+            min_rbf_mean=min_rbf_mean,
+            max_rbf_mean=max_rbf_mean,
+            num_rbf=num_rbf,
+        )
+
+        if model_type == "ligand_mpnn":
+            common_kwargs["num_context_atoms"] = num_context_atoms
+            common_kwargs["num_context_encoding_layers"] = num_context_encoding_layers
+
+        self.model = model_cls(**common_kwargs)
+
+    def forward(
+        self,
+        X: torch.Tensor,
+        S: Optional[torch.Tensor] = None,
+        residue_mask: Optional[torch.Tensor] = None,
+        designed_residue_mask: Optional[torch.Tensor] = None,
+        chain_labels: Optional[torch.Tensor] = None,
+        R_idx: Optional[torch.Tensor] = None,
+        temperature: float = 0.1,
+        **kwargs,
+    ) -> MPNNModelOutput:
+        """
+        Run ProteinMPNN / LigandMPNN sequence design.
+
+        Args:
+            X: Backbone atom coordinates [B, L, num_atoms, 3].
+                For ProteinMPNN: num_atoms=4 (N, CA, C, O).
+            S: Ground-truth sequence tokens [B, L] (optional, for teacher forcing).
+            residue_mask: Valid residue mask [B, L] (default: all valid).
+            designed_residue_mask: Which residues to design [B, L] (default: all).
+            chain_labels: Chain identifiers [B, L] (default: single chain).
+            R_idx: Residue indices [B, L] (default: 0..L-1).
+            temperature: Sampling temperature (default: 0.1).
+
+        Returns:
+            MPNNModelOutput with sequence logits and sampled indices.
+        """
+        B, L = X.shape[0], X.shape[1]
+        device = X.device
+
+        if S is None:
+            S = torch.zeros(B, L, dtype=torch.long, device=device)
+        if residue_mask is None:
+            residue_mask = torch.ones(B, L, dtype=torch.bool, device=device)
+        if designed_residue_mask is None:
+            designed_residue_mask = torch.ones(B, L, dtype=torch.bool, device=device)
+        if chain_labels is None:
+            chain_labels = torch.zeros(B, L, dtype=torch.long, device=device)
+        if R_idx is None:
+            R_idx = torch.arange(L, device=device).unsqueeze(0).expand(B, -1)
+
+        # Atom mask: mark all atoms as valid based on coordinate presence
+        X_m = (X.abs().sum(dim=-1) > 0).float()  # [B, L, num_atoms]
+
+        network_input = {
+            "X": X,
+            "X_m": X_m,
+            "S": S,
+            "R_idx": R_idx,
+            "chain_labels": chain_labels,
+            "residue_mask": residue_mask,
+            "designed_residue_mask": designed_residue_mask,
+            "temperature": temperature,
+            **kwargs,
+        }
+
+        output = self.model(network_input)
+
+        logits = output["decoder_features"]["logits"]  # [B, L, n_vocab]
+        S_sampled = output["decoder_features"].get(
+            "S_sampled", logits.argmax(dim=-1)
+        )
+
+        return MPNNModelOutput(
+            sequence_logits=logits,
+            sequence_indices=S_sampled,
+            decoder_features=output["decoder_features"],
+        )

pyproject.toml

@@ -0,0 +1,66 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "modular-rfdiffusion"
+version = "0.1.0"
+description = "Modular Diffusers Pipeline for RFDiffusion protein structure generation"
+readme = "README.md"
+license = {text = "Apache-2.0"}
+requires-python = ">=3.8"
+authors = [
+    {name = "Dhruv Nair"}
+]
+keywords = ["diffusion", "protein", "rfdiffusion", "deep-learning", "pytorch"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+]
+
+dependencies = [
+    "torch>=2.0.0",
+    "numpy>=1.21.0",
+    "diffusers>=0.25.0",
+    "huggingface-hub>=0.20.0",
+    "scipy>=1.7.0",
+    "hydra-core>=1.0.0",
+    "omegaconf>=2.0.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0.0",
+    "pytest-cov>=4.0.0",
+    "black>=23.0.0",
+    "ruff>=0.1.0",
+    "mypy>=1.0.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/DN6/modular-diffusers"
+Repository = "https://github.com/DN6/modular-diffusers"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["*"]
+
+[tool.black]
+line-length = 119
+target-version = ["py38", "py39", "py310", "py311"]
+
+[tool.ruff]
+line-length = 119
+target-version = "py38"
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I", "N"]
+ignore = ["E501"]

scheduler/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+from .model import RFDiffusionScheduler

scheduler/config.json

@@ -0,0 +1,16 @@
+{
+  "_class_name": "RFDiffusionScheduler",
+  "_diffusers_version": "0.37.0.dev0",
+  "auto_map": {
+    "AutoModel": "model.RFDiffusionScheduler"
+  },
+  "num_timesteps": 200,
+  "sigma_data": 16.0,
+  "s_min": 4e-4,
+  "s_max": 160.0,
+  "p": 7.0,
+  "gamma_0": 0.6,
+  "gamma_min": 1.0,
+  "noise_scale": 1.003,
+  "step_scale": 1.5
+}

scheduler/model.py

@@ -0,0 +1,152 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+"""
+RFDiffusion3 Scheduler.
+
+A thin diffusers-compatible wrapper around the foundry EDM noise schedule
+and stochastic sampling logic from `rfd3.model.inference_sampler`.
+"""
+
+from typing import Optional
+
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+
+# Reuse the original noise schedule and sampling config directly
+from rfd3.model.inference_sampler import SampleDiffusionWithMotif
+
+
+class RFDiffusionScheduler(ConfigMixin):
+    """
+    Diffusers-compatible scheduler wrapping the foundry EDM sampler.
+
+    Delegates noise schedule construction and sampling parameters to
+    `rfd3.model.inference_sampler.SampleDiffusionWithMotif`.
+    """
+
+    config_name = "config.json"
+
+    @register_to_config
+    def __init__(
+        self,
+        num_timesteps: int = 200,
+        sigma_data: float = 16.0,
+        s_min: float = 4e-4,
+        s_max: float = 160.0,
+        p: float = 7.0,
+        gamma_0: float = 0.6,
+        gamma_min: float = 1.0,
+        noise_scale: float = 1.003,
+        step_scale: float = 1.5,
+    ):
+        # Instantiate the foundry sampler with matching parameters
+        self._sampler = SampleDiffusionWithMotif(
+            num_timesteps=num_timesteps,
+            sigma_data=sigma_data,
+            s_min=s_min,
+            s_max=s_max,
+            p=p,
+            gamma_0=gamma_0,
+            gamma_min=gamma_min,
+            noise_scale=noise_scale,
+            step_scale=step_scale,
+        )
+
+    @property
+    def sampler(self) -> SampleDiffusionWithMotif:
+        return self._sampler
+
+    def get_noise_schedule(self, device: torch.device = None) -> torch.Tensor:
+        """
+        Construct the EDM noise schedule using the foundry implementation.
+
+        Returns:
+            torch.Tensor: Noise schedule [num_timesteps] from high to low noise.
+        """
+        return self._sampler._construct_inference_noise_schedule(
+            device=device or torch.device("cpu")
+        )
+
+    def get_initial_noise_level(self, device: torch.device = None) -> torch.Tensor:
+        """Get the first (largest) noise level from the schedule."""
+        return self.get_noise_schedule(device=device)[0]
+
+    def step(
+        self,
+        xyz_pred: torch.Tensor,
+        xyz_noisy: torch.Tensor,
+        c_t_minus_1: torch.Tensor,
+        c_t: torch.Tensor,
+        motif_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Perform one Euler denoising step matching the foundry sampler.
+
+        The foundry ``sample_diffusion_like_af3`` does NOT clamp motif
+        coordinates after the Euler update — it relies on noise injection
+        having zeroed epsilon for motif atoms so the model's delta is ~0
+        there.  We replicate that behaviour here.
+
+        Args:
+            xyz_pred: Model's denoised prediction X_denoised_L [B, L, 3]
+            xyz_noisy: Noise-injected coordinates X_noisy_L [B, L, 3]
+            c_t_minus_1: Previous noise level
+            c_t: Next (lower) noise level
+            motif_mask: Boolean mask for fixed positions (True = fixed) [L]
+                        (unused — kept for API compatibility)
+
+        Returns:
+            Updated coordinates X_L [B, L, 3]
+        """
+        gamma = self._sampler.gamma_0 if c_t > self._sampler.gamma_min else 0.0
+        t_hat = c_t_minus_1 * (gamma + 1.0)
+
+        delta_L = (xyz_noisy - xyz_pred) / t_hat
+        d_t = c_t - t_hat
+        xyz_next = xyz_noisy + self._sampler.step_scale * d_t * delta_L
+
+        return xyz_next
+
+    def add_noise(
+        self,
+        xyz: torch.Tensor,
+        c_t_minus_1: torch.Tensor,
+        c_t: torch.Tensor,
+        motif_mask: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Inject stochastic noise before the model call, matching the foundry sampler.
+
+        Args:
+            xyz: Current coordinates X_L [B, L, 3]
+            c_t_minus_1: Previous noise level
+            c_t: Current (next lower) noise level
+            motif_mask: Boolean mask for fixed positions (True = fixed) [L]
+
+        Returns:
+            Tuple of (noisy coordinates X_noisy_L, t_hat scalar)
+        """
+        gamma = self._sampler.gamma_0 if c_t > self._sampler.gamma_min else 0.0
+        t_hat = c_t_minus_1 * (gamma + 1.0)
+
+        noise_std = self._sampler.noise_scale * torch.sqrt(t_hat**2 - c_t_minus_1**2)
+        epsilon = noise_std * torch.randn_like(xyz)
+
+        if motif_mask is not None:
+            epsilon[:, motif_mask] = 0.0
+
+        xyz_noisy = xyz + epsilon
+        return xyz_noisy, t_hat

transformer/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+from .model_rfdiffusion import RFDiffusionTransformerModel, RFDiffusionTransformerOutput

transformer/config.json

@@ -0,0 +1,22 @@
+{
+  "_class_name": "RFDiffusionTransformerModel",
+  "_diffusers_version": "0.37.0.dev0",
+  "auto_map": {
+    "AutoModel": "model_rfdiffusion.RFDiffusionTransformerModel"
+  },
+  "c_s": 384,
+  "c_z": 128,
+  "c_atom": 128,
+  "c_atompair": 16,
+  "c_token": 768,
+  "c_t_embed": 256,
+  "sigma_data": 16.0,
+  "n_pairformer_block": 2,
+  "n_diffusion_block": 18,
+  "n_atom_encoder_block": 3,
+  "n_atom_decoder_block": 3,
+  "n_head": 16,
+  "n_pairformer_head": 16,
+  "n_recycle": 2,
+  "p_drop": 0.0
+}

transformer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1fe11cd4fe730654e0659d5218228a9d27bdcdac4c880faad993feaa532e99eb
+size 672279561

transformer/model_rfdiffusion.py

@@ -0,0 +1,297 @@
+# Copyright 2025 Dhruv Nair. 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.
+
+"""
+RFDiffusion3 Transformer model.
+
+A thin diffusers-compatible wrapper around the existing foundry RFD3
+model components. Reuses all layer implementations from `rfd3.model.layers.*`
+and `foundry.model.layers.*` directly, adding only the ModelMixin/ConfigMixin
+interface needed for diffusers ModularPipeline integration.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+# Reuse all components from the foundry/rfd3 package directly
+from rfd3.model.RFD3_diffusion_module import RFD3DiffusionModule
+from rfd3.model.layers.encoders import TokenInitializer
+
+
+@dataclass
+class RFDiffusionTransformerOutput:
+    """Output class for RFDiffusion transformer."""
+
+    xyz: torch.Tensor
+    single: torch.Tensor
+    pair: torch.Tensor
+    sequence_logits: Optional[torch.Tensor] = None
+    sequence_indices: Optional[torch.Tensor] = None
+
+
+class RFDiffusionTransformerModel(ModelMixin, ConfigMixin):
+    """
+    Diffusers-compatible wrapper around the foundry RFD3 model.
+
+    This wraps `rfd3.model.RFD3_diffusion_module.RFD3DiffusionModule` and
+    `rfd3.model.layers.encoders.TokenInitializer` to provide a diffusers
+    ModelMixin/ConfigMixin interface. All actual model logic is delegated
+    to the foundry implementation.
+
+    State dict keys match the foundry checkpoint format via the
+    `token_initializer.*` and `diffusion_module.*` prefixes.
+    """
+
+    config_name = "config.json"
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        c_s: int = 384,
+        c_z: int = 128,
+        c_atom: int = 128,
+        c_atompair: int = 16,
+        c_token: int = 768,
+        c_t_embed: int = 256,
+        sigma_data: float = 16.0,
+        n_pairformer_block: int = 2,
+        n_diffusion_block: int = 18,
+        n_atom_encoder_block: int = 3,
+        n_atom_decoder_block: int = 3,
+        n_head: int = 16,
+        n_pairformer_head: int = 16,
+        n_recycle: int = 2,
+        p_drop: float = 0.0,
+    ):
+        super().__init__()
+
+        # ── Shared sub-configs matching rfd3_net.yaml exactly ────────────
+        cross_attention_block = {
+            "n_head": 4,
+            "c_model": c_atom,
+            "dropout": 0.0,
+            "kq_norm": True,
+        }
+        downcast_cfg = {
+            "method": "cross_attention",
+            "cross_attention_block": cross_attention_block,
+        }
+        upcast_cfg = {
+            "method": "cross_attention",
+            "n_split": 3,
+            "cross_attention_block": cross_attention_block,
+        }
+
+        # ── TokenInitializer (foundry) ──────────────────────────────────
+        token_1d_features = {
+            "ref_motif_token_type": 3,
+            "restype": 32,
+            "ref_plddt": 1,
+            "is_non_loopy": 1,
+        }
+        atom_1d_features = {
+            "ref_atom_name_chars": 256,
+            "ref_element": 128,
+            "ref_charge": 1,
+            "ref_mask": 1,
+            "ref_is_motif_atom_with_fixed_coord": 1,
+            "ref_is_motif_atom_unindexed": 1,
+            "has_zero_occupancy": 1,
+            "ref_pos": 3,
+            "ref_atomwise_rasa": 3,
+            "active_donor": 1,
+            "active_acceptor": 1,
+            "is_atom_level_hotspot": 1,
+        }
+
+        self.token_initializer = TokenInitializer(
+            c_s=c_s,
+            c_z=c_z,
+            c_atom=c_atom,
+            c_atompair=c_atompair,
+            relative_position_encoding={"r_max": 32, "s_max": 2},
+            n_pairformer_blocks=n_pairformer_block,
+            pairformer_block={
+                "attention_pair_bias": {"n_head": n_pairformer_head, "kq_norm": True},
+            },
+            downcast=downcast_cfg,
+            token_1d_features=token_1d_features,
+            atom_1d_features=atom_1d_features,
+            atom_transformer={
+                "n_blocks": 0,
+                "atom_transformer_block": {
+                    "n_head": 4,
+                    "kq_norm": True,
+                    "no_residual_connection_between_attention_and_transition": False,
+                    "dropout": 0.0,
+                    "n_attn_seq_neighbours": 4,
+                    "n_attn_keys": 128,
+                },
+            },
+        )
+
+        # ── RFD3DiffusionModule (foundry) ───────────────────────────────
+        self.diffusion_module = RFD3DiffusionModule(
+            c_atom=c_atom,
+            c_atompair=c_atompair,
+            c_token=c_token,
+            c_s=c_s,
+            c_z=c_z,
+            c_t_embed=c_t_embed,
+            sigma_data=sigma_data,
+            f_pred="edm",
+            n_attn_seq_neighbours=2,
+            n_attn_keys=128,
+            n_recycle=n_recycle,
+            atom_attention_encoder={
+                "n_blocks": n_atom_encoder_block,
+                "atom_transformer_block": {
+                    "n_head": 4,
+                    "kq_norm": True,
+                    "no_residual_connection_between_attention_and_transition": False,
+                    "dropout": 0.0,
+                },
+            },
+            diffusion_token_encoder={
+                "sigma_data": sigma_data,
+                "n_pairformer_blocks": n_pairformer_block,
+                "pairformer_block": {
+                    "attention_pair_bias": {"n_head": n_pairformer_head, "kq_norm": True},
+                },
+                "use_distogram": True,
+                "use_self": True,
+                "use_sinusoidal_distogram_embedder": False,
+            },
+            diffusion_transformer={
+                "n_block": n_diffusion_block,
+                "diffusion_transformer_block": {
+                    "n_head": n_head,
+                    "kq_norm": True,
+                    "no_residual_connection_between_attention_and_transition": False,
+                    "dropout": p_drop,
+                },
+            },
+            atom_attention_decoder={
+                "n_blocks": n_atom_decoder_block,
+                "upcast": upcast_cfg,
+                "downcast": downcast_cfg,
+                "atom_transformer_block": {
+                    "n_head": 4,
+                    "kq_norm": True,
+                    "no_residual_connection_between_attention_and_transition": False,
+                    "dropout": p_drop,
+                },
+            },
+            downcast=downcast_cfg,
+        )
+
+    @property
+    def sigma_data(self) -> float:
+        return self.diffusion_module.sigma_data
+
+    def forward(
+        self,
+        xyz_noisy: torch.Tensor,
+        t: torch.Tensor,
+        f: Optional[dict] = None,
+        motif_mask: Optional[torch.Tensor] = None,
+        n_recycle: Optional[int] = None,
+        **kwargs,
+    ) -> RFDiffusionTransformerOutput:
+        """
+        Forward pass delegated to the foundry RFD3DiffusionModule.
+
+        Args:
+            xyz_noisy: Noisy atom coordinates [B, L, 3]
+            t: Noise level / timestep [B]
+            f: Feature dictionary (as expected by foundry). If None, a minimal
+               feature dict is constructed from xyz_noisy and motif_mask.
+            motif_mask: Mask for fixed motif atoms [L] (used when f is None)
+            n_recycle: Number of recycling iterations
+
+        Returns:
+            RFDiffusionTransformerOutput with denoised coordinates and predictions
+        """
+        B, L, _ = xyz_noisy.shape
+
+        # If caller provides a full feature dict, use the native foundry path
+        if f is not None:
+            initializer_outputs = self.token_initializer(f)
+            outs = self.diffusion_module(
+                X_noisy_L=xyz_noisy,
+                t=t,
+                f=f,
+                n_recycle=n_recycle,
+                **initializer_outputs,
+            )
+            return RFDiffusionTransformerOutput(
+                xyz=outs["X_L"],
+                single=torch.zeros(1),  # not directly exposed by foundry
+                pair=torch.zeros(1),
+                sequence_logits=outs.get("sequence_logits_I"),
+                sequence_indices=outs.get("sequence_indices_I"),
+            )
+
+        # Simplified path: construct minimal feature dict and call dm.forward()
+        # For unconditional generation, each residue has 1 atom (CA), so L = I
+        device = xyz_noisy.device
+        dtype = xyz_noisy.dtype
+
+        if motif_mask is None:
+            motif_mask = torch.zeros(L, dtype=torch.bool, device=device)
+        else:
+            motif_mask = motif_mask.to(device)
+
+        # Construct minimal feature dict with all keys required by foundry
+        f = {
+            "atom_to_token_map": torch.arange(L, device=device),  # 1:1 atom-to-token
+            "unindexing_pair_mask": torch.zeros(L, L, dtype=torch.bool, device=device),
+            "is_ca": torch.ones(L, dtype=torch.bool, device=device),
+            "is_motif_atom_with_fixed_coord": motif_mask,
+            "is_motif_token_with_fully_fixed_coord": motif_mask,
+        }
+
+        # Zero-initialized TokenInitializer outputs (no conditioning features)
+        Q_L_init = torch.zeros(L, self.config.c_atom, device=device, dtype=dtype)
+        C_L = torch.zeros(L, self.config.c_atom, device=device, dtype=dtype)
+        P_LL = torch.zeros(L, L, self.config.c_atompair, device=device, dtype=dtype)
+        S_I = torch.zeros(L, self.config.c_s, device=device, dtype=dtype)
+        Z_II = torch.zeros(L, L, self.config.c_z, device=device, dtype=dtype)
+
+        outs = self.diffusion_module(
+            X_noisy_L=xyz_noisy,
+            t=t,
+            f=f,
+            Q_L_init=Q_L_init,
+            C_L=C_L,
+            P_LL=P_LL,
+            S_I=S_I,
+            Z_II=Z_II,
+            n_recycle=n_recycle,
+        )
+
+        return RFDiffusionTransformerOutput(
+            xyz=outs["X_L"],
+            single=torch.zeros(1),
+            pair=torch.zeros(1),
+            sequence_logits=outs.get("sequence_logits_I"),
+            sequence_indices=outs.get("sequence_indices_I"),
+        )