DDIM Sampling #
Reverse a finite diffusion schedule with a caller-provided epsilon predictor. The predictor receives both the schedule index and current tensor, so image time channels and other conditioning remain independent of the sampling algorithm.
def
TorchLean.diffusion.ddim
{shape : Shape}
{T : ℕ}
(predict : Fin T → Tensor Float shape → IO (Tensor Float shape))
(alphaBars : Tensor Float [T])
(initial : Tensor Float shape)
(postprocess : Tensor Float shape → Tensor Float shape := fun (reconstruction : Tensor Float shape) =>
reconstruction.clamp (-1) 1)
(denominatorFloor : Float := 1e-12)
(start : Option (Fin T) := none)
:
Run deterministic DDIM updates down through timestep zero.
By default, start at the last schedule entry. Pass start := some index to resume from a
particular timestep. An empty schedule leaves the input unchanged and never calls the predictor.
The schedule coefficients must describe the same forward process used to train predict.
The endpoint convention is cumulative alpha one before the first timestep.
Every update uses the supplied reconstruction postprocessor and denominator floor.