From scratch implementation of several diffusion architectures; DDPM, ADM, and DiT

Examples

Here are examples of how to trigger a training run:

ADM DDPM style CIFAR10 32x32 images:

python3 main.py --device cuda --dataset cifar10 --total_steps 800000 --batch_size 128 --grad_accum 1

DiT on CelebA 32x32 images with 4 patches:

python3 main.py --device cuda --dataset celeb_small --total_steps 800000 --batch_size 64 --grad_accum 2 --model_name dit --patch_size 4

Weights & Biases

Authenticate once with uv run wandb login, then configure the run under trainer:

wandb:
	project: diffusion
	entity: null
	name: null
	log_frequency: 100

Omit the block or set wandb: null to disable logging. The trainer logs accumulation-averaged loss and learning rate, generated preview grids, and the best completed FID score. Use WANDB_MODE=offline to record a run without network access.

To-Do:

• Training warmup
• For eval metrics for validation set; FID, sFID, IS, precision/recall
• LSUN dataset variant

There is probably more, but these are top of mind