Skip to content

README.md

Latest commit

 

History

History
399 lines (305 loc) · 10.5 KB

README.md

File metadata and controls

399 lines (305 loc) · 10.5 KB
Raw

AleDiffuser 🎨

Custom (from-scratch) implementation of Denoising Diffusion Probabilistic Models (DDPM) with U-Net architecture, built entirely from the ground up for deep learning research and experimentation.

🌟 Overview

AleDiffuser is a complete, educational implementation of diffusion models that demonstrates the inner workings of state-of-the-art generative AI. Every component—from the U-Net architecture to the diffusion process—is implemented without relying on high-level abstractions, making it perfect for learning and experimentation.

Key Features

  • ✨ Complete DDPM Implementation - Full forward and reverse diffusion processes
  • 🏗️ Custom U-Net Architecture - ResNet blocks, spatial attention, and skip connections
  • 🎯 Conditional Generation - Class-conditional image synthesis with classifier-free guidance
  • 📊 Flexible Configuration - YAML/JSON config system for reproducible experiments
  • 🔧 Modular Design - Clean, extensible codebase for research and development
  • 🚀 Production Ready - Training scripts, checkpointing, and evaluation utilities

🏛️ Architecture

DDPM (Denoising Diffusion Probabilistic Model)

The core diffusion model implements the mathematical framework from Ho et al. 2020:

Forward Process (Noise Addition):

q(x_t | x_0) = 𝓝(x_t; √ᾱ_t x_0, (1 - ᾱ_t)I)

Reverse Process (Denoising):

p_θ(x_{t-1} | x_t) = 𝓝(x_{t-1}; μ_θ(x_t, t), Σ_θ(x_t, t))

U-Net Architecture

Custom implementation featuring:

  • Encoder Path: Progressive downsampling with ResNet blocks
  • Bottleneck: Self-attention layers for global context
  • Decoder Path: Symmetric upsampling with skip connections
  • Time Embedding: Sinusoidal positional encoding for timestep information
  • Conditional Embedding: Optional class conditioning for controlled generation
Input (3×32×32)
    ↓
Conv + ResNet Blocks (128)
    ↓
Downsample + Attention (256)
    ↓
Downsample + Attention (512)
    ↓
Bottleneck + Self-Attention
    ↓
Upsample + Attention (256)
    ↓
Upsample + Attention (128)
    ↓
Output (3×32×32)

📦 Installation

# Clone the repository
git clone https://github.com/AlejandroParedesLT/alediffuser
cd alediffuser

# Create virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install dependencies
pip install torch torchvision matplotlib pyyaml numpy

Requirements

  • Python 3.8+
  • PyTorch 2.0+
  • torchvision
  • matplotlib
  • PyYAML
  • numpy

🚀 Quick Start

1. Prepare Your Configuration

Create a config file configs/cifar10_experiment.yaml:

# Model Architecture
TIME: 1000
in_channels: 3
out_channels: 3
num_classes: 10
recommended_steps: [1, 2, 2, 2]
recommended_attn_step_indexes: [1, 2, 3]
p_cond: 0.2

# Training Hyperparameters
EPOCHS: 100
batch_size: 128
lr: 0.001
weight_decay: 0.001
beta1: 0.9
beta2: 0.999

# Learning Rate Schedule
alpha_max: 0.0006
alpha_min: 0.000006
T_w: 1000
T_c: 15000

# Output
PATH_TO_SAVE_MODEL: './checkpoints/exp1_T1000_epochs100.pth'
ckpt_path: './outputs'
prefix_name_experiment: 'cifar10_ddpm'

# Hardware
device: 'cuda'
dtype: 'float32'
mixed_precision: false
seed: 42

2. Train the Model

# Basic training
python train.py --config configs/cifar10_experiment.yaml

# With overrides
python train.py --config configs/cifar10_experiment.yaml \
    --batch_size 256 \
    --learning_rate 0.0005 \
    --num_epochs 200

# Debug mode
python train.py --config configs/cifar10_experiment.yaml --debug

# Resume from checkpoint
python train.py --config configs/cifar10_experiment.yaml \
    --resume_from checkpoints/checkpoint_epoch_50.pth

3. Generate Samples

import torch
from alediffuser.models.ddpm import DDPM
from alediffuser.models.unet import UNet

# Load trained model
device = 'cuda' if torch.cuda.is_available() else 'cpu'

model = DDPM(
    T=1000,
    p_cond=0.2,
    eps_model=UNet(
        in_channels=3,
        out_channels=3,
        T=1001,
        num_classes=10,
        device=device,
        dtype=torch.float32
    ),
    device=device,
    dtype=torch.float32
)

model.load_state_dict(torch.load('checkpoints/model.pth'))

# Generate samples for each class
samples = model.sample(
    n_samples=10,
    size=(3, 32, 32),
    classes=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
    w=0.2  # Classifier-free guidance weight
)

# Save results
from torchvision.utils import save_image, make_grid
grid = make_grid(samples, nrow=10)
save_image(grid, 'generated_samples.png')

📁 Project Structure

