novelkathor/DDIM_Image_Generation

1"""Inference: load a checkpoint and generate samples / trajectory / interp grid.2 3Usage:4    # 16 random faces with DDIM 50 steps5    python3 sample.py --ckpt checkpoints/stage-256_best.pt --num 16 --steps 506 7    # save denoising trajectory as a GIF8    python3 sample.py --ckpt checkpoints/stage-256_best.pt --trajectory \9        --num 4 --steps 50 --out samples/traj.gif10 11    # interpolate between two random latents (8 frames, slerp)12    python3 sample.py --ckpt checkpoints/stage-256_best.pt --interpolate 8 \13        --out samples/interp.png14 15    # DDPM-1000 vs DDIM-50 side-by-side16    python3 sample.py --ckpt ... --compare-ddpm --num 417"""18from __future__ import annotations19 20import argparse21import os22from typing import Optional23 24os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")25 26import torch27 28from config import Config29from models.unet import UNet30from models.diffusion import GaussianDiffusion, EMA31from utils.visualize import (save_image_grid, trajectory_to_gif,32                             interpolate_latents, make_grid)33from PIL import Image34 35 36# ---------------------------------------------------------------------------37def load_run(ckpt_path: str, device: torch.device, prefer_ema: bool = True):38    payload = torch.load(ckpt_path, map_location=device)39    cfg_dict = payload["config"]40    cfg = Config(**cfg_dict)41    model = UNet(42        image_size=cfg.image_size,43        in_channels=cfg.in_channels,44        base_channels=cfg.base_channels,45        channel_mults=cfg.channel_mults,46        num_res_blocks=cfg.num_res_blocks,47        attn_resolutions=cfg.attn_resolutions,48        time_embed_dim=cfg.time_embed_dim,49        dropout=cfg.dropout,50    ).to(device)51    if prefer_ema and payload.get("ema") is not None:52        model.load_state_dict(payload["ema"], strict=True)53        print("[sample] loaded EMA weights")54    else:55        model.load_state_dict(payload["model"], strict=True)56        print("[sample] loaded raw weights")57    model.eval()58    diffusion = GaussianDiffusion(59        timesteps=cfg.timesteps, beta_start=cfg.beta_start,60        beta_end=cfg.beta_end, schedule=cfg.beta_schedule,61    ).to(device)62    return cfg, model, diffusion63 64 65# ---------------------------------------------------------------------------66def parse_args():67    p = argparse.ArgumentParser()68    p.add_argument("--ckpt", required=True)69    p.add_argument("--num", type=int, default=16)70    p.add_argument("--steps", type=int, default=50)71    p.add_argument("--eta", type=float, default=0.0)72    p.add_argument("--seed", type=int, default=None)73    p.add_argument("--out", type=str, default=None)74    p.add_argument("--no-ema", action="store_true")75    p.add_argument("--device", type=str, default=None)76    # mode flags77    p.add_argument("--trajectory", action="store_true",78                   help="save denoising trajectory as a GIF")79    p.add_argument("--interpolate", type=int, default=0,80                   help="number of interpolation frames between two latents")81    p.add_argument("--compare-ddpm", action="store_true",82                   help="generate DDIM-N vs DDPM-T side-by-side comparison")83    return p.parse_args()84 85 86# ---------------------------------------------------------------------------87def main():88    args = parse_args()89    device = torch.device(args.device or ("mps" if torch.backends.mps.is_available() else "cpu"))90    cfg, model, diffusion = load_run(args.ckpt, device, prefer_ema=not args.no_ema)91    print(f"[sample] image_size={cfg.image_size} run={cfg.run_name} device={device}")92 93    if args.seed is not None:94        torch.manual_seed(args.seed)95 96    shape = (args.num, cfg.in_channels, cfg.image_size, cfg.image_size)97    out_dir = cfg.sample_dir98    os.makedirs(out_dir, exist_ok=True)99 100    # ---- interpolation -------------------------------------------------101    if args.interpolate > 0:102        n = args.interpolate103        z1 = torch.randn(1, *shape[1:], device=device)104        z2 = torch.randn(1, *shape[1:], device=device)105        latents = interpolate_latents(z1.cpu(), z2.cpu(), num_steps=n).squeeze(1).to(device)106        # latents shape: (n, C, H, W). One sampling pass per frame.107        with torch.no_grad():108            samples = diffusion.ddim_sample(109                model, (n, *shape[1:]), num_steps=args.steps, eta=args.eta,110                x_T=latents, device=device,111            )112        out = args.out or os.path.join(out_dir, f"interp_{n}.png")113        save_image_grid(samples.cpu(), out, nrow=n)114        print(f"[sample] interpolation saved -> {out}")115        return116 117    # ---- trajectory GIF ------------------------------------------------118    if args.trajectory:119        x_T = torch.randn(shape, device=device)120        with torch.no_grad():121            _, traj = diffusion.ddim_sample(122                model, shape, num_steps=args.steps, eta=args.eta,123                x_T=x_T, device=device,124                return_trajectory=True, trajectory_stride=1,125            )126        out = args.out or os.path.join(out_dir, f"traj_{args.steps}.gif")127        trajectory_to_gif(traj, out, fps=10)128        print(f"[sample] trajectory saved -> {out}")129        return130 131    # ---- DDIM vs DDPM comparison --------------------------------------132    if args.compare_ddpm:133        x_T = torch.randn(shape, device=device)134        with torch.no_grad():135            ddim = diffusion.ddim_sample(model, shape, num_steps=args.steps,136                                         eta=args.eta, x_T=x_T.clone(), device=device)137            ddpm = diffusion.ddim_sample(model, shape, num_steps=cfg.timesteps,138                                         eta=1.0, x_T=x_T.clone(), device=device)139        # stack as 2 rows140        side = torch.cat([ddim.cpu(), ddpm.cpu()], dim=0)141        out = args.out or os.path.join(out_dir, f"compare_ddim{args.steps}_vs_ddpm.png")142        save_image_grid(side, out, nrow=args.num)143        print(f"[sample] comparison saved -> {out}  (top: DDIM-{args.steps}, bottom: DDPM-{cfg.timesteps})")144        return145 146    # ---- default: simple grid -----------------------------------------147    with torch.no_grad():148        samples = diffusion.ddim_sample(149            model, shape, num_steps=args.steps, eta=args.eta, device=device,150        )151    out = args.out or os.path.join(out_dir, f"samples_n{args.num}_s{args.steps}.png")152    save_image_grid(samples.cpu(), out)153    print(f"[sample] grid saved -> {out}")154 155 156if __name__ == "__main__":157    main()158