Denoising Diffusion Probabilistic Model (DDPM) Implementation in PyTorch

This implementation builds a minimal Denoising Diffusion Probabilistic Model (DDPM) from scratch. It includes a linear-beta forward diffusion utility, a reverse sampling loop, a time-conditioned UNet denoising network, and a test script that forward-diffuses a batch, runs a forward pass, and samples from Gaussian noise.

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from tqdm import tqdm

class SimpleDiffusionModel:
    def __init__(self, n_steps=1000, beta_start=1e-4, beta_end=0.02):
        self.n_steps = n_steps
        
        self.beta = torch.linspace(beta_start, beta_end, n_steps)
        self.alpha = 1 - self.beta
        self.alpha_bar = torch.cumprod(self.alpha, dim=0)
        
        # Pre-compute values for sampling
        self.sqrt_alpha_bar = torch.sqrt(self.alpha_bar)
        self.sqrt_one_minus_alpha_bar = torch.sqrt(1 - self.alpha_bar)

    def forward_diffusion(self, x_0, t):
        noise = torch.randn_like(x_0)
        sqrt_alpha_bar_t = self.sqrt_alpha_bar[t]
        sqrt_one_minus_alpha_bar_t = self.sqrt_one_minus_alpha_bar[t]
        
        sqrt_alpha_bar_t = sqrt_alpha_bar_t.view(-1, 1, 1, 1)
        sqrt_one_minus_alpha_bar_t = sqrt_one_minus_alpha_bar_t.view(-1, 1, 1, 1)
        
        # Apply forward diffusion formula
        x_t = sqrt_alpha_bar_t * x_0 + sqrt_one_minus_alpha_bar_t * noise
        return x_t, noise

    def sample(self, model, shape, device, guidance_scale=1.0):
        x = torch.randn(shape).to(device)
        
        self.beta = self.beta.to(device)
        self.alpha = self.alpha.to(device)
        self.alpha_bar = self.alpha_bar.to(device)

        # Iterate through reverse diffusion process
        for t in tqdm(reversed(range(self.n_steps)), desc='Sampling'):
            t_tensor = torch.tensor([t], device=device).repeat(shape[0])
            
            # Get model prediction
            with torch.no_grad():
                predicted_noise = model(x, t_tensor)
            
            if t > 0:
                noise = torch.randn_like(x)
            else:
                noise = torch.zeros_like(x)
            
            alpha_t = self.alpha[t].view(1, 1, 1, 1)
            alpha_bar_t = self.alpha_bar[t].view(1, 1, 1, 1)
            beta_t = self.beta[t].view(1, 1, 1, 1)
            
            # Apply reverse diffusion formula with proper broadcasting
            x = (1 / torch.sqrt(alpha_t)) * (
                x - (beta_t / (torch.sqrt(1 - alpha_bar_t))) * predicted_noise
            ) + torch.sqrt(beta_t) * noise
            
        return x

class SimpleUNet(nn.Module):
    def __init__(self, input_channels=1, hidden_dim=64):
        super().__init__()

        self.time_embed = nn.Sequential(
            nn.Linear(1, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
        )

        self.enc1 = nn.Conv2d(input_channels, hidden_dim, 3, padding=1)
        self.enc2 = nn.Conv2d(hidden_dim, hidden_dim * 2, 3, padding=1, stride=2)
        
        self.middle = nn.Sequential(
            nn.Conv2d(hidden_dim * 2 + hidden_dim, hidden_dim * 4, 3, padding=1),
            nn.SiLU(),
            nn.Conv2d(hidden_dim * 4, hidden_dim * 2, 3, padding=1),
        )

        self.dec1 = nn.ConvTranspose2d(hidden_dim * 4, hidden_dim * 2, 4, padding=1, stride=2)
        self.dec2 = nn.Conv2d(hidden_dim * 3, input_channels, 3, padding=1)
        
    def forward(self, x, t):
        t_emb = self.time_embed(t.float().view(-1, 1))
        
        e1 = F.silu(self.enc1(x))  
        e2 = F.silu(self.enc2(e1))  
        
        t_emb = t_emb.view(-1, t_emb.shape[1], 1, 1)
        t_emb = F.interpolate(t_emb, size=e2.shape[2:], mode='nearest')
        
        middle = self.middle(torch.cat([e2, t_emb], dim=1))
        
        d1 = F.silu(self.dec1(torch.cat([middle, e2], dim=1)))  # Upsample and concatenate
        d2 = self.dec2(torch.cat([d1, e1], dim=1))  # Final convolution
        
        return d2

# Test the implementation
def test_diffusion_model():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    diffusion = SimpleDiffusionModel(n_steps=100)
    model = SimpleUNet().to(device)

    batch_size = 2
    channels = 1
    height = 32
    width = 32
    
    x_0 = torch.randn(batch_size, channels, height, width).to(device)
    t = torch.tensor([50]).repeat(batch_size)
    
    print("\nTesting shapes before forward diffusion:")
    print(f"Input shape: {x_0.shape}")
    print(f"Timestep shape: {t.shape}")

    x_t, noise = diffusion.forward_diffusion(x_0, t)
    
    print("\nShapes after forward diffusion:")
    print(f"Noised output shape: {x_t.shape}")
    print(f"Added noise shape: {noise.shape}")

    print("\nTesting model forward pass...")
    model_output = model(x_t, t.to(device))
    print(f"Model output shape: {model_output.shape}")
    assert model_output.shape == x_0.shape, f"Model output shape {model_output.shape} doesn't match input shape {x_0.shape}"
    
    print("\nTesting sampling...")
    sample_shape = (batch_size, channels, height, width)
    samples = diffusion.sample(model, sample_shape, device)
    print(f"Generated sample shape: {samples.shape}")
    assert samples.shape == sample_shape, f"Sample shape {samples.shape} doesn't match expected shape {sample_shape}"
    
    return {
        'input': x_0,
        'noised': x_t,
        'noise': noise,
        'model_output': model_output,
        'sample': samples
    }

if __name__ == "__main__":
    test_results = test_diffusion_model()