Expose n_recycle as inference sampler parameter

models/rfd3/README.md

@@ -10,8 +10,8 @@ both are described in more detail below.
   <img src="docs/.assets/overview.png" alt="All-atom design with RFD3">
 </p>
 
-> [!NOTE]
-> Looking for information about configuration options? See a through list of possible inputs in our [external documentation](https://rosettacommons.github.io/foundry/models/rfd3/input.html).
+> [!IMPORTANT]
+> **Looking for the InputSpecification?** The complete reference for all input fields, CLI arguments, and inference sampler options is in [`docs/input.md`](./docs/input.md) ([external docs](https://rosettacommons.github.io/foundry/models/rfd3/input.html)). This covers everything from contig strings and conditioning options to diffusion sampler parameters like `num_timesteps` and `n_recycle`.
 
 ## Getting Started
 1. Install RFdiffusion3. 

models/rfd3/configs/inference_engine/rfdiffusion3.yaml

@@ -37,6 +37,9 @@ inference_sampler:
   inference_noise_scaling_factor: 1.0
   allow_realignment: False
 
+  # Recycling
+  n_recycle: null  # Override model default n_recycle for inference (default from checkpoint: 2)
+
   # Diffusion args:
   num_timesteps: 200
   step_scale: 1.5  # 1.5 - 1.0   |   Higher values lead to less diverse, more designable, structures

models/rfd3/docs/input.md

@@ -79,9 +79,11 @@ rfd3 design out_dir=<path/to/outdir> inputs=<path/to/inputs>
 ### Other CLI Options:
 - `json_keys_subset` — Allows the user to extract only a subset of the JSON keys provided in the `inputs` file (default: null).
 - `inference_sampler` —
+  - `num_timesteps` — Number of diffusion denoising timesteps (default: 200). Controls how many steps the reverse diffusion process takes. More steps can improve quality at the cost of runtime.
+  - `n_recycle` — Number of recycling iterations per diffusion step (default: null, uses the model checkpoint default of 2). Recycling allows the network to refine its predictions iteratively within each denoising step.
   - `kind` — Change this value to `symmetry` (default: default) to turn on symmetry mode for the inference sampler.
-  - `cfg_features` — The values specified (options are `active donor, active_acceptor, or ref_atomwise_rasa`) are set to 0 for classifier-free guidance. Classifier-free guidance is how the diffusion model can steer the calculation towards a condition without training a separate classifier.
-  - `use_classifier_free_guidance` — If set to `True`, RFD3 can use classifier-free guidance to guide the system towards a condition without training a separate classifier (default: `Fasle`).
+  - `cfg_features` — The values specified (options are `active_donor, active_acceptor, or ref_atomwise_rasa`) are set to 0 for classifier-free guidance. Classifier-free guidance is how the diffusion model can steer the calculation towards a condition without training a separate classifier.
+  - `use_classifier_free_guidance` — If set to `True`, RFD3 can use classifier-free guidance to guide the system towards a condition without training a separate classifier (default: `False`).
   - `cfg_t_max` — The maximum time to apply classifier-free guidance to the inference run (default: null).
   - `cfg_scale` — Controls the influence of the classifier-free guidance adjustment (default: 1.5).
   - `center_option`: Specifies how to center the coordinates during the inference run to ensure that structures are alined around a specific point. Options include: 

models/rfd3/src/rfd3/model/RFD3_diffusion_module.py

@@ -272,7 +272,7 @@ def forward_with_recycle(
         **kwargs,
     ):
         if not self.training:
-            n_recycle = self.n_recycle
+            n_recycle = n_recycle if n_recycle is not None else self.n_recycle
         else:
             assert n_recycle is not None
 

models/rfd3/src/rfd3/model/inference_sampler.py

@@ -51,6 +51,9 @@ class SampleDiffusionConfig:
     cfg_scale: float = 2.0
     cfg_t_max: float | None = None
 
+    # Recycling
+    n_recycle: int | None = None  # Override model default n_recycle for inference
+
 
 class SampleDiffusionWithMotif(SampleDiffusionConfig):
     """Diffusion sampler that supports optional motif alignment."""
@@ -245,6 +248,7 @@ def sample_diffusion_like_af3(
                     P_LL=None,  # Not used in chunked mode
                     chunked_pairwise_embedder=chunked_embedder,
                     initializer_outputs=other_outputs,
+                    n_recycle=self.n_recycle,
                     **other_outputs,
                 )
                 toc = time.time()
@@ -257,6 +261,7 @@ def sample_diffusion_like_af3(
                     X_noisy_L=X_noisy_L,
                     t=t_hat.tile(D),
                     f=f,
+                    n_recycle=self.n_recycle,
                     **initializer_outputs,
                 )
 
@@ -278,6 +283,7 @@ def sample_diffusion_like_af3(
                     X_noisy_L=X_noisy_L_stripped,  # modify X
                     t=t_hat.tile(D),
                     f=f_ref,  # modified f
+                    n_recycle=self.n_recycle,
                     **ref_initializer_outputs,
                 )
 
@@ -474,6 +480,7 @@ def sample_diffusion_like_af3(
                     P_LL=None,  # Not used in chunked mode
                     chunked_pairwise_embedder=chunked_embedder,
                     initializer_outputs=other_outputs,
+                    n_recycle=self.n_recycle,
                     **other_outputs,
                 )
                 toc = time.time()
@@ -486,6 +493,7 @@ def sample_diffusion_like_af3(
                     X_noisy_L=X_noisy_L,
                     t=t_hat.tile(D),
                     f=f,
+                    n_recycle=self.n_recycle,
                     **initializer_outputs,
                 )
             # apply symmetry to X_denoised_L