Denoising Diffusion Probabilistic Model (DDPM) Implementation in PyTorch

This implementation builds a minimal Denoising Diffusion Probabilistic Model (DDPM) from scratch. It includes a UNet denoiser with linear time conditioning, a linear-beta forward diffusion utility, an iterative reverse sampling loop, and a test script that verifies noise addition, model output, and sampling shapes.

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

class UNet(nn.Module):
    """
    A UNet for the denoising model.
    """
    def __init__(self, input_channels=1, hidden_dims=64):
        super().__init__()

        self.init_conv = nn.Conv2d(input_channels, hidden_dims, 3, padding=1)

        # Downsampling 
        self.down1 = nn.Conv2d(hidden_dims, hidden_dims*2, 4, stride=2, padding=1)
        self.down2 = nn.Conv2d(hidden_dims*2, hidden_dims*4, 4, stride=2, padding=1)

        # Time embedding
        self.time_mlp = nn.Sequential(
            nn.Linear(1, hidden_dims*4),
            nn.ReLU(),
            nn.Linear(hidden_dims*4, hidden_dims*4)
        )

        # Upsampling 
        self.up1 = nn.ConvTranspose2d(hidden_dims*4, hidden_dims*2, 4, stride=2, padding=1)
        self.up2 = nn.ConvTranspose2d(hidden_dims*2, hidden_dims, 4, stride=2, padding=1)

        self.final_conv = nn.Conv2d(hidden_dims, input_channels, 3, padding=1)

    def forward(self, x, t):
        x = self.init_conv(x)
        x1 = F.relu(x)
        
        x2 = F.relu(self.down1(x1))
        x3 = F.relu(self.down2(x2))

        t = t.float().unsqueeze(-1)
        t = self.time_mlp(t)
        t = t.view(-1, t.shape[1], 1, 1).expand(-1, -1, x3.shape[2], x3.shape[3])
        x3 = x3 + t

        x = F.relu(self.up1(x3))
        x = F.relu(self.up2(x))
        
        return self.final_conv(x)
    
class DDPM:
    def __init__(self, n_steps=1000, beta_start=1e-4, beta_end=0.02):
        """   
        Args:
            n_steps: Number of diffusion steps
            beta_start: Starting value for noise schedule
            beta_end: Ending value for noise schedule
        """
        self.n_steps = n_steps
        self.beta_start = beta_start
        self.beta_end = beta_end

        # noise schedule
        self.betas = torch.linspace(beta_start, beta_end, n_steps)
        self.alphas = 1 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)

        # Pre-calculate values for inference
        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1 - self.alphas_cumprod)

    def add_noise(self, x_0, t):
        noise = torch.randn_like(x_0)
        sqrt_alpha_cumprod = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
        sqrt_one_minus_alpha_cumprod = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)
        
        noisy_image = sqrt_alpha_cumprod * x_0 + sqrt_one_minus_alpha_cumprod * noise
        return noisy_image, noise
    
    def sample(self, model, n_samples, img_size, device):
        model.eval()
        with torch.no_grad():
            x = torch.randn(n_samples, 1, img_size, img_size).to(device)

            # denoise
            for t in tqdm(reversed(range(self.n_steps)), desc='Sampling'):
                t_batch = torch.ones(n_samples, dtype=torch.long).to(device) * t
                predicted_noise = model(x, t_batch)
                
                alpha = self.alphas[t]
                alpha_cumprod = self.alphas_cumprod[t]
                beta = self.betas[t]
                
                if t > 0:
                    noise = torch.randn_like(x)
                else:
                    noise = torch.zeros_like(x)
                    
                x = (1 / torch.sqrt(alpha)) * (
                    x - ((1 - alpha) / torch.sqrt(1 - alpha_cumprod)) * predicted_noise
                ) + torch.sqrt(beta) * noise

        model.train()
        return x
    
def test_ddpm():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = UNet().to(device)
    ddpm = DDPM(n_steps=1000)

    print("Testing forward noise addition")
    x_0 = torch.randn(4, 1, 32, 32).to(device)  
    t = torch.randint(0, 1000, (4,)).to(device)  
    noisy_images, noise = ddpm.add_noise(x_0, t)
    assert noisy_images.shape == x_0.shape, f"Shape mismatch: {noisy_images.shape} vs {x_0.shape}"
    print("✓ Forward noise addition test passed")

    print("\nTTesting model forward pass")
    predicted_noise = model(noisy_images, t)
    assert predicted_noise.shape == noise.shape, f"Shape mismatch: {predicted_noise.shape} vs {noise.shape}"
    print("✓ Model forward pass test passed")

    print("\nTesting sampling process")
    samples = ddpm.sample(model, n_samples=2, img_size=32, device=device)
    assert samples.shape == (2, 1, 32, 32), f"Wrong sample shape: {samples.shape}"
    assert not torch.isnan(samples).any(), "Samples contain NaN values"
    print("✓ Sampling test passed")

    print("\nAll tests passed successfully!")

if __name__ == "__main__":
    test_ddpm()