Stable Diffusion (DDPM): Minimal PyTorch Implementation

This implementation builds a minimal Denoising Diffusion Probabilistic Model (DDPM), the core of the Stable Diffusion family. It includes sinusoidal time embeddings, a U-Net with skip connections and time-conditioned features, a forward noising process, and an ancestral reverse denoising process for sampling.

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math

class SinusoidalPositionEmbeddings(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, time):
        device = time.device
        half_dim = self.dim // 2
        embeddings = math.log(10000) / (half_dim - 1)
        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
        embeddings = time[:, None] * embeddings[None, :]
        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
        return embeddings

class UNet(nn.Module):
    def __init__(self, in_channels=3, time_dim=256):
        super().__init__()

        self.time_mlp = nn.Sequential(
            SinusoidalPositionEmbeddings(time_dim),
            nn.Linear(time_dim, time_dim),
            nn.GELU(),
            nn.Linear(time_dim, time_dim)
        )
        
        self.conv_in = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)

        # Downsampling path
        self.down1 = nn.Sequential(
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.down2 = nn.Sequential(
            nn.Conv2d(128, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )

        # Middle blocks
        self.mid_block1 = nn.Sequential(
            nn.Conv2d(256, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(),
        )

        # Time embedding projection
        self.time_proj = nn.Conv2d(time_dim, 512, 1)
        
        self.mid_block2 = nn.Sequential(
            nn.Conv2d(512, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
        )

        # Upsampling path
        self.up1 = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 2, stride=2),
            nn.BatchNorm2d(128),
            nn.ReLU()
        )
        self.up2 = nn.Sequential(
            nn.ConvTranspose2d(256, 64, 2, stride=2),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )

        self.conv_out = nn.Conv2d(128, in_channels, kernel_size=3, padding=1)

    def forward(self, x, t):
        t_emb = self.time_mlp(t)  
        
        x1 = self.conv_in(x)

        # Downsample
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        
        # Middle block
        x3 = self.mid_block1(x3)
        
        # Project time embedding and add to features
        t_emb = self.time_proj(t_emb.unsqueeze(-1).unsqueeze(-1))
        x3 = x3 + t_emb
        
        x3 = self.mid_block2(x3)
        
        # Upsample
        x = self.up1(x3)
        x = torch.cat([x, x2], dim=1)
        x = self.up2(x)
        x = torch.cat([x, x1], dim=1)

        return self.conv_out(x)
    
class DiffusionModel:
    def __init__(self, timesteps=1000, beta_start=1e-4, beta_end=0.02):
        self.timesteps = timesteps
        
        self.beta = torch.linspace(beta_start, beta_end, timesteps)
        self.alpha = 1 - self.beta
        self.alpha_bar = torch.cumprod(self.alpha, dim=0)
        
        self.unet = UNet()

    def add_noise(self, x, t):
        sqrt_alpha_bar = torch.sqrt(self.alpha_bar[t])[:, None, None, None]
        sqrt_one_minus_alpha_bar = torch.sqrt(1 - self.alpha_bar[t])[:, None, None, None]
        epsilon = torch.randn_like(x)
        
        return sqrt_alpha_bar * x + sqrt_one_minus_alpha_bar * epsilon, epsilon
    
    def remove_noise(self, x, t, noise):
        alpha_t = self.alpha[t][:, None, None, None]
        alpha_bar_t = self.alpha_bar[t][:, None, None, None]
        beta_t = self.beta[t][:, None, None, None]
        
        factor = (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t)
        noise_pred = self.unet(x, t.float())
        mean = (1 / torch.sqrt(alpha_t)) * (x - factor * noise_pred)
        
        if t[0] > 0:  # Check first timestep value
            variance = torch.sqrt(beta_t) * torch.randn_like(x)
            return mean + variance
        return mean
    
    def sample(self, batch_size=1, img_size=32):
        device = next(self.unet.parameters()).device
        
        # random noise
        x = torch.randn(batch_size, 3, img_size, img_size).to(device)
        
        for t in reversed(range(self.timesteps)):
            t_batch = torch.full((batch_size,), t, device=device)
            x = self.remove_noise(x, t_batch, None)
        
        return x
    
def test_diffusion_model():
    import math  
    
    model = DiffusionModel(timesteps=100)

    print("Testing noise addition")
    x = torch.randn(1, 3, 32, 32)
    t = torch.tensor([50])
    noised_x, noise = model.add_noise(x, t)
    assert noised_x.shape == x.shape, f"Expected shape {x.shape}, got {noised_x.shape}"
    print("✓ Noise addition shape test passed")

    print("\nTesting U-Net forward pass")
    noise_pred = model.unet(noised_x, t.float())
    assert noise_pred.shape == x.shape, f"Expected shape {x.shape}, got {noise_pred.shape}"
    print("✓ U-Net forward pass shape test passed")

    print("\nTesting sampling process")
    samples = model.sample(batch_size=2, img_size=32)
    assert samples.shape == (2, 3, 32, 32), f"Expected shape (2, 3, 32, 32), got {samples.shape}"
    assert torch.isfinite(samples).all(), "Generated samples contain NaN or infinite values"
    print("✓ Sampling test passed")

if __name__ == "__main__":
    test_diffusion_model()