Module ml::diffusion
Diffusion: noise schedules (EDM/Karras, DDPM linear) and samplers (DDPM ancestral, DDIM, DPM-Solver, DPM-Solver++).
Operates on plain std::Tensor samples (1D or 2D — a sampler step is just elementwise arithmetic over whatever shape the caller's data/ noise tensors already are; nothing here is image-shape-specific). Forward-only orchestration around a caller-supplied denoiser callback, like the rest of std/ml's newer additions (see attention.h).
Convention: DDPM/DDIM/DPM-Solver/DPM-Solver++ all use the VP (variance-preserving) parametrization: alpha_t = sqrt(alphaBar_t), sigma_t = sqrt(1 - alphaBar_t), alpha_t^2 + sigma_t^2 = 1. EDM's Karras schedule is a separate, VE-flavored (variance-exploding) convention used by a different family of samplers not implemented here — it's included as a standalone schedule generator (edm_karras_sigmas) rather than force-fit into the VP sampler functions below.
Scope note: DPM-Solver-v3 (Zheng et al. 2023) is NOT implemented — its distinguishing idea is an affine reparametrization of the ODE using "empirical model statistics" (EMS) calibrated from a real trained model over real data, which isn't something this environment has access to; a version that skipped the calibration step wouldn't actually implement the paper's contribution, so it's left out rather than faked. std/ml/ROADMAP.md tracks this.
Structs
| DpmMidpoint | Internal: the order-2 solvers' midpoint sample and noise level (pp: the x0-prediction form), and the update from the midpoint prediction. |
Traits
| DenoiseModel | A denoiser the order-2 solvers call at their midpoint: an eps-prediction model for dpm_solver_2_step, an x0-prediction one for dpm_solver_pp_2_step. Any struct with this method conforms. |
Functions
| edm_karras_sigmas | sigma(i) = (sigmaMax^(1/rho) + i/(numSteps-1)*(sigmaMin^(1/rho)-sigmaMax^(1/rho)))^rho for i in [0,numSteps), followed by sigma(numSteps)=0 (the final, noise-free point). 'outSigmas' must have numSteps+1 slots (caller- allocated) — the standard Karras et al. sigma sequence used by EDM/ k-diffusion-style samplers (e.g. Heun/Euler ancestral, not reimplemented here — this schedule alone is the deferred-scope-item std/ml/ROADMAP.md previously listed). |
| ddpm_linear_schedule | beta_t = betaStart + (betaEnd-betaStart)*t/(numSteps-1), alpha_t = 1-beta_t, alphaBar_t = prod_{s<=t} alpha_s. 'outBeta'/'outAlphaBar' must each have numSteps slots (caller-allocated). |
| ddpm_sampler_step | x_{t-1} = (1/sqrt(1-beta_t)) * (x_t - beta_t/sqrt(1-alphaBar_t) * epsPred) + sqrt(beta_t) * noise Pass a zero tensor for 'noise' to get the deterministic (posterior- mean-only) step — useful for verification, matching how the rest of std/ml avoids baking in an RNG (see rnn.h's header comment). |
| ddim_sampler_step | x0Pred = (x_t - sqrt(1-alphaBar_t)*epsPred) / sqrt(alphaBar_t) x_{t-1} = sqrt(alphaBar_prev)*x0Pred + sqrt(1-alphaBar_prev)*epsPred |
| dpm_solver_1_step | in log-SNR space. lambda_t = log(alpha_t/sigma_t); order-1 (dpm_solver_1_step) is the paper's eq. 4.3 first-order update. Note this is NOT the same update as ddim_sampler_step — that equivalence holds for the data-prediction form below (dpm_solver_pp_1_step), not this eps-prediction one; see diffusion.sc's header comment. |
| dpm_solver_2_step | Order-2 (midpoint variant, r1=1/2): evaluates the eps-model a second time at the log-SNR midpoint between the current and next step for a higher-accuracy update: model.predict(x_s, sigma_s) at the midpoint sample / noise level, returning the predicted eps. |
| dpm_solver_pp_1_step | Same exponential-integrator idea as DPM-Solver, reformulated around a predicted x0 (the denoised sample) instead of a predicted eps (the noise). Order-1 is provably identical to ddim_sampler_step — see diffusion.sc's header comment for the derivation. |
| dpm_solver_pp_2_step | Order-2 midpoint variant, x0-prediction form: model.predict(x_s, sigma_s) at the log-SNR midpoint returns a predicted x0 (not eps). |