yuanchenyang/smalldiffusion

Simple and readable code for training and sampling from diffusion models

What it solves

smalldiffusion 是一个轻量级库，旨在简化扩散模型和流模型的训练与采样过程。它消除了从头实现这些模型的复杂性，提供了一个简洁的核心（代码行数不足 100 行），便于研究人员和开发者在实验新模型或采样器时进行扩展。

How it works

库通过三个主要组件的交互来工作：

  • Data：使用 PyTorch 的 Dataset 和 DataLoader 进行训练。内置支持 2D 玩具数据集、像素空间图像以及潜在空间图像数据集。
  • Model：支持多种架构，如 MLP、U‑Net 和 Diffusion Transformer（DiT）。可选择不同的参数化方式，包括 score‑prediction、flow‑prediction 和 data‑prediction（x0）。
  • Schedule：管理训练和采样时使用的噪声水平（sigmas）。提供多种内置调度器，如 LogLinear、DDPM、LDM（Stable Diffusion 风格）以及 Flow matching 调度器。

采样由统一的 samples() 生成器处理，通过调整 gam、mu 等超参数即可支持 DDPM、DDIM 以及其他加速采样算法。

Who it’s for

适用于希望在小到中等规模数据集（如 CIFAR‑10、FashionMNIST）上进行实验、训练，或从预训练的 Hugging Face diffusers 模型进行采样的研究人员和开发者。

Highlights

  • Minimalist Core：训练和采样逻辑少于 100 行代码。
  • Versatile Model Support：包含 U‑Net 与 Diffusion Transformer（DiT）的实现。
  • Flexible Parameterizations：支持 score、flow 和 data 预测目标。
  • Pretrained Model Integration：提供包装器，可将库的采样器与预训练的 Stable Diffusion 模型一起使用。
  • High Performance：在 ImageNet 256×256 上使用 flow matching 进行示例训练，可实现约 27 的无条件 FID，性能出色。