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 -R -K cuda=true build
CUDA_VISIBLE_DEVICES=0 lake -R -K cuda=true exe torchlean diffusion --device cuda \
  --dataset imagenet64 --n-total 800 --steps 1000 --hidden-c 8 --T 100 --beta-end 0.12 \
  --log data/examples/diffusion_trainlog.json \
  --reference-ppm data/examples/diffusion_reference.ppm \
  --noisy-ppm data/examples/diffusion_noisy.ppm \
  --reconstruct-ppm data/examples/diffusion_reconstruct.ppm \
  --sample-ppm data/examples/diffusion_sample.ppm

CIFAR run:

python3 scripts/datasets/download_example_data.py --cifar10
lake -R -K cuda=true exe torchlean diffusion --device cuda --dataset cifar10 --n-total 1 \
  --steps 1 --hidden-c 2 --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.defaultLogPath : System.FilePath

Default JSON loss-curve path for this command.

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

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

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

Cropped CIFAR height for the compact runnable diffusion example.

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

Cropped CIFAR width for the compact runnable diffusion example.

theorem NN.Examples.Models.Generative.Diffusion.instNeZeroNatCifarCropHeight : NeZero cifarCropHeight

theorem NN.Examples.Models.Generative.Diffusion.instNeZeroNatCifarCropWidth : NeZero cifarCropWidth

@[reducible, inline]

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

Clean image batch dimensions $x_0$: NCHW with the fixed command batch size.

@[reducible, inline]

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

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

def NN.Examples.Models.Generative.Diffusion.config (c h w hiddenChannels : ℕ) : TorchLean.nn.models.Diffusion.NoisePredictor.Config 2

Architecture of the epsilon predictor for a particular image size.

A kernel radius of one means every convolution uses a 3 x 3 same-padding kernel.

def NN.Examples.Models.Generative.Diffusion.model (c h w hiddenChannels : ℕ) : TorchLean.nn.Builder (TorchLean.nn.Sequential (input c h w) (output c h w))

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

Every 3×3 convolution preserves the spatial extent, and the residual blocks mix information between neighboring pixels while remaining small enough for the runnable CUDA check.

def NN.Examples.Models.Generative.Diffusion.cifarBatch (batchSample : TorchLean.Sample.Batch Float batchSize RealData.CifarImage RealData.CifarTarget) : Except String (TorchLean.Tensor Float (output RealData.cifarChannels cifarCropHeight cifarCropWidth))

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.imageNet64Batch (batchSample : TorchLean.Sample.Batch Float batchSize RealData.ImageNet64Image RealData.ImageNet64Target) : TorchLean.Tensor Float (output RealData.imagenet64Channels RealData.imagenet64Height RealData.imagenet64Width)

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

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

def NN.Examples.Models.Generative.Diffusion.loadCifar (xPath yPath : System.FilePath) (nRows seed : ℕ) : IO ((count : ℕ) × TorchLean.Tensor Float ((output RealData.cifarChannels cifarCropHeight cifarCropWidth).prependDim count))

Load CIFAR-10 batches as a tensor of clean diffusion minibatches.

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

def NN.Examples.Models.Generative.Diffusion.loadImageNet64 (xPath yPath : System.FilePath) (nRows seed : ℕ) : IO ((count : ℕ) × TorchLean.Tensor Float ((output RealData.imagenet64Channels RealData.imagenet64Height RealData.imagenet64Width).prependDim count))

Load ImageNet64-style batches as a tensor of clean diffusion minibatches.

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.conditionedPredictor {c h w T : ℕ} (predict : TorchLean.Tensor Float (input c h w) → IO (TorchLean.Tensor Float (output c h w))) (index : Fin T) (sample : TorchLean.Tensor Float (output c h w)) : IO (TorchLean.Tensor Float (output c h w))

Adapt image-and-time conditioning to the dataset-independent DDIM sampler.

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

Diffusion command-line options after parsing.

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

• training : TorchLean.CLI.Training.OptimizerOptions

  Optimizer, step, batching, and logging controls.

• schedule : Support.DiffusionScheduleFlags

  Diffusion timestep and beta schedule.

• artifacts : Support.ImageArtifactFlags

  Optional generated and reconstructed image paths.

• hiddenChannels : ℕ

  Hidden channel width of the epsilon predictor.

@[instance_reducible]

instance NN.Examples.Models.Generative.Diffusion.instReprOptions : Repr Options

def NN.Examples.Models.Generative.Diffusion.instReprOptions.repr : Options → ℕ → Std.Format

def NN.Examples.Models.Generative.Diffusion.Training.run {c h w : ℕ} [NeZero c] [NeZero h] [NeZero w] (runtime : TorchLean.Runtime.Config) (loadBatches : IO ((count : ℕ) × TorchLean.Tensor Float ((output c h w).prependDim count))) (config : Options) (schedule : TorchLean.diffusion.Schedule config.schedule.T) : 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.

def NN.Examples.Models.Generative.Diffusion.train {c h w : ℕ} [NeZero c] [NeZero h] [NeZero w] (runtime : TorchLean.Runtime.Config) (loadBatches : IO ((count : ℕ) × TorchLean.Tensor Float ((output c h w).prependDim count))) (config : Options) : IO Runtime.Training.Curve

Train the diffusion example after rejecting an empty timestep schedule.

def NN.Examples.Models.Generative.Diffusion.Options.parse (args : List String) : Except String (Options × 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.Options.sourceNotes (datasetName : String) (data : Support.NpyDataFlags) : Array String

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

def NN.Examples.Models.Generative.Diffusion.Options.logNotes (config : Options) (dataset : String) (runtime : TorchLean.Runtime.Config) (sourceNotes : Array String := #[]) : Array String

TrainLog note fields shared by all diffusion dataset branches.

def NN.Examples.Models.Generative.Diffusion.writeLog (log : Runtime.Training.LogDestination) (dataset : String) (sourceNotes : Array String) (config : Options) (runtime : TorchLean.Runtime.Config) (curve : Runtime.Training.Curve) : IO Unit

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

def NN.Examples.Models.Generative.Diffusion.runDataset {c h w : ℕ} [NeZero c] [NeZero h] [NeZero w] (runtime : TorchLean.Runtime.Config) (args : List String) (datasetName : String) (parseData : List String → Except String (Support.NpyDataFlags × List String)) (loadBatches : System.FilePath → System.FilePath → ℕ → ℕ → IO ((count : ℕ) × TorchLean.Tensor Float ((output c h w).prependDim count))) : 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 a positive hidden-channel count, train the epsilon predictor, then write the same curve log.

def NN.Examples.Models.Generative.Diffusion.runImageNet64 (runtime : TorchLean.Runtime.Config) (args : List String) : IO Unit

Run the ImageNet64 branch with shape-specialized model construction.

def NN.Examples.Models.Generative.Diffusion.runCifar10 (runtime : TorchLean.Runtime.Config) (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/typed-graph settings first; the remaining arguments select the dataset branch and diffusion training configuration.