[RFD3] Feed MPNN real backbone atoms and the schema it validates

README.md

@@ -20,7 +20,7 @@ import torch
 from diffusers import ModularPipeline
 
 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)
+pipe.load_components(device_map="cuda", torch_dtype=torch.float32, trust_remote_code=True)
 
 state = pipe(contigs="100")
 print(state.output.xyz.shape)  # [1, 100, 3]
@@ -34,7 +34,6 @@ The active workflow is selected automatically based on which inputs you provide:
 |----------|---------------|-----------|
 | `structure_only` | `contigs` | RFdiffusion3 |
 | `structure_and_sequence` | `contigs`, `temperature` | RFdiffusion3 → MPNN |
-| `motif_structure_and_sequence` | `contigs`, `input_xyz`, `temperature` | Motif-conditioned RFdiffusion3 → MPNN |
 
 ### Structure Only
 
@@ -67,18 +66,12 @@ Three MPNN variants are available:
 
 ### Motif-Conditioned Design
 
-Passing `input_xyz` enables motif conditioning — fix specific residues in place while designing the rest:
+Not supported yet. Motif contigs such as `"A10-25/50"` and the `input_xyz` argument raise a
+`ValueError`.
 
-```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,
-)
-```
+Conditioning on a motif requires per-atom element, atom-name and occupancy annotations for the
+fixed residues, which the feature pipeline derives from a reference structure. A bare coordinate
+tensor cannot supply them, so this needs a structure input rather than `input_xyz`.
 
 ### Full Design Pipeline
 
@@ -94,7 +87,7 @@ 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)
+design_pipe.load_components(device_map="cuda", torch_dtype=torch.float32, trust_remote_code=True)
 
 mpnn = AutoModel.from_pretrained("dn6/RFDiffusion-3", subfolder="mpnn", trust_remote_code=True)
 design_pipe.update_components(mpnn=mpnn)
@@ -104,7 +97,7 @@ designed_sequence = state.mpnn_output.designed_sequence
 
 # 2. Validate the fold with RF3
 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)
