From 9960c688cf47c60f57a3048eeb9c7f119e295dee Mon Sep 17 00:00:00 2001 From: David El Malih Date: Thu, 8 Jan 2026 10:25:37 +0100 Subject: [PATCH] docs: improve docstring scheduling_consistency_decoder.py --- .../scheduling_consistency_decoder.py | 54 +++++++++++++++---- 1 file changed, 44 insertions(+), 10 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_consistency_decoder.py b/src/diffusers/schedulers/scheduling_consistency_decoder.py index 767fa9157f59..2b0e70f5f828 100644 --- a/src/diffusers/schedulers/scheduling_consistency_decoder.py +++ b/src/diffusers/schedulers/scheduling_consistency_decoder.py @@ -71,6 +71,22 @@ class ConsistencyDecoderSchedulerOutput(BaseOutput): class ConsistencyDecoderScheduler(SchedulerMixin, ConfigMixin): + """ + A scheduler for the consistency decoder used in Stable Diffusion pipelines. + + This scheduler implements a two-step denoising process using consistency models for decoding latent representations + into images. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, *optional*, defaults to `1024`): + The number of diffusion steps to train the model. + sigma_data (`float`, *optional*, defaults to `0.5`): + The standard deviation of the data distribution. Used for computing the skip and output scaling factors. + """ + order = 1 @register_to_config @@ -78,7 +94,7 @@ def __init__( self, num_train_timesteps: int = 1024, sigma_data: float = 0.5, - ): + ) -> None: betas = betas_for_alpha_bar(num_train_timesteps) alphas = 1.0 - betas @@ -98,8 +114,18 @@ def __init__( def set_timesteps( self, num_inference_steps: Optional[int] = None, - device: Union[str, torch.device] = None, - ): + device: Optional[Union[str, torch.device]] = None, + ) -> None: + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`, *optional*): + The number of diffusion steps used when generating samples with a pre-trained model. Currently, only + `2` inference steps are supported. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ if num_inference_steps != 2: raise ValueError("Currently more than 2 inference steps are not supported.") @@ -111,7 +137,15 @@ def set_timesteps( self.c_in = self.c_in.to(device) @property - def init_noise_sigma(self): + def init_noise_sigma(self) -> torch.Tensor: + """ + Return the standard deviation of the initial noise distribution. + + Returns: + `torch.Tensor`: + The initial noise sigma value from the precomputed `sqrt_one_minus_alphas_cumprod` at the first + timestep. + """ return self.sqrt_one_minus_alphas_cumprod[self.timesteps[0]] def scale_model_input(self, sample: torch.Tensor, timestep: Optional[int] = None) -> torch.Tensor: @@ -146,20 +180,20 @@ def step( Args: model_output (`torch.Tensor`): The direct output from the learned diffusion model. - timestep (`float`): + timestep (`float` or `torch.Tensor`): The current timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. generator (`torch.Generator`, *optional*): - A random number generator. + A random number generator for reproducibility. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a - [`~schedulers.scheduling_consistency_models.ConsistencyDecoderSchedulerOutput`] or `tuple`. + [`~schedulers.scheduling_consistency_decoder.ConsistencyDecoderSchedulerOutput`] or `tuple`. Returns: - [`~schedulers.scheduling_consistency_models.ConsistencyDecoderSchedulerOutput`] or `tuple`: - If return_dict is `True`, - [`~schedulers.scheduling_consistency_models.ConsistencyDecoderSchedulerOutput`] is returned, otherwise + [`~schedulers.scheduling_consistency_decoder.ConsistencyDecoderSchedulerOutput`] or `tuple`: + If `return_dict` is `True`, + [`~schedulers.scheduling_consistency_decoder.ConsistencyDecoderSchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the sample tensor. """ x_0 = self.c_out[timestep] * model_output + self.c_skip[timestep] * sample