dew.sampling.solvers

One reverse step each, from t to t_next, given the model’s denoising at t.

A solver is a value. What it needs between steps travels in its state; init builds it from x_T, a concrete time grid, the process and the walk’s root key; step threads it through sample’s scan. The rates of the sampling schedule come from process; a solver that needs another evaluation of the model (Heun’s corrector, RK4’s stages, KDPM2’s midpoint) calls denoise. A solver that integrates dx / dsigma = eps says so by refusing a schedule whose alpha is not one.

The solvers named after Diffusers 0.34.0 schedulers reproduce their arithmetic; tests/test_samplers.py holds their trajectories and trajectory gradients against the fixtures tools/diffusers_reference.py records. Process supplies the time grid. Initializing with a concrete grid and process checks each algorithm’s endpoint domain before the compiled scan. Finite endpoint limits are verified separately in tools/diffusers_limits_reference.py; DEIS history, UniPC epsilon correction, and non-++ SDE noise can survive a zero-sigma target. Undefined limits raise rather than substitute an update.

init’s key is the walk’s own key, the one sample folds per step. Every solver draws its per-step noise from the folded key it is handed, so the root key is unused except by DPMSolverSDE, whose source noise sampler is one Brownian tree over the whole trajectory and needs a state its steps share.

Name	Summary
DDIM	DDIM (Song et al.
DDPM	Exact ancestral sampler for the reverse diffusion SDE.
DEIS	DEIS (Zhang and Chen 2023, arXiv 2204.13902) in its log-rho multistep form, Diffusers 0.34.0’s DEISMultistepScheduler: the exponential integrator of eps with the polynomial-in-log(rho) interpolation of the last outputs, rho = sigma / alpha, integrated in closed form over the step.
KDPM2	k-diffusion’s DPM-Solver-2 (sample_dpm_2), the update of Diffusers 0.34.0’s KDPM2DiscreteScheduler, and with ancestral its sample_dpm_2_ancestral and KDPM2AncestralDiscreteScheduler.
LMS	Linear multistep over dx/dsigma = (x - x_0) / sigma, k-diffusion’s sample_lms and Diffusers 0.34.0’s LMSDiscreteScheduler: the last order derivatives interpolated by the Lagrange polynomial through their sigmas and integrated over the step, in closed form where Diffusers quadratures; the order grows with the history.
PNDM	PNDM (Liu et al.
RK4	Classical Runge-Kutta over dx/dsigma = eps, on a variance exploding schedule; the stages at half steps read the model at the time the schedule maps that sigma back to.
TCD	Trajectory consistency sampling (Zheng et al.
Consistency	Multistep consistency sampling (Song et al.
DPMSolverMultistep	DPM-Solver (Lu et al.
DPMSolverSDE	Diffusers 0.34.0’s DPMSolverSDEScheduler, k-diffusion’s sample_dpmpp_sde midpoint solver over a Brownian tree.
DPMSolverSinglestep	Diffusers 0.34.0’s grouped DPM-Solver updates from each group’s anchor.
Euler	The DDIM update written as an Euler step of the probability flow ODE.
EulerAncestral	Euler with the ancestral noise injection of k-diffusion (get_ancestral_step, eta 1).
Heun	Heun’s second order method (Karras et al.
MultiStepDPM	A third order multistep integrator of dx/dsigma = eps on a variance exploding schedule, from finite differences of the last three eps.
Solver	A step of a sampler, and whatever it carries between steps.
UniPC	UniPC (Zhao et al.

DDIM

dataclass source

class DDIM(eta: float = 0.0)

DDIM (Song et al. 2021); eta is the stochasticity, 0 deterministic and 1 DDPM-like.

Diffusers 0.34.0’s DDIMScheduler limits the clean prediction under clip_sample or thresholding and keeps the model’s own output as its epsilon, so the direction term is the unlimited one. That pairing is SourceLimitedPrediction’s, in the process’s conversion.

DDIM.init

def init(x, times, process, *, key)

DDIM.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

DDPM

dataclass source

class DDPM(variance: Literal['small', 'large'] = 'small')

Exact ancestral sampler for the reverse diffusion SDE.

One step draws from the forward posterior q(x_s | x_t, x_0) for x_t = alpha_t x_0 + sigma_t eps, written in signal and noise rates so it holds for any schedule and any step stride. The posterior mean is alpha_s x_0 + alpha_t sigma_s^2 / (alpha_s sigma_t) eps and its variance is sigma_s^2 (1 - alpha_t^2 sigma_s^2 / (alpha_s^2 sigma_t^2)).

variance is which of Diffusers 0.34.0’s fixed DDPMScheduler posterior variances the draw takes. "small" is that posterior’s own, written in rates and so defined on any schedule. "large" is the forward step’s beta, 1 - alpha_t^2 / alpha_s^2, the wider Glide choice: that is a variance-preserving statement, and it is zero wherever alpha is one, so a variance-exploding grid is refused rather than sampled without noise.

Neither draws on the step whose own time is the schedule’s zero: x_t is the least noised state the schedule holds there, and the source gates its draw on that time the same way. Elsewhere the wide variance at a terminal alpha of one is exactly sigma_t, which is what the source’s own current_beta_t reduces to.

DDPM.init

def init(x, times, process, *, key)

DDPM.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

DEIS

dataclass source

class DEIS(order: int = 2, lower_order_final: bool = True)

DEIS (Zhang and Chen 2023, arXiv 2204.13902) in its log-rho multistep form, Diffusers 0.34.0’s DEISMultistepScheduler: the exponential integrator of eps with the polynomial-in-log(rho) interpolation of the last outputs, rho = sigma / alpha, integrated in closed form over the step. The first order is DPM-Solver’s. Orders grow with the history; lower_order_final is the same short-walk taper as DPMSolverMultistep’s. At sigma=0 the integrated log-rho basis retains its history terms. An alpha=0 source contributes a node at infinite rho; that node’s weight vanishes in subsequent finite-interval integrals.

DEIS.init

def init(x, times, process, *, key)

DEIS.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

KDPM2

dataclass source

class KDPM2(ancestral: bool = False)

k-diffusion’s DPM-Solver-2 (sample_dpm_2), the update of Diffusers 0.34.0’s KDPM2DiscreteScheduler, and with ancestral its sample_dpm_2_ancestral and KDPM2AncestralDiscreteScheduler.

An Euler step to the geometric midpoint of sigma_t and the target level, the model read there, and the step from x taken with that midpoint derivative. The target is sigma_s, or under ancestral the sigma_down of k-diffusion’s ancestral step with sigma_up of fresh noise added after. The midpoint’s time comes from the schedule’s t_of_sigma, so this integrates a GeneralizedNoiseScheduler.

KDPM2.init

def init(x, times, process, *, key)

KDPM2.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

LMS

dataclass source

class LMS(order: int = 4)

Linear multistep over dx/dsigma = (x - x_0) / sigma, k-diffusion’s sample_lms and Diffusers 0.34.0’s LMSDiscreteScheduler: the last order derivatives interpolated by the Lagrange polynomial through their sigmas and integrated over the step, in closed form where Diffusers quadratures; the order grows with the history. Integrates a GeneralizedNoiseScheduler.

LMS.init

def init(x, times, process, *, key)

LMS.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

PNDM

dataclass source

class PNDM(skip_prk_steps: bool = False)

PNDM (Liu et al. 2022, arXiv 2202.09778), Diffusers 0.34.0’s PNDMScheduler: Adams-Bashforth over eps with DDIM as the transfer, fourth order once four outputs are in hand. The warmup is the paper’s, three pseudo Runge-Kutta steps of four model evaluations each (the stages at the interval’s midpoint and end), which seed the history with each step’s first eps; under skip_prk_steps, the PLMS form Stable Diffusion runs, the first step is a predictor-corrector pair (an Euler step, eps re-read at its end, the step retaken with the mean) and the orders grow from there. The schedule owns the transfer stride: ordinary native grids use their adjacent interval, while published integer grids retain their fixed training stride. Transfers cannot start at alpha = 0.

PNDM.init

def init(x, times, process, *, key)

PNDM.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

RK4

dataclass source

class RK4()

Classical Runge-Kutta over dx/dsigma = eps, on a variance exploding schedule; the stages at half steps read the model at the time the schedule maps that sigma back to.

RK4.init

def init(x, times, process, *, key)

RK4.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

TCD

dataclass source

class TCD(eta: float = 0.3)

Trajectory consistency sampling (Zheng et al. 2024, arXiv 2402.19159), Diffusers 0.34.0’s TCDScheduler: the deterministic DDIM step lands at (1 - eta) t_next, and the forward process noises it up to t_next with fresh noise, the paper’s gamma-sampling with gamma = eta. At eta 0 it is DDIM; at eta 1 the step goes through the clean end of the schedule. A tabulated schedule reads the intermediate time at its truncated index, as the reference floors it. The last step of a walk lands on the grid’s end itself and takes no noise.

TCD.init

def init(x, times, process, *, key)

TCD.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

Consistency

dataclass source

class Consistency()

Multistep consistency sampling (Song et al. 2023, Algorithm 1), the update of Diffusers 0.34.0’s LCMScheduler: the clean prediction is noised again to the next level with fresh noise, x_s = alpha_s x_0 + sigma_s z, and the last step keeps x_0 as it is. A latent consistency model’s x_0 is the consistency function’s output, which ConsistencyBoundary reads out of the model’s prediction.

Consistency.init

def init(x, times, process, *, key)

Consistency.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

DPMSolverMultistep

dataclass source

class DPMSolverMultistep(
    order: int = 2,
    algorithm: Algorithm = 'dpmsolver++',
    solver_type: Literal['midpoint', 'heun'] = 'midpoint',
    lower_order_final: bool = True,
    euler_at_final: bool = False,
)

DPM-Solver (Lu et al. 2022, arXiv 2206.00927) and DPM-Solver++ (arXiv 2211.01095) as multistep integrators in lambda = log(alpha) - log(sigma): the four algorithms, three orders and two second-order forms of Diffusers 0.34.0’s DPMSolverMultistepScheduler, with its defaults.

dpmsolver++ and sde-dpmsolver++ integrate the clean prediction, the other two eps; the sde- forms add the noise term of the SDE solver. With h = lambda_t - lambda_s0 over the step and D0, D1, D2 the finite differences of the last outputs in lambda, the deterministic dpmsolver++ step is

x_t = (sigma_t / sigma_s0) x - alpha_t (e^-h - 1) D0 + c D1 + c_2 D2

and _dpm_terms holds each algorithm’s coefficients as Diffusers writes them. The first step has no history and is first order, the second at most second. lower_order_final is Diffusers’ rule verbatim, which acts only in a walk under 15 steps: first order on the last step and at most second on the one before. euler_at_final makes the last step first order whatever the length. A zero target forces the clean limit for deterministic and ++ updates. The non-++ SDE first-order limit also retains alpha_tsigma_s/alpha_s(noise-eps); it requires a finite source alpha and a final order reduction. That endpoint is an equation-limit extension: Diffusers refuses a literal zero-terminal non-++ config. EDM uses the ++ algorithms over its existing process.

DPM-Solver++ 2M with no order taper is order=2, algorithm=“dpmsolver++”, solver_type=“midpoint”, lower_order_final=False, euler_at_final=False.

DPMSolverMultistep.init

def init(x, times, process, *, key)

DPMSolverMultistep.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

DPMSolverSDE

dataclass source

class DPMSolverSDE(depth: int = MAX_BROWNIAN_DEPTH, seed: int | None = None)

Diffusers 0.34.0’s DPMSolverSDEScheduler, k-diffusion’s sample_dpmpp_sde midpoint solver over a Brownian tree.

Each interval takes two ancestral first-order steps from its own start: one to the geometric midpoint of sigma_t and sigma_s, which the model is read at, and one to sigma_s with that midpoint’s clean prediction. Both steps go down to k-diffusion’s sigma_down and add sigma_up of noise, and both draw that noise from one Brownian path over the trajectory’s sigma interval: the first over [sigma_t, sigma_mid] and the second over [sigma_t, sigma_s], so the two are correlated exactly as nested increments of one path. The source’s sampler transforms sigma with the identity even though its own steps integrate -log(sigma), so the interval widths are sigma differences.

depth resolves the root interval to (sigma_max - sigma_min) / 2**depth: on a published VP table’s span of about 14.6 the default reaches 8.7e-7, at or inside the reference tree’s own 1e-6 tolerance, and it is also where a float32 position runs out of mantissa, so no deeper descent tells two sigmas apart. A zero-sigma target has no ancestral step and lands on the clean prediction.

The root interval is the schedule’s own positive sigma domain, not the extremes of the grid handed to init: the source builds its tree from all the positive sigmas it prepared, so a continuation that walks a suffix of that grid keeps the path the same key gives the whole one. A grid whose only interval lands on sigma zero leaves that domain a single point, which the source also prepares and never queries.

seed is the source’s noise_sampler_seed: with it the tree’s entropy is the checkpoint’s rather than the caller’s, so every walk over the same grid integrates one fixed path however the sampling key changes. It seeds this bridge, not the reference tree, because a Torch seed does not name a JAX stream; what carries over is the contract, a path independent of the walk’s key.

DPMSolverSDE.init

def init(x, times, process, *, key)

DPMSolverSDE.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

DPMSolverSinglestep

dataclass source

class DPMSolverSinglestep(
    order: int = 2,
    algorithm: Literal['dpmsolver++', 'dpmsolver', 'sde-dpmsolver++'] = 'dpmsolver++',
    solver_type: Literal['midpoint', 'heun'] = 'midpoint',
    lower_order_final: bool = False,
)

Diffusers 0.34.0’s grouped DPM-Solver updates from each group’s anchor.

A group of k model evaluations completes one k-th order update. Orders repeat [1, 2] or [1, 2, 3]; an incomplete final group uses lower order, as set_timesteps does in the reference. lower_order_final also lowers the final complete group, and a zero-sigma target forces final order 1. Source-domain checks use that effective order list.

At alpha=0, clean-prediction midpoint and deterministic Heun groups have finite limits. Noise-prediction groups above order 1 and third-order SDE Heun groups diverge; initialization rejects those source/grid pairs.

DPMSolverSinglestep.order_list

def order_list(steps: int) -> list[int]

Reference groups, completing an uneven final group at lower order.

DPMSolverSinglestep.init

def init(x, times, process, *, key)

DPMSolverSinglestep.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

Euler

dataclass source

class Euler()

The DDIM update written as an Euler step of the probability flow ODE. On a variance exploding schedule it is dx/dsigma = eps.

Euler.init

def init(x, times, process, *, key)

Euler.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

EulerAncestral

dataclass source

class EulerAncestral()

Euler with the ancestral noise injection of k-diffusion (get_ancestral_step, eta 1). The step goes down to sigma_down, and sigma_up of fresh noise brings the marginal back to sigma_s. Integrates a GeneralizedNoiseScheduler.

EulerAncestral.init

def init(x, times, process, *, key)

EulerAncestral.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

Heun

dataclass source

class Heun()

Heun’s second order method (Karras et al. 2022, Algorithm 2): an Euler step, the derivative re-evaluated at its end, and the average of the two.

Diffusers 0.34.0’s HeunDiscreteScheduler limits the clean prediction of both stages under clip_sample; that limit belongs to the process’s conversion, SourceLimitedPrediction, so both evaluations here read the limited prediction without the solver knowing about it.

Heun.init

def init(x, times, process, *, key)

Heun.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

MultiStepDPM

dataclass source

class MultiStepDPM()

A third order multistep integrator of dx/dsigma = eps on a variance exploding schedule, from finite differences of the last three eps.

MultiStepDPM.init

def init(x, times, process, *, key)

MultiStepDPM.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

Solver

class source

class Solver(Protocol[StateT])

A step of a sampler, and whatever it carries between steps.

StateT is that carried value: nothing for a one-step solver, the previous model outputs for a multi-step one. It is a type parameter, so a solver’s own state type is checked at its call sites.

Solver.init

def init(x, times, process, *, key) -> StateT

Prepare state and check endpoint domains on the concrete time grid.

sample() materializes this grid at compile time, so validation adds no host callbacks to the compiled step. key is the walk’s root key; a solver whose source draws one correlated path over the whole trajectory keeps that path’s state, and every other solver ignores it and draws from the per-step key step is handed.

Every argument is on every solver because this is the surface sample calls. x sizes the carried history (LMS, MultiStepDPM, DEIS, UniPC and the DPM-Solvers), times and process check the grid’s endpoints and count its steps (DDPM, Consistency, DPMSolverSDE, DEIS, UniPC and the DPM-Solvers), and key seeds the Brownian tree of DPMSolverSDE alone. A one-step solver reads none of them and answers ().

Solver.step

def step(
    x,
    t,
    t_next,
    denoised,
    eps,
    state,
    key,
    process,
    denoise,
    /,
) -> tuple[jax.Array, StateT]

x at t_next from x at t and the model’s (denoised, eps) at t. sample passes every argument by position, so a solver over another algebra names the pair for what it reads (the discrete one takes log-probabilities where a Gaussian one takes eps).

UniPC

dataclass source

class UniPC(
    order: int = 2,
    solver_type: Literal['bh1', 'bh2'] = 'bh2',
    predict_x0: bool = True,
    lower_order_final: bool = True,
    disable_corrector: tuple[int, ...] = (),
)

UniPC (Zhao et al. 2023, arXiv 2302.04867), Diffusers 0.34.0’s UniPCMultistepScheduler: a unified predictor and corrector in lambda whose weights solve a small linear system over the history’s positions. Each step first corrects the sample the last predictor produced, with the model output just read there and that predictor’s order, then predicts the next sample from the corrected one. solver_type picks B(h) as h (bh1) or e^h - 1 (bh2); predict_x0 integrates the clean prediction, otherwise eps. lower_order_final caps the order by the steps remaining, Diffusers’ rule, so the last step is first order; disable_corrector names the step indices whose predictor’s output is not corrected. The lowered clean-prediction terminal is alpha_t*x_0; epsilon prediction uses the corrected sample with the pre-correction epsilon. Initialization rejects an unlowered higher-order zero target. At an alpha=0 source, bh1 and epsilon correctors diverge unless the first correction is disabled. The finite infinite-node weights use the same Vandermonde system with column scaling.

UniPC.init

def init(x, times, process, *, key)

UniPC.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)