Add missing Args: entries to scheduler docstrings (#14354)

Add missing Args entries to scheduler docstrings

Sixteen scheduler methods had Args: blocks that omitted parameters present
in the signature. Adds the missing entries:

- noise, on the DPMSolver single/multistep and cosine-multistep update
  methods (required by the sde-* algorithm types)
- mu, on DPMSolverMultistep/Singlestep set_timesteps
- device, on CogVideoXDDIM and DDIMInverse set_timesteps
- sigma / sigma_next / sigma_before, on the Helios UniPC methods

Also corrects the rescale_zero_terminal_snr entry in scheduling_dpm_cogvideox,
which documented a `betas` parameter the function does not take (it takes
alphas_cumprod).

The scheduling_dpmsolver_multistep_inverse change is propagated by
make fix-copies, not hand-edited.

Documentation only; no behavior change.

Author: 4ktLuffy

src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py

@@ -459,6 +459,9 @@ def dpm_solver_first_order_update(
                 The direct output from the learned diffusion model.
             sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:
@@ -497,6 +500,9 @@ def multistep_dpm_solver_second_order_update(
                 The direct outputs from learned diffusion model at current and latter timesteps.
             sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:

src/diffusers/schedulers/scheduling_ddim_cogvideox.py

@@ -285,6 +285,8 @@ def set_timesteps(
         Args:
             num_inference_steps (`int`):
                 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.
         """
 
         if num_inference_steps > self.config.num_train_timesteps:

src/diffusers/schedulers/scheduling_ddim_inverse.py

@@ -275,6 +275,8 @@ def set_timesteps(
         Args:
             num_inference_steps (`int`):
                 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.
         """
 
         if num_inference_steps > self.config.num_train_timesteps:

src/diffusers/schedulers/scheduling_dpm_cogvideox.py

@@ -106,8 +106,8 @@ def rescale_zero_terminal_snr(alphas_cumprod):
     Rescales betas to have zero terminal SNR Based on https://huggingface.co/papers/2305.08891 (Algorithm 1)
 
     Args:
-        betas (`torch.Tensor`):
-            the betas that the scheduler is being initialized with.
+        alphas_cumprod (`torch.Tensor`):
+            the cumulative product of alphas that the scheduler is being initialized with.
 
     Returns:
         `torch.Tensor`: rescaled betas with zero terminal SNR

src/diffusers/schedulers/scheduling_dpmsolver_multistep.py

@@ -379,6 +379,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.
+            mu (`float`, *optional*):
+                Exponent for the dynamic time shift. Requires `use_dynamic_shifting` and a `time_shift_type` of
+                `"exponential"`; when passed, `flow_shift` is set to `exp(mu)`.
             timesteps (`list[int]`, *optional*):
                 Custom timesteps used to support arbitrary timesteps schedule. If `None`, timesteps will be generated
                 based on the `timestep_spacing` attribute. If `timesteps` is passed, `num_inference_steps` and `sigmas`
@@ -931,6 +934,9 @@ def multistep_dpm_solver_second_order_update(
                 The direct outputs from learned diffusion model at current and latter timesteps.
             sample (`torch.Tensor`, *optional*):
                 A current instance of a sample created by the diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:

src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py

@@ -787,6 +787,9 @@ def multistep_dpm_solver_second_order_update(
                 The direct outputs from learned diffusion model at current and latter timesteps.
             sample (`torch.Tensor`, *optional*):
                 A current instance of a sample created by the diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:

src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py

@@ -342,6 +342,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.
+            mu (`float`, *optional*):
+                Exponent for the dynamic time shift. Requires `use_dynamic_shifting` and a `time_shift_type` of
+                `"exponential"`; when passed, `flow_shift` is set to `exp(mu)`.
             timesteps (`list[int]`, *optional*):
                 Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
                 timestep spacing strategy of equal spacing between timesteps schedule is used. If `timesteps` is
@@ -776,6 +779,9 @@ def dpm_solver_first_order_update(
                 The previous discrete timestep in the diffusion chain.
             sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:
@@ -841,6 +847,9 @@ def singlestep_dpm_solver_second_order_update(
                 The previous discrete timestep in the diffusion chain.
             sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:
@@ -952,6 +961,9 @@ def singlestep_dpm_solver_third_order_update(
                 The previous discrete timestep in the diffusion chain.
             sample (`torch.Tensor`):
                 A current instance of a sample created by diffusion process.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:
@@ -1076,6 +1088,9 @@ def singlestep_dpm_solver_update(
                 A current instance of a sample created by diffusion process.
             order (`int`):
                 The solver order at this step.
+            noise (`torch.Tensor`, *optional*):
+                Random noise used by the stochastic (`sde-*`) solver variants. Required when `algorithm_type` is set to
+                one of them, and unused otherwise.
 
         Returns:
             `torch.Tensor`:

src/diffusers/schedulers/scheduling_helios.py

@@ -386,6 +386,8 @@ def convert_model_output(
                 The current discrete timestep in the diffusion chain.
             sample (`torch.Tensor`):
                 A current instance of a sample created by the diffusion process.
+            sigma (`torch.Tensor`, *optional*):
+                The sigma of the current step in the noise schedule.
 
         Returns:
             `torch.Tensor`:
@@ -470,6 +472,10 @@ def multistep_uni_p_bh_update(
                 A current instance of a sample created by the diffusion process.
             order (`int`):
                 The order of UniP at this timestep (corresponds to the *p* in UniPC-p).
+            sigma (`torch.Tensor`, *optional*):
+                The sigma of the current step in the noise schedule.
+            sigma_next (`torch.Tensor`, *optional*):
+                The sigma of the next step in the noise schedule.
 
         Returns:
             `torch.Tensor`:
@@ -607,6 +613,10 @@ def multistep_uni_c_bh_update(
                 The generated sample after the last predictor `x_{t}`.
             order (`int`):
                 The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`.
+            sigma_before (`torch.Tensor`, *optional*):
+                The sigma of the previous step in the noise schedule.
+            sigma (`torch.Tensor`, *optional*):
+                The sigma of the current step in the noise schedule.
 
         Returns:
             `torch.Tensor`: