sample.py

import torch
import torch.nn as nn
from torchvision.utils import save_image
from unet import UNet
from linear_alpha import DDPMScheduler
from tqdm import tqdm

def generate_samples():
    num_images=16
    time_steps=1000

    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # Initialize Model
    model=UNet(in_channels=1, out_channels=1, time_emb_dim=32).to(device)

    # Load the checkpoint file
    try:
        model.load_state_dict(torch.load("ddpm_mnist_epoch_20.pth", map_location=device))
    except FileNotFoundError:
        print("Warning: 'ddpm_mnist_epoch_20.pth' not found. Ensure you have trained the model first.")
    
    model.eval()

    scheduler=DDPMScheduler(num_train_timesteps=time_steps)
    
    # Move scheduler tensors to device once for efficiency
    scheduler.betas = scheduler.betas.to(device)
    scheduler.alphas = scheduler.alphas.to(device)
    scheduler.alphas_cumprod = scheduler.alphas_cumprod.to(device)

    # Start from pure gaussian noise: [16, 1, 28, 28]
    x=torch.randn(num_images, 1, 28, 28, device=device)

    print("Running reverse diffusion Langevin loop...")
    with torch.no_grad():
        for t in reversed(tqdm(range(time_steps), desc="Sampling")):
            # Create a batch tensor filled with current timestep value
            t_tensor=torch.full((num_images,), t, device=device, dtype=torch.long)

            # Predict the noise inside the image
            predicted_noise=model(x, t_tensor)

            # Extract constants
            beta = scheduler.betas[t]
            alpha = scheduler.alphas[t]
            alpha_hat = scheduler.alphas_cumprod[t]

            # Calculate the variance factor (sigma_t)
            if t > 0:
                noise_jitter = torch.randn_like(x)
                sigma = torch.sqrt(beta)
            else:
                noise_jitter = torch.zeros_like(x)
                sigma = 0.0

            # Execute the reverse process step equation
            x = (1 / torch.sqrt(alpha)) * (x - ((beta / torch.sqrt(1 - alpha_hat)) * predicted_noise)) + (sigma * noise_jitter)
    
    # Scale from [-1, 1] back to [0, 1] for saving
    x=(x+1.0)/2.0
    x=torch.clamp(x, 0.0, 1.0)

    save_image(x, "ddpm_generated_digits.png", nrow=4)
    print("Sampling complete! Your images are saved as 'ddpm_generated_digits.png'.")

if __name__ == "__main__":
    generate_samples()