"""Proof-of-life DDPM on CIFAR-10. Trains a small U-Net to predict noise (Ho et al. 2020, simplified). Saves sample grids and checkpoints to diffusion/out/ as it goes. Run: python diffusion/train_diffusion.py [--base 128 --epochs 300 --out out_big --resume] """ import argparse import copy import math import os import time import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms from torchvision.utils import save_image # ---------------- config ---------------- DATA_ROOT = os.path.join(os.path.dirname(__file__), "..", "cifar-10-python") BATCH_SIZE = 128 LR = 2e-4 T = 1000 # diffusion timesteps DDIM_STEPS = 50 # sampling steps DEVICE = "cuda" # ---------------- model ---------------- class TimeEmbedding(nn.Module): def __init__(self, dim): super().__init__() self.dim = dim self.mlp = nn.Sequential(nn.Linear(dim, dim * 4), nn.SiLU(), nn.Linear(dim * 4, dim * 4)) def forward(self, t): half = self.dim // 2 freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half) args = t[:, None].float() * freqs[None] emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1) return self.mlp(emb) class ResBlock(nn.Module): def __init__(self, in_ch, out_ch, emb_dim): super().__init__() self.norm1 = nn.GroupNorm(8, in_ch) self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1) self.emb_proj = nn.Linear(emb_dim, out_ch) self.norm2 = nn.GroupNorm(8, out_ch) self.dropout = nn.Dropout(0.1) self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1) self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity() def forward(self, x, emb): h = self.conv1(F.silu(self.norm1(x))) h = h + self.emb_proj(F.silu(emb))[:, :, None, None] h = self.conv2(self.dropout(F.silu(self.norm2(h)))) return h + self.skip(x) class Attention(nn.Module): def __init__(self, ch): super().__init__() self.norm = nn.GroupNorm(8, ch) self.qkv = nn.Conv2d(ch, ch * 3, 1) self.proj = nn.Conv2d(ch, ch, 1) def forward(self, x): b, c, h, w = x.shape q, k, v = self.qkv(self.norm(x)).reshape(b, 3, c, h * w).unbind(1) attn = torch.softmax(q.transpose(1, 2) @ k / math.sqrt(c), dim=-1) out = (v @ attn.transpose(1, 2)).reshape(b, c, h, w) return x + self.proj(out) class UNet(nn.Module): """Small U-Net: 32 -> 16 -> 8, attention at 16 and 8.""" def __init__(self, base=64): super().__init__() chs = [base, base * 2, base * 2] # per resolution level emb_dim = base * 4 self.time_emb = TimeEmbedding(base) self.stem = nn.Conv2d(3, base, 3, padding=1) # down self.d1a = ResBlock(base, chs[0], emb_dim) self.d1b = ResBlock(chs[0], chs[0], emb_dim) self.down1 = nn.Conv2d(chs[0], chs[0], 3, stride=2, padding=1) self.d2a = ResBlock(chs[0], chs[1], emb_dim) self.d2b = ResBlock(chs[1], chs[1], emb_dim) self.attn2 = Attention(chs[1]) self.down2 = nn.Conv2d(chs[1], chs[1], 3, stride=2, padding=1) self.d3a = ResBlock(chs[1], chs[2], emb_dim) self.d3b = ResBlock(chs[2], chs[2], emb_dim) self.attn3 = Attention(chs[2]) # middle self.mid1 = ResBlock(chs[2], chs[2], emb_dim) self.mid_attn = Attention(chs[2]) self.mid2 = ResBlock(chs[2], chs[2], emb_dim) # up self.u3a = ResBlock(chs[2] * 2, chs[2], emb_dim) self.u3b = ResBlock(chs[2], chs[1], emb_dim) self.uattn3 = Attention(chs[1]) self.up2 = nn.ConvTranspose2d(chs[1], chs[1], 4, stride=2, padding=1) self.u2a = ResBlock(chs[1] * 2, chs[1], emb_dim) self.u2b = ResBlock(chs[1], chs[0], emb_dim) self.uattn2 = Attention(chs[0]) self.up1 = nn.ConvTranspose2d(chs[0], chs[0], 4, stride=2, padding=1) self.u1a = ResBlock(chs[0] * 2, chs[0], emb_dim) self.u1b = ResBlock(chs[0], chs[0], emb_dim) self.out_norm = nn.GroupNorm(8, chs[0]) self.out = nn.Conv2d(chs[0], 3, 3, padding=1) def forward(self, x, t): emb = self.time_emb(t) h1 = self.stem(x) h1 = self.d1b(self.d1a(h1, emb), emb) # 32 h2 = self.down1(h1) h2 = self.attn2(self.d2b(self.d2a(h2, emb), emb)) # 16 h3 = self.down2(h2) h3 = self.attn3(self.d3b(self.d3a(h3, emb), emb)) # 8 m = self.mid2(self.mid_attn(self.mid1(h3, emb)), emb) u = self.u3b(self.u3a(torch.cat([m, h3], 1), emb), emb) u = self.uattn3(u) u = self.up2(u) u = self.u2b(self.u2a(torch.cat([u, h2], 1), emb), emb) u = self.uattn2(u) u = self.up1(u) u = self.u1b(self.u1a(torch.cat([u, h1], 1), emb), emb) return self.out(F.silu(self.out_norm(u))) # ---------------- diffusion ---------------- class Diffusion: def __init__(self, T, device): self.T = T beta = torch.linspace(1e-4, 0.02, T, device=device) self.alpha_bar = torch.cumprod(1 - beta, dim=0) def add_noise(self, x0, t, noise): ab = self.alpha_bar[t][:, None, None, None] return ab.sqrt() * x0 + (1 - ab).sqrt() * noise @torch.no_grad() def ddim_sample(self, model, n, steps, device): ts = torch.linspace(self.T - 1, 0, steps, device=device).long() x = torch.randn(n, 3, 32, 32, device=device) for i in range(steps): t = ts[i].expand(n) eps = model(x, t) ab_t = self.alpha_bar[ts[i]] x0 = (x - (1 - ab_t).sqrt() * eps) / ab_t.sqrt() x0 = x0.clamp(-1, 1) if i + 1 < steps: ab_prev = self.alpha_bar[ts[i + 1]] x = ab_prev.sqrt() * x0 + (1 - ab_prev).sqrt() * eps else: x = x0 return x # ---------------- training ---------------- def main(): ap = argparse.ArgumentParser() ap.add_argument("--base", type=int, default=64, help="U-Net base channels") ap.add_argument("--epochs", type=int, default=50) ap.add_argument("--out", default="out", help="output dir name (under diffusion/)") ap.add_argument("--sample-every", type=int, default=5) ap.add_argument("--ema-decay", type=float, default=0.999) ap.add_argument("--resume", action="store_true", help="resume from checkpoint.pt in out dir") args = ap.parse_args() out_dir = os.path.join(os.path.dirname(__file__), args.out) os.makedirs(out_dir, exist_ok=True) torch.backends.cudnn.benchmark = True tfm = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.5] * 3, [0.5] * 3), # -> [-1, 1] ]) ds = datasets.CIFAR10(DATA_ROOT, train=True, transform=tfm, download=True) dl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True, persistent_workers=True, drop_last=True) model = UNet(base=args.base).to(DEVICE) ema = copy.deepcopy(model).requires_grad_(False) opt = torch.optim.Adam(model.parameters(), lr=LR) scaler = torch.amp.GradScaler() diff = Diffusion(T, DEVICE) start_epoch = 1 ckpt_path = os.path.join(out_dir, "checkpoint.pt") if args.resume and os.path.exists(ckpt_path): ck = torch.load(ckpt_path, map_location=DEVICE, weights_only=True) model.load_state_dict(ck["model"]) ema.load_state_dict(ck["ema"]) opt.load_state_dict(ck["opt"]) start_epoch = ck["epoch"] + 1 print(f"resumed from epoch {ck['epoch']}", flush=True) n_params = sum(p.numel() for p in model.parameters()) print(f"params: {n_params/1e6:.1f}M | batches/epoch: {len(dl)}", flush=True) step = 0 for epoch in range(start_epoch, args.epochs + 1): t0 = time.time() running = 0.0 for x, _ in dl: x = x.to(DEVICE, non_blocking=True) t = torch.randint(0, T, (x.size(0),), device=DEVICE) noise = torch.randn_like(x) xt = diff.add_noise(x, t, noise) with torch.amp.autocast("cuda"): loss = F.mse_loss(model(xt, t), noise) opt.zero_grad(set_to_none=True) scaler.scale(loss).backward() scaler.step(opt) scaler.update() with torch.no_grad(): for pe, pm in zip(ema.parameters(), model.parameters()): pe.lerp_(pm, 1 - args.ema_decay) for be, bm in zip(ema.buffers(), model.buffers()): be.copy_(bm) running += loss.item() step += 1 dt = time.time() - t0 print(f"epoch {epoch:3d}/{args.epochs} | loss {running/len(dl):.4f} | " f"{dt:.0f}s | {len(dl)/dt:.1f} it/s", flush=True) if epoch % args.sample_every == 0 or epoch == args.epochs: ema.eval() imgs = diff.ddim_sample(ema, 64, DDIM_STEPS, DEVICE) save_image(imgs * 0.5 + 0.5, os.path.join(out_dir, f"samples_epoch{epoch:03d}.png"), nrow=8) torch.save({"model": model.state_dict(), "ema": ema.state_dict(), "opt": opt.state_dict(), "epoch": epoch}, ckpt_path) print(f" -> saved samples_epoch{epoch:03d}.png + checkpoint", flush=True) if __name__ == "__main__": main()