Top AI Repos — open-source AI, indexed and scored
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Simple and readable code for training and sampling from diffusion models
| Date | Stars |
|---|---|
| 2026-07-31 | 777 |
| 2026-08-02 | 777 |
| 2026-08-06 | 777 |
Today
— stars today
This week
— stars this week
This month
— stars this month
Momentum
0.0
growth rate 0.00%/day
# smalldiffusion [![Tutorial blog post][blog-img]][blog-url] [![Paper link][arxiv-img]][arxiv-url] [![Open in Colab][colab-img]][colab-url] [![Pypi project][pypi-img]][pypi-url] [![Build Status][build-img]][build-url] A lightweight diffusion library for training and sampling from diffusion and flow models. Features: - Designed for ease of experimentation when training new models or developing new samplers - Dataset support: [2D toy datasets](#toy-models), [pixel](#u-net-models) and [latent-space](#dit-on-imagenet-with-flow-matching) image datasets - Example training code (with close to SOTA FID): [FashionMNIST](/examples/fashion_mnist_dit.py), [CIFAR10](/examples/cifar_unet.py), [Imagenet](/examples/imagenet_dit.py) - Models: [MLP](/src/smalldiffusion/model.py), [U-Net](/src/smalldiffusion/model_unet.py) and [DiT](/src/smalldiffusion/model_dit.py) - Supports multiple parameterizations: [score-, flow- or data-prediction](#model) - [Small but extensible core][diffusion-py]: less than 100 lines of code for training and sampling To install from [pypi][pypi-url]: ``` pip install smalldiffusion ``` For local development with `uv`: ``` uv sync --extra dev --extra test --extra examples uv run pytest ``` ## Toy models To train and sample from the `Swissroll` toy dataset in 10 lines of code (see [examples/toyexample.ipynb](/examples/toyexample.ipynb) for a detailed guide): ```python 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) ``` Results on various toy datasets: <p align="center"> <img src="https://raw.githubusercontent.com/yuanchenyang/smalldiffusion/main/imgs/toy_models.png" width=100%> </p> ### Conditional training and sampling with classifier-free guidance We can also train conditional diffusion models and sample from them using [classifier-free guidance][cfg-paper]. In [examples/cond_tree_model.ipynb](/examples/cond_tree_model.ipynb), samples from each class in the 2D tree dataset are represented with a different color. <p align="center"> <img src="https://raw.githubusercontent.com/yuanchenyang/smalldiffusion/main/imgs/cfg.png" width=100%> </p> ## Diffusion transformer We provide [a concise implementation][model-code] of the diffusion transformer introduced in [[Peebles and Xie 2022]][dit-paper]. ### DiT on ImageNet with flow matching We provide [an example script](/examples/imagenet_dit.py) for training a DiT-B/2 model on ImageNet 256×256 using the flow matching formulation in the latent space of [Stable Diffusion's VAE](https://huggingface.co/stabilityai/sd-vae-ft-mse). The script trains on precomputed VAE latents and supports multi-GPU training via `accelerate`: ``` uv run accelerate config uv run accelerate launch examples/imagenet_dit.py ``` After training for 400k steps (~10 hours on 8 GPUs), the model achieves an **unconditional FID of around 27**, compared to 33 for [SiT](https://github.com/willisma/sit) and 43 for [DiT](https://github.com/facebookresearch/dit). <p align="center"> <img src="https://raw.githubusercontent.com/yuanchenyang/smalldiffusion/main/imgs/imagenet_samples.png" width=100%> </p> ### FashionMNIST dataset To train a diffusion transformer model on the FashionMNIST dataset and generate a batch of samples (after first running `uv run accelerate config`): ``` uv run accelerate launch examples/fashion_mnist_dit.py ``` With the provided default parameters and training on a single GPU for around 2 hours, the model can achieve a **[FID score](https://paperswithcode.com/sota/image-generation-on-fashion-mnist)
Excerpt of 12,593 characters
Read on GitHub
75
Would you bet a product on this? Bounded 0–100 and slow moving.
matched fp:fd3acaacbc3cbe68, topic:diffusion-models