---
license: mit
pipeline_tag: unconditional-image-generation
tags:
- diffusion
- rectified-flow
- patch-diffusion
- anime
---
# Waifu Diffusion

A 130M-parameter diffusion model trained on 10,000 anime faces (90% monochrome) using **rectified flow**, **patch diffusion**, and **CIELAB color space decoupling**.

## Model Details

- **Architecture**: Diffusion Transformer (DiT-B) with Vision RoPE
- **Parameters**: 130M
- **Training Data**: 10k anime faces (80×80), 90% corrupted to grayscale
- **Training Steps**: 1280 epochs × batch 256
- **Sampling**: 50-step Euler integration

### Versions

| Model | Details |
|-------|---------|
| `waifu_diffusion_1280_bs256.safetensors` | Full training (1280 epochs, bs=256) |
| `waifu_diffusion_128_bs32.safetensors` | Shallow trained version (128 epochs, bs=32) |

## Quick Start

```python
import torch
from safetensors.torch import load_file
from skimage import color
import numpy as np

# Load model
model = JiT(
    input_size=80,
    patch_size=4,
    in_channels=3,
    hidden_size=768,
    depth=12,
    num_heads=12,
    num_classes=1
)
state_dict = load_file("waifu_diffusion_1280_bs256.safetensors")
model.load_state_dict(state_dict)
model.eval()

# Generate
device = "cuda"
model.to(device)

with torch.no_grad():
    xt = torch.randn((1, 3, 80, 80), device=device)
    y = torch.zeros(1, dtype=torch.long, device=device)
    
    for step in range(50):
        t = torch.tensor(step / 50, device=device)
        pred_x1 = model(xt, t, y, top_idx=0, left_idx=0)
        v = (pred_x1 - xt) / max(1.0 - step / 50, 1e-2)
        xt = xt + v / 50

# Convert CIELAB → RGB
lab = torch.clamp(pred_x1[0], -1, 1).cpu().numpy()
L = (lab[0] + 1) * 50
a = lab[1] * 128
b = lab[2] * 128
rgb = color.lab2rgb(np.stack([L, a, b], axis=-1))
```

## Key Techniques

- **Rectified Flow**: Straight-line paths from noise to data (50 steps vs. 1000s for DDPM)
- **CIELAB Decoupling**: Separate luminance from color; mask gradients on monochrome → learn structure from all 10k, color from 1k
- **Patch Diffusion**: Random 40×80 px crops act as data augmentation; effectively 10k → ~50k samples
- **Vision RoPE**: 2D rotary embeddings for spatial consistency across patches

## Links

- **GitHub**: https://github.com/ruwwww/waifu_diffusion
- **Blog Post**: [Training a Waifu Diffusion Model](https://ruwwww.github.io/al-folio/blog/2026/waifu-diffusion/)

## Citation

```bibtex
@misc{waifu_diffusion_2026,
  author = {Abdurrahman Izzuddin Al Faruq},
  title = {Training a Waifu Diffusion Model with Patch Diffusion and Rectified Flow},
  year = {2026},
  url = {https://github.com/ruwwww/waifu_diffusion}
}
```

## License

MIT