Skip to content

Commit e276555

Browse files
delmalihstevhliu
delmalih
and
stevhliu
authored
docs: improve docstring scheduling_sasolver.py (#14148)
* Improve docstring scheduling sasolver * Update src/diffusers/schedulers/scheduling_sasolver.py Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update src/diffusers/schedulers/scheduling_sasolver.py Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update src/diffusers/schedulers/scheduling_sasolver.py Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * Update src/diffusers/schedulers/scheduling_sasolver.py Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --------- Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
1 parent 284419b commit e276555

1 file changed

Lines changed: 108 additions & 23 deletions

File tree

‎src/diffusers/schedulers/scheduling_sasolver.py‎

Lines changed: 108 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
# The codebase is modified based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
1717

1818
import math
19-
from typing import Callable, Literal
19+
from typing import Callable, List, Literal
2020

2121
import numpy as np
2222
import torch
@@ -109,10 +109,10 @@ class SASolverScheduler(SchedulerMixin, ConfigMixin):
109109
corrector_order (`int`, defaults to 2):
110110
The corrector order which can be `1` or `2` or `3` or '4'. It is recommended to use `corrector_order=2` for
111111
guided sampling, and `corrector_order=3` for unconditional sampling.
112-
prediction_type (`str`, defaults to `epsilon`, *optional*):
113-
Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
114-
`sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
115-
Video](https://huggingface.co/papers/2210.02303) paper).
112+
prediction_type (`"epsilon"`, `"sample"`, `"v_prediction"`, or `"flow_prediction"`, defaults to `"epsilon"`):
113+
Prediction type of the scheduler function. `epsilon` predicts the noise of the diffusion process, `sample`
114+
directly predicts the noisy sample, `v_prediction` predicts the velocity (see section 2.4 of [Imagen
115+
Video](https://huggingface.co/papers/2210.02303) paper), and `flow_prediction` predicts the flow.
116116
tau_func (`Callable`, *optional*):
117117
Stochasticity during the sampling. Default in init is `lambda t: 1 if t >= 200 and t <= 800 else 0`.
118118
SA-Solver will sample from vanilla diffusion ODE if tau_func is set to `lambda t: 0`. SA-Solver will sample
@@ -139,6 +139,10 @@ class SASolverScheduler(SchedulerMixin, ConfigMixin):
139139
use_beta_sigmas (`bool`, *optional*, defaults to `False`):
140140
Whether to use beta sigmas for step sizes in the noise schedule during the sampling process. Refer to [Beta
141141
Sampling is All You Need](https://huggingface.co/papers/2407.12173) for more information.
142+
use_flow_sigmas (`bool`, *optional*, defaults to `False`):
143+
Whether to use flow sigmas for step sizes in the noise schedule during the sampling process.
144+
flow_shift (`float`, *optional*, defaults to `1.0`):
145+
The shift value for the timestep schedule for flow matching.
142146
lambda_min_clipped (`float`, defaults to `-inf`):
143147
Clipping threshold for the minimum value of `lambda(t)` for numerical stability. This is critical for the
144148
cosine (`squaredcos_cap_v2`) noise schedule.
@@ -271,7 +275,7 @@ def set_timesteps(self, num_inference_steps: int = None, device: str | torch.dev
271275
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
272276
273277
Args:
274-
num_inference_steps (`int`):
278+
num_inference_steps (`int`, *optional*):
275279
The number of diffusion steps used when generating samples with a pre-trained model.
276280
device (`str` or `torch.device`, *optional*):
277281
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
@@ -451,7 +455,7 @@ def _sigma_to_alpha_sigma_t(self, sigma):
451455
return alpha_t, sigma_t
452456

453457
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
454-
def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
458+
def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
455459
"""
456460
Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative
457461
Models](https://huggingface.co/papers/2206.00364).
@@ -576,7 +580,7 @@ def convert_model_output(
576580
self,
577581
model_output: torch.Tensor,
578582
*args,
579-
sample: torch.Tensor = None,
583+
sample: torch.Tensor | None = None,
580584
**kwargs,
581585
) -> torch.Tensor:
582586
"""
@@ -663,9 +667,24 @@ def convert_model_output(
663667

664668
return epsilon
665669

666-
def get_coefficients_exponential_negative(self, order, interval_start, interval_end):
667-
"""
668-
Calculate the integral of exp(-x) * x^order dx from interval_start to interval_end
670+
def get_coefficients_exponential_negative(
671+
self, order: int, interval_start: torch.Tensor, interval_end: torch.Tensor
672+
) -> torch.Tensor:
673+
r"""
674+
Compute the integral of $\exp(-x) \cdot x^{order}$ over the interval `[interval_start, interval_end]`. This is
675+
used as a building block of the SA-Solver algorithm (negative exponential case).
676+
677+
Args:
678+
order (`int`):
679+
The order of the polynomial factor $x^{order}$ in the integrand. Must be in `[0, 1, 2, 3]`.
680+
interval_start (`torch.Tensor`):
681+
The lower bound of the integration interval.
682+
interval_end (`torch.Tensor`):
683+
The upper bound of the integration interval.
684+
685+
Returns:
686+
`torch.Tensor`:
687+
The value of the integral $\int_{interval_start}^{interval_end} \exp(-x) \cdot x^{order} \, dx$.
669688
"""
670689
assert order in [0, 1, 2, 3], "order is only supported for 0, 1, 2 and 3"
671690

@@ -687,9 +706,27 @@ def get_coefficients_exponential_negative(self, order, interval_start, interval_
687706
- (interval_end**3 + 3 * interval_end**2 + 6 * interval_end + 6)
688707
)
689708

690-
def get_coefficients_exponential_positive(self, order, interval_start, interval_end, tau):
691-
"""
692-
Calculate the integral of exp(x(1+tau^2)) * x^order dx from interval_start to interval_end
709+
def get_coefficients_exponential_positive(
710+
self, order: int, interval_start: torch.Tensor, interval_end: torch.Tensor, tau: torch.Tensor
711+
) -> torch.Tensor:
712+
r"""
713+
Compute the integral of $\exp(x(1+\tau^2)) \cdot x^{order}$ over the interval `[interval_start, interval_end]`.
714+
This is used as a building block of the SA-Solver algorithm (positive exponential case).
715+
716+
Args:
717+
order (`int`):
718+
The order of the polynomial factor $x^{order}$ in the integrand. Must be in `[0, 1, 2, 3]`.
719+
interval_start (`torch.Tensor`):
720+
The lower bound of the integration interval.
721+
interval_end (`torch.Tensor`):
722+
The upper bound of the integration interval.
723+
tau (`torch.Tensor`):
724+
The stochasticity coefficient $\tau$ that scales the exponent in the integrand.
725+
726+
Returns:
727+
`torch.Tensor`:
728+
The value of the integral $\int_{interval_start}^{interval_end} \exp(x(1+\tau^2)) \cdot x^{order} \,
729+
dx$.
693730
"""
694731
assert order in [0, 1, 2, 3], "order is only supported for 0, 1, 2 and 3"
695732

@@ -731,9 +768,23 @@ def get_coefficients_exponential_positive(self, order, interval_start, interval_
731768
/ ((1 + tau**2) ** 4)
732769
)
733770

734-
def lagrange_polynomial_coefficient(self, order, lambda_list):
735-
"""
736-
Calculate the coefficient of lagrange polynomial
771+
def lagrange_polynomial_coefficient(self, order: int, lambda_list: List[torch.Tensor]) -> List[List[torch.Tensor]]:
772+
r"""
773+
Compute the coefficients of the Lagrange polynomial of order `order` that interpolates the points $(x_i, y_i) =
774+
(\lambda_i, i)$ for $i \in [0, order]`. The returned coefficients describe the polynomial in its expanded form
775+
$\sum_{k=0}^{order} c_k \cdot x^k$.
776+
777+
Args:
778+
order (`int`):
779+
The order of the Lagrange polynomial. Must be in `[0, 1, 2, 3]`.
780+
lambda_list (`List[torch.Tensor]`):
781+
A list of `order + 1` $\lambda$ values (i.e. the $x$-coordinates of the interpolation nodes). Must
782+
satisfy `len(lambda_list) == order + 1`.
783+
784+
Returns:
785+
`List[List[torch.Tensor]]`:
786+
A list of `order + 1` coefficient lists — one for each output channel — where each inner list contains
787+
the polynomial coefficients $c_0, c_1, \ldots, c_{order}$.
737788
"""
738789

739790
assert order in [0, 1, 2, 3]
@@ -840,7 +891,35 @@ def lagrange_polynomial_coefficient(self, order, lambda_list):
840891
],
841892
]
842893

843-
def get_coefficients_fn(self, order, interval_start, interval_end, lambda_list, tau):
894+
def get_coefficients_fn(
895+
self,
896+
order: int,
897+
interval_start: torch.Tensor,
898+
interval_end: torch.Tensor,
899+
lambda_list: List[torch.Tensor],
900+
tau: torch.Tensor,
901+
) -> List[torch.Tensor]:
902+
r"""
903+
Compute the gradient coefficients used by the SA-Predictor and SA-Corrector. The coefficients are obtained by
904+
combining the Lagrange polynomial interpolation with the closed-form exponential integrals.
905+
906+
Args:
907+
order (`int`):
908+
The order of the solver. Must be in `[1, 2, 3, 4]` and equal to `len(lambda_list)`.
909+
interval_start (`torch.Tensor`):
910+
The lower bound $\lambda_{s_0}$ of the integration interval.
911+
interval_end (`torch.Tensor`):
912+
The upper bound $\lambda_t$ of the integration interval.
913+
lambda_list (`List[torch.Tensor]`):
914+
The $\lambda$ values of the past `order` timesteps used as the interpolation nodes of the Lagrange
915+
polynomial.
916+
tau (`torch.Tensor`):
917+
The stochasticity coefficient $\tau$ at the current step.
918+
919+
Returns:
920+
`List[torch.Tensor]`:
921+
A list of `order` gradient coefficients, one per past model output.
922+
"""
844923
assert order in [1, 2, 3, 4]
845924
assert order == len(lambda_list), "the length of lambda list must be equal to the order"
846925
coefficients = []
@@ -876,12 +955,15 @@ def stochastic_adams_bashforth_update(
876955
Args:
877956
model_output (`torch.Tensor`):
878957
The direct output from the learned diffusion model at the current timestep.
879-
prev_timestep (`int`):
880-
The previous discrete timestep in the diffusion chain.
881958
sample (`torch.Tensor`):
882959
A current instance of a sample created by the diffusion process.
960+
noise (`torch.Tensor`):
961+
A noise tensor of the same shape as `sample`, drawn from a standard normal distribution. Used to inject
962+
stochasticity into the update.
883963
order (`int`):
884964
The order of SA-Predictor at this timestep.
965+
tau (`torch.Tensor`):
966+
The stochasticity coefficient tau at the current step.
885967
886968
Returns:
887969
`torch.Tensor`:
@@ -1004,14 +1086,17 @@ def stochastic_adams_moulton_update(
10041086
Args:
10051087
this_model_output (`torch.Tensor`):
10061088
The model outputs at `x_t`.
1007-
this_timestep (`int`):
1008-
The current timestep `t`.
10091089
last_sample (`torch.Tensor`):
10101090
The generated sample before the last predictor `x_{t-1}`.
1091+
last_noise (`torch.Tensor`):
1092+
The noise tensor injected by the last predictor at the previous step. Used to ensure consistent
1093+
stochasticity between the predictor and the corrector.
10111094
this_sample (`torch.Tensor`):
10121095
The generated sample after the last predictor `x_{t}`.
10131096
order (`int`):
10141097
The order of SA-Corrector at this step.
1098+
tau (`torch.Tensor`):
1099+
The stochasticity coefficient tau at the current step.
10151100
10161101
Returns:
10171102
`torch.Tensor`:
@@ -1195,7 +1280,7 @@ def step(
11951280
A current instance of a sample created by the diffusion process.
11961281
generator (`torch.Generator`, *optional*):
11971282
A random number generator.
1198-
return_dict (`bool`):
1283+
return_dict (`bool`, defaults to `True`):
11991284
Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
12001285
12011286
Returns:

0 commit comments

Comments
 (0)