Diffusion Training Example

Runnable torchlean diffusion example.

This is the maintained diffusion command. It supports two real-data modes:

• --dataset imagenet64 (default): user-provided ImageNet/Imagenette/Tiny-ImageNet-style images converted to (N,3,64,64) .npy tensors.
• --dataset cifar10: prepared CIFAR-10 (N,3,32,32) arrays.

The command is one public entrypoint, but the implementation keeps separate typed branches because Lean tracks image height and width in the tensor type.

Why unconditional samples are still rough

The default epsilon predictor is a small same-resolution residual CNN with a broadcast time channel. That is enough to validate real image loading, CUDA training, logging, reconstruction diagnostics, and DDIM replay from Lean. It is deliberately not advertised as a high-fidelity image generator: good unconditional samples need a full U-Net with multiscale skips, richer timestep embeddings, EMA, more training, more timesteps, and runtime support that avoids eager-autograd buffer blow-up for wider models.

Examples

Prepare ImageNet-style data:

python3 scripts/datasets/torchlean_data_convert.py image-folder \
  --input /path/to/imagenet/train \
  --x-output data/real/imagenet64/imagenet64_train_X.npy \
  --y-output data/real/imagenet64/imagenet64_train_y.npy \
  --height 64 --width 64 --labels-from-dirs --limit 800

Train on ImageNet64 and save visual artifacts:

lake build -R -K cuda=true
CUDA_VISIBLE_DEVICES=0 lake exe -K cuda=true torchlean diffusion --cuda --fast-kernels \
  --dataset imagenet64 --n-total 800 --steps 1000 --hidden-c 8 --T 100 --beta-end 0.12 \
  --reference-ppm data/model_zoo/diffusion_reference.ppm \
  --noisy-ppm data/model_zoo/diffusion_noisy.ppm \
  --reconstruct-ppm data/model_zoo/diffusion_reconstruct.ppm \
  --sample-ppm data/model_zoo/diffusion_sample.ppm

CIFAR smoke path:

python3 scripts/datasets/download_example_data.py --cifar10
lake exe torchlean diffusion --dataset cifar10 --cuda --fast-kernels --steps 200

def NN.Examples.Models.Generative.Diffusion.exeName : String

def NN.Examples.Models.Generative.Diffusion.defaultLogJson : System.FilePath

def NN.Examples.Models.Generative.Diffusion.batch : ℕ

@[reducible, inline]

abbrev NN.Examples.Models.Generative.Diffusion.x0Shape (c h w : ℕ) : Shape

@[reducible, inline]

abbrev NN.Examples.Models.Generative.Diffusion.xInShape (c h w : ℕ) : Shape

def NN.Examples.Models.Generative.Diffusion.cfgFor (c h w hiddenC : ℕ) : API.nn.models.EpsConvNetConfig

def NN.Examples.Models.Generative.Diffusion.mkModel (c h w hiddenC : ℕ) [NeZero c] [NeZero h] [NeZero w] (h_hiddenC : hiddenC ≠ 0) : API.nn.M (API.nn.Sequential (xInShape c h w) (x0Shape c h w))

Build the default epsilon predictor for a specific typed image shape.

We use the residual CNN from NN.API.Models.Diffusion: it is still small enough for tutorial-scale CUDA runs, but the skip paths train much better than the plain convolution chain. The plain epsConvNet remains in the API as the smaller baseline; this example uses the residual default so the documented command matches the maintained training path.

def NN.Examples.Models.Generative.Diffusion.toDiffusionRange {c h w : ℕ} (x01 : Spec.Tensor Float (x0Shape c h w)) : Spec.Tensor Float (x0Shape c h w)

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

The input is already NCHW because the dataset converter and RealData loaders enforce that layout.

def NN.Examples.Models.Generative.Diffusion.cifarBatchX0 (batchSample : API.sample.Batch Float batch RealData.CifarImage RealData.CifarTarget) : Spec.Tensor Float (x0Shape RealData.cifarChannels RealData.cifarHeight RealData.cifarWidth)

def NN.Examples.Models.Generative.Diffusion.imageNet64BatchX0 (batchSample : API.sample.Batch Float batch RealData.ImageNet64Image RealData.ImageNet64Target) : Spec.Tensor Float (x0Shape RealData.imagenet64Channels RealData.imagenet64Height RealData.imagenet64Width)

def NN.Examples.Models.Generative.Diffusion.loadCifarX0Batches (xPath yPath : System.FilePath) (nRows seed : ℕ) : IO (List (Spec.Tensor Float (x0Shape RealData.cifarChannels RealData.cifarHeight RealData.cifarWidth)))

def NN.Examples.Models.Generative.Diffusion.loadImageNet64X0Batches (xPath yPath : System.FilePath) (nRows seed : ℕ) : IO (List (Spec.Tensor Float (x0Shape RealData.imagenet64Channels RealData.imagenet64Height RealData.imagenet64Width)))

