Upload folder using huggingface_hub

README.md

@@ -0,0 +1,140 @@
+# Protein Structure Prediction with Diffusers
+
+A [diffusers](https://github.com/huggingface/diffusers) `ModularPipeline` wrapper for [RosettaFold3](https://doi.org/10.1101/2025.08.14.670328) (RF3) — a diffusion-based protein structure prediction model that predicts 3D atomic coordinates from amino acid sequences.
+
+RF3 relies on [Foundry](https://github.com/RosettaCommons/foundry) for its underlying implementation 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
+pip install rc-foundry[all]
+pip install diffusers
+```
+
+### Running with Diffusers
+
+```python
+import torch
+from diffusers import ModularPipeline
+
+pipe = ModularPipeline.from_pretrained("dn6/RosettaFold-3", trust_remote_code=True)
+pipe.load_components(device_map="cuda", torch_dtype=torch.bfloat16, trust_remote_code=True)
+
+state = pipe(sequence="MKVLSEGDPWRK...")
+print(state.output.xyz.shape)  # [D, L, 3]
+```
+
+## Workflows
+
+| Workflow | Trigger inputs | What runs |
+|----------|---------------|-----------|
+| `fold` | `sequence` | Full structure prediction (recycling trunk + diffusion) |
+
+### Fold a Sequence
+
+```python
+state = pipe(sequence="MKVLSEGDPWRK...", output_type="cif.gz", output_path="prediction")
+print(state.output.atom_array)
+```
+
+### Full Design Pipeline
+
+RF3 is typically used as a validation step after backbone design with [RFdiffusion3](https://huggingface.co/dn6/RFDiffusion-3):
+
+```
+RFD3 (design backbone) → MPNN (design sequence) → RF3 (validate fold)
+```
+
+```python
+import torch
+from diffusers import AutoModel, ModularPipeline
+
+# 1. Design a backbone + sequence
+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)
+designed_sequence = state.mpnn_output.designed_sequence
+
+# 2. Validate the fold
+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")
+```
+
+## Customizing Workflows
+
+```python
+# Inspect the pipeline structure
+print(pipe.blocks)
+
+# Add a custom block
+from diffusers.modular_pipelines import ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, OutputParam
+
+class ComputeRadiusOfGyration(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
+
+pipe._blocks.sub_blocks.insert("rog", ComputeRadiusOfGyration(), index=3)
+```
+
+## Output Types
+
+| `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 |
+
+```python
+# CIF output with AtomArray
+state = pipe(sequence="MKVLSEG...", output_type="cif.gz", output_path="fold_0")
+atom_array = state.output.atom_array
+
+# Denoising trajectory
+trajectory = state.output.trajectory_stack
+
+# PDB output
+state = pipe(sequence="MKVLSEG...", output_type="pdb", output_path="fold_0.pdb")
+```
+
+## Model Architecture
+
+RF3 is a diffusion model with the same EDM noise schedule as RFdiffusion3 (200 steps), but conditioned on sequence/MSA/template representations from a large recycling trunk:
+
+| Component | Subfolder | Description |
+|-----------|-----------|-------------|
+| `transformer` | `transformer/` | `RF3TransformerModel` (366M params) — FeatureInitializer + Recycler (48 pairformer blocks) + DiffusionModule (24 transformer blocks) + DistogramHead |
+| `scheduler` | `scheduler/` | `RF3Scheduler` (EDM schedule, gamma_0=0.8) |
+
+## Citation
+
+```bibtex
+@article{corley2025accelerating,
+    author = {Corley, Nathaniel and Mathis, Simon and Krishna, Rohith and Bauer, Magnus S and Thompson, Tuscan R and Ahern, Woody and Kazman, Maxwell W and Brent, Rafael I and Didi, Kieran and Kubaney, Andrew and others},
+    title = {Accelerating biomolecular modeling with AtomWorks and RF3},
+    journal = {bioRxiv},
+    year = {2025},
+}
+```

__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2025 Dhruv Nair. All rights reserved.
+# Licensed under the Apache License, Version 2.0
+
+from .transformer import RF3TransformerModel, RF3TransformerOutput
+from .scheduler import RF3Scheduler
+from .modular_blocks import (
+    RF3AutoBeforeDenoiseStep,
+    RF3AutoBlocks,
+    RF3AutoDecodeStep,
+    RF3AutoDenoiseStep,
+)
+from .before_denoise import (
+    RF3InputStep,
+    RF3PrepareLatentsStep,
+    RF3RecyclingStep,
+    RF3SetTimestepsStep,
+)
+from .denoise import RF3DenoiseStep
+from .decoders import RF3DecodeStep, RF3PipelineOutput

before_denoise.py

@@ -0,0 +1,223 @@
+# Copyright 2025 Dhruv Nair. All rights reserved.
+# Licensed under the Apache License, Version 2.0
+
+"""
+Pre-denoising steps for RF3: input processing, timestep setup, recycling trunk, latent preparation.
+"""
+
+from typing import 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 RF3InputStep(ModularPipelineBlocks):
+    """Parse sequence input and prepare feature dict for RF3."""
+
+    model_name = "rf3"
+
+    @property
+    def description(self) -> str:
+        return "Parse sequence and optional MSA/template inputs for structure prediction."
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam("sequence", required=True, type_hint=str, description="Amino acid sequence (one-letter codes)"),
+            InputParam("f", type_hint=dict, description="Pre-built feature dict (overrides sequence)"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam("f", type_hint=dict, description="Feature dictionary for RF3"),
+            OutputParam("L", type_hint=int, description="Sequence length (num atoms)"),
+            OutputParam("I", type_hint=int, description="Num tokens"),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        block_state = self.get_block_state(state)
+
+        f = block_state.f
+        sequence = block_state.sequence
+
+        if f is None:
+            # Build minimal feature dict from sequence
+            L = len(sequence)
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+            # Map sequence to restype indices
+            AA_ORDER = "ARNDCQEGHILKMFPSTWYV"
+            restype = torch.zeros(L, 32, device=device)
+            for i, aa in enumerate(sequence):
+                idx = AA_ORDER.find(aa)
+                if idx >= 0:
+                    restype[i, idx] = 1.0
+                else:
+                    restype[i, 20] = 1.0  # unknown
+
+            f = {
+                "restype": restype,
+                "atom_to_token_map": torch.arange(L, device=device),
+                "is_ca": torch.ones(L, dtype=torch.bool, device=device),
+                "ref_pos": torch.zeros(L, 3, device=device),
+                "ref_charge": torch.zeros(L, device=device),
+                "ref_mask": torch.ones(L, device=device),
+                "ref_element": torch.zeros(L, 128, device=device),
+                "ref_atom_name_chars": torch.zeros(L, 4, 64, device=device),
+            }
+        else:
+            L = f.get("ref_element", f.get("restype")).shape[0]
+
+        block_state.f = f
+        block_state.L = L
+        block_state.I = L  # token count = atom count for CA-only
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class RF3SetTimestepsStep(ModularPipelineBlocks):
+    """Set up EDM noise schedule for RF3."""
+
+    model_name = "rf3"
+
+    @property
+    def description(self) -> str:
+        return "Construct EDM noise schedule for RF3 diffusion sampling."
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [ComponentSpec("scheduler", description="RF3 EDM scheduler")]
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam("num_inference_steps", default=None, type_hint=int),
+            InputParam("L", required=True, type_hint=int),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam("noise_schedule", type_hint=torch.Tensor),
+            OutputParam("num_inference_steps", type_hint=int),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        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:
+            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 RF3RecyclingStep(ModularPipelineBlocks):
+    """Run the recycling trunk (pairformer + MSA + templates)."""
+
+    model_name = "rf3"
+
+    @property
+    def description(self) -> str:
+        return "Run RF3 recycling trunk to produce single/pair representations."
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [ComponentSpec("transformer", description="RF3 transformer model")]
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam("f", required=True, type_hint=dict),
+            InputParam("n_recycles", default=None, type_hint=int),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam("single", type_hint=torch.Tensor, description="Single representation [I, c_s]"),
+            OutputParam("pair", type_hint=torch.Tensor, description="Pair representation [I, I, c_z]"),
+            OutputParam("s_inputs", type_hint=torch.Tensor, description="Input embeddings [I, c_s_inputs]"),
+            OutputParam("distogram", type_hint=torch.Tensor, description="Distogram prediction [I, I, bins]"),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        block_state = self.get_block_state(state)
+
+        f = block_state.f
+        n_recycles = block_state.n_recycles
+
+        if hasattr(components, "transformer") and components.transformer is not None:
+            output = components.transformer(f=f, n_recycles=n_recycles)
+            block_state.single = output.single
+            block_state.pair = output.pair
+            block_state.distogram = output.distogram
+            block_state.s_inputs = None  # populated inside forward
+        else:
+            # Placeholder when no model loaded
+            block_state.single = None
+            block_state.pair = None
+            block_state.distogram = None
+            block_state.s_inputs = None
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class RF3PrepareLatentsStep(ModularPipelineBlocks):
+    """Prepare initial noised coordinates for diffusion sampling."""
+
+    model_name = "rf3"
+
+    @property
+    def description(self) -> str:
+        return "Sample initial Gaussian noise scaled by the first noise schedule value."
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam("generator", type_hint=torch.Generator),
+            InputParam("diffusion_batch_size", default=5, type_hint=int),
+            InputParam("L", required=True, type_hint=int),
+            InputParam("noise_schedule", required=True, type_hint=torch.Tensor),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam("xyz", type_hint=torch.Tensor, description="Initial noised coords [D, L, 3]"),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state):
+        block_state = self.get_block_state(state)
+
+        L = block_state.L
+        noise_schedule = block_state.noise_schedule
+        D = block_state.diffusion_batch_size or 5
+        generator = block_state.generator
+
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        c0 = noise_schedule[0]
+        xyz = c0 * torch.randn((D, L, 3), device=device, generator=generator)
+
+        block_state.xyz = xyz
+
+        self.set_block_state(state, block_state)
+        return components, state

decoders.py

@@ -0,0 +1,169 @@
+# Copyright 2025 Dhruv Nair. All rights reserved.
+# Licensed under the Apache License, Version 2.0
+
+"""
+Decode step for RF3 — converts denoised coordinates to output structures.
+Supports tensor, PDB, and CIF (via AtomWorks) output formats.
+"""
+
+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_ORDER = "ARNDCQEGHILKMFPSTWYV"
+AA_NAMES_3 = [
+    "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, sequence: Optional[str] = None) -> AtomArray:
+    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 sequence:
+        arr.res_name = np.array([
+            AA_NAMES_3[AA_ORDER.find(aa)] if aa in AA_ORDER else "UNK"
+            for aa in sequence
+        ])
+    else:
+        arr.res_name = np.full(L, "ALA")
+    return arr
+
+
+def _build_atom_array_stack(xyz: torch.Tensor, sequence: Optional[str] = None) -> AtomArrayStack:
+    template = _build_atom_array(xyz[0], sequence)
+    arr_stack = stack([template for _ in range(xyz.shape[0])])
+    arr_stack.coord = xyz.detach().cpu().float().numpy()
+    return arr_stack
+
+
+@dataclass
+class RF3PipelineOutput:
+    """Output class for RF3 pipeline."""
+
+    xyz: torch.Tensor
+    atom_array: Optional[AtomArray] = None
+    atom_array_stack: Optional[AtomArrayStack] = None
+    trajectory_stack: Optional[AtomArrayStack] = None
+    distogram: Optional[torch.Tensor] = None
+    sequence: Optional[str] = None
+    pdb_string: Optional[str] = None
+    trajectory: Optional[List[torch.Tensor]] = None
+
+
+class RF3DecodeStep(ModularPipelineBlocks):
+    """
+    Decode step for RF3.
+
+    Supported ``output_type`` values: ``"tensor"``, ``"pdb"``, ``"cif"``, ``"cif.gz"``.
+    """
+
+    model_name = "rf3"
+
+    @property
+    def description(self) -> str:
+        return "Convert predicted coordinates to output format (tensor/PDB/CIF)."
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam("output_type", default="tensor", type_hint=str),
+            InputParam("output_path", type_hint=str),
+            InputParam("xyz", required=True, type_hint=torch.Tensor),
+            InputParam("sequence", type_hint=str),
+            InputParam("distogram", type_hint=torch.Tensor),
+            InputParam("trajectory", type_hint=List[torch.Tensor]),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam("output", type_hint=RF3PipelineOutput),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        xyz = block_state.xyz
+        sequence = block_state.sequence
+        distogram = block_state.distogram
+        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
+
+        if output_type in ("cif", "cif.gz"):
+            atom_array = _build_atom_array(xyz[0], sequence)
+            if xyz.shape[0] > 1:
+                atom_array_stack = _build_atom_array_stack(xyz, sequence)
+            if trajectory:
+                traj_coords = torch.stack([t[0] for t in trajectory])
+                template = _build_atom_array(traj_coords[0], sequence)
+                trajectory_stack = stack([template for _ in range(traj_coords.shape[0])])
+                trajectory_stack.coord = traj_coords.detach().cpu().float().numpy()
+
+        if output_type == "pdb":
+            pdb_string = self._coords_to_pdb(xyz[0], sequence)
+
+        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)
+            elif output_type == "pdb" and pdb_string:
+                with open(output_path, "w") as f:
+                    f.write(pdb_string)
+
+        output = RF3PipelineOutput(
+            xyz=xyz,
+            atom_array=atom_array,
+            atom_array_stack=atom_array_stack,
+            trajectory_stack=trajectory_stack,
+            distogram=distogram,
+            sequence=sequence,
+            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, sequence: Optional[str] = None) -> str:
+        xyz_np = xyz.cpu().numpy()
+        L = xyz_np.shape[0]
+        lines = []
+        for i in range(L):
+            aa = sequence[i] if sequence and i < len(sequence) else "A"
+            aa3 = AA_NAMES_3[AA_ORDER.find(aa)] if aa in AA_ORDER else "UNK"
+            x, y, z = xyz_np[i, :]
+            lines.append(
+                f"ATOM  {i+1:5d}  CA  {aa3: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)

denoise.py

@@ -0,0 +1,135 @@
+# Copyright 2025 Dhruv Nair. All rights reserved.
+# Licensed under the Apache License, Version 2.0
+
+"""
+Denoising loop for RF3.
+
+Same EDM stochastic sampling as RFD3, but conditioned on trunk
+representations (single S_I, pair Z_II) from the recycling step.
+"""
+
+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 RF3DenoiseStep(ModularPipelineBlocks):
+    """
+    Iterative denoising step for RF3.
+
+    Uses trunk representations from the recycling step as conditioning
+    for the diffusion module at each denoising step.
+    """
+
+    model_name = "rf3"
+
+    @property
+    def description(self) -> str:
+        return "Iteratively denoise protein structure conditioned on sequence/MSA representations."
+
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec("transformer", description="RF3 transformer (provides diffusion_module)"),
+            ComponentSpec("scheduler", description="RF3 EDM scheduler"),
+        ]
+
+    @property
+    def inputs(self) -> List[InputParam]:
+        return [
+            InputParam("xyz", required=True, type_hint=torch.Tensor, description="Initial noised coords [D, L, 3]"),
+            InputParam("noise_schedule", required=True, type_hint=torch.Tensor),
+            InputParam("f", required=True, type_hint=dict, description="Feature dictionary"),
+            InputParam("single", type_hint=torch.Tensor, description="Trunk single repr [I, c_s]"),
+            InputParam("pair", type_hint=torch.Tensor, description="Trunk pair repr [I, I, c_z]"),
+            InputParam("s_inputs", type_hint=torch.Tensor, description="Input embeddings [I, c_s_inputs]"),
+            InputParam("callback", type_hint=Callable),
+            InputParam("callback_steps", default=1, type_hint=int),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> List[OutputParam]:
+        return [
+            OutputParam("xyz", type_hint=torch.Tensor, description="Denoised coords [D, L, 3]"),
+            OutputParam("trajectory", type_hint=List[torch.Tensor]),
+        ]
+
+    @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
+        f = block_state.f
+        single = block_state.single
+        pair = block_state.pair
+        s_inputs = block_state.s_inputs
+        callback = block_state.callback
+        callback_steps = block_state.callback_steps or 1
+
+        X_L = xyz.clone()
+        D = X_L.shape[0]
+        device = X_L.device
+        noise_schedule = noise_schedule.to(device)
+
+        trajectory = []
+
+        has_transformer = hasattr(components, "transformer") and components.transformer is not None
+        has_scheduler = hasattr(components, "scheduler") and components.scheduler is not None
+
+        for step_num in range(len(noise_schedule) - 1):
+            c_t_minus_1 = noise_schedule[step_num]
+            c_t = noise_schedule[step_num + 1]
+
+            # Noise injection
+            if has_scheduler:
+                X_noisy, t_hat = components.scheduler.add_noise(X_L, c_t_minus_1, c_t)
+            else:
+                X_noisy = X_L
+                t_hat = c_t_minus_1
+
+            # Model forward (diffusion module conditioned on trunk)
+            if has_transformer:
+                t_batch = (t_hat.to(device).expand(D) if isinstance(t_hat, torch.Tensor)
+                          else torch.full((D,), t_hat, device=device))
+
+                outs = components.transformer.diffusion_module(
+                    X_noisy_L=X_noisy,
+                    t=t_batch,
+                    f=f,
+                    S_inputs_I=s_inputs,
+                    S_trunk_I=single,
+                    Z_trunk_II=pair,
+                )
+                X_denoised = outs if isinstance(outs, torch.Tensor) else outs.get("X_L", outs)
+            else:
+                X_denoised = X_noisy
+
+            # Euler step
+            if has_scheduler:
+                X_L = components.scheduler.step(
+                    xyz_pred=X_denoised, xyz_noisy=X_noisy,
+                    c_t_minus_1=c_t_minus_1, c_t=c_t,
+                )
+            else:
+                delta = (X_noisy - X_denoised) / (t_hat + 1e-8)
+                d_t = c_t - t_hat
+                X_L = X_noisy + d_t * delta
+
+            trajectory.append(X_denoised.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.trajectory = trajectory
+
+        self.set_block_state(state, block_state)
+        return components, state

modular_blocks.py

@@ -0,0 +1,91 @@
+# Copyright 2025 Dhruv Nair. All rights reserved.
+# Licensed under the Apache License, Version 2.0
+
+from diffusers.utils import logging
+from diffusers.modular_pipelines import AutoPipelineBlocks, SequentialPipelineBlocks
+
+from .before_denoise import (
+    RF3InputStep,
+    RF3PrepareLatentsStep,
+    RF3RecyclingStep,
+    RF3SetTimestepsStep,
+)
+from .decoders import RF3DecodeStep
+from .denoise import RF3DenoiseStep
+
+
+logger = logging.get_logger(__name__)
+
+
+class RF3BeforeDenoiseStep(SequentialPipelineBlocks):
+    """Sequential block for pre-denoising: input → timesteps → recycling → latents."""
+
+    block_classes = [
+        RF3InputStep,
+        RF3SetTimestepsStep,
+        RF3RecyclingStep,
+        RF3PrepareLatentsStep,
+    ]
+    block_names = ["input", "set_timesteps", "recycling", "prepare_latents"]
+
+    @property
+    def description(self):
+        return (
+            "Before denoise step:\n"
+            " - `RF3InputStep` parses sequence and builds feature dict\n"
+            " - `RF3SetTimestepsStep` constructs EDM noise schedule\n"
+            " - `RF3RecyclingStep` runs trunk recycler (pairformer + MSA + templates)\n"
+            " - `RF3PrepareLatentsStep` samples initial noised coordinates\n"
+        )
+
+
+class RF3AutoBeforeDenoiseStep(AutoPipelineBlocks):
+    block_classes = [RF3BeforeDenoiseStep]
+    block_names = ["fold"]
+    block_trigger_inputs = [None]
+
+    @property
+    def description(self):
+        return "Before denoise step for RF3 structure prediction."
+
+
+class RF3AutoDenoiseStep(AutoPipelineBlocks):
+    block_classes = [RF3DenoiseStep]
+    block_names = ["denoise"]
+    block_trigger_inputs = [None]
+
+    @property
+    def description(self) -> str:
+        return "Denoise step for RF3 structure prediction."
+
+
+class RF3AutoDecodeStep(AutoPipelineBlocks):
+    block_classes = [RF3DecodeStep]
+    block_names = ["decode"]
+    block_trigger_inputs = [None]
+
+    @property
+    def description(self):
+        return "Decode step for RF3 — coordinates to tensor/PDB/CIF."
+
+
+class RF3AutoBlocks(SequentialPipelineBlocks):
+    """Full RF3 structure prediction pipeline."""
+
+    block_classes = [
+        RF3AutoBeforeDenoiseStep,
+        RF3AutoDenoiseStep,
+        RF3AutoDecodeStep,
+    ]
+    block_names = [
+        "before_denoise",
+        "denoise",
+        "decoder",
+    ]
+
+    @property
+    def description(self):
+        return (
+            "Modular pipeline for protein structure prediction using RF3.\n"
+            "Provide `sequence` to predict a protein's 3D structure."
+        )

modular_config.json

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

modular_model_index.json

@@ -0,0 +1,34 @@
+{
+    "_blocks_class_name": "RF3AutoBlocks",
+    "_class_name": "ModularPipeline",
+    "_diffusers_version": "0.37.0.dev0",
+    "transformer": [
+        null,
+        null,
+        {
+            "pretrained_model_name_or_path": "dn6/RosettaFold-3",
+            "subfolder": "transformer",
+            "type_hint": [
+                "diffusers",
+                "AutoModel"
+            ],
+            "revision": null,
+            "variant": null
+        }
+    ],
+    "scheduler": [
+        null,
+        null,
+        {
+            "pretrained_model_name_or_path": "dn6/RosettaFold-3",
+            "subfolder": "scheduler",
+            "type_hint": [
+                "diffusers",
+                "AutoModel"
+            ],
+            "revision": null,
+            "variant": null,
+            "default_creation_method": "from_config"
+        }
+    ]
+}

scheduler/__init__.py

@@ -0,0 +1,4 @@
+# Copyright 2025 Dhruv Nair. All rights reserved.
+# Licensed under the Apache License, Version 2.0
+
+from .model import RF3Scheduler

scheduler/config.json

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

scheduler/model.py

@@ -0,0 +1,90 @@
+# Copyright 2025 Dhruv Nair. All rights reserved.
+# Licensed under the Apache License, Version 2.0
+
+"""
+RF3 Scheduler.
+
+A diffusers-compatible wrapper around the foundry EDM noise schedule
+for RF3. Same schedule formula as RFD3 but with gamma_0=0.8 (vs 0.6).
+"""
+
+from typing import Optional
+
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+
+from rf3.diffusion_samplers.inference_sampler import SampleDiffusion
+
+
+class RF3Scheduler(ConfigMixin):
+    """
+    Diffusers-compatible scheduler wrapping the foundry RF3 EDM sampler.
+    """
+
+    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.8,
+        gamma_min: float = 1.0,
+        noise_scale: float = 1.003,
+        step_scale: float = 1.5,
+    ):
+        self._sampler = SampleDiffusion(
+            num_timesteps=num_timesteps,
+            min_t=0,
+            max_t=1,
+            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,
+            solver="af3",
+        )
+
+    @property
+    def sampler(self) -> SampleDiffusion:
+        return self._sampler
+
+    def get_noise_schedule(self, device: torch.device = None) -> torch.Tensor:
+        """Construct the EDM noise schedule."""
+        return self._sampler._construct_inference_noise_schedule(
+            device=device or torch.device("cpu")
+        )
+
+    def add_noise(
+        self,
+        xyz: torch.Tensor,
+        c_t_minus_1: torch.Tensor,
+        c_t: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Inject stochastic noise before the model call."""
+        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)
+        return xyz + epsilon, t_hat
+
+    def step(
+        self,
+        xyz_pred: torch.Tensor,
+        xyz_noisy: torch.Tensor,
+        c_t_minus_1: torch.Tensor,
+        c_t: torch.Tensor,
+    ) -> torch.Tensor:
+        """Perform one Euler denoising step."""
+        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 = (xyz_noisy - xyz_pred) / t_hat
+        d_t = c_t - t_hat
+        return xyz_noisy + self._sampler.step_scale * d_t * delta

transformer/__init__.py

@@ -0,0 +1,4 @@
+# Copyright 2025 Dhruv Nair. All rights reserved.
+# Licensed under the Apache License, Version 2.0
+
+from .model import RF3TransformerModel, RF3TransformerOutput

transformer/config.json

@@ -0,0 +1,25 @@
+{
+  "_class_name": "RF3TransformerModel",
+  "_diffusers_version": "0.37.0.dev0",
+  "auto_map": {
+    "AutoModel": "model.RF3TransformerModel"
+  },
+  "c_s": 384,
+  "c_z": 128,
+  "c_atom": 128,
+  "c_atompair": 16,
+  "c_s_inputs": 449,
+  "c_token": 768,
+  "sigma_data": 16.0,
+  "n_pairformer_blocks": 48,
+  "n_diffusion_blocks": 24,
+  "n_atom_encoder_blocks": 3,
+  "n_atom_decoder_blocks": 3,
+  "n_msa_blocks": 4,
+  "n_template_blocks": 2,
+  "n_head": 16,
+  "n_pairformer_head": 16,
+  "n_recycles": 10,
+  "distogram_bins": 65,
+  "p_drop": 0.25
+}

transformer/diffusion_pytorch_model.safetensors

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

transformer/model.py

@@ -0,0 +1,286 @@
+# Copyright 2025 Dhruv Nair. All rights reserved.
+# Licensed under the Apache License, Version 2.0
+
+"""
+RF3 (RosettaFold3) Transformer model.
+
+A diffusers-compatible wrapper around the foundry RF3 model components.
+Reuses FeatureInitializer, Recycler, DiffusionModule, and DistogramHead
+from ``rf3.model.*`` directly, adding only the ModelMixin/ConfigMixin
+interface needed for diffusers ModularPipeline integration.
+
+RF3 is structurally similar to RFD3 but adds a trunk recycler (48
+pairformer blocks + MSA + templates) for sequence-conditioned folding.
+"""
+
+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 rf3.model.RF3_structure import DiffusionModule, DistogramHead, Recycler
+from rf3.model.layers.pairformer_layers import FeatureInitializer
+
+
+@dataclass
+class RF3TransformerOutput:
+    """Output class for RF3 transformer."""
+
+    xyz: torch.Tensor                                       # [D, L, 3]
+    distogram: Optional[torch.Tensor] = None                # [I, I, bins]
+    single: Optional[torch.Tensor] = None                   # [I, c_s]
+    pair: Optional[torch.Tensor] = None                     # [I, I, c_z]
+    trajectory_noisy: Optional[list] = None                 # list of [D, L, 3]
+    trajectory_denoised: Optional[list] = None              # list of [D, L, 3]
+
+
+class RF3TransformerModel(ModelMixin, ConfigMixin):
+    """
+    Diffusers-compatible wrapper around the foundry RF3 model.
+
+    Wraps FeatureInitializer, Recycler, DiffusionModule, and DistogramHead
+    to provide a diffusers ModelMixin/ConfigMixin interface.
+
+    State dict keys match the foundry checkpoint format via the
+    ``feature_initializer.*``, ``recycler.*``, ``diffusion_module.*``,
+    and ``distogram_head.*`` 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_s_inputs: int = 449,
+        c_token: int = 768,
+        sigma_data: float = 16.0,
+        n_pairformer_blocks: int = 48,
+        n_diffusion_blocks: int = 24,
+        n_atom_encoder_blocks: int = 3,
+        n_atom_decoder_blocks: int = 3,
+        n_msa_blocks: int = 4,
+        n_template_blocks: int = 2,
+        n_head: int = 16,
+        n_pairformer_head: int = 16,
+        n_recycles: int = 10,
+        distogram_bins: int = 65,
+        p_drop: float = 0.25,
+    ):
+        super().__init__()
+
+        # ── FeatureInitializer ──────────────────────────────────────────
+        self.feature_initializer = FeatureInitializer(
+            c_s=c_s,
+            c_z=c_z,
+            c_atom=c_atom,
+            c_atompair=c_atompair,
+            c_s_inputs=c_s_inputs,
+            input_feature_embedder={
+                "features": ["restype", "profile", "deletion_mean"],
+                "atom_attention_encoder": {
+                    "c_token": c_s,
+                    "c_atom_1d_features": 389,
+                    "c_tokenpair": c_z,
+                    "use_inv_dist_squared": True,
+                    "atom_1d_features": [
+                        "ref_pos", "ref_charge", "ref_mask",
+                        "ref_element", "ref_atom_name_chars",
+                    ],
+                    "atom_transformer": {
+                        "n_queries": 32,
+                        "n_keys": 128,
+                        "diffusion_transformer": {
+                            "n_block": 3,
+                            "diffusion_transformer_block": {
+                                "n_head": 4,
+                                "no_residual_connection_between_attention_and_transition": True,
+                                "kq_norm": True,
+                            },
+                        },
+                    },
+                },
+            },
+            relative_position_encoding={"r_max": 32, "s_max": 2},
+        )
+
+        # ── Recycler (trunk) ───────────────────────────────────────────
+        self.recycler = Recycler(
+            c_s=c_s,
+            c_z=c_z,
+            n_pairformer_blocks=n_pairformer_blocks,
+            pairformer_block={
+                "p_drop": p_drop,
+                "triangle_multiplication": {"d_hidden": 128},
+                "triangle_attention": {"n_head": 4, "d_hidden": 32},
+                "attention_pair_bias": {"n_head": n_head},
+            },
+            template_embedder={
+                "n_block": n_template_blocks,
+                "raw_template_dim": 108,
+                "c": 64,
+                "p_drop": p_drop,
+            },
+            msa_module={
+                "n_block": n_msa_blocks,
+                "c_m": 64,
+                "p_drop_msa": 0.15,
+                "p_drop_pair": p_drop,
+                "msa_subsample_embedder": {
+                    "num_sequences": 1024,
+                    "dim_raw_msa": 34,
+                    "c_s_inputs": c_s_inputs,
+                    "c_msa_embed": 64,
+                },
+                "outer_product": {
+                    "c_msa_embed": 64,
+                    "c_outer_product": 32,
+                    "c_out": c_z,
+                },
+                "msa_pair_weighted_averaging": {
+                    "n_heads": 8,
+                    "c_weighted_average": 32,
+                    "c_msa_embed": 64,
+                    "c_z": c_z,
+                    "separate_gate_for_every_channel": True,
+                },
+                "msa_transition": {"n": 4, "c": 64},
+                "triangle_multiplication_outgoing": {
+                    "d_pair": c_z, "d_hidden": 128, "bias": True,
+                },
+                "triangle_multiplication_incoming": {
+                    "d_pair": c_z, "d_hidden": 128, "bias": True,
+                },
+                "triangle_attention_starting": {
+                    "d_pair": c_z, "n_head": 4, "d_hidden": 32, "p_drop": 0.0,
+                },
+                "triangle_attention_ending": {
+                    "d_pair": c_z, "n_head": 4, "d_hidden": 32, "p_drop": 0.0,
+                },
+                "pair_transition": {"n": 4, "c": c_z},
+            },
+        )
+
+        # ── DiffusionModule ────────────────────────────────────────────
+        self.diffusion_module = DiffusionModule(
+            sigma_data=sigma_data,
+            c_atom=c_atom,
+            c_atompair=c_atompair,
+            c_token=c_token,
+            c_s=c_s,
+            c_z=c_z,
+            diffusion_conditioning={
+                "c_s_inputs": c_s_inputs,
+                "c_t_embed": 256,
+                "relative_position_encoding": {"r_max": 32, "s_max": 2},
+            },
+            atom_attention_encoder={
+                "c_tokenpair": c_z,
+                "c_atom_1d_features": 389,
+                "use_inv_dist_squared": True,
+                "atom_1d_features": [
+                    "ref_pos", "ref_charge", "ref_mask",
+                    "ref_element", "ref_atom_name_chars",
+                ],
+                "atom_transformer": {
+                    "n_queries": 32,
+                    "n_keys": 128,
+                    "diffusion_transformer": {
+                        "n_block": n_atom_encoder_blocks,
+                        "diffusion_transformer_block": {
+                            "n_head": 4,
+                            "no_residual_connection_between_attention_and_transition": True,
+                            "kq_norm": True,
+                        },
+                    },
+                },
+                "broadcast_trunk_feats_on_1dim_old": False,
+                "use_chiral_features": True,
+                "no_grad_on_chiral_center": False,
+            },
+            diffusion_transformer={
+                "n_block": n_diffusion_blocks,
+                "diffusion_transformer_block": {
+                    "n_head": n_head,
+                    "no_residual_connection_between_attention_and_transition": True,
+                    "kq_norm": True,
+                },
+            },
+            atom_attention_decoder={
+                "atom_transformer": {
+                    "n_queries": 32,
+                    "n_keys": 128,
+                    "diffusion_transformer": {
+                        "n_block": n_atom_decoder_blocks,
+                        "diffusion_transformer_block": {
+                            "n_head": 4,
+                            "no_residual_connection_between_attention_and_transition": True,
+                            "kq_norm": True,
+                        },
+                    },
+                },
+            },
+        )
+
+        # ── DistogramHead ──────────────────────────────────────────────
+        self.distogram_head = DistogramHead(c_z=c_z, bins=distogram_bins)
+
+        self._n_recycles = n_recycles
+
+    def forward(
+        self,
+        f: dict,
+        n_recycles: Optional[int] = None,
+        diffusion_batch_size: int = 1,
+        coord_atom_lvl_to_be_noised: Optional[torch.Tensor] = None,
+    ) -> RF3TransformerOutput:
+        """
+        Forward pass: recycling trunk → diffusion sampling.
+
+        Args:
+            f: Feature dictionary (sequence, MSA, templates, atom features).
+            n_recycles: Number of recycling iterations (default: config value).
+            diffusion_batch_size: Number of diffusion samples.
+            coord_atom_lvl_to_be_noised: Initial coordinates for partial diffusion.
+
+        Returns:
+            RF3TransformerOutput with predicted coordinates and distogram.
+        """
+        n_recycles = n_recycles or self._n_recycles
+
+        # Pre-recycle: initialize features
+        initialized = self.feature_initializer(f)
+        S_inputs_I = initialized["S_inputs_I"]
+        S_I = initialized.get("S_init_I", initialized.get("S_I"))
+        Z_II = initialized.get("Z_init_II", initialized.get("Z_II"))
+
+        # Recycling trunk
+        for i in range(n_recycles):
+            ctx = torch.no_grad() if i < n_recycles - 1 else torch.enable_grad()
+            with ctx:
+                recycled = self.recycler(
+                    S_I=S_I,
+                    Z_II=Z_II,
+                    S_inputs_I=S_inputs_I,
+                    f=f,
+                )
+                S_I = recycled["S_I"]
+                Z_II = recycled["Z_II"]
+
+        # Distogram prediction
+        distogram = self.distogram_head(Z_II)
+
+        return RF3TransformerOutput(
+            xyz=torch.zeros(1),  # placeholder — filled by sampler in denoise step
+            distogram=distogram,
+            single=S_I,
+            pair=Z_II,
+        )