mini-diffusion — a denoising diffusion model trained from scratch, overnight, on one laptop

A DDPM written from scratch in plain PyTorch and trained in a single night on a MacBook Pro (Apple M5 Pro, 20-core GPU, 24 GB unified memory, macOS 26.6) through the MPS backend. No cloud, no CUDA, no rented A100 — the entire run is nine hours on a machine you could carry to a café.

Every number on this page was measured on that machine, not quoted from a paper. Its companion project, a 30M-parameter transformer built the same way, is mini-llm.

64 faces that do not exist. Unconditional, DDIM-50, EMA weights, one fixed seed. FID-tv 30.12.

Lessons 9 and 10 of part 2 call Stable Diffusion through diffusers. That teaches the shape of the pipeline but not the thing itself. Here the U-Net, the noise schedule, the four samplers and the training loop are written out, and diffusers is never imported. What runs overnight is the same algorithm as DDPM (Ho et al. 2020) and DDIM (Song et al. 2021), at a size that fits one night on one Mac.

The same six attributes at five guidance scales — identical noise and identical attributes in every row, w rising 1.0 → 1.5 → 2.0 → 3.0 → 5.0 top to bottom. This is the mechanism behind the guidance-scale slider in Stable Diffusion, implemented in about ten lines:

Results, figures and the FID curve: RESULTS.md (generated automatically). How the samples evolved: samples_progress.md.

Everything is plain PyTorch

No diffusers, no accelerate, no lightning, no miniai. One dependency was added to the environment — pyarrow, because HuggingFace ships CelebA as parquet and reading a columnar file format is not what "from scratch" is about. FID uses the Inception network that ships with torchvision, since training an Inception classifier is a different project.

Setup

Python 3.11, and a GPU that PyTorch can reach. Everything below was run on Apple Silicon through MPS; CUDA should work unchanged, since nothing in the code is Metal-specific — but the memory guard rails in train.py and bench.py are calibrated against torch.mps.recommended_max_memory() and would need adjusting.

python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

Just want to generate faces?

The trained weights are on the Hub, so you can skip both the dataset and the overnight run:

python -c "
from huggingface_hub import hf_hub_download
import shutil, os
os.makedirs('ckpt', exist_ok=True)
shutil.copy(hf_hub_download('vous99/mini-diffusion', 'best.pt'), 'ckpt/best.pt')
"
python sample.py

Weights: https://huggingface.co/vous99/mini-diffusion (142 MB). Sampling needs no dataset — only retraining does.

What it cannot do

The grid at the top is the good half. The honest version:

• Roughly 15–20% of samples collapse into noise. This is the dominant remaining defect and it is visible in the very first grid on this page — count the broken frames.
• Colour casts. Some faces come out tinted green or blue. The cause is not guidance (verified: the mean RGB barely moves as w goes 1.0 → 3.0) but the schedule: even cosine ends at abar_T ≈ 2.4e-9 rather than 0, so the model never trains on pure noise yet is handed pure noise at sampling time. With epsilon-prediction that surfaces as a per-image brightness bias. It faded substantially over training but did not disappear; v-prediction with zero terminal SNR is the proper fix.
• 64×64 only. 128px would quadruple the FLOPs at the most expensive level for the same parameter count — about 4 epochs instead of 17, and strictly worse samples.
• No text conditioning. The condition is six binary flags, not a prompt. There is no text encoder and none is planned: cross-attention over six booleans would be ceremony.
• Aligned frontal portraits only. CelebA is centred faces. Profiles, several people in frame, hands and full bodies are out of distribution.
• It cannot draw a named celebrity. Beyond having no text encoder, CelebA ships identities as integers and never publishes the names, and 16 photographs per identity would not be enough anyway.
• This is 4.4% of the compute the DDIM paper spent on CelebA-64 (2.83M images against 64M). The register is a good 2016 GAN, not "indistinguishable from a photograph".

Licence. The weights inherit CelebA's: non-commercial research use only. This is a study artifact, not a product component. It generates faces of people who do not exist — do not use it to impersonate anyone or to produce anything presented as a real photograph of a real person.

How to run

# 1. Build the dataset: ~10 GB of downloads, ~1 minute of decoding. Idempotent.
python prepare_data.py

# 2. Check the machine and pick a configuration. Prints images/s and memory per config.
python bench.py

# 3. Train. The step count is derived from the measured speed, not hard-coded.
nohup python -u train.py --hours 9 > train.log 2>&1 &
nohup ./watchdog.sh > watchdog.log 2>&1 &          # restarts training if it dies
nohup ./finish_report.sh > finish_report.log 2>&1 &  # builds RESULTS.md when it ends

