From f49348fbf8619122744fee97704dfb6a383128f0 Mon Sep 17 00:00:00 2001 From: jbutch Date: Sat, 4 Apr 2026 15:09:54 -0700 Subject: [PATCH 1/2] 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) --- models/rfd3/README.md | 4 ++-- models/rfd3/configs/inference_engine/rfdiffusion3.yaml | 3 +++ models/rfd3/docs/input.md | 2 ++ models/rfd3/src/rfd3/model/RFD3_diffusion_module.py | 2 +- models/rfd3/src/rfd3/model/inference_sampler.py | 8 ++++++++ 5 files changed, 16 insertions(+), 3 deletions(-) diff --git a/models/rfd3/README.md b/models/rfd3/README.md index ecb4f34c..58514781 100644 --- a/models/rfd3/README.md +++ b/models/rfd3/README.md @@ -10,8 +10,8 @@ both are described in more detail below. All-atom design with RFD3

-> [!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. diff --git a/models/rfd3/configs/inference_engine/rfdiffusion3.yaml b/models/rfd3/configs/inference_engine/rfdiffusion3.yaml index c5a6a72b..cd468d5f 100644 --- a/models/rfd3/configs/inference_engine/rfdiffusion3.yaml +++ b/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 diff --git a/models/rfd3/docs/input.md b/models/rfd3/docs/input.md index 9253abc5..fd173f29 100644 --- a/models/rfd3/docs/input.md +++ b/models/rfd3/docs/input.md @@ -79,6 +79,8 @@ rfd3 design out_dir= 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`). diff --git a/models/rfd3/src/rfd3/model/RFD3_diffusion_module.py b/models/rfd3/src/rfd3/model/RFD3_diffusion_module.py index b6512d02..b77371bd 100644 --- a/models/rfd3/src/rfd3/model/RFD3_diffusion_module.py +++ b/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 diff --git a/models/rfd3/src/rfd3/model/inference_sampler.py b/models/rfd3/src/rfd3/model/inference_sampler.py index ccc638dd..ef3f1440 100644 --- a/models/rfd3/src/rfd3/model/inference_sampler.py +++ b/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 From 2a6cd12fa4012ce8861d944e980e03f6b3bd09df Mon Sep 17 00:00:00 2001 From: Jasper Butcher <66851659+Ubiquinone-dot@users.noreply.github.com> Date: Sat, 4 Apr 2026 15:11:13 -0700 Subject: [PATCH 2/2] Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- models/rfd3/docs/input.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/rfd3/docs/input.md b/models/rfd3/docs/input.md index fd173f29..101e34a0 100644 --- a/models/rfd3/docs/input.md +++ b/models/rfd3/docs/input.md @@ -82,8 +82,8 @@ rfd3 design out_dir= inputs= - `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: