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 modest

The default epsilon predictor is a compact 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. High-fidelity unconditional samples require more machinery: 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 \
  --log data/model_zoo/diffusion_trainlog.json \
  --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 run:

python3 scripts/datasets/download_example_data.py --cifar10
lake exe -K cuda=true torchlean diffusion --cuda --dataset cifar10 --n-total 1 --steps 1 --hidden-c 1 --T 2

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

CLI subcommand name used in terminal banners and error messages.

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

Default JSON loss-curve path for this command.

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

Static minibatch size used by both CIFAR-10 and ImageNet64 typed branches.

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

Cropped CIFAR height for the compact runnable diffusion example.

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

Cropped CIFAR width for the compact runnable diffusion example.

theorem NN.Examples.Models.Generative.Diffusion.instNeZeroNatCifarTinyH : NeZero cifarTinyH

theorem NN.Examples.Models.Generative.Diffusion.instNeZeroNatCifarTinyW : NeZero cifarTinyW

@[reducible, inline]

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

Clean image batch shape x₀: NCHW with the fixed command batch size.

@[reducible, inline]

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

Epsilon-model input shape: image channels plus one broadcast timestep channel.

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

Shape-level configuration for the epsilon predictor.

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

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

We use the plain compact epsilon CNN from the public diffusion model API. The residual denoiser stays available in the API for larger opt-in experiments, but the runnable command should remain a quick CUDA quick check.

def NN.Examples.Models.Generative.Diffusion.cifarBatchX0 (batchSample : TorchLean.Sample.Batch Float batch RealData.CifarImage RealData.CifarTarget) : TorchLean.Tensor.T Float (x0Shape RealData.cifarChannels cifarTinyH cifarTinyW)

Convert one typed CIFAR minibatch into diffusion-space clean images.

The loader returns images in [0,1]; diffusion training uses [-1,1], so this function performs the range conversion after Lean has established the CIFAR NCHW shape.

def NN.Examples.Models.Generative.Diffusion.imageNet64BatchX0 (batchSample : TorchLean.Sample.Batch Float batch RealData.ImageNet64Image RealData.ImageNet64Target) : TorchLean.Tensor.T Float (x0Shape RealData.imagenet64Channels RealData.imagenet64Height RealData.imagenet64Width)

Convert one typed ImageNet64 minibatch into diffusion-space clean images.

This mirrors cifarBatchX0 but keeps the ImageNet64 height/width/channel constants in the type.

def NN.Examples.Models.Generative.Diffusion.loadCifarX0Batches (xPath yPath : System.FilePath) (nRows seed : ℕ) : IO (List (TorchLean.Tensor.T Float (x0Shape RealData.cifarChannels cifarTinyH cifarTinyW)))

Load CIFAR-10 batches as a finite list of clean diffusion images.

The function validates the .npy paths, builds a typed Data.batchLoader, drops incomplete final batches, and returns NCHW tensors already mapped into [-1,1].

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

Load ImageNet64-style batches as a finite list of clean diffusion images.

The converter accepts ImageNet/Imagenette/Tiny-ImageNet-style folders ahead of time; this Lean path only consumes the prepared .npy arrays and keeps the tensor shapes explicit.

def NN.Examples.Models.Generative.Diffusion.reverseDdim {c h w : ℕ} (predict : TorchLean.Tensor.T Float (xInShape c h w) → IO (TorchLean.Tensor.T Float (x0Shape c h w))) (alphaBars : Array Float) (T : ℕ) (xStart : TorchLean.Tensor.T Float (x0Shape c h w)) : IO (TorchLean.Tensor.T Float (x0Shape c h w))

Run deterministic DDIM reverse steps from a starting noisy image.

This is used for unconditional sample artifacts: start from Gaussian noise, repeatedly ask the model for ε̂, and apply the DDIM previous-step formula.

def NN.Examples.Models.Generative.Diffusion.reverseDdimFrom {c h w : ℕ} (predict : TorchLean.Tensor.T Float (xInShape c h w) → IO (TorchLean.Tensor.T Float (x0Shape c h w))) (alphaBars : Array Float) (T tStart : ℕ) (xStart : TorchLean.Tensor.T Float (x0Shape c h w)) : IO (TorchLean.Tensor.T Float (x0Shape c h w))