def NN.Examples.Models.Generative.Diffusion.randomEps {c h w : ℕ} (seed step : ℕ) : Spec.Tensor Float (x0Shape c h w)

def NN.Examples.Models.Generative.Diffusion.mkNoisedSample {c h w : ℕ} (alphaBars : Array Float) (T : ℕ) (x0 : Spec.Tensor Float (x0Shape c h w)) (seed step : ℕ) : API.sample.Supervised Float (xInShape c h w) (x0Shape c h w)

def NN.Examples.Models.Generative.Diffusion.reverseDdim {c h w : ℕ} (opts : Runtime.Autograd.Torch.Options) (model : API.nn.Sequential (xInShape c h w) (x0Shape c h w)) (params : Runtime.Autograd.Torch.ParamList Float (Runtime.Autograd.TorchLean.NN.Seq.paramShapes model)) (alphaBars : Array Float) (T : ℕ) (xStart : Spec.Tensor Float (x0Shape c h w)) : IO (Spec.Tensor Float (x0Shape c h w))

def NN.Examples.Models.Generative.Diffusion.reverseDdimFrom {c h w : ℕ} (opts : Runtime.Autograd.Torch.Options) (model : API.nn.Sequential (xInShape c h w) (x0Shape c h w)) (params : Runtime.Autograd.Torch.ParamList Float (Runtime.Autograd.TorchLean.NN.Seq.paramShapes model)) (alphaBars : Array Float) (T tStart : ℕ) (xStart : Spec.Tensor Float (x0Shape c h w)) : IO (Spec.Tensor Float (x0Shape c h w))

Reverse DDIM from a chosen timestep for reconstruction diagnostics.

This is intentionally separate from unconditional sampling. It lets us corrupt a real image to a moderate timestep, denoise from there, and check whether reconstruction improves over the noisy input.

structure NN.Examples.Models.Generative.Diffusion.TrainConfig : Type

• steps : ℕ
• logEvery : ℕ
• lr : Float
• T : ℕ
• hiddenC : ℕ
• betaStart : Float
• betaEnd : Float
• reconstructStep? : Option ℕ
• samplePpm? : Option System.FilePath
• referencePpm? : Option System.FilePath
• noisyPpm? : Option System.FilePath
• reconstructPpm? : Option System.FilePath

def NN.Examples.Models.Generative.Diffusion.trainCurveFloat {c h w : ℕ} [NeZero c] [NeZero h] [NeZero w] (opts : Runtime.Autograd.Torch.Options) (loadBatches : IO (List (Spec.Tensor Float (x0Shape c h w)))) (cfg : TrainConfig) (h_hiddenC : cfg.hiddenC ≠ 0) : IO Runtime.Training.Curve

Shared training loop for both CIFAR-10 and ImageNet64 branches.

The loop optimizes epsilon prediction and can emit four visual artifacts:

• reference-ppm: clean evaluation image,
• noisy-ppm: clean image after forward diffusion to reconstruct-step,
• reconstruct-ppm: DDIM denoising from that timestep,
• sample-ppm: unconditional DDIM sample from Gaussian noise.

inductive NN.Examples.Models.Generative.Diffusion.DatasetChoice : Type

• imagenet64 : DatasetChoice
• cifar10 : DatasetChoice

def NN.Examples.Models.Generative.Diffusion.instReprDatasetChoice.repr : DatasetChoice → ℕ → Std.Format

@[implicit_reducible]

instance NN.Examples.Models.Generative.Diffusion.instReprDatasetChoice : Repr DatasetChoice

def NN.Examples.Models.Generative.Diffusion.instBEqDatasetChoice.beq : DatasetChoice → DatasetChoice → Bool

@[implicit_reducible]

instance NN.Examples.Models.Generative.Diffusion.instBEqDatasetChoice : BEq DatasetChoice

def NN.Examples.Models.Generative.Diffusion.DatasetChoice.parse (args : List String) : Except String (DatasetChoice × List String)

def NN.Examples.Models.Generative.Diffusion.parseTrainConfig (args : List String) : Except String (TrainConfig × Runtime.Training.LogDestination × List String)

def NN.Examples.Models.Generative.Diffusion.writeTrainingLog (log : Runtime.Training.LogDestination) (dataset : String) (sourceNotes : Array String) (cfg : TrainConfig) (opts : Runtime.Autograd.Torch.Options) (curve : Runtime.Training.Curve) : IO Unit

def NN.Examples.Models.Generative.Diffusion.runImageNet64 (opts : Runtime.Autograd.Torch.Options) (args : List String) : IO Unit

def NN.Examples.Models.Generative.Diffusion.runCifar10 (opts : Runtime.Autograd.Torch.Options) (args : List String) : IO Unit

def NN.Examples.Models.Generative.Diffusion.main (args : List String) : IO UInt32