smalldiffusion is a diffusion library small enough to read end to end

A core you can read in one sitting

smalldiffusion is described as a lightweight diffusion library for training and sampling from diffusion and flow models, and the README gives the reason for the name: it is designed for ease of experimentation when training new models or developing new samplers.

The claim that defines the project is a line count. The core is described as small but extensible, less than 100 lines of code for training and sampling. That is the whole pitch, and it is a pitch aimed at a specific person: the researcher who wants to change one thing about a diffusion model and does not want to spend a week reading a framework first.

The project is a single Python package, smalldiffusion, with the source under src/smalldiffusion, examples under examples, tests under tests and a datasets directory. The package metadata puts it at version 0.5.0 with the development status of Beta, requires Python 3.10 or later, and lists six runtime dependencies: accelerate, numpy, torchvision, torch, tqdm and einops. The author and maintainer is Chenyang Yuan.

The documentation strategy is a blog post rather than a manual. The README links a tutorial blog post, an arXiv paper, a Colab notebook, the PyPI project and a build badge, and it tells you to see the accompanying tutorial for a detailed introduction to diffusion models and the notation used in the code. That is a reasonable choice for a teaching-focused library, and it also means the code assumes you already know the vocabulary the code uses.

Everything turns on data, model and schedule

The how-to-use section states the shape of the whole system in one sentence: the core of smalldiffusion depends on the interaction between data, model and schedule objects.

Data is ordinary PyTorch. The library uses Dataset and DataLoader batches, and by default training_loop expects each batch to be a tensor of data only. To drop labels or otherwise map items, you wrap the underlying dataset with MappedDataset(dataset, fn) so each __getitem__ returns what the loop should see. For classifier-free guidance style conditional training, you pass conditional=True to training_loop, and the batch then carries the labels as well.

The model is a callable that takes the pose and identity inputs and returns what you need. Three are provided: an MLP, a U-Net and a diffusion transformer, the last described as a concise implementation of the diffusion transformer introduced by Peebles and Xie. The same training code trains all three, which is the point of the parameterisation support: score-prediction, flow-prediction and data-prediction are all supported, so you can switch the objective without rewriting the loop.

The schedule is the noise schedule, and it appears in the code as a value with a name, for example ScheduleLogLinear with a number of steps, a sigma_min and a sigma_max, or ScheduleLDM for latent diffusion. The schedule is also what produces the sigmas you sample along, which is why the sampler takes a schedule rather than a step count.

Three objects, one interaction. That is a small enough API to hold in your head, which is exactly the point of the project.

Ten lines train a model on a toy dataset

The Swissroll example is the best introduction to the library, because the whole training and sampling loop fits in ten lines of code.

from torch.utils.data import DataLoader
from smalldiffusion import Swissroll, TimeInputMLP, ScheduleLogLinear, training_loop, samples

dataset  = Swissroll(np.pi/2, 5*np.pi, 100)
loader   = DataLoader(dataset, batch_size=2048)
model    = TimeInputMLP(hidden_dims=(16,128,128,128,128,16))
schedule = ScheduleLogLinear(N=200, sigma_min=0.005, sigma_max=10)
trainer  = training_loop(loader, model, schedule, epochs=15000)
losses   = [ns.loss.item() for ns in trainer]
*xt, x0  = samples(model, schedule.sample_sigmas(20), gam=2)

Every stage is visible. The dataset is a Swissroll with a start angle, an end angle and a noise level. The model is a time-input MLP with a small hidden layer stack. The schedule is log-linear over 200 steps with a noise range. The training loop takes the loader, the model and the schedule and returns a trainer whose elements carry a loss, so the loss curve is one list comprehension. And the sampler takes the model and a list of sigmas from the schedule and returns the intermediate states and the final sample, with a parameter called gam that controls the sampler.

Two things are worth noticing. The listing does not import numpy even though the dataset call uses np.pi, so the snippet assumes numpy is already in scope; and training_loop returns an iterable of per-step state objects rather than a model, which is a deliberate shape: the loop is a generator you can inspect, not a black box that returns a checkpoint.

A second example covers conditional training with classifier-free guidance on a two-dimensional tree dataset, where each class is drawn in a different colour.

Example scripts carry the ImageNet and CIFAR-10 results

The examples directory is where the library's claims are demonstrated, and the file names describe the range.

There is a DiT script for ImageNet at 256 by 256 using the flow matching formulation in the latent space of Stable Diffusion's VAE, training on precomputed VAE latents and supporting multi-GPU training via accelerate. The README reports that after 400k steps, around 10 hours on 8 GPUs, the model achieves an unconditional FID of around 27, compared to 33 for SiT and 43 for DiT.

The FashionMNIST script trains a diffusion transformer and, with the provided default parameters on a single GPU for around two hours, the README reports an FID score of around 5 to 6. The CIFAR-10 U-Net example is reported at an unconditional FID of around 3 to 4.

