Merge branch 'main' into add-joyimage-edit-plus

Author: dg845

src/diffusers/schedulers/scheduling_lcm.py

@@ -40,9 +40,9 @@ class LCMSchedulerOutput(BaseOutput):
         prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
             Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
             denoising loop.
-        pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
-            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
-            `pred_original_sample` can be used to preview progress or for guidance.
+        denoised (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images, *optional*):
+            The predicted denoised sample `(x_{0})` based on the model output from the current timestep. `denoised` can
+            be used to preview progress or for guidance.
     """
 
     prev_sample: torch.Tensor
@@ -150,45 +150,45 @@ class LCMScheduler(SchedulerMixin, ConfigMixin):
     functionality via the [`SchedulerMixin.save_pretrained`] and [`~SchedulerMixin.from_pretrained`] functions.
 
     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.00085`):
             The starting `beta` value of inference.
-        beta_end (`float`, defaults to 0.02):
+        beta_end (`float`, defaults to `0.012`):
             The final `beta` value.
-        beta_schedule (`str`, defaults to `"linear"`):
+        beta_schedule (`str`, defaults to `"scaled_linear"`):
             The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
             `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
-        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`.
-        original_inference_steps (`int`, *optional*, defaults to 50):
+        original_inference_steps (`int`, defaults to `50`):
             The default number of inference steps used to generate a linearly-spaced timestep schedule, from which we
             will ultimately take `num_inference_steps` evenly spaced timesteps to form the final timestep schedule.
-        clip_sample (`bool`, defaults to `True`):
+        clip_sample (`bool`, defaults to `False`):
             Clip the predicted sample for numerical stability.
-        clip_sample_range (`float`, defaults to 1.0):
+        clip_sample_range (`float`, defaults to `1.0`):
             The maximum magnitude for sample clipping. Valid only when `clip_sample=True`.
         set_alpha_to_one (`bool`, defaults to `True`):
             Each diffusion step uses the alphas product value at that step and at the previous one. For the final step
             there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,
             otherwise it uses the alpha value at step 0.
-        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.
-        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).
         thresholding (`bool`, defaults to `False`):
             Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
             as Stable Diffusion.
-        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+        dynamic_thresholding_ratio (`float`, defaults to `0.995`):
             The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
-        sample_max_value (`float`, defaults to 1.0):
+        sample_max_value (`float`, defaults to `1.0`):
             The threshold value for dynamic thresholding. Valid only when `thresholding=True`.
         timestep_spacing (`str`, defaults to `"leading"`):
             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.
-        timestep_scaling (`float`, defaults to 10.0):
+        timestep_scaling (`float`, defaults to `10.0`):
             The factor the timesteps will be multiplied by when calculating the consistency model boundary conditions
             `c_skip` and `c_out`. Increasing this will decrease the approximation error (although the approximation
             error at the default of `10.0` is already pretty small).
@@ -306,13 +306,24 @@ def _init_step_index(self, timestep: float | torch.Tensor) -> None:
             self._step_index = self._begin_index
 
     @property
-    def step_index(self):
+    def step_index(self) -> int | None:
+        """
+        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) -> 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 for the scheduler, or `None` if not set.
         """
         return self._begin_index
 
@@ -337,6 +348,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -
                 The input sample.
             timestep (`int`, *optional*):
                 The current timestep in the diffusion chain.
+
         Returns:
             `torch.Tensor`:
                 A scaled input sample.
@@ -390,11 +402,11 @@ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
     def set_timesteps(
         self,
         num_inference_steps: int | None = None,
-        device: str | torch.device = None,
+        device: str | torch.device | None = None,
         original_inference_steps: int | None = None,
         timesteps: list[int] | None = None,
-        strength: int = 1.0,
-    ):
+        strength: float = 1.0,
+    ) -> None:
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
 
@@ -413,6 +425,12 @@ def set_timesteps(
                 Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
                 timestep spacing strategy of equal spacing between timesteps on the training/distillation timestep
                 schedule is used. If `timesteps` is passed, `num_inference_steps` must be `None`.
+            strength (`float`, defaults to `1.0`):
+                Strength factor used to generate a partial timestep schedule (e.g. for image-to-image). A value of
+                `1.0` uses the full schedule.
+
+        Returns:
+            `None`
         """
         # 0. Check inputs
         if num_inference_steps is None and timesteps is None:
@@ -527,7 +545,20 @@ def set_timesteps(
         self._step_index = None
         self._begin_index = None
 
-    def get_scalings_for_boundary_condition_discrete(self, timestep):
+    def get_scalings_for_boundary_condition_discrete(self, timestep: int) -> tuple[float, float]:
+        """
+        Computes the boundary condition scalings (`c_skip` and `c_out`) for the given discrete timestep, as used in the
+        Latent Consistency Model.
+
+        Args:
+            timestep (`int`):
+                The discrete timestep for which to compute the scalings.
+
+        Returns:
+            `tuple[float, float]`:
+                A tuple containing `c_skip` (scaling for the input sample) and `c_out` (scaling for the predicted
+                denoised sample).
+        """
         self.sigma_data = 0.5  # Default: 0.5
         scaled_timestep = timestep * self.config.timestep_scaling
 
@@ -550,16 +581,17 @@ def step(
         Args:
             model_output (`torch.Tensor`):
                 The direct output from learned diffusion model.
-            timestep (`float`):
+            timestep (`int`):
                 The current discrete 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.
-            return_dict (`bool`, *optional*, defaults to `True`):
+                A random number generator for reproducible sampling.
+            return_dict (`bool`, defaults to `True`):
                 Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`.
+
         Returns:
-            [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`:
+            [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`:
                 If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a
                 tuple is returned where the first element is the sample tensor.
         """
@@ -709,7 +741,14 @@ def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: tor
         velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
         return velocity
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """
+        Returns the number of train timesteps.
+
+        Returns:
+            `int`:
+                The number of diffusion steps used to train the model.
+        """
         return self.config.num_train_timesteps
 
     # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.previous_timestep