From 1886e6456c24e26d4e22659bbc1d5e6f91c4154c Mon Sep 17 00:00:00 2001 From: Pierre Chapuis Date: Wed, 29 May 2024 15:50:49 +0200 Subject: [PATCH] add a docstring for set_inference_steps This explains the relation between first_step and strength, as shown by @holwech here: https://github.com/finegrain-ai/refiners/discussions/374 --- src/refiners/foundationals/latent_diffusion/model.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/refiners/foundationals/latent_diffusion/model.py b/src/refiners/foundationals/latent_diffusion/model.py index 4bb51e8..11d8450 100644 --- a/src/refiners/foundationals/latent_diffusion/model.py +++ b/src/refiners/foundationals/latent_diffusion/model.py @@ -33,6 +33,15 @@ class LatentDiffusionModel(fl.Module, ABC): self.classifier_free_guidance = classifier_free_guidance def set_inference_steps(self, num_steps: int, first_step: int = 0) -> None: + """Set the steps of the diffusion process. + + Args: + num_steps: The number of inference steps. + first_step: The first inference step, used for image-to-image diffusion. + You may be used to setting a float in `[0, 1]` called `strength` instead, + which is an abstraction for this. The first step is + `round((1 - strength) * (num_steps - 1))`. + """ self.solver = self.solver.rebuild(num_inference_steps=num_steps, first_inference_step=first_step) @staticmethod