These are the project's own reported numbers, taken from the README, and the comparison figures for SiT and DiT are the reference point the README uses. What is worth noticing is the shape of the evidence: the numbers come from example scripts that you run, with the FID evaluation left to you, rather than from a benchmark harness inside the library.

The model code itself is split into three modules referenced from the README, an MLP model, a U-Net model and a DiT model, and the same training code is used across them, which is the clearest demonstration of the library's central claim.

Sampling from a pretrained model through a diffusers wrapper

The sampler is not tied to models trained with this library. The README states that smalldiffusion's sampler works with any pretrained diffusion model and supports DDPM, DDIM as well as accelerated sampling algorithms.

The examples directory contains a wrapper for any pretrained huggingface diffusers latent diffusion model, and the example is short:

from diffusers_wrapper import ModelLatentDiffusion
from smalldiffusion import ScheduleLDM, samples

schedule = ScheduleLDM(1000)
model    = ModelLatentDiffusion('stabilityai/stable-diffusion-2-1-base')
model.set_text_condition('An astronaut riding a horse')
*xts, x0 = samples(model, schedule.sample_sigmas(50))
decoded  = model.decode_latents(x0)

The flow is the same one the toy example uses, with a different model object. The schedule is the latent diffusion one, the model sets a text condition, the sampler walks the schedule's sigmas and returns the final latent, and the wrapper decodes the latents back to pixels.

That is the practical payoff of the design. Because the sampler only needs a model that is callable and a schedule, you can point the same sampler at a Stable Diffusion checkpoint without retraining anything, and the README notes that it is easy to experiment with different sampler parameters and sampling schedules, demonstrated in another example script.

The gam parameter visible in the toy example is the knob to turn there, and the README shows several example images of tweaking it.

Installing, running the tests and launching the examples

The package is on PyPI and installs with:

pip install smalldiffusion

For local development the README recommends uv, and the two commands are:

uv sync --extra dev --extra test --extra examples
uv run pytest

The extras are declared in the project metadata. dev is twine, test is pytest together with a pinned diffusers release, and examples is the larger set: diffusers, transformers, datasets, jupyter, matplotlib and torch_ema. The how-to-use section also mentions uv sync --all-extras, or the make install-local target, and the subset form for a lighter setup.

The Makefile is small and shows the intended loop:

build:
	uv build

test:
	uv run pytest

test-slow:
	uv run pytest --run_slow

The pytest configuration registers a run_slow marker, and the marker description says tests marked slow are skipped unless the flag is passed, which is why the Makefile has a separate test-slow target.

The example scripts are launched through accelerate, after configuring it once:

uv run accelerate config
uv run accelerate launch examples/imagenet_dit.py

The same pattern applies to the FashionMNIST and CIFAR-10 scripts, and the tests use the built-in uv run pytest entry point, with the build backend being uv_build rather than setuptools.

Where smalldiffusion is the wrong tool

The scope is the problem, and it is a deliberate trade.

First, this is a research and teaching library, not a training service. The examples show you how to launch a run with accelerate, but the production concerns of a training job, checkpointing and resuming, evaluation loops, experiment tracking, data versioning and scheduling, are not part of the package. The Makefile has build, test, test-slow, publish and install-local targets, and that is the whole operational surface.

Second, the package is a beta. The metadata says Development Status 4 - Beta, the version is 0.5.0, and the repository has no GitHub releases at all, so there is no changelog to read and no version to pin except the PyPI one. The last push was on 2026-04-08, more than six months ago, so the branch has not been receiving changes even though the project is not archived. Anyone building on it is building on a snapshot.

Third, the reported FID numbers are the project's own claims, produced by example scripts whose evaluation code is not in the library. They are a reasonable indication of what the code can do and not a substitute for running the experiment on your own data and your own evaluation.

Fourth, the library deliberately omits things a production pipeline needs. There is no built-in mixed precision or gradient checkpointing policy, no EMA handling beyond the torch_ema dependency listed under the examples extra, and no data pipeline beyond wrapping a PyTorch Dataset. The README also points to the accompanying tutorial for the notation, which is another prerequisite for reading the code well.

What it is excellent for is the thing the README names: ease of experimentation when training new models or developing new samplers, with a core you can read in an evening.

Editorial conclusion

Researchers who want to understand or modify a diffusion training loop should adopt smalldiffusion, because a core you can read in one sitting is the fastest way to see how data, model and schedule interact, and the examples reach from a toy dataset to ImageNet-scale training. It is the wrong tool for a production training service, since checkpointing, evaluation and distributed training live in the example scripts rather than the library, and it is the wrong tool for a beta status package in a pinned production dependency. Before adopting, read the accompanying tutorial for the notation the code uses, confirm the parameterisation your task needs is supported, and pin version 0.5.0 rather than tracking a branch that has not been pushed to since 2026-04-08.

Frequently asked questions

Official sources

1. Issues
2. License: MIT
3. Project website
4. README
5. yuanchenyang/smalldiffusion on GitHub