Assignment 4 – Diffusion + Energy models on top of A1/A2/A3

Adds two CIFAR-10 generators to the FastAPI service:

• DDPM diffusion model — small UNet with sinusoidal time embedding, linear β schedule (app/diffusion.py).
• Energy-based model — CNN energy net trained with contrastive divergence, sampled via Langevin dynamics (app/energy.py).

All previous endpoints from Assignments 1–3 (embeddings, similarity, CIFAR-10 classifier, MNIST DCGAN) still work.

Endpoints

Method	Path	Purpose
GET	/	service metadata + which models are loaded
GET	/health	liveness probe
POST	/embed	word → 300-d spaCy vector (A1)
POST	/similarity	cosine similarity between two words (A1)
POST	/classify	upload image → CIFAR-10 class (A2)
POST	/generate	N MNIST digits from DCGAN as base64 PNGs (A3)
GET	/generate/image	single MNIST digit as image/png (A3)
POST	/generate/diffusion	N CIFAR samples from the DDPM as base64 PNGs
GET	/generate/diffusion/image	single DDPM sample as image/png
POST	/generate/energy	N CIFAR samples from the EBM (Langevin) as base64 PNGs
GET	/generate/energy/image	single EBM sample as image/png
GET	/docs	Swagger UI

Model architecture (A4 additions)

Diffusion (app/diffusion.py)

• UNet — 2 downsampling stages (32×32 → 16×16 → 8×8 bottleneck), 2.66M params, GroupNorm+SiLU residual blocks, time embedding injected as an additive shift on channels
• Sinusoidal time embedding — implements the paired-sine/cosine formula from Vaswani et al. (see SOLUTIONS_PART2.md Q1/Q2)
• Schedule — linear β from 1e-4 to 0.02 over T=200 timesteps
• Training objective — MSE(ε, ε̂_θ(x_t, t)) (Ho et al. 2020, Alg. 1)
• Sampling — reverse-diffusion loop (Ho et al. 2020, Alg. 2) in p_sample_loop

Energy-based model (app/energy.py)

• EnergyNet — 4 stride-2 conv blocks with SiLU → scalar E(x), 393K params
• Training — contrastive divergence: positives from real CIFAR, negatives from Langevin dynamics samples of the current model, plus an L2-on-energies stabiliser (Du & Mordatch 2019)
• Sampling — x_{k+1} = x_k - (η/2)·∇_x E(x_k) + σ·noise, run for 60 steps by default. Note the input has requires_grad_(True) — this is the "gradient descent on the input" pattern the assignment highlights.

Train the models

pip install -r requirements.txt
python -m spacy download en_core_web_md

# ~4 min on M-series Mac (MPS), ~15 min on CPU
python -m app.train_diffusion    # writes app/diffusion_weights.pth

# ~2 min on MPS
python -m app.train_energy       # writes app/energy_weights.pth

Both training scripts prefer the fast.ai-format CIFAR-10 folder at data/cifar10/train/<class>/*.png (much faster to download than the standard torchvision tarball); if that folder is missing they fall back to torchvision.datasets.CIFAR10(download=True).

The API loads each checkpoint lazily on the first request to its endpoint, so cold-start latency is a few seconds for diffusion (multi-step reverse process) and a fraction of a second for Langevin.

Run with Docker

docker compose up --build

Service on http://localhost:8000 — visit /docs for the interactive Swagger.

Try it

# 1 diffusion sample as PNG
curl "http://localhost:8000/generate/diffusion/image?seed=42" -o diffusion.png

# 4 diffusion samples as base64 JSON
curl -X POST http://localhost:8000/generate/diffusion \
     -H "Content-Type: application/json" \
     -d '{"n": 4, "seed": 7}'

# 1 energy-model sample with 100 Langevin steps
curl "http://localhost:8000/generate/energy/image?seed=42&steps=100" -o ebm.png

# 4 EBM samples with custom step size + noise
curl -X POST http://localhost:8000/generate/energy \
     -H "Content-Type: application/json" \
     -d '{"n": 4, "seed": 7, "steps": 80, "step_size": 10.0, "noise_scale": 0.005}'

Tests

pytest -v

10 tests covering all endpoints. Diffusion/energy tests degrade to a 503-tolerant assertion if the checkpoint is missing (so the suite still passes on a fresh clone before training runs).

Deliverables for Part 2 / Part 3

• SOLUTIONS_PART2.md — diffusion theory (sinusoidal embeddings, positional-encoding comparison, UNet bottleneck arithmetic, MSE loss)
• SOLUTIONS_PART3.md — PyTorch autograd questions (basic gradients, weights, .detach, gradient accumulation), each with a verified code run

File map

app/
  main.py                # FastAPI routes (all assignments)
  models.py              # Pydantic request/response schemas
  embeddings.py          # A1: spaCy
  classifier.py, cnn.py, train.py           # A2: CIFAR-10 CNN
  gan.py, generator_service.py, train_gan.py, generator_weights.pth   # A3: MNIST DCGAN
  diffusion.py, diffusion_service.py, train_diffusion.py, diffusion_weights.pth
  energy.py,    energy_service.py,    train_energy.py,    energy_weights.pth
tests/test_api.py        # pytest for every endpoint
Dockerfile, docker-compose.yml, requirements.txt
SOLUTIONS_PART2.md, SOLUTIONS_PART3.md