inference.py

import torch
import deepinv
from pathlib import Path
from torchvision.utils import save_image, make_grid
from tqdm import tqdm

device = "cuda" if torch.cuda.is_available() else "cpu"
image_size = 32
timesteps = 1000
n_samples = 32

checkpoint_path = "./trained_diffusion_model.pth"

# Load model
model = deepinv.models.DiffUNet(in_channels=1, out_channels=1, pretrained=None).to(device)
model.load_state_dict(torch.load(checkpoint_path, map_location=device))
model.eval()

# Diffusion schedule
beta_start, beta_end = 1e-4, 0.02
betas = torch.linspace(beta_start, beta_end, timesteps, device=device)
alphas = 1.0 - betas
alphas_cumprod = torch.cumprod(alphas, dim=0)
sqrt_recip_alphas = torch.sqrt(1.0 / alphas)
sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod)
sqrt_betas = torch.sqrt(betas)

# Sampling
with torch.no_grad():
    x = torch.randn(n_samples, 1, image_size, image_size, device=device)

    for t in tqdm(reversed(range(timesteps)), total=timesteps, desc="Sampling", ncols=100):
        t_tensor = torch.full((n_samples,), t, device=device, dtype=torch.long)
        predicted_noise = model(x, t_tensor, type_t="timestep")

        alpha = alphas[t]
        alpha_cumprod = alphas_cumprod[t]
        beta = betas[t]

        # Reverse diffusion step
        x = (
            sqrt_recip_alphas[t]
            * (x - (beta / torch.sqrt(1 - alpha_cumprod)) * predicted_noise)
        )

        if t > 0:
            noise = torch.randn_like(x)
            x += sqrt_betas[t] * noise

# Denormalize and save
x = (x.clamp(-1, 1) + 1) / 2.0  # map [-1,1] -> [0,1]
grid = make_grid(x, nrow=8)
save_image(grid, "samples.png")

print("Saved generated samples to samples.png")