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..101e34a0 100644 --- a/models/rfd3/docs/input.md +++ b/models/rfd3/docs/input.md @@ -79,9 +79,11 @@ 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`). + - `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: 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