This post documents a small, didactic project I built to understand diffusion models end-to-end. It’s a from-scratch, minimal implementation of schedules, samplers, and a compact U-Net that learns to draw a valid path on a 2D maze by diffusing only the path channel while conditioning on fixed walls/start/end.
System at a glance
- Data contract. On disk: NumPy arrays
[4, H, W]in[0,1]ordered[walls, start, end, solution]. The loader emits:-
condition=[3, H, W]in[0,1](walls/start/end) -
target=[1, H, W]scaled to[-1,1](path to diffuse)
-
- Model. Compact U-Net (GroupNorm + SiLU, residual conv blocks) with sinusoidal time embeddings.
in_channels=4(noisy path + 3 condition),out_channels=1(predicted noise ε). - Diffusion core. Linear and cosine schedules implemented in-house with precomputed
α_t,\bar{α}_t, posterior variance, forward factors, and DDIM coefficients for a uniform timestep ladder. - Samplers. DDPM (ancestral) and DDIM (fast, η=0 deterministic or η>0 stochastic), with classifier-free guidance via two passes (cond/uncond) and linear mixing.
- Training. Sample
t, add noise only to the path channel, concatenate condition (optionally dropped for CFG training), predict ε with MSE loss. AdamW + grad clip + EMA. - Inference. Solver prefers EMA weights, uses cosine schedule and DDIM by default, expects images with black=walls, green=start, red=end, white=free. Returns raw
[-1,1]and thresholded[0,1]path maps, plus validation. - Evaluation. Validity (no wall overlap; start/end on path; BFS 4-connectivity), plus IoU/F1 vs. ground truth and optimality gap by path length.
Note on channels: Inference must use the same in_channels as training. If a checkpoint was trained with, say, 7 channels (e.g., extra aux features), construct the inference U-Net with in_channels=7 or weights won’t load.How it works (briefly)How it works (brief)How it works
As a compact learning project, this codebase made diffusion mechanics tangible: ε-prediction training, schedule design, sampler behavior, and guidance trade-offs. If you run the quick experiments above, you’ll see how each choice-cosine vs. linear, steps, CFG scale, η-maps to validity, quality, and speed in a way that’s easy to reason about and extend.