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
This commit is contained in:
Pierre Chapuis 2024-05-29 15:50:49 +02:00
parent ddcd84c740
commit 1886e6456c

View file

@ -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