From 0b78e1dc0eb6f6fece6784b92e649e0de26e3275 Mon Sep 17 00:00:00 2001 From: David El Malih Date: Mon, 17 Nov 2025 17:34:03 +0100 Subject: [PATCH 1/3] Enhance type hints and docstrings in LMSDiscreteScheduler class Updated type hints for function parameters and return types to improve code clarity and maintainability. Enhanced docstrings for several methods, providing clearer descriptions of their functionality and expected arguments. Notable changes include specifying Literal types for certain parameters and ensuring consistent return type annotations across the class. --- .../schedulers/scheduling_lms_discrete.py | 82 +++++++++++++------ 1 file changed, 57 insertions(+), 25 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index 573678b100ba..3b51eff91699 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -99,15 +99,14 @@ class LMSDiscreteScheduler(SchedulerMixin, ConfigMixin): methods the library implements for all schedulers such as loading and saving. Args: - num_train_timesteps (`int`, defaults to 1000): + num_train_timesteps (`int`, defaults to `1000`): The number of diffusion steps to train the model. - beta_start (`float`, defaults to 0.0001): + beta_start (`float`, defaults to `0.0001`): The starting `beta` value of inference. - beta_end (`float`, defaults to 0.02): + beta_end (`float`, defaults to `0.02`): The final `beta` value. - 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`. + beta_schedule (`"linear"`, `"scaled_linear"`, or `"squaredcos_cap_v2"`, defaults to `"linear"`): + The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. trained_betas (`np.ndarray`, *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`): @@ -118,14 +117,14 @@ class LMSDiscreteScheduler(SchedulerMixin, ConfigMixin): use_beta_sigmas (`bool`, *optional*, 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 (`"epsilon"`, `"sample"`, or `"v_prediction"`, 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://imagen.research.google/video/paper.pdf) paper). - timestep_spacing (`str`, defaults to `"linspace"`): + timestep_spacing (`"linspace"`, `"leading"`, or `"trailing"`, defaults to `"linspace"`): The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. - steps_offset (`int`, defaults to 0): + steps_offset (`int`, defaults to `0`): An offset added to the inference steps, as required by some model families. """ @@ -138,13 +137,13 @@ def __init__( num_train_timesteps: int = 1000, beta_start: float = 0.0001, beta_end: float = 0.02, - beta_schedule: str = "linear", + beta_schedule: Literal["linear", "scaled_linear", "squaredcos_cap_v2"] = "linear", trained_betas: Optional[Union[np.ndarray, List[float]]] = None, use_karras_sigmas: Optional[bool] = False, use_exponential_sigmas: Optional[bool] = False, use_beta_sigmas: Optional[bool] = False, - prediction_type: str = "epsilon", - timestep_spacing: str = "linspace", + prediction_type: Literal["epsilon", "sample", "v_prediction"] = "epsilon", + timestep_spacing: Literal["linspace", "leading", "trailing"] = "linspace", steps_offset: int = 0, ): if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1: @@ -183,7 +182,15 @@ def __init__( self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @property - def init_noise_sigma(self): + def init_noise_sigma(self) -> Union[float, torch.Tensor]: + """ + The standard deviation of the initial noise distribution. + + Returns: + `float` or `torch.Tensor`: + The standard deviation of the initial noise distribution, computed based on the maximum sigma value and + the timestep spacing configuration. + """ # standard deviation of the initial noise distribution if self.config.timestep_spacing in ["linspace", "trailing"]: return self.sigmas.max() @@ -191,21 +198,29 @@ def init_noise_sigma(self): return (self.sigmas.max() ** 2 + 1) ** 0.5 @property - def step_index(self): + def step_index(self) -> Optional[int]: """ - The index counter for current timestep. It will increase 1 after each scheduler step. + The index counter for current timestep. It will increase by 1 after each scheduler step. + + Returns: + `int` or `None`: + The current step index, or `None` if not initialized. """ return self._step_index @property - def begin_index(self): + def begin_index(self) -> Optional[int]: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + + Returns: + `int` or `None`: + The begin index for the scheduler, or `None` if not 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. @@ -239,14 +254,21 @@ def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.T self.is_scale_input_called = True return sample - def get_lms_coefficient(self, order, t, current_order): + def get_lms_coefficient(self, order: int, t: int, current_order: int) -> float: """ Compute the linear multistep coefficient. Args: - order (): - t (): - current_order (): + order (`int`): + The order of the linear multistep method. + t (`int`): + The current timestep index. + current_order (`int`): + The current order for which to compute the coefficient. + + Returns: + `float`: + The computed linear multistep coefficient. """ def lms_derivative(tau): @@ -261,7 +283,7 @@ def lms_derivative(tau): return integrated_coeff - def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None): + def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None) -> None: """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -367,7 +389,7 @@ def _init_step_index(self, timestep: Union[float, torch.Tensor]) -> None: self._step_index = self._begin_index # 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. @@ -405,7 +427,17 @@ def _sigma_to_t(self, sigma, log_sigmas): # copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras def _convert_to_karras(self, in_sigmas: torch.Tensor) -> torch.Tensor: - """Constructs the noise schedule of Karras et al. (2022).""" + """ + Construct the noise schedule of Karras et al. (2022). + + Args: + in_sigmas (`torch.Tensor`): + The input sigma values to be converted. + + Returns: + `torch.Tensor`: + The converted sigma values following the Karras noise schedule. + """ sigma_min: float = in_sigmas[-1].item() sigma_max: float = in_sigmas[0].item() @@ -629,5 +661,5 @@ def add_noise( noisy_samples = original_samples + noise * sigma return noisy_samples - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps From b19fd52c582cb62dfa6da6b6232834d7a104c919 Mon Sep 17 00:00:00 2001 From: David El Malih Date: Thu, 20 Nov 2025 09:42:08 +0100 Subject: [PATCH 2/3] docs: Add specific paper reference to `_convert_to_karras` docstring. --- src/diffusers/schedulers/scheduling_lms_discrete.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index 3b51eff91699..09d53df872d6 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -428,7 +428,8 @@ def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: # copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras def _convert_to_karras(self, in_sigmas: torch.Tensor) -> torch.Tensor: """ - Construct the noise schedule of Karras et al. (2022). + Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative + Models](https://huggingface.co/papers/2206.00364). Args: in_sigmas (`torch.Tensor`): From 4bdd818d27ac408b2d1f64226c17d58b16ea20d0 Mon Sep 17 00:00:00 2001 From: David El Malih Date: Thu, 20 Nov 2025 21:16:29 +0100 Subject: [PATCH 3/3] Refactor `_convert_to_karras` docstring in DPMSolverSDEScheduler to include detailed descriptions and a specific paper reference, enhancing clarity and documentation consistency. --- .../schedulers/scheduling_dpmsolver_sde.py | 15 +++++++++++++-- .../schedulers/scheduling_lms_discrete.py | 1 - 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_dpmsolver_sde.py b/src/diffusers/schedulers/scheduling_dpmsolver_sde.py index e22954d4e6ea..ef89feb1cad6 100644 --- a/src/diffusers/schedulers/scheduling_dpmsolver_sde.py +++ b/src/diffusers/schedulers/scheduling_dpmsolver_sde.py @@ -488,9 +488,20 @@ def _sigma_to_t(self, sigma, log_sigmas): t = t.reshape(sigma.shape) return t - # copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras + # Copied from diffusers.schedulers.scheduling_lms_discrete.LMSDiscreteScheduler._convert_to_karras def _convert_to_karras(self, in_sigmas: torch.Tensor) -> torch.Tensor: - """Constructs the noise schedule of Karras et al. (2022).""" + """ + Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative + Models](https://huggingface.co/papers/2206.00364). + + Args: + in_sigmas (`torch.Tensor`): + The input sigma values to be converted. + + Returns: + `torch.Tensor`: + The converted sigma values following the Karras noise schedule. + """ sigma_min: float = in_sigmas[-1].item() sigma_max: float = in_sigmas[0].item() diff --git a/src/diffusers/schedulers/scheduling_lms_discrete.py b/src/diffusers/schedulers/scheduling_lms_discrete.py index 09d53df872d6..d0766eed1b66 100644 --- a/src/diffusers/schedulers/scheduling_lms_discrete.py +++ b/src/diffusers/schedulers/scheduling_lms_discrete.py @@ -425,7 +425,6 @@ def _sigma_to_t(self, sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: t = t.reshape(sigma.shape) return t - # copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras def _convert_to_karras(self, in_sigmas: torch.Tensor) -> torch.Tensor: """ Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative