Diffusion Model Helpers (API)

Config-style diffusion model constructors plus reusable, dataset-independent DDPM/DDIM helpers.

The runnable examples decide where data comes from (CIFAR-10, ImageNet-style folders, synthetic artifacts). The definitions here are shape-parametric and can be reused by tests, examples, and future proof-facing specifications.

structure NN.API.nn.models.EpsConvNetConfig : Type

Configuration for a minimal epsilon-predictor conv net.

• batch : ℕ
• dataC : ℕ

  Data channels (e.g. 3 for RGB). The model input has one extra channel for time.

• h : ℕ
• w : ℕ
• hiddenC : ℕ

  Hidden channel width.

@[reducible, inline]

abbrev NN.API.nn.models.epsConvNetInShape (cfg : EpsConvNetConfig) : Shape

Epsilon-predictor input shape, with one extra channel carrying the diffusion time.

@[reducible, inline]

abbrev NN.API.nn.models.epsConvNetOutShape (cfg : EpsConvNetConfig) : Shape

Epsilon-predictor output shape matching the denoised data channels.

def NN.API.nn.models.epsConvNet (cfg : EpsConvNetConfig) (h_batch : cfg.batch ≠ 0 := by decide) (h_dataC : cfg.dataC ≠ 0 := by decide) (h_inC : cfg.dataC + 1 ≠ 0 := by decide) (h_h : cfg.h ≠ 0 := by decide) (h_w : cfg.w ≠ 0 := by decide) (h_hiddenC : cfg.hiddenC ≠ 0 := by decide) : M (Sequential (epsConvNetInShape cfg) (epsConvNetOutShape cfg))

Build a minimal epsilon-predictor conv net: conv -> relu -> conv -> relu -> conv -> relu -> conv.

This stays compact enough for the eager CUDA example while giving the CIFAR trainer more denoising capacity than a bare two-layer network.

def NN.API.nn.models.epsResidualConvNet (cfg : EpsConvNetConfig) (h_batch : cfg.batch ≠ 0 := by decide) (h_dataC : cfg.dataC ≠ 0 := by decide) (h_inC : cfg.dataC + 1 ≠ 0 := by decide) (h_h : cfg.h ≠ 0 := by decide) (h_w : cfg.w ≠ 0 := by decide) (h_hiddenC : cfg.hiddenC ≠ 0 := by decide) : M (Sequential (epsConvNetInShape cfg) (epsConvNetOutShape cfg))

Build a stronger same-resolution residual epsilon predictor.

Architecture:

stem conv -> relu -> residual block -> relu -> residual block -> relu -> output conv

Each residual block has shape hiddenC×H×W -> hiddenC×H×W and computes x + conv(relu(conv(x))). This compact residual denoiser omits U-Net downsampling, upsampling, and multi-scale skip concatenation. It is still a useful compact architecture because residual paths make the denoising problem much easier than a plain conv chain while staying within the eager CUDA memory envelope used by examples.

def NN.API.diffusion.toMinusOneOne {batch c h w : ℕ} (x01 : Spec.Tensor Float (Tensor.Shape.NCHW batch c h w)) : Spec.Tensor Float (Tensor.Shape.NCHW batch c h w)

Map image tensors from [0,1] into the standard diffusion training range [-1,1].

def NN.API.diffusion.randomEps {batch c h w : ℕ} (seed step : ℕ) : Spec.Tensor Float (Tensor.Shape.NCHW batch c h w)

Deterministic Gaussian epsilon tensor for an NCHW diffusion shape.

The (seed, step) pair is turned into the runtime RNG key, so examples and artifact generation can reproduce the same noising path without ambient randomness.

def NN.API.diffusion.linearBeta (T : ℕ) (betaStart betaEnd : Float) (t : ℕ) : Float

Linear beta schedule value at timestep t.

def NN.API.diffusion.alphaBarsLinear (T : ℕ) (betaStart betaEnd : Float) : Array Float

Compute cumulative products ᾱ_t = ∏_{s≤t} (1 - β_s) for a linear beta schedule.

These values connect clean data x₀, noised data x_t, and the epsilon target used by DDPM-style training.

def NN.API.diffusion.appendTimeChannel {batch c h w : ℕ} (x : Spec.Tensor Float (Tensor.Shape.NCHW batch c h w)) (tNorm : Float) : Spec.Tensor Float (Tensor.Shape.NCHW batch (c + 1) h w)

Append a constant time channel to an NCHW image batch.

The epsilon predictor consumes (data channels + 1) channels: noisy image channels plus a scalar timestep broadcast over spatial positions.

def NN.API.diffusion.noisedSampleFromEps {batch c h w : ℕ} (alphaBars : Array Float) (T : ℕ) (x0 eps : Spec.Tensor Float (Tensor.Shape.NCHW batch c h w)) (step : ℕ) : SupervisedSample Float (Tensor.Shape.NCHW batch (c + 1) h w) (Tensor.Shape.NCHW batch c h w)

Build an epsilon-prediction training sample from explicit noise.

The caller supplies eps, usually from the runtime RNG. Keeping randomness outside this helper makes the transformation reusable:

x_t = sqrt(ᾱ_t) * x₀ + sqrt(1 - ᾱ_t) * eps, target eps.

def NN.API.diffusion.noisedSample {batch c h w : ℕ} (alphaBars : Array Float) (T : ℕ) (x0 : Spec.Tensor Float (Tensor.Shape.NCHW batch c h w)) (seed step : ℕ) : SupervisedSample Float (Tensor.Shape.NCHW batch (c + 1) h w) (Tensor.Shape.NCHW batch c h w)

Build a deterministic epsilon-prediction training sample.

This is the common DDPM training step used by examples: draw reproducible Gaussian noise from (seed, step), corrupt x₀, and use that same noise as the target.

def NN.API.diffusion.ddimPrev {batch c h w : ℕ} (abPrev ab : Float) (x_t epsHat : Spec.Tensor Float (Tensor.Shape.NCHW batch c h w)) : Spec.Tensor Float (Tensor.Shape.NCHW batch c h w)

One deterministic DDIM reverse update (η = 0).

Given x_t, predicted epsilon, and adjacent schedule values, this estimates x₀ and remixes it to the previous timestep.

We clamp the intermediate x₀ estimate to the training image range [-1,1]. This is the standard "clipped denoised" stabilizer used by many DDPM/DDIM samplers: without it, a compact model can drive one color channel far outside the data range and the final PPM exporter merely clips the damage into saturated color blobs.

def NN.API.diffusion.writeFirstRgbNchwPpm {batch c h w : ℕ} (path : System.FilePath) (x : Spec.Tensor Float (Tensor.Shape.NCHW batch c h w)) : IO Unit

Write the first image in an RGB NCHW batch as an ASCII PPM.

This dependency-free writer emits portable image artifacts for examples and rendered diagnostics.