Small enough to read, serious enough for ImageNet

What it does

smalldiffusion is a compact Python library for training and sampling diffusion and flow models. It wraps standard architectures—U-Net, DiT, and simple MLPs—behind a unified interface that handles score, flow, and data-prediction objectives. The library also bundles schedulers, classifier-free guidance, and a thin wrapper for experimenting with pretrained Hugging Face diffusers models such as Stable Diffusion.

The interesting bit

The core training and sampling logic lives in fewer than 100 lines of code, yet the provided examples scale from 2D toy datasets all the way to ImageNet 256×256. That is a deliberate design bet: when the foundation is tiny, the experiments built on top are easier to reason about and modify.

Key highlights

• Core training and sampling loop in under 100 lines.
• Supports multiple prediction objectives via composable wrappers: PredX0, PredV, and PredFlow for data, score, and flow matching.
• Example scripts report near-SOTA FID scores: ~3–4 on CIFAR-10, ~5–6 on FashionMNIST, and ~27 unconditional FID on ImageNet with a DiT-B/2.
• Includes built-in schedulers for DDPM, DDIM, flow matching, and latent diffusion, plus a concise DiT implementation.
• Compatible with Hugging Face Accelerate for multi-GPU training and with pretrained Stable Diffusion models for sampling experiments.

Verdict

Researchers and tinkerers who want to understand or extend diffusion mechanics without fighting a monolithic framework will feel at home. If you are looking for a batteries-included production pipeline with a built-in model zoo, this is not it.