update class docstrings

Author: natolambert

src/diffusers/schedulers/scheduling_ddim.py

@@ -70,10 +70,12 @@ class DDIMScheduler(SchedulerMixin, ConfigMixin):
             `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
         trained_betas (`np.ndarray`, optional): TODO
         timestep_values (`np.ndarray`, optional): TODO
-        clip_sample (`bool`, default `True`): option to clip predicted sample between -1 and 1 for numerical stability
+        clip_sample (`bool`, default `True`):
+            option to clip predicted sample between -1 and 1 for numerical stability.
         set_alpha_to_one (`bool`, default `True`):
-            if alpha for final step is 1 or the final alpha of the "non-previous" one
-        tensor_format (`str`): whether the scheduler expects pytorch or numpy arrays
+            if alpha for final step is 1 or the final alpha of the "non-previous" one.
+        tensor_format (`str`): whether the scheduler expects pytorch or numpy arrays.
+
     """
 
     @register_to_config

src/diffusers/schedulers/scheduling_ddpm.py

@@ -67,12 +67,14 @@ class DDPMScheduler(SchedulerMixin, ConfigMixin):
         beta_schedule (`str`):
             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 (): TODO
+        trained_betas (`np.ndarray`, optional): TODO
         variance_type (`str`):
             options to clip the variance used when adding noise to the denoised sample. Choose from `fixed_small`,
             `fixed_small_log`, `fixed_large`, `fixed_large_log`, `learned` or `learned_range`.
-        clip_sample (`bool`, default `True`): option to clip predicted sample between -1 and 1 for numerical stability
-        tensor_format (`str`): whether the scheduler expects pytorch or numpy arrays
+        clip_sample (`bool`, default `True`):
+            option to clip predicted sample between -1 and 1 for numerical stability.
+        tensor_format (`str`): whether the scheduler expects pytorch or numpy arrays.
+
     """
 
     @register_to_config

src/diffusers/schedulers/scheduling_karras_ve.py

@@ -65,7 +65,7 @@ class KarrasVeScheduler(SchedulerMixin, ConfigMixin):
             A reasonable range is [0, 10].
         s_max (`float`): the end value of the sigma range where we add noise.
             A reasonable range is [0.2, 80].
-        tensor_format (`str`): whether the scheduler expects pytorch or numpy arrays
+        tensor_format (`str`): whether the scheduler expects pytorch or numpy arrays.
 
     """
 

src/diffusers/schedulers/scheduling_lms_discrete.py

@@ -24,6 +24,26 @@
 
 
 class LMSDiscreteScheduler(SchedulerMixin, ConfigMixin):
+    """
+    Linear Multistep Scheduler for discrete beta schedules. Based on the original k-diffusion implementation by
+    Katherine Crowson:
+    https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181
+
+    Args:
+        num_train_timesteps (`int`): number of diffusion steps used to train the model.
+        beta_start (`float`): the starting `beta` value of inference.
+        beta_end (`float`): the final `beta` value.
+        beta_schedule (`str`):
+            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): TODO
+            options to clip the variance used when adding noise to the denoised sample. Choose from `fixed_small`,
+            `fixed_small_log`, `fixed_large`, `fixed_large_log`, `learned` or `learned_range`.
+        timestep_values (`np.ndarry`, optional): TODO
+        tensor_format (`str`): whether the scheduler expects pytorch or numpy arrays.
+
+    """
+
     @register_to_config
     def __init__(
         self,
@@ -35,12 +55,8 @@ def __init__(
         timestep_values: Optional[np.ndarray] = None,
         tensor_format: str = "pt",
     ):
-        """
-        Linear Multistep Scheduler for discrete beta schedules. Based on the original k-diffusion implementation by
-        Katherine Crowson:
-        https://github.com/crowsonkb/k-diffusion/blob/481677d114f6ea445aa009cf5bd7a9cdee909e47/k_diffusion/sampling.py#L181
-        """
-
+        if trained_betas is not None:
+            self.betas = np.asarray(trained_betas)
         if beta_schedule == "linear":
             self.betas = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32)
         elif beta_schedule == "scaled_linear":

src/diffusers/schedulers/scheduling_pndm.py

@@ -67,11 +67,12 @@ class PNDMScheduler(SchedulerMixin, ConfigMixin):
         beta_schedule (`str`):
             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 (): TODO
+        trained_betas (`np.ndarray`, optional): TODO
         tensor_format (`str`): whether the scheduler expects pytorch or numpy arrays
         skip_prk_steps (`bool`):
             allows the scheduler to skip the Runge-Kutta steps that are defined in the original paper as being required
             before plms steps; defaults to `False`.
+
     """
 
     @register_to_config

src/diffusers/schedulers/scheduling_sde_ve.py

@@ -53,7 +53,7 @@ class ScoreSdeVeScheduler(SchedulerMixin, ConfigMixin):
         sigma_min (`float`):
                 initial noise scale for sigma sequence in sampling procedure. The minimum sigma should mirror the
                 distribution of the data.
-        sigma_max (`float`): TODO
+        sigma_max (`float`): maximum value used for the range of continuous timesteps passed into the model.
         sampling_eps (`float`): the end value of sampling, where timesteps decrease progessively from 1 to
         epsilon.
         correct_steps (`int`): number of correction steps performed on a produced sample.

src/diffusers/schedulers/scheduling_utils.py

@@ -38,6 +38,9 @@ class SchedulerOutput(BaseOutput):
 
 
 class SchedulerMixin:
+    """
+    Mixin containing common functions for the schedulers.
+    """
 
     config_name = SCHEDULER_CONFIG_NAME
     ignore_for_config = ["tensor_format"]