+fold_pipe.load_components(device_map="cuda", torch_dtype=torch.float32, trust_remote_code=True)
 
 state = fold_pipe(sequence=designed_sequence, output_type="cif.gz", output_path="prediction")
 ```

before_denoise.py

@@ -23,6 +23,65 @@ from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, In
 
 logger = logging.get_logger(__name__)
 
+# Feature widths the checkpoint's token initializer was trained with, from
+# rfd3/configs/model/components/rfd3_net.yaml. They set the input width of the embedding layers,
+# so a mismatch raises on layer shape rather than silently degrading the design.
+_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,
+}
+
+# Every token is padded to this many atom slots, so the coordinate tensors the model sees are
+# atom level with L = _N_ATOMS_PER_TOKEN * n_residues.
+_N_ATOMS_PER_TOKEN = 14
+
+
+def build_design_features(length: int, diffusion_batch_size: int, sigma_data: float) -> dict:
+    """
+    Build the foundry feature dict for an unconditional design of `length` residues.
+
+    The token initializer embeds real chemical and positional features; there is no meaningful
+    zero substitute for them, and no API in foundry that turns a length into features in one
+    call. This mirrors what `rfd3.engine.RFD3InferenceEngine` does, minus the hydra app and the
+    checkpoint, by driving the same specification and transform pipeline directly.
+
+    Returns:
+        The transformed example, carrying `feats` and `coord_atom_lvl_to_be_noised`.
+    """
+    from rfd3.inference.input_parsing import DesignInputSpecification
+    from rfd3.transforms.pipelines import build_atom14_base_pipeline
+
+    # `length` and `contig` are mutually exclusive in the specification: passing `contig` without
+    # a structure input fails validation, so unconditional designs go through `length`.
+    spec = DesignInputSpecification(length=str(length))
+    data = spec.to_pipeline_input(example_id=f"rfd3_{length}")
+
+    pipeline = build_atom14_base_pipeline(
+        is_inference=True,
+        diffusion_batch_size=diffusion_batch_size,
+        sigma_data=sigma_data,
+        central_atom="CB",
+        n_atoms_per_token=_N_ATOMS_PER_TOKEN,
+        generate_conformers=True,
+        provide_reference_conformer_when_unmasked=True,
+        ground_truth_conformer_policy="IGNORE",
+        use_element_for_atom_names_of_atomized_tokens=True,
+        token_1d_features=_TOKEN_1D_FEATURES,
+        atom_1d_features=_ATOM_1D_FEATURES,
+    )
+    return pipeline(data)
+
 
 def parse_contig_string(contig_str: str) -> Tuple[int, List[Tuple[int, int]]]:
     """
@@ -100,34 +159,45 @@ class RFDiffusionInputStep(ModularPipelineBlocks):
             ),
         ]
 
+    @property
+    def expected_components(self) -> List[ComponentSpec]:
+        return [
+            ComponentSpec("scheduler", description="RFDiffusion3 EDM scheduler"),
+        ]
+
     @property
     def intermediate_outputs(self) -> List[OutputParam]:
         return [
+            OutputParam(
+                "f",
+                type_hint=dict,
+                description="Foundry feature dict consumed by the token initializer",
+            ),
+            OutputParam(
+                "coord_atom_lvl_to_be_noised",
+                type_hint=torch.Tensor,
+                description="Reference atom-level coordinates [D, L_atom, 3]",
+            ),
             OutputParam(
                 "motif_mask",
                 type_hint=torch.Tensor,
-                description="Boolean mask for motif (fixed) positions",
+                description="Atom-level boolean mask for motif (fixed) positions [L_atom]",
             ),
             OutputParam(
-                "motif_xyz",
+                "motif_token_mask",
                 type_hint=torch.Tensor,
-                description="Coordinates for motif residues",
+                description="Residue-level boolean mask for motif (fixed) positions [L]",
             ),
             OutputParam(
                 "L",
                 type_hint=int,
-                description="Total length of the protein being designed",
+                description="Total length of the protein being designed, in residues",
             ),
             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):
@@ -149,20 +219,33 @@ class RFDiffusionInputStep(ModularPipelineBlocks):
 
         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
-
+        # Motif conditioning needs a reference structure so the transform pipeline can build
+        # per-atom features for the fixed residues. A coordinate tensor alone cannot supply the
+        # element, atom-name and occupancy annotations those features are derived from.
+        if motif_ranges:
+            raise ValueError(
+                f"Motif-conditioned contigs are not supported yet, got `contigs={contig_str!r}`. "
+                "Pass a plain design length such as `contigs='100'`."
+            )
         if input_xyz is not None:
-            motif_xyz = input_xyz
-        else:
-            motif_xyz = None
+            raise ValueError(
+                "`input_xyz` is not supported yet. Pass a plain design length such as "
+                "`contigs='100'`."
+            )
+
+        batch_size = 1
+        example = build_design_features(
+            length=L,
+            diffusion_batch_size=batch_size,
+            sigma_data=components.scheduler.config.sigma_data,
+        )
 
-        block_state.motif_mask = motif_mask
-        block_state.motif_xyz = motif_xyz
+        block_state.f = example["feats"]
+        block_state.coord_atom_lvl_to_be_noised = example["coord_atom_lvl_to_be_noised"]
+        block_state.motif_mask = example["feats"]["is_motif_atom_with_fixed_coord"]
+        block_state.motif_token_mask = example["feats"]["is_motif_token_with_fully_fixed_coord"]
         block_state.L = L
-        block_state.batch_size = 1
-        block_state.dtype = torch.float32
+        block_state.batch_size = batch_size
 
         self.set_block_state(state, block_state)
         return components, state
@@ -216,11 +299,15 @@ class RFDiffusionSetTimestepsStep(ModularPipelineBlocks):
     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)
+        # A linear stand-in for the EDM schedule silently changes the sampler, so require
+        # the real one rather than degrading.
+        if components.scheduler is None:
+            raise ValueError(
+                "`scheduler` is not loaded. Call `load_components(trust_remote_code=True)` on the "
+                "pipeline before calling it."
+            )
+
+        noise_schedule = components.scheduler.get_noise_schedule()
 
         block_state.noise_schedule = noise_schedule
         block_state.num_inference_steps = len(noise_schedule)
@@ -261,52 +348,44 @@ class RFDiffusionPrepareLatentsStep(ModularPipelineBlocks):
             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("f", required=True, type_hint=dict),
+            InputParam("coord_atom_lvl_to_be_noised", required=True, type_hint=torch.Tensor),
             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]"),
+            OutputParam("xyz", type_hint=torch.Tensor, description="Initial noised coordinates [D, L_atom, 3]"),
+            OutputParam("initializer_outputs", type_hint=dict, description="Embedded conditioning, reused every step"),
         ]
 
     @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
-
+        device = components.transformer.device
+
+        # The feature dict is built on CPU and is read on every denoising step, so move it once.
+        f = {k: v.to(device) if torch.is_tensor(v) else v for k, v in block_state.f.items()}
+        coord = block_state.coord_atom_lvl_to_be_noised.to(device)
+        motif_mask = f["is_motif_atom_with_fixed_coord"]
+
+        # Matches rfd3.model.inference_sampler._get_initial_structure:
+        #   noise = c0 * randn(D, L, 3); noise[..., is_motif, :] = 0; X_L = noise + coord
+        c0 = noise_schedule[0].to(device)
+        L_atom = coord.shape[-2]
+        noise = c0 * torch.randn((D, L_atom, 3), device=device, generator=generator)
+        noise[..., motif_mask, :] = 0.0
+        xyz = noise + coord
+
+        block_state.f = f
         block_state.xyz = xyz
+        block_state.initializer_outputs = components.transformer.encode_conditioning(f)
 
         self.set_block_state(state, block_state)
         return components, state

denoise.py

@@ -80,15 +80,18 @@ class RFDiffusionDenoiseStep(ModularPipelineBlocks):
                 type_hint=int,
                 description="Frequency of callback invocation",
             ),
-            InputParam("xyz", required=True, type_hint=torch.Tensor, description="Initial noised coordinates [D, L, 3]"),
+            InputParam("xyz", required=True, type_hint=torch.Tensor, description="Initial noised coordinates [D, L_atom, 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"),
+            InputParam("f", required=True, type_hint=dict, description="Foundry feature dict"),
+            InputParam("initializer_outputs", required=True, type_hint=dict, description="Embedded conditioning"),
         ]
 
     @property
     def intermediate_outputs(self) -> List[OutputParam]:
         return [
             OutputParam("xyz", type_hint=torch.Tensor, description="Denoised coordinates [D, L, 3]"),
+            OutputParam("xyz_backbone", type_hint=torch.Tensor, description="Backbone N/CA/C/O coordinates [D, L, 4, 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"),
@@ -98,11 +101,27 @@ class RFDiffusionDenoiseStep(ModularPipelineBlocks):
 
     @torch.no_grad()
     def __call__(self, components: ModularPipeline, state: PipelineState) -> PipelineState:
+        # Both components are required. Falling back to a hand-rolled Euler step when the
+        # scheduler is missing silently swaps out the EDM sampler and yields coordinates that
+        # look finite but are not a protein backbone.
+        if components.transformer is None:
+            raise ValueError(
+                "`transformer` is not loaded. Call `load_components(trust_remote_code=True)` on the "
+                "pipeline before calling it."
+            )
+        if components.scheduler is None:
+            raise ValueError(
+                "`scheduler` is not loaded. Call `load_components(trust_remote_code=True)` on the "
+                "pipeline before calling it."
+            )
+
         block_state = self.get_block_state(state)
 
         xyz = block_state.xyz
         noise_schedule = block_state.noise_schedule
         motif_mask = block_state.motif_mask
+        f = block_state.f
+        initializer_outputs = block_state.initializer_outputs
 
         n_recycle = block_state.n_recycle
         callback = block_state.callback
@@ -123,9 +142,6 @@ class RFDiffusionDenoiseStep(ModularPipelineBlocks):
         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):
@@ -133,61 +149,65 @@ class RFDiffusionDenoiseStep(ModularPipelineBlocks):
             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
+            X_noisy_L, t_hat = components.scheduler.add_noise(
+                X_L, c_t_minus_1, c_t, motif_mask=motif_mask
+            )
 
             # 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
+            # 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,
+                f=f,
+                initializer_outputs=initializer_outputs,
+                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
 
             # 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_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,
+            )
 
             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
+        # The sampler runs on padded atom-level coordinates (14 slots per residue). Downstream
+        # blocks and the documented output are one point per residue, so collapse to CA here.
+        is_ca = f["is_ca"]
+        block_state.xyz = X_L[:, is_ca]
+
+        # MPNN wants N, CA, C and O per residue. They are present in the padded tensor, so select
+        # them instead of synthesising them from CA: foundry's `BACKBONE_ATOM_NAMES` is
+        # ["N", "CA", "C", "O"], which is also the CCD atom order for amino acids.
+        is_backbone = f["is_backbone"]
+        n_residues = int(is_ca.sum())
+        n_backbone = int(is_backbone.sum())
+        if n_backbone != 4 * n_residues:
+            raise ValueError(
+                f"Expected 4 backbone atoms per residue, got {n_backbone} for {n_residues} residues. "
+                "The padded atom layout does not match the N/CA/C/O ordering MPNN assumes."
+            )
+        block_state.xyz_backbone = X_L[:, is_backbone].reshape(X_L.shape[0], n_residues, 4, 3)
         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
+        block_state.trajectory = [step[:, is_ca] for step in X_denoised_L_traj]
 
         self.set_block_state(state, block_state)
         return components, state

modular_blocks.py

@@ -175,8 +175,12 @@ class MPNNSequenceDesignStep(ModularPipelineBlocks):
                 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.",
+                "xyz_backbone", required=True, type_hint=torch.Tensor,
+                description="Backbone N/CA/C/O coordinates [B, L, 4, 3]",
+            ),
+            InputParam(
+                "motif_token_mask", type_hint=torch.Tensor,
+                description="Residue-level mask for fixed/motif positions [L]. True = fixed sequence.",
             ),
             InputParam(
                 "sequence_indices", type_hint=torch.Tensor,
@@ -214,7 +218,8 @@ class MPNNSequenceDesignStep(ModularPipelineBlocks):
         block_state = self.get_block_state(state)
 
         xyz = block_state.xyz
-        motif_mask = block_state.motif_mask
+        xyz_backbone = block_state.xyz_backbone
+        motif_mask = block_state.motif_token_mask
         known_seq = block_state.sequence_indices
         temperature = block_state.temperature or 0.1
         output_type = block_state.output_type or "tensor"
@@ -227,7 +232,7 @@ class MPNNSequenceDesignStep(ModularPipelineBlocks):
 
         if has_mpnn:
             sequence_logits, sequence_indices = self._run_mpnn(
-                components.mpnn, xyz, motif_mask, known_seq, temperature,
+                components.mpnn, xyz_backbone, motif_mask, known_seq, temperature,
             )
         else:
             if known_seq is not None:
@@ -264,32 +269,30 @@ class MPNNSequenceDesignStep(ModularPipelineBlocks):
         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)
+    def _run_mpnn(self, mpnn, xyz_backbone, motif_mask, known_seq, temperature):
+        """Run the MPNNModel wrapper on the sampler's N/CA/C/O coordinates."""
+        B = xyz_backbone.shape[0]
+        out_device = xyz_backbone.device
 
         if motif_mask is not None:
             designed_mask = ~motif_mask.unsqueeze(0).expand(B, -1)
         else:
             designed_mask = None
 
+        # MPNN is attached with `update_components` and is often left wherever it loaded, which
+        # need not be the device the sampler ran on. It is a 1.7M parameter model, so following it
+        # is cheaper than making the caller place it. Only the call arguments move; everything used
+        # afterwards stays on the sampler's device.
+        mpnn_device = next(mpnn.parameters()).device
         output = mpnn(
-            X=X, S=known_seq, designed_residue_mask=designed_mask, temperature=temperature,
+            X=xyz_backbone.to(mpnn_device),
+            S=None if known_seq is None else known_seq.to(mpnn_device),
+            designed_residue_mask=None if designed_mask is None else designed_mask.to(mpnn_device),
+            temperature=temperature,
         )
 
-        logits = output.sequence_logits
-        indices = output.sequence_indices
+        logits = output.sequence_logits.to(out_device)
+        indices = output.sequence_indices.to(out_device)
 
         if motif_mask is not None and known_seq is not None:
             indices[:, motif_mask] = known_seq[:, motif_mask]

modular_model_index.json

@@ -30,8 +30,7 @@
                 "AutoModel"
             ],
             "revision": null,
-            "variant": null,
-            "default_creation_method": "from_config"
+            "variant": null
         }
     ]
 }

mpnn/model_mpnn.py

@@ -39,6 +39,27 @@ MODEL_CLASSES = {
 }
 
 
+# Scalar decoding settings, matching MPNN_PER_INPUT_INFERENCE_DEFAULTS in foundry's
+# mpnn/utils/inference.py. foundry's feature-aggregation transform normally fills these in;
+# calling the network directly means supplying them here.
+_INFERENCE_DECODE_SETTINGS = {
+    "structure_noise": 0.0,
+    "decode_type": "auto_regressive",
+    "causality_pattern": "auto_regressive",
+    "initialize_sequence_embedding_with_ground_truth": False,
+    "atomize_side_chains": False,
+    "features_to_return": None,
+    "repeat_sample_num": None,
+}
+
+_OPTIONAL_CONDITIONING_KEYS = (
+    "bias",
+    "pair_bias",
+    "symmetry_equivalence_group",
+    "symmetry_weight",
+)
+
+
 @dataclass
 class MPNNModelOutput:
     """Output from the MPNN model wrapper."""
@@ -130,7 +151,7 @@ class MPNNModel(ModelMixin, ConfigMixin):
             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).
+            temperature: Sampling temperature, scalar or per-residue [B, L] (default: 0.1).
 
         Returns:
             MPNNModelOutput with sequence logits and sampled indices.
@@ -152,7 +173,12 @@ class MPNNModel(ModelMixin, ConfigMixin):
         # 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 = {
+        # foundry reads every tensor from `network_input["input_features"]`, and wants a
+        # per-residue temperature rather than a scalar.
+        if not torch.is_tensor(temperature):
+            temperature = torch.full((B, L), float(temperature), device=device)
+
+        input_features = {
             "X": X,
             "X_m": X_m,
             "S": S,
@@ -160,11 +186,16 @@ class MPNNModel(ModelMixin, ConfigMixin):
             "chain_labels": chain_labels,
             "residue_mask": residue_mask,
             "designed_residue_mask": designed_residue_mask,
+            "mask_for_loss": residue_mask,
             "temperature": temperature,
+            **_INFERENCE_DECODE_SETTINGS,
             **kwargs,
         }
+        # The network checks these keys are present but accepts None for all of them.
+        for key in _OPTIONAL_CONDITIONING_KEYS:
+            input_features.setdefault(key, None)
 
-        output = self.model(network_input)
+        output = self.model({"input_features": input_features})
 
         logits = output["decoder_features"]["logits"]  # [B, L, n_vocab]
         S_sampled = output["decoder_features"].get(

mpnn_ligand/model_mpnn.py

@@ -39,6 +39,27 @@ MODEL_CLASSES = {
 }
 
 
+# Scalar decoding settings, matching MPNN_PER_INPUT_INFERENCE_DEFAULTS in foundry's
+# mpnn/utils/inference.py. foundry's feature-aggregation transform normally fills these in;
+# calling the network directly means supplying them here.
+_INFERENCE_DECODE_SETTINGS = {
+    "structure_noise": 0.0,
+    "decode_type": "auto_regressive",
+    "causality_pattern": "auto_regressive",
+    "initialize_sequence_embedding_with_ground_truth": False,
+    "atomize_side_chains": False,
+    "features_to_return": None,
+    "repeat_sample_num": None,
+}
+
+_OPTIONAL_CONDITIONING_KEYS = (
+    "bias",
+    "pair_bias",
+    "symmetry_equivalence_group",
+    "symmetry_weight",
+)
+
+
 @dataclass
 class MPNNModelOutput:
     """Output from the MPNN model wrapper."""
@@ -130,7 +151,7 @@ class MPNNModel(ModelMixin, ConfigMixin):
             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).
+            temperature: Sampling temperature, scalar or per-residue [B, L] (default: 0.1).
 
         Returns:
             MPNNModelOutput with sequence logits and sampled indices.
@@ -152,7 +173,12 @@ class MPNNModel(ModelMixin, ConfigMixin):
         # 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 = {
+        # foundry reads every tensor from `network_input["input_features"]`, and wants a
+        # per-residue temperature rather than a scalar.
+        if not torch.is_tensor(temperature):
+            temperature = torch.full((B, L), float(temperature), device=device)
+
+        input_features = {
             "X": X,
             "X_m": X_m,
             "S": S,
@@ -160,11 +186,16 @@ class MPNNModel(ModelMixin, ConfigMixin):
             "chain_labels": chain_labels,
             "residue_mask": residue_mask,
             "designed_residue_mask": designed_residue_mask,
+            "mask_for_loss": residue_mask,
             "temperature": temperature,
+            **_INFERENCE_DECODE_SETTINGS,
             **kwargs,
         }
+        # The network checks these keys are present but accepts None for all of them.
+        for key in _OPTIONAL_CONDITIONING_KEYS:
+            input_features.setdefault(key, None)
 
-        output = self.model(network_input)
+        output = self.model({"input_features": input_features})
 
         logits = output["decoder_features"]["logits"]  # [B, L, n_vocab]
         S_sampled = output["decoder_features"].get(

mpnn_soluble/model_mpnn.py

@@ -39,6 +39,27 @@ MODEL_CLASSES = {
 }
 
 
+# Scalar decoding settings, matching MPNN_PER_INPUT_INFERENCE_DEFAULTS in foundry's
+# mpnn/utils/inference.py. foundry's feature-aggregation transform normally fills these in;
+# calling the network directly means supplying them here.
+_INFERENCE_DECODE_SETTINGS = {
+    "structure_noise": 0.0,
+    "decode_type": "auto_regressive",
+    "causality_pattern": "auto_regressive",
+    "initialize_sequence_embedding_with_ground_truth": False,
+    "atomize_side_chains": False,
+    "features_to_return": None,
+    "repeat_sample_num": None,
+}
+
+_OPTIONAL_CONDITIONING_KEYS = (
+    "bias",
+    "pair_bias",
+    "symmetry_equivalence_group",
+    "symmetry_weight",
+)
+
+
 @dataclass
 class MPNNModelOutput:
     """Output from the MPNN model wrapper."""
@@ -130,7 +151,7 @@ class MPNNModel(ModelMixin, ConfigMixin):
             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).
+            temperature: Sampling temperature, scalar or per-residue [B, L] (default: 0.1).
 
         Returns:
             MPNNModelOutput with sequence logits and sampled indices.
@@ -152,7 +173,12 @@ class MPNNModel(ModelMixin, ConfigMixin):
         # 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 = {
+        # foundry reads every tensor from `network_input["input_features"]`, and wants a
+        # per-residue temperature rather than a scalar.
+        if not torch.is_tensor(temperature):
+            temperature = torch.full((B, L), float(temperature), device=device)
+
+        input_features = {
             "X": X,
             "X_m": X_m,
             "S": S,
@@ -160,11 +186,16 @@ class MPNNModel(ModelMixin, ConfigMixin):
             "chain_labels": chain_labels,
             "residue_mask": residue_mask,
             "designed_residue_mask": designed_residue_mask,
+            "mask_for_loss": residue_mask,
             "temperature": temperature,
+            **_INFERENCE_DECODE_SETTINGS,
             **kwargs,
         }
+        # The network checks these keys are present but accepts None for all of them.
+        for key in _OPTIONAL_CONDITIONING_KEYS:
+            input_features.setdefault(key, None)
 
-        output = self.model(network_input)
+        output = self.model({"input_features": input_features})
 
         logits = output["decoder_features"]["logits"]  # [B, L, n_vocab]
         S_sampled = output["decoder_features"].get(

scheduler/model.py

@@ -24,12 +24,19 @@ from typing import Optional
 import torch
 
 from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
 
 # Reuse the original noise schedule and sampling config directly
 from rfd3.model.inference_sampler import SampleDiffusionWithMotif
 
+# `ComponentSpec.load` only strips `dtype` for a `type_hint` that is not a torch module, and
+# modular_model_index.json records this scheduler as `AutoModel`, so the loader forwards its
+# weight-placement kwargs here. None of them are constructor arguments, and letting
+# `register_to_config` capture `dtype` would write a torch.dtype into config.json.
+_LOADER_ONLY_KWARGS = ("dtype", "torch_dtype", "device_map", "variant", "trust_remote_code", "low_cpu_mem_usage")
 
-class RFDiffusionScheduler(ConfigMixin):
+
+class RFDiffusionScheduler(SchedulerMixin, ConfigMixin):
     """
     Diffusers-compatible scheduler wrapping the foundry EDM sampler.
 
@@ -65,6 +72,13 @@ class RFDiffusionScheduler(ConfigMixin):
             step_scale=step_scale,
         )
 
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path=None, subfolder=None, **kwargs):
+        for key in _LOADER_ONLY_KWARGS:
+            kwargs.pop(key, None)
+
+        return super().from_pretrained(pretrained_model_name_or_path, subfolder=subfolder, **kwargs)
+
     @property
     def sampler(self) -> SampleDiffusionWithMotif:
         return self._sampler

transformer/model_rfdiffusion.py

@@ -207,12 +207,25 @@ class RFDiffusionTransformerModel(ModelMixin, ConfigMixin):
     def sigma_data(self) -> float:
         return self.diffusion_module.sigma_data
 
+    def encode_conditioning(self, f: dict) -> dict:
+        """
+        Embed the feature dict once, before the denoising loop.
+
+        The token initializer does not depend on the noise level, so foundry runs it once per
+        design and reuses the result for every step. Doing it inside `forward` would repeat the
+        pairformer stack on all 200 steps.
+
+        Returns:
+            The `Q_L_init` / `C_L` / `P_LL` / `S_I` / `Z_II` kwargs for `forward`.
+        """
+        return self.token_initializer(f)
+
     def forward(
         self,
         xyz_noisy: torch.Tensor,
         t: torch.Tensor,
-        f: Optional[dict] = None,
-        motif_mask: Optional[torch.Tensor] = None,
+        f: dict,
+        initializer_outputs: dict,
         n_recycle: Optional[int] = None,
         **kwargs,
     ) -> RFDiffusionTransformerOutput:
@@ -220,77 +233,37 @@ class RFDiffusionTransformerModel(ModelMixin, ConfigMixin):
         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)
+            xyz_noisy: Noisy atom coordinates [D, L, 3], atom level, L = 14 * n_residues
+            t: Noise level / timestep [D]
+            f: Foundry feature dictionary, from `build_design_features`
+            initializer_outputs: Output of `encode_conditioning`
             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,
+        # foundry builds the per-atom noise level with a hardcoded `.float()`
+        # (rfd3/model/RFD3_diffusion_module.py), so the scaled coordinates it feeds to the first
+        # linear layer are float32 no matter what dtype the caller supplies. Half-precision weights
+        # cannot consume them. Fail here with the remedy rather than deep inside foundry.
+        weight_dtype = next(self.parameters()).dtype
+        if weight_dtype != torch.float32:
+            raise ValueError(
+                f"{self.__class__.__name__} only runs in float32, but its weights are {weight_dtype}. "
+                "Reload the pipeline with `torch_dtype=torch.float32`."
             )
-            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,
+            **initializer_outputs,
         )
 
         return RFDiffusionTransformerOutput(
             xyz=outs["X_L"],
-            single=torch.zeros(1),
+            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"),