[RFD3] Build real conditioning features instead of zero tensors

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,40 @@ 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(
-                "motif_mask",
+                "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="Boolean mask for motif (fixed) positions",
+                description="Reference atom-level coordinates [D, L_atom, 3]",
             ),
             OutputParam(
-                "motif_xyz",
+                "motif_mask",
                 type_hint=torch.Tensor,
-                description="Coordinates for motif residues",
+                description="Atom-level boolean mask for motif (fixed) positions [L_atom]",
             ),
             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 +214,32 @@ 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.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 +293,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 +342,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,9 +80,11 @@ 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
@@ -98,11 +100,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 +141,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 +148,52 @@ 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]
         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_model_index.json

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

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"),