alediffuser/
├── alediffuser/
│   ├── core/
│   │   └── modules.py           # Basic building blocks (ResNet, Attention, etc.)
│   ├── models/
│   │   ├── ddpm.py              # DDPM implementation
│   │   └── unet.py              # U-Net architecture
│   └── training/
│       ├── data/
│       │   └── cifar10.py       # Data loading utilities
│       └── train.py             # Training loop
├── configs/
│   └── cifar10_experiment.yaml  # Example configuration
├── train.py                     # Main training script
├── requirements.txt
└── README.md

🔬 Components Breakdown

Core Modules (alediffuser/core/modules.py)

PositionalEmbedding

Sinusoidal positional encoding for timestep information:

PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

ResNet Block

Residual block with:

  • Group normalization
  • SiLU/Swish activation
  • Dropout regularization
  • Time embedding injection
  • Skip connections

SpatialSelfAttention

Multi-head self-attention mechanism for capturing global dependencies in feature maps.

DownSample / UpSample

Spatial resolution changes with optional convolution for learnable downsampling/upsampling.

DDPM Model (alediffuser/models/ddpm.py)

Training:

  • Samples random timestep t
  • Adds noise according to schedule: x_t = √ᾱ_t·x_0 + √(1-ᾱ_t)·ε
  • Predicts noise ε_θ(x_t, t, c) with optional conditioning
  • Computes MSE loss between predicted and actual noise

Sampling:

  • Starts from pure noise x_T ~ 𝓝(0, I)
  • Iteratively denoises: x_{t-1} = μ_θ(x_t, t) + σ_t·z
  • Supports classifier-free guidance for conditional generation

U-Net (alediffuser/models/unet.py)

Encoder:

  • Initial convolution: 3 → 128 channels
  • 4 resolution levels with ResNet blocks
  • Self-attention at specified levels
  • Progressive channel increase: 128 → 256 → 512 → 1024

Bottleneck:

  • ResNet → Attention → ResNet
  • Maximum feature compression

Decoder:

  • Symmetric to encoder with skip connections
  • Progressive upsampling and channel reduction
  • Attention at corresponding levels

⚙️ Configuration Options

Model Parameters

Parameter Description Default
TIME Number of diffusion timesteps 1000
in_channels Input image channels 3
out_channels Output image channels 3
num_classes Number of conditional classes 10
recommended_steps Channel multipliers per level [1, 2, 2, 2]
recommended_attn_step_indexes Levels with attention [1, 2, 3]
p_cond Unconditional training probability 0.2

Training Hyperparameters

Parameter Description Default
EPOCHS Training epochs 100
batch_size Batch size 128
lr Learning rate 0.001
weight_decay L2 regularization 0.001
beta1, beta2 Adam optimizer betas 0.9, 0.999

Learning Rate Schedule

Parameter Description
alpha_max Peak learning rate
alpha_min Minimum learning rate
T_w Warmup iterations
T_c Total iterations

🔧 Advanced Usage

Custom Datasets

from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
    def __init__(self, data_path):
        # Your data loading logic
        pass
    
    def __getitem__(self, idx):
        # Return (image, label) tuple
        pass

# In your training script
custom_data = DataLoader(
    CustomDataset('path/to/data'),
    batch_size=128,
    shuffle=True
)

Custom Noise Schedules

# In DDPM.__init__
# Linear schedule (default)
beta_schedule = torch.linspace(1e-4, 0.02, T+1)

# Cosine schedule
beta_schedule = self.cosine_beta_schedule(T+1)

def cosine_beta_schedule(self, timesteps, s=0.008):
    steps = timesteps + 1
    x = torch.linspace(0, timesteps, steps)
    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return torch.clip(betas, 0.0001, 0.9999)

Classifier-Free Guidance

Adjust the guidance weight w during sampling:

# w = 0: Unconditional generation
# w > 0: Stronger conditioning
# w < 0: Negative conditioning (avoid class)

samples = model.sample(
    n_samples=10,
    size=(3, 32, 32),
    classes=[5] * 10,  # Generate 10 images of class 5
    w=0.5  # Moderate guidance
)

📊 Results

CIFAR-10 Samples (After 100 Epochs)

[Include generated sample images here]

Training Metrics

  • Final Validation Loss: ~0.015
  • Training Time: ~8 hours on RTX 3090
  • FID Score: [Add your FID score]

🤝 Contributing

Contributions are welcome! Areas for improvement:

  • Additional noise schedules (cosine, exponential)
  • Latent diffusion support
  • Multi-GPU training
  • Evaluation metrics (FID, IS)
  • Additional datasets (ImageNet, custom)
  • Web interface for interactive generation

📚 References

This implementation is based on:

  1. DDPM: Denoising Diffusion Probabilistic Models - Ho et al., 2020
  2. U-Net: U-Net: Convolutional Networks for Biomedical Image Segmentation - Ronneberger et al., 2015
  3. Classifier-Free Guidance: Classifier-Free Diffusion Guidance - Ho & Salimans, 2022

📄 License

MIT License - feel free to use this code for learning, research, or production!

🙏 Acknowledgments

Built from scratch with inspiration from the diffusion models research community. Special thanks to the authors of the original DDPM paper for their groundbreaking work.

Made with by Alejandro Paredes La Torre