smalldiffusion

https://github.com/yuanchenyang/smalldiffusion
Admin
smalldiffusion is a lightweight library for training and sampling from diffusion and flow models,...
Tokens:19,509
Snippets:127
Trust Score:9.2
License:MIT
Update:1 week ago
Show doc for...
Tokens:
Raw
### Configure and launch DiT training on ImageNet with accelerate

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/README.md

Run these commands to configure accelerate and start training the DiT-B/2 model on ImageNet 256x256 using flow matching. The script trains on precomputed VAE latents and supports multi-GPU training. After 400k steps (~10 hours on 8 GPUs), the model achieves an unconditional FID of around 27.

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

--------------------------------

### Install smalldiffusion via pip

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/README.md

Install the smalldiffusion package from PyPI using pip.

```bash
pip install smalldiffusion
```

--------------------------------

### Sample from Stable Diffusion

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/README.md

Example of using the ModelLatentDiffusion wrapper to sample from a pretrained Hugging Face latent diffusion model.

```python
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)
```

--------------------------------

### Set up local development with uv

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/README.md

Sync development dependencies and run tests using uv.

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

--------------------------------

### Launch DiT training on FashionMNIST with accelerate

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/README.md

Run this command to train a diffusion transformer on FashionMNIST and generate a batch of samples. Requires first running 'uv run accelerate config'. With default parameters on a single GPU for around 2 hours, the model achieves an FID score of around 5-6.

```bash
uv run accelerate launch examples/fashion_mnist_dit.py
```

--------------------------------

### Train U-Net on Fashion-MNIST

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/README.md

Command to launch the training script for a U-Net model on the Fashion-MNIST dataset.

```bash
uv run accelerate launch examples/fashion_mnist_unet.py
```

--------------------------------

### Import libraries and define plotting/moving average helpers

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/cond_tree_model.ipynb

Sets up the environment with necessary imports and helper functions for plotting batches and computing moving averages. Includes autoreload magic for development.

```python
import numpy as np
import matplotlib.pyplot as plt 
import torch
from torch.utils.data import DataLoader
from smalldiffusion import TreeDataset, training_loop, samples, ScheduleLogLinear, ConditionalMLP

def plot_batch(batch, ax=None, **kwargs):
    batch = batch.cpu().numpy()
    ax = ax or plt
    return ax.scatter(batch[:,0], batch[:,1], marker='.', **kwargs)
    
def moving_average(x, w):
    return np.convolve(x, np.ones(w), 'valid') / w    

%load_ext autoreload
%autoreload 2
```

--------------------------------

### Sample from conditional model using Classifier-Free Guidance

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/cond_tree_model.ipynb

Defines a function to sample from the model for each class using the samples function with configurable CFG scale and other parameters. Then calls it to generate and plot samples for all classes.

```python
def sample_and_plot(model, schedule, num_classes, ax=None,
                    N_sample=20, sample_batch_size=100,
                    gam=1.5, cfg_scale=1.0, mu=0.0):
    for c in range(num_classes):
        *_, x0 = samples(
            model, schedule.sample_sigmas(N_sample),
            batchsize=sample_batch_size, gam=gam, mu=mu, cfg_scale=cfg_scale,
            cond=torch.tensor([c for _ in range(sample_batch_size)])
        )
        plot_batch(x0, ax=ax, c=[c]*sample_batch_size, cmap='jet', vmin=0, vmax=num_classes)

sample_and_plot(model, schedule, dataset.total_leaves)
```

```text
Result:
<Figure size 640x480 with 1 Axes>
```

--------------------------------

### Train U-Net on CIFAR-10

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/README.md

Command to launch the training script for a U-Net model on the CIFAR-10 dataset, achieving an unconditional FID of around 3-4.

```bash
uv run accelerate launch examples/cifar_unet.py
```

--------------------------------

### Train and sample from Swissroll toy dataset

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/README.md

Trains a TimeInputMLP model on the Swissroll dataset and generates samples using a log-linear noise schedule. The training loop runs for 15000 epochs and samples are drawn with 20 sigmas and gam=2.

```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)
```

--------------------------------

### Sample from the diffusion model

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/toyexample.ipynb

Samples 1500 points using 20 sampling steps with gam=2 and mu=0 (gradient estimation sampler). Comments show how to switch to DDPM (gam=1, mu=0.5) or DDIM (gam=1, mu=0). Plots the generated batch.

```python
# For DDPM sampling, change to gam=1, mu=0.5
# For DDIM sampling, change to gam=1, mu=0
*xts, x0 = samples(model, schedule.sample_sigmas(20), batchsize=1500, gam=2, mu=0)
plot_batch(x0)
```

--------------------------------

### Import libraries and define helper functions

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/toyexample.ipynb

Imports numpy, torch, matplotlib, DataLoader, and smalldiffusion components. Defines plot_batch to scatter plot a batch of 2D points and moving_average for smoothing loss curves.

```python
import numpy as np
import torch
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from smalldiffusion import (
    TimeInputMLP, ScheduleLogLinear, training_loop, samples,
    DatasaurusDozen, Swissroll
)

def plot_batch(batch):
    batch = batch.cpu().numpy()
    plt.scatter(batch[:,0], batch[:,1], marker='.')

def moving_average(x, w):
    return np.convolve(x, np.ones(w), 'valid') / w
```

--------------------------------

