DDPM from scratch — Oxford Flowers 102 at 64×64

A complete Denoising Diffusion Probabilistic Model (Ho et al. 2020) implemented from scratch in PyTorch — no diffusers, no pretrained pipelines. Forward process, noise schedules, U-Net, training objective, and both DDPM and DDIM samplers are all written out and commented, math first. Everything runs on free tools: the dataset auto-downloads, training fits in a single free Colab/Kaggle GPU session, and logging is plain CSV + matplotlib (no paid tracking).

Why Oxford Flowers 102?

• Free and frictionless — auto-downloads through torchvision.datasets.Flowers102 (~330 MB, no Kaggle API key or account needed).
• Small but not tiny — ~8,200 images: enough diversity to learn from, small enough that an epoch takes seconds on a free GPU.
• Niche and visually diagnosable — fine-grained flower species with saturated colors and strong radial structure. At 64×64 you can tell at a glance whether the model is learning (petal shapes, green backgrounds, flower-centered composition) or producing mush. That makes it ideal for a project whose acceptance test is "does the sample grid look like the training domain?"

Project structure

File	What it does
dataset.py	Loads Flowers 102, center-crops to 64×64, scales to [-1, 1]
diffusion.py	Noise schedules (linear + cosine), closed-form q(x_t|x_0), DDPM ancestral sampling, DDIM sampling
unet.py	~19M-param U-Net: residual blocks, sinusoidal timestep embeddings injected into every block, self-attention at 16×16 and the 8×8 bottleneck
train.py	Training loop (Algorithm 1): random t, add noise, MSE on predicted ε. AMP, EMA weights, checkpointing/resume, CSV loss log + matplotlib curve, periodic sample grids
sample.py	Generates a 4×4 grid PNG from a checkpoint (the acceptance test)
app.py	Gradio demo: Generate button + denoising-steps slider (speed/quality tradeoff)
fid.py	Optional FID score via torchmetrics (free)

The mechanics (what's implemented from scratch)

1. Forward process — cosine β-schedule (linear also available) and the closed-form jump to any noise level: x_t = √ᾱ_t·x₀ + √(1−ᾱ_t)·ε, ε ~ N(0, I). See diffusion.py:q_sample.
2. Training objective — ε-prediction with the simplified loss L = ‖ε − ε_θ(x_t, t)‖² at uniformly random t. See diffusion.py:p_losses.
3. Reverse process — DDPM ancestral sampling with the posterior mean μ_θ = (x_t − β_t/√(1−ᾱ_t)·ε_θ)/√α_t and variance β̃_t (p_sample_loop, 1000 network calls), plus DDIM for 10–100-step sampling (ddim_sample).
4. U-Net — residual blocks with time-embedding injection, sinusoidal timestep embeddings, self-attention at low resolutions. See unet.py.

Quickstart

pip install -r requirements.txt

# Train (~2.5–3 h on a free Colab T4; see budget below)
python train.py --steps 30000

# Acceptance test: 4x4 grid of generated flowers -> results/samples_grid.png
python sample.py --ckpt checkpoints/final.pt

# Interactive demo with a denoising-steps slider
python app.py --ckpt checkpoints/final.pt        # add --share on Colab

On Colab: Runtime → Change runtime type → T4 GPU, then

!git clone <this repo> && cd <repo> && pip install -r requirements.txt
!python train.py --steps 30000

If the session disconnects, resume with python train.py --steps 30000 --resume checkpoints/last.pt (keep checkpoints/ on Drive if you want belt-and-braces safety).

Training budget (free tier)

Designed to finish inside one free GPU session:

Knob	Value	Why
Resolution	64×64	4× cheaper than 128×128 per step
U-Net	~19M params (base 64, mults 1-2-2-4)	A fraction of the paper's model; plenty for 64×64
Batch / steps	128 / 30,000 (~470 epochs)	Fits T4 memory with AMP
Precision	AMP (float16)	~2× throughput on T4

Expected wall-clock: ~2.5–3 hours on a Colab free-tier T4 (~3 it/s), faster on Kaggle's P100/T4×2. Recognizable flowers appear by ~10k steps (~1 h); 30k steps gives clearly better texture and color coherence. For a quicker end-to-end sanity run, --image-size 32 --steps 10000 finishes in under an hour.

Training on a local GPU instead

Works fine on any NVIDIA card — two gotchas:

• Install the CUDA build of PyTorch. On Windows, plain pip install torch gives a CPU-only wheel (x.y.z+cpu) and training will say "no GPU found" even though you have one. Instead: pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
• Match the batch size to your VRAM. The default --batch-size 128 is sized for a 16 GB T4 and peaks at ~11 GB. On an 8 GB card (e.g. RTX 3060 Ti) the driver silently spills into system RAM and training runs ~10× slower — use --batch-size 64 (~5.7 GB peak, measured 3.7 it/s on a 3060 Ti). python train.py --steps 40000 --batch-size 64 finishes in ~3 hours and sees a comparable number of images to the Colab recipe.

Evaluation

• Early vs late checkpoints — train.py saves a 16-image grid to samples/step_XXXXXX.png every 2,500 steps. Put step_002500.png (colored blobs) next to step_030000.png (recognizable flowers) to see learning progress directly.
• Loss curve — logged to results/loss_log.csv and plotted to results/loss_curve.png with matplotlib. Expect a fast drop to ~0.05 then a slow grind downward — sample quality keeps improving even after the loss looks flat (most of the loss mass is at high-noise timesteps).
• FID (optional) — python fid.py --ckpt checkpoints/final.pt computes FID against the real dataset using torchmetrics' free implementation. Skip it if GPU quota is tight; the grids are the primary evidence. This checkpoint (40k steps): FID 48.95 (2,048 samples/side, DDIM 100-step) — see MODEL_CARD.md for the full writeup.

Acceptance test

python sample.py --ckpt checkpoints/final.pt --n 16 --out results/samples_grid.png

produces a 4×4 grid that should show clearly flower-like images — petal structures, flower-appropriate palettes (pinks/purples/yellows), green/dark backgrounds — not noise.

Trained for 40,000 steps on an RTX 3060 Ti; results below. See MODEL_CARD.md for full training/eval details, skills demonstrated, and FID score.

Design notes

• Cosine schedule over linear: at 64×64 the linear schedule destroys nearly all signal in the first half of the steps; cosine spreads information destruction evenly and trains noticeably better (Nichol & Dhariwal 2021).
• EMA weights (decay 0.9995) are used for all sampling — raw weights jitter with each minibatch; the EMA average samples much cleaner.
• DDIM by default for sampling — deterministic (η=0) DDIM at 100 steps is ~10× faster than the full 1000-step ancestral sampler with minimal quality loss; the Gradio slider lets you explore this tradeoff live.