Denoising Diffusion Implicit Model (DDIM) Implementation in PyTorch

This implementation builds a Denoising Diffusion Implicit Model (DDIM) sampler from scratch. It includes a UNet denoiser with GELU time conditioning, a linear-β noise schedule, an eta-controlled reverse process that can run on a sub-sequence of timesteps for accelerated sampling, and a test script that verifies noise addition, model output, sampling shapes, and deterministic reproducibility.

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from tqdm import tqdm
from typing import List, Optional, Tuple

class UNet(nn.Module):
    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.GELU(),
            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):
        t = t.float().unsqueeze(-1)
        t = self.time_mlp(t)
        
        x = self.init_conv(x)
        x1 = F.gelu(x)

        x2 = F.gelu(self.down1(x1))
        x3 = F.gelu(self.down2(x2))
        
        t = t.view(-1, t.shape[1], 1, 1).expand(-1, -1, x3.shape[2], x3.shape[3])
        x3 = x3 + t

        x = F.gelu(self.up1(x3))
        x = F.gelu(self.up2(x))
        
        return self.final_conv(x)
    
class DDIM:
    def __init__(
        self,
        n_steps: int = 1000,
        n_sampling_steps: int = 50,
        beta_start: float = 1e-4,
        beta_end: float = 0.02,
        eta: float = 0.0, 
        seed: Optional[int] = None
    ):
        """    
        Args:
            n_steps: Total number of diffusion steps
            n_sampling_steps: Number of steps to use for sampling (can be < n_steps)
            beta_start: Starting value for noise schedule
            beta_end: Ending value for noise schedule
            eta: Controls the stochasticity (0 for deterministic, 1 for stochastic)
            seed: Random seed 
        """
        self.n_steps = n_steps
        self.n_sampling_steps = n_sampling_steps
        self.eta = eta
        self.seed = seed

        # 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)

        self.sampling_timesteps = torch.linspace(
            0, n_steps - 1, n_sampling_steps, dtype=torch.long
        )

    def set_seed(self, seed: Optional[int] = None):
        if seed is not None:
            torch.manual_seed(seed)
            if torch.cuda.is_available():
                torch.cuda.manual_seed_all(seed)
            np.random.seed(seed)

    def add_noise(self, x_0: torch.Tensor, t: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        noise = torch.randn_like(x_0)
        sqrt_alpha_cumprod = torch.sqrt(self.alphas_cumprod[t]).view(-1, 1, 1, 1)
        sqrt_one_minus_alpha_cumprod = torch.sqrt(1 - self.alphas_cumprod[t]).view(-1, 1, 1, 1)
        
        return sqrt_alpha_cumprod * x_0 + sqrt_one_minus_alpha_cumprod * noise, noise
    
    def sample(
        self,
        model: nn.Module,
        n_samples: int,
        img_size: int,
        device: torch.device,
        clip_denoised: bool = True,
        seed: Optional[int] = None
    ) -> torch.Tensor:
        """
        Generate samples using the DDIM sampling process.
        """
        if seed is not None or self.seed is not None:
            self.set_seed(seed if seed is not None else self.seed)

        model.eval()
        with torch.no_grad():
            x = torch.randn(n_samples, 1, img_size, img_size).to(device)

            # denoise
            for i in tqdm(reversed(range(self.n_sampling_steps)), desc='DDIM Sampling'):
                timestep = self.sampling_timesteps[i]
                timestep_next = self.sampling_timesteps[i-1] if i > 0 else torch.tensor([-1])
                
                t_batch = torch.ones(n_samples, dtype=torch.long, device=device) * timestep
                
                predicted_noise = model(x, t_batch)

                # alpha values for current and next timestep
                alpha_cumprod_t = self.alphas_cumprod[timestep]
                alpha_cumprod_t_next = self.alphas_cumprod[timestep_next] if i > 0 else torch.tensor(1.0)
                
                x_0_predicted = (
                    x - torch.sqrt(1 - alpha_cumprod_t).view(-1, 1, 1, 1) * predicted_noise
                ) / torch.sqrt(alpha_cumprod_t).view(-1, 1, 1, 1)
                
                if clip_denoised:
                    x_0_predicted = torch.clamp(x_0_predicted, -1, 1)

                # Calculate sigma for noise
                sigma_t = self.eta * torch.sqrt(
                    (1 - alpha_cumprod_t_next) / (1 - alpha_cumprod_t) *
                    (1 - alpha_cumprod_t / alpha_cumprod_t_next)
                )

                c1 = torch.sqrt(alpha_cumprod_t_next)
                c2 = torch.sqrt(1 - alpha_cumprod_t_next - sigma_t**2)
                
                if i > 0:
                    noise = torch.randn_like(x) if self.eta > 0 else 0
                    x = c1.view(-1, 1, 1, 1) * x_0_predicted + \
                        c2.view(-1, 1, 1, 1) * predicted_noise + \
                        sigma_t.view(-1, 1, 1, 1) * noise
                else:
                    x = x_0_predicted

        model.train()
        return x
    
def test_ddim():
    torch.manual_seed(42)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(42)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = UNet().to(device)
    ddim = DDIM(n_steps=1000, n_sampling_steps=50, eta=0.0, seed=42)

    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 = ddim.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("\nTesting 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 = ddim.sample(model, n_samples=2, img_size=32, device=device, seed=42)
    assert samples.shape == (2, 1, 32, 32), f"Wrong sample shape: {samples.shape}"
    assert not torch.isnan(samples).any(), "Samples contain NaN values"

    samples2 = ddim.sample(model, n_samples=2, img_size=32, device=device, seed=42)
    assert torch.allclose(samples, samples2, atol=1e-5), "Deterministic sampling failed"
    print("✓ Sampling test passed")

    print("\nAll tests passed successfully!")

if __name__ == "__main__":
    test_ddim()