### Define and plot tree dataset

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/cond_tree_model.ipynb

Creates a TreeDataset with branching factor 3, depth 3, and 50 samples per path. Loads it into a DataLoader and plots the first batch colored by class label.

```python
dataset = TreeDataset(branching_factor=3, depth=3, num_samples_per_path=50)
loader = DataLoader(dataset, batch_size=len(dataset)*2, shuffle=True)
x, labels = next(iter(loader))
plot_batch(x, c=labels, cmap='jet')
plt.show()
```

```text
Result:
<Figure size 640x480 with 1 Axes>
```

--------------------------------

### Define log-linear sigma schedule

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/toyexample.ipynb

Creates a ScheduleLogLinear with 200 steps, sigma_min=0.01, sigma_max=10, and plots the sigma values on a log scale.

```python
schedule = ScheduleLogLinear(N=200, sigma_min=0.01, sigma_max=10)
plt.plot(schedule.sigmas)
plt.xlabel('$t$')
plt.ylabel('$\sigma_t$')
plt.yscale('log')
```

--------------------------------

### Varying CFG scale pushes samples towards leaves

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/cond_tree_model.ipynb

Creates a 1x4 subplot grid and for each CFG scale (0,1,2,3) calls sample_and_plot to generate samples, setting axis limits and titles. Shows how increasing CFG scale concentrates samples near the leaves.

```python
lim = 1.2
fig, axs = plt.subplots(1, 4, figsize=(20,5.5))
cfg_scales = [0, 1, 2, 3]

for ax, cfg_scale in zip(axs, cfg_scales):    
    ax.set_ylim((-lim,lim))
    ax.set_xlim((-lim,lim))
    sample_and_plot(model, schedule, dataset.total_leaves, ax=ax, cfg_scale=cfg_scale)
    ax.axis('off')
    ax.set_title(f'CFG Scale = {cfg_scale}', fontsize=15, pad=20)
```

```text
Result:
<Figure size 2000x550 with 4 Axes>
```

--------------------------------

### Load and plot Datasaurus Dozen dataset

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/toyexample.ipynb

Loads the 'dino' dataset from the Datasaurus Dozen TSV file and creates a DataLoader with batch size 2130. The comment shows alternative datasets and the Swissroll option. Plots the first batch.

```python
# Try replacing dataset with 'dino', 'bullseye', 'h_lines', 'x_shape', etc. 
dataset = DatasaurusDozen(csv_file='../datasets/DatasaurusDozen.tsv', dataset='dino')
# Or use the SwissRoll dataset
# dataset = Swissroll(np.pi/2, 5*np.pi, 100)
loader = DataLoader(dataset, batch_size=2130)
plot_batch(next(iter(loader)))
```

--------------------------------

### Train conditional diffusion model

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/cond_tree_model.ipynb

Initializes a ConditionalMLP model with specified hidden dimensions and conditional dimension, sets up a log-linear noise schedule, and runs the training loop for 30000 epochs. Collects the loss values from each training step.

```python
model = ConditionalMLP(dim=2, hidden_dims=(16, 128, 256, 128, 16), cond_dim=4, num_classes=dataset.total_leaves)
schedule = ScheduleLogLinear(N=200, sigma_min=0.01, sigma_max=10)
trainer = training_loop(loader, model, schedule, epochs=30000, lr=1e-3, conditional=True)
losses = [ns.loss.item() for ns in trainer]
```

```python
plt.plot(moving_average(losses, 100))
plt.show()
```

```text
Result:
<Figure size 640x480 with 1 Axes>
```

--------------------------------

### Train the diffusion model

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/toyexample.ipynb

Runs the training loop for 15000 epochs with learning rate 1e-3, collecting the loss values from each iteration.

```python
trainer = training_loop(loader, model, schedule, epochs=15000, lr=1e-3)
losses = [ns.loss.item() for ns in trainer]
```

--------------------------------

### Training loop usage

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/README.md

Iterate over the training_loop generator to access training step locals such as loss. The loop yields a namespace with attributes like loss, optimizer, x0, and sigma.

```python
for ns in training_loop(loader, model, schedule):
    print(ns.loss.item())
```

--------------------------------

### Define TimeInputMLP model

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/toyexample.ipynb

Creates a TimeInputMLP with hidden dimensions (16,128,128,128,128,16) and prints the model architecture. The input is 4-dimensional (sigma embeddings concatenated with x) and output is 2-dimensional predicted noise.

```python
model = TimeInputMLP(hidden_dims=(16,128,128,128,128,16))
print(model)
```

--------------------------------

### Plot training loss

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/toyexample.ipynb

Plots the moving average of the training losses with a window of 100 to smooth the curve.

```python
plt.plot(moving_average(losses, 100))
plt.show()
```

--------------------------------

### Plot sigma embeddings

Source: https://github.com/yuanchenyang/smalldiffusion/blob/main/examples/toyexample.ipynb

Computes and plots the sine and cosine components of the sigma embeddings using get_sigma_embeds, showing how the schedule is encoded.

```python
from smalldiffusion.model import get_sigma_embeds
sx, sy = get_sigma_embeds(len(schedule), schedule.sigmas).T
plt.plot(sx, label='$\sin(\log(\sigma_t)/2)$')
plt.plot(sy, label='$\cos(\log(\sigma_t)/2)$')
plt.xlabel('$t$')
plt.legend()
plt.show()
```
Tokens:
Raw