Look at what it learned

python sample.py                                  # 8x8 grid
python sample.py --attrs "Male=1,Eyeglasses=1"    # conditional
python sample.py --guidance-sweep                 # the CFG figure
python sample.py --sampler-comparison             # DDPM/DDIM/Heun
python evaluate.py --memorisation                 # did it copy?
python fid.py --n 10000                           # the headline number

Or do it in a browser

python serve.py     # playground at http://127.0.0.1:8891

Tick attributes, drag the guidance slider, and watch each face condense out of noise — the page streams snapshots of the reverse process as it runs, rather than waiting for the final image. Hold the seed fixed and move one control to see what that control actually does. Standard library only, bound to localhost.

Files

File	Purpose
schedule.py	cosine and linear noise schedules, noisify, the eps/x0/v conversions, stratified timesteps
model.py	the U-Net: ResBlocks, self-attention, time embedding, attribute conditioning, zero-init output
samplers.py	DDPM ancestral, DDIM with eta, Karras+Heun, DDIM inversion, classifier-free guidance, slerp
prepare_data.py	CelebA parquet → data/*_images.bin uint8 memmaps + attributes + meta.json
train.py	the training loop, EMA, time budget, sanity checks, checkpoints, progress grids
bench.py	configuration sweep with a memory ceiling and a per-config deadline
sample.py	generate from a checkpoint; guidance and sampler sweeps
evaluate.py	L_simple on train/val/test, loss by timestep bin, the memorisation study
fid.py	Inception features, cached reference statistics, FID-tv
make_report.py	every figure plus RESULTS.md
serve.py	browser playground: live attribute and guidance controls, streamed denoising
checkpoint.py, images.py	loading a checkpoint; turning tensors into PNG grids
watchdog.sh, finish_report.sh	keep the overnight run alive; build the report when it ends

Measurements on the training machine — MacBook Pro (Apple M5 Pro, 20-core GPU, 24 GB unified memory), MPS

From bench.py, on the real U-Net at 64×64 with bf16 autocast:

Configuration	s/step	images/s	GPU memory	epochs per 9 h
batch 64	0.774	82.7	47%	16.5
batch 64 × accum 2	1.422	90.0	47%	17.9
batch 96	1.054	91.0	59%	18.1
batch 128	1.407	91.0	81%	18.1
batch 192	—	skipped, would swap	116%	—

Throughput saturates above batch 96, so the choice is made on memory, not speed. 64 × 2 keeps the effective batch at 128 — matching the DDIM CelebA-64 reference, so its learning rate and EMA transfer unchanged — while holding 47% of the memory budget instead of 81%.

Variant at batch 96	images/s	verdict
baseline	90.9
channels_last	65.8	28% slower — do not use on MPS
fp32, no autocast	75.9	bf16 is worth 20%
no attention	97.4	attention at 16²/8² costs 7%

Forward-only throughput is 265 images/s, which is what makes the evaluation affordable: an 8×8 grid through DDIM-50 takes 12 s, 10,000 samples take 31 minutes, and 10,000 samples through DDPM-1000 would take 10 hours — which is why that last one is never done.

Actual training runs at ~87 images/s rather than the benchmark's 90: the difference is the batch fetch (15 ms) and the EMA update, neither of which the benchmark includes.

What changed compared to the DDPM paper, and why

DDPM (Ho et al. 2020)	Here	Why
linear beta schedule	cosine (Nichol & Dhariwal)	at 64×64 a linear schedule destroys the image early and the last fifth of the timesteps teaches almost nothing
t ~ Uniform per sample	stratified over the batch	same expectation, much lower gradient variance — the currency a 22,000-step run is short of
EMA decay 0.9999	0.999	0.9999 averages over 10,000 steps, nearly half our entire run; the average would still carry the early garbage
dropout 0.1	0.0	DDPM ran ~2000 epochs over CIFAR and genuinely overfit; we run 17 epochs over 163k images
plain Adam	AdamW, weight_decay=0	same thing, explicitly
unconditional	6 CelebA attributes + CFG	one model, cond_dropout=0.15, serves both branches; guidance is the largest quality lever an undertrained model has
ConvTranspose2d upsample	nearest + conv	avoids checkerboard artifacts, which a short run has no time to learn its way out of

The attributes are Male, Smiling, Young, Eyeglasses, Blond_Hair, Bangs — two balanced, one skewed (Young, 75% positive, which exercises the null path), one rare and dramatic (Eyeglasses, 7%, the best guidance demo), and two that move different visual axes.

An honest note on the conditioning: Stable Diffusion injects its condition through cross-attention to a sequence of text tokens. Six fixed flags are not a sequence, so adding their projection to the time embedding is the correct analogue — the same mechanism, in that the network learns a conditional score, with a simpler carrier. Cross-attention over six booleans would be ceremony.

Sanity checks that run before every training run

A diffusion model fails silently in a specific way: get the coefficient broadcast wrong and the loss still falls beautifully while the model learns nothing. These run in milliseconds.

• Schedule invariants — alphas_cumprod strictly decreasing, abar_0 > 0.99, abar_T < 1e-2, sqrt(abar)² + sqrt(1-abar)² = 1, SNR strictly decreasing.
• Noisify round-trip on the device — reconstruct x_0 from x_t and eps; max error must be under 1e-3. This is the direct analogue of mini-LLM's "y is x shifted by one token".
• Per-sample spread — if the coefficients broadcast along the wrong axis, every image in the batch is noised identically, and that is checked explicitly.
• Transfer integrity — pixel range, per-image variance above zero, no two identical samples, and md5 checksums of eight reference thumbnails against meta.json.
• Loss at initialisation is 1.0000 — the prettiest one. The output convolution is zero-init, so the model predicts no noise, so the loss is exactly E||eps||² = 1. That single number confirms the zero-init, that the target is eps and not x0 or v, that pixels are scaled to [-1,1], that the reduction is mean, and that autocast is not mangling the loss.

Pitfalls we already hit

MPS does not raise when it runs out of memory — it swaps. torch.mps.recommended_max_memory() is 17.8 GB on an M5 Pro with 24 GB of unified memory. Go past it and Metal pages GPU allocations to disk: everything becomes about a hundred times slower while looking perfectly alive. During planning, one benchmark configuration spent 17 minutes failing to finish 11 steps while swap grew by 6 GB. bench.py therefore refuses any configuration that is past 85% of the budget after warm-up, and train.py logs the driver allocation into the CSV every step.

The number that matters is driver_allocated_memory(), not current_allocated_memory(). Live tensors are only 0.45 GB at batch 128; the caching allocator's pool is 14.4 GB. It is the pool that hits the cliff.

This is almost certainly what caused the unexplained 23K → 3.2K tokens/s collapse partway through the mini-LLM run — the same machine, the same silent failure.

non_blocking=True still corrupts batches on MPS. Same trap as in mini-LLM: the async copy starts from ordinary unpinned memory, the source is freed before it runs, and garbage arrives on the device. Banned here, with the assertions to catch it if it ever creeps back.

channels_last is 28% slower. Reasonable-sounding advice from CUDA that does not transfer: MPSGraph manages its own internal layouts and the extra transposes are pure cost. Measured, not assumed.

Do not set PYTORCH_ENABLE_MPS_FALLBACK. With it unset, an unsupported operation raises loudly. With it set, the operation silently round-trips to the CPU and you discover the slowdown in the morning. train.py asserts it is unset.

Pure noise out of the sampler at step 200 is not a bug. Early on the model has only learned the easy regime — at high t the input already is the noise, so predicting it is nearly free, and that alone drives the loss from 1.00 to 0.26 without any generative ability. Denoising a real image at t=100 produced a near-perfect face at that point, which is how we confirmed the plumbing was right and simply waited.

Where to go next

• Min-SNR-γ=5 loss weighting — the paper reports 3.4× faster convergence on ImageNet-64, which is exactly what a compute-starved run wants. Implemented nowhere yet; the first experiment to run.
• v-prediction with zero terminal SNR. Even the cosine schedule ends at abar_T ≈ 2e-9 rather than 0, so the model never trains on pure noise but is handed pure noise at sampling time. With ε-prediction this shows up as a mean-brightness bias.
• RL post-training (DDPO / DRaFT) on top of this sampler, with a reward from CLIP or from an attribute classifier. That is a separate camp artifact, and an 18M model with DDIM-20 makes it genuinely feasible locally.
• Latent diffusion, reusing the pretrained VAE already in the HuggingFace cache, to reach 128–256 px on the same compute budget.

---

Licence

Code: MIT — see LICENSE. Use it for anything, including commercially.

The weights are a separate matter. They are trained on CelebA, which its authors release for non-commercial research only, and no licence file in this repository can widen that. Retrain on your own data if you need weights you can ship.