Expose n_recycle as inference sampler parameter (#259)

* Expose n_recycle as inference sampler parameter and improve docs

The n_recycle parameter was previously hardcoded in the diffusion module
config and not overridable at inference time. This exposes it through the
inference sampler so users can control recycling iterations via CLI
(e.g. inference_sampler.n_recycle=3). Also adds num_timesteps and
n_recycle to the "Other CLI Options" docs section, and makes the
InputSpecification reference more prominent in the README.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

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

@@ -270,7 +270,7 @@ class RFD3DiffusionModule(nn.Module):
         **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 @@ class SampleDiffusionWithMotif(SampleDiffusionConfig):
                     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 @@ class SampleDiffusionWithMotif(SampleDiffusionConfig):
                     X_noisy_L=X_noisy_L,
                     t=t_hat.tile(D),
                     f=f,
+                    n_recycle=self.n_recycle,
                     **initializer_outputs,
                 )
 
@@ -278,6 +283,7 @@ class SampleDiffusionWithMotif(SampleDiffusionConfig):
                     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 @@ class SampleDiffusionWithSymmetry(SampleDiffusionWithMotif):
                     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 @@ class SampleDiffusionWithSymmetry(SampleDiffusionWithMotif):
                     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