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，效能優異。