From 8366d2676dc246074b0b46ff249fc913047eb7c9 Mon Sep 17 00:00:00 2001 From: David El Malih Date: Sat, 18 Jul 2026 09:24:27 +0200 Subject: [PATCH] Improve docstring scheduling k dpm 2 ancestral discrete --- .../scheduling_k_dpm_2_ancestral_discrete.py | 82 +++++++++++++------ 1 file changed, 57 insertions(+), 25 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py b/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py index 706190e291a4..0a7bffb01a86 100644 --- a/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py +++ b/src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py @@ -119,17 +119,17 @@ class KDPM2AncestralDiscreteScheduler(SchedulerMixin, ConfigMixin): beta_schedule (`str`, defaults to `"linear"`): The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from `linear` or `scaled_linear`. - trained_betas (`np.ndarray`, *optional*): + trained_betas (`np.ndarray` or `list[float]`, *optional*): Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - use_karras_sigmas (`bool`, *optional*, defaults to `False`): + use_karras_sigmas (`bool`, defaults to `False`): Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, the sigmas are determined according to a sequence of noise levels {σi}. - use_exponential_sigmas (`bool`, *optional*, defaults to `False`): + use_exponential_sigmas (`bool`, defaults to `False`): Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. - use_beta_sigmas (`bool`, *optional*, defaults to `False`): + use_beta_sigmas (`bool`, defaults to `False`): Whether to use beta sigmas for step sizes in the noise schedule during the sampling process. Refer to [Beta Sampling is All You Need](https://huggingface.co/papers/2407.12173) for more information. - prediction_type (`str`, defaults to `epsilon`, *optional*): + prediction_type (`str`, defaults to `"epsilon"`): Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process), `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen Video](https://huggingface.co/papers/2210.02303) paper). @@ -157,7 +157,7 @@ def __init__( prediction_type: str = "epsilon", timestep_spacing: str = "linspace", steps_offset: int = 0, - ): + ) -> None: if self.config.use_beta_sigmas and not is_scipy_available(): raise ImportError("Make sure to install scipy if you want to use beta sigmas.") if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1: @@ -187,29 +187,43 @@ def __init__( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property - def init_noise_sigma(self): - # standard deviation of the initial noise distribution + def init_noise_sigma(self) -> torch.Tensor: + """ + The standard deviation of the initial noise distribution. + + Returns: + `torch.Tensor`: + The standard deviation of the initial noise distribution. + """ if self.config.timestep_spacing in ["linspace", "trailing"]: return self.sigmas.max() return (self.sigmas.max() ** 2 + 1) ** 0.5 @property - def step_index(self): + def step_index(self) -> int | None: """ - The index counter for current timestep. It will increase 1 after each scheduler step. + The index counter for the current timestep. It increases by 1 after each scheduler step. + + Returns: + `int` or `None`: + The current step index, or `None` if it has not been initialized. """ return self._step_index @property - def begin_index(self): + def begin_index(self) -> int | None: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + + Returns: + `int` or `None`: + The begin index, or `None` if it has not been set. """ return self._begin_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index - def set_begin_index(self, begin_index: int = 0): + def set_begin_index(self, begin_index: int = 0) -> None: """ Sets the begin index for the scheduler. This function should be run from pipeline before the inference. @@ -231,7 +245,7 @@ def scale_model_input( Args: sample (`torch.Tensor`): The input sample. - timestep (`int`, *optional*): + timestep (`float` or `torch.Tensor`): The current timestep in the diffusion chain. Returns: @@ -254,7 +268,7 @@ def set_timesteps( num_inference_steps: int, device: str | torch.device | None = None, num_train_timesteps: int | None = None, - ): + ) -> None: """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -263,6 +277,9 @@ def set_timesteps( The number of diffusion steps used when generating samples with a pre-trained model. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + num_train_timesteps (`int`, *optional*): + The number of diffusion steps used to train the model. If `None`, `self.config.num_train_timesteps` is + used. """ self.num_inference_steps = num_inference_steps @@ -349,7 +366,7 @@ def set_timesteps( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t - def _sigma_to_t(self, sigma, log_sigmas): + def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: """ Convert sigma values to corresponding timestep values through interpolation. @@ -386,7 +403,7 @@ def _sigma_to_t(self, sigma, log_sigmas): return t # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras - def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor: + def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: """ Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative Models](https://huggingface.co/papers/2206.00364). @@ -508,7 +525,14 @@ def _convert_to_beta( return sigmas @property - def state_in_first_order(self): + def state_in_first_order(self) -> bool: + """ + Whether the scheduler is currently performing a first-order update. + + Returns: + `bool`: + `True` if the scheduler is performing a first-order update, otherwise `False`. + """ return self.sample is None # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.index_for_timestep @@ -565,27 +589,28 @@ def step( sample: torch.Tensor | np.ndarray, generator: torch.Generator | None = None, return_dict: bool = True, - ) -> KDPM2AncestralDiscreteSchedulerOutput | tuple: + ) -> KDPM2AncestralDiscreteSchedulerOutput | tuple[torch.Tensor, torch.Tensor]: """ Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). Args: - model_output (`torch.Tensor`): + model_output (`torch.Tensor` or `np.ndarray`): The direct output from learned diffusion model. - timestep (`float`): + timestep (`float` or `torch.Tensor`): The current discrete timestep in the diffusion chain. - sample (`torch.Tensor`): + sample (`torch.Tensor` or `np.ndarray`): A current instance of a sample created by the diffusion process. generator (`torch.Generator`, *optional*): A random number generator. - return_dict (`bool`): + return_dict (`bool`, defaults to `True`): Whether or not to return a [`~schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteSchedulerOutput`] or tuple. Returns: - [`~schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteSchedulerOutput`] or `tuple`: - If return_dict is `True`, + [`~schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteSchedulerOutput`] or + `tuple[torch.Tensor, torch.Tensor]`: + If `return_dict` is `True`, [`~schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteSchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the sample tensor. """ @@ -714,5 +739,12 @@ def add_noise( noisy_samples = original_samples + noise * sigma return noisy_samples - def __len__(self): + def __len__(self) -> int: + """ + Return the number of training timesteps configured for the scheduler. + + Returns: + `int`: + The number of training timesteps. + """ return self.config.num_train_timesteps