Reverse DDIM from a chosen timestep for reconstruction diagnostics.

This reconstruction path is separate from unconditional sampling. It corrupts a real image to a moderate timestep, denoises from there, and checks whether reconstruction improves over the noisy input.

structure NN.Examples.Models.Generative.Diffusion.DiffusionOptionsextends NN.API.Common.ModelTrainFlags, NN.API.Common.DiffusionScheduleFlags, NN.API.Common.ImageArtifactFlags : Type

Diffusion command-line options after parsing.

The inherited pieces make the CLI shape explicit: ordinary training flags come from ModelZoo, diffusion math lives in ModelZoo.DiffusionScheduleFlags, visual outputs live in ModelZoo.ImageArtifactFlags, and the epsilon-network width is the model-specific knob.

• steps : ℕ
• batchSize : ℕ
• log : Runtime.Training.LogDestination
• logPath : System.FilePath
• cudaMemWatch : ℕ
• lr : Float
• T : ℕ
• betaStart : Float
• betaEnd : Float
• reconstructStep? : Option ℕ
• samplePpm? : Option System.FilePath
• referencePpm? : Option System.FilePath
• noisyPpm? : Option System.FilePath
• reconstructPpm? : Option System.FilePath
• hiddenC : ℕ

  Hidden channel width of the epsilon predictor.

def NN.Examples.Models.Generative.Diffusion.instReprDiffusionOptions.repr : DiffusionOptions → ℕ → Std.Format

@[implicit_reducible]

instance NN.Examples.Models.Generative.Diffusion.instReprDiffusionOptions : Repr DiffusionOptions

def NN.Examples.Models.Generative.Diffusion.trainCurveFloat {c h w : ℕ} [NeZero c] [NeZero h] [NeZero w] (opts : TorchLean.Options) (loadBatches : IO (List (TorchLean.Tensor.T Float (x0Shape c h w)))) (cfg : DiffusionOptions) (h_hiddenC : cfg.hiddenC ≠ 0) : IO TorchLean.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.

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

Parse diffusion-specific training flags after runtime/device flags and dataset flags.

The shared parser handles --steps, --log, and --cuda-mem-watch; this parser handles diffusion schedule parameters, model width, and optional PPM artifact paths.

def NN.Examples.Models.Generative.Diffusion.DiffusionOptions.ensureValid (cfg : DiffusionOptions) : IO Unit

Reject unsupported diffusion hyperparameters before shape-specialized execution begins.

def NN.Examples.Models.Generative.Diffusion.DiffusionOptions.sourceNotes (datasetName : String) (data : TorchLean.ModelZoo.NpyDataFlags) : Array String

Dataset/source note fields shared by the CIFAR-10 and ImageNet64 branches.

def NN.Examples.Models.Generative.Diffusion.DiffusionOptions.logNotes (cfg : DiffusionOptions) (dataset : String) (opts : TorchLean.Options) (sourceNotes : Array String := #[]) : Array String

TrainLog note fields shared by all diffusion dataset branches.

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

Write the diffusion loss curve plus dataset, schedule, model, and artifact metadata.

def NN.Examples.Models.Generative.Diffusion.runTypedDataset {c h w : ℕ} [NeZero c] [NeZero h] [NeZero w] (opts : TorchLean.Options) (args : List String) (datasetName : String) (parseData : List String → Except String (TorchLean.ModelZoo.NpyDataFlags × List String)) (loadBatches : System.FilePath → System.FilePath → ℕ → ℕ → IO (List (TorchLean.Tensor.T Float (x0Shape c h w)))) : IO Unit

Run one typed diffusion dataset branch.

The CIFAR-10 and ImageNet64 commands differ in their shape-level loader and default .npy paths, but after parsing those inputs they follow the same command flow: parse training flags, reject unused args, require hiddenC > 0, train the epsilon predictor, then write the same curve log.

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

Run the ImageNet64 branch with shape-specialized model construction.

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

Run the CIFAR-10 branch with shape-specialized model construction.

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

Executable entrypoint for diffusion training.

The runtime parser selects CPU/CUDA and eager/compiled settings first; the remaining arguments select the dataset branch and diffusion training configuration.