Spaces:
Runtime error
Runtime error
| # sr_diffusion.py | |
| # ----------------------------------------------------------- | |
| # Super-Resolution Diffusion for CIFAR-10 (32×32 → 128×128) | |
| # Stage-2 after your text-to-image model | |
| # ----------------------------------------------------------- | |
| import os, math, argparse, random, glob | |
| from typing import Optional | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.utils.data import DataLoader | |
| from torchvision import datasets, transforms, utils as vutils | |
| from tqdm import tqdm | |
| # ------------------ utils ------------------ | |
| def set_seed(seed: int = 0): | |
| random.seed(seed) | |
| torch.manual_seed(seed) | |
| torch.cuda.manual_seed_all(seed) | |
| def _raw_module(m): | |
| return getattr(m, "_orig_mod", m) | |
| # ------------------ time embedding ------------------ | |
| class SinusoidalTimeEmbedding(nn.Module): | |
| def __init__(self, dim): | |
| super().__init__() | |
| self.dim = dim | |
| def forward(self, t): | |
| half = self.dim // 2 | |
| emb = torch.log(torch.tensor(10000.0)) / (half - 1) | |
| emb = torch.exp(torch.arange(half, device=t.device) * -emb) | |
| emb = t[:, None].float() * emb[None, :] | |
| emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) | |
| if self.dim % 2 == 1: | |
| emb = F.pad(emb, (0,1)) | |
| return emb | |
| # ------------------ building blocks ------------------ | |
| class ResBlock(nn.Module): | |
| def __init__(self, dim_in, dim_out, cond_dim): | |
| super().__init__() | |
| self.norm1 = nn.GroupNorm(32, dim_in) | |
| self.conv1 = nn.Conv2d(dim_in, dim_out, 3, padding=1) | |
| self.norm2 = nn.GroupNorm(32, dim_out) | |
| self.conv2 = nn.Conv2d(dim_out, dim_out, 3, padding=1) | |
| self.emb = nn.Sequential( | |
| nn.SiLU(), | |
| nn.Linear(cond_dim, dim_out * 2) | |
| ) | |
| self.act = nn.SiLU() | |
| self.skip = nn.Conv2d(dim_in, dim_out, 1) if dim_in != dim_out else nn.Identity() | |
| def forward(self, x, cond): | |
| h = self.conv1(self.act(self.norm1(x))) | |
| scale, shift = self.emb(cond).chunk(2, dim=1) | |
| scale, shift = scale[:, :, None, None], shift[:, :, None, None] | |
| h = h * (1 + scale) + shift | |
| h = self.conv2(self.act(self.norm2(h))) | |
| return h + self.skip(x) | |
| class Down(nn.Module): | |
| def __init__(self, dim_in, dim_out, cond_dim): | |
| super().__init__() | |
| self.block1 = ResBlock(dim_in, dim_out, cond_dim) | |
| self.block2 = ResBlock(dim_out, dim_out, cond_dim) | |
| self.down = nn.Conv2d(dim_out, dim_out, 4, stride=2, padding=1) | |
| def forward(self, x, cond): | |
| x = self.block1(x, cond) | |
| x = self.block2(x, cond) | |
| x = self.down(x) | |
| return x | |
| class Up(nn.Module): | |
| def __init__(self, dim_in, dim_out, cond_dim): | |
| super().__init__() | |
| self.block1 = ResBlock(dim_in, dim_out, cond_dim) | |
| self.block2 = ResBlock(dim_out, dim_out, cond_dim) | |
| self.up = nn.ConvTranspose2d(dim_out, dim_out, 4, stride=2, padding=1) | |
| def forward(self, x, cond): | |
| x = self.block1(x, cond) | |
| x = self.block2(x, cond) | |
| x = self.up(x) | |
| return x | |
| class UNetSR(nn.Module): | |
| def __init__(self, in_ch=3, base=64, cond_dim=256): | |
| super().__init__() | |
| self.in_conv = nn.Conv2d(in_ch, base, 3, padding=1) | |
| self.down1 = Down(base, base*2, cond_dim) | |
| self.down2 = Down(base*2, base*4, cond_dim) | |
| self.mid1 = ResBlock(base*4, base*4, cond_dim) | |
| self.mid2 = ResBlock(base*4, base*4, cond_dim) | |
| self.up2 = Up(base*4, base*2, cond_dim) | |
| self.up1 = Up(base*2, base, cond_dim) | |
| self.out_norm = nn.GroupNorm(8, base) | |
| self.out_conv = nn.Conv2d(base, in_ch, 3, padding=1) | |
| self.act = nn.SiLU() | |
| def forward(self, x, cond): | |
| x0 = self.in_conv(x) | |
| d1 = self.down1(x0, cond) | |
| d2 = self.down2(d1, cond) | |
| m = self.mid1(d2, cond) | |
| m = self.mid2(m, cond) | |
| u2 = self.up2(m + d2, cond) | |
| u1 = self.up1(u2 + d1, cond) | |
| h = self.act(self.out_norm(u1)) | |
| return self.out_conv(h) | |
| # ------------------ diffusion ------------------ | |
| class Diffusion(nn.Module): | |
| def __init__(self, model: nn.Module, T=1000, schedule="cosine"): | |
| super().__init__() | |
| self.model = model | |
| self.T = T | |
| self.register_buffer("betas", self.make_betas(T, schedule)) | |
| alphas = 1.0 - self.betas | |
| alphas_cumprod = torch.cumprod(alphas, dim=0) | |
| self.register_buffer("alphas_cumprod", alphas_cumprod) | |
| self.register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod)) | |
| self.register_buffer("sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod)) | |
| from diffusers.schedulers.scheduling_ddpm import DDPMScheduler | |
| self.scheduler = DDPMScheduler(num_train_timesteps=self.T, beta_schedule="squaredcos_cap_v2") | |
| def make_betas(T, schedule): | |
| if schedule == "linear": | |
| return torch.linspace(1e-4, 0.02, T) | |
| s = 0.008 | |
| steps = T + 1 | |
| x = torch.linspace(0, T, steps) | |
| alphas_cum = torch.cos(((x / T) + s) / (1 + s) * math.pi * 0.5) ** 2 | |
| alphas_cum = alphas_cum / alphas_cum[0] | |
| betas = 1 - (alphas_cum[1:] / alphas_cum[:-1]) | |
| return betas.clamp(max=0.999) | |
| def sample(self, shape, steps=1000, cond=None): | |
| device = next(self.parameters()).device | |
| x = torch.randn(shape, device=device) | |
| self.scheduler.set_timesteps(steps) | |
| for t in self.scheduler.timesteps: | |
| eps_pred = self.model(x, torch.full((shape[0],), t, device=device, dtype=torch.long), cond) | |
| x = self.scheduler.step(eps_pred, t, x).prev_sample | |
| return x | |
| # ------------------ Wrapper ------------------ | |
| class CondUNetSR(nn.Module): | |
| def __init__(self, base=64, time_dim=128, cond_dim=256): | |
| super().__init__() | |
| self.time_emb = SinusoidalTimeEmbedding(time_dim) | |
| self.time_mlp = nn.Sequential( | |
| nn.Linear(time_dim, cond_dim), nn.SiLU(), nn.Linear(cond_dim, cond_dim) | |
| ) | |
| # learnable projection for conditioning image | |
| self.cond_proj = nn.Sequential( | |
| nn.Conv2d(3, 16, 3, stride=2, padding=1), | |
| nn.SiLU(), | |
| nn.Conv2d(16, 32, 3, stride=2, padding=1), | |
| nn.SiLU(), | |
| nn.AdaptiveAvgPool2d(1), | |
| nn.Flatten(), | |
| nn.Linear(32, cond_dim) | |
| ) | |
| self.net = UNetSR(in_ch=3, base=base, cond_dim=cond_dim) | |
| def forward(self, x, t, cond_img=None): | |
| t_emb = self.time_mlp(self.time_emb(t)) # [B, 256] | |
| if cond_img is not None: | |
| cond_emb = self.cond_proj(cond_img) # [B, 256] | |
| cond = t_emb + cond_emb | |
| else: | |
| cond = t_emb | |
| return self.net(x, cond) | |
| # ------------------ Trainer ------------------ | |
| class SRTrainer: | |
| def __init__(self, args): | |
| self.args = args | |
| set_seed(args.seed) | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| tfm_hr = transforms.Compose([ | |
| transforms.Resize(args.hr_size), | |
| transforms.ToTensor(), | |
| transforms.Lambda(lambda x: x * 2 - 1), | |
| ]) | |
| self.ds = datasets.CIFAR10(args.data_root, train=True, transform=tfm_hr, download=True) | |
| self.loader = DataLoader(self.ds, batch_size=args.bs, shuffle=True, | |
| num_workers=args.workers, drop_last=True) | |
| self.model = CondUNetSR(base=args.base, time_dim=128, cond_dim=args.cond_dim).to(self.device) | |
| self.diff = Diffusion(self.model, T=args.T, schedule=args.schedule).to(self.device) | |
| self.opt = torch.optim.AdamW(self.model.parameters(), lr=args.lr) | |
| os.makedirs(args.outdir, exist_ok=True) | |
| # preview set | |
| test_ds = datasets.CIFAR10(args.data_root, train=False, | |
| transform=transforms.ToTensor(), download=True) | |
| self.sample_imgs = torch.stack([test_ds[i][0] for i in range(16)]).to(self.device) * 2 - 1 | |
| # auto resume | |
| self.start_epoch = 0 | |
| if args.resume: | |
| self.start_epoch = self.try_resume() | |
| # ---- checkpoints ---- | |
| def save_ckpt(self, epoch): | |
| ckpt = { | |
| "model": _raw_module(self.model).state_dict(), | |
| "opt": self.opt.state_dict(), | |
| "epoch": epoch | |
| } | |
| path = os.path.join(self.args.outdir, f"ckpt_epoch{epoch}.pt") | |
| torch.save(ckpt, path) | |
| print(f"[ckpt] saved {path}") | |
| def load_ckpt(self, path): | |
| ckpt = torch.load(path, map_location=self.device) | |
| self.model.load_state_dict(ckpt["model"]) | |
| self.opt.load_state_dict(ckpt["opt"]) | |
| print(f"[ckpt] loaded {path} (epoch {ckpt['epoch']})") | |
| return ckpt["epoch"] | |
| def try_resume(self): | |
| ckpts = sorted(glob.glob(os.path.join(self.args.outdir, "ckpt_epoch*.pt"))) | |
| if len(ckpts) == 0: | |
| print("[resume] no checkpoints found") | |
| return 0 | |
| last_ckpt = ckpts[-1] | |
| epoch = self.load_ckpt(last_ckpt) | |
| return epoch | |
| # ---- training ---- | |
| def train(self): | |
| self.model.train() | |
| for epoch in range(self.start_epoch, self.args.epochs): | |
| losses = [] | |
| pbar = tqdm(self.loader, desc=f"Epoch {epoch+1}/{self.args.epochs}") | |
| for hr, _ in pbar: | |
| hr = hr.to(self.device) | |
| lr = F.interpolate(hr, size=32, mode="bilinear") # low-res input | |
| t = torch.randint(0, self.args.T, (hr.size(0),), device=self.device) | |
| noise = torch.randn_like(hr) | |
| a_t = self.diff.sqrt_alphas_cumprod[t][:, None, None, None] | |
| am1_t = self.diff.sqrt_one_minus_alphas_cumprod[t][:, None, None, None] | |
| hr_noisy = a_t * hr + am1_t * noise | |
| pred = self.model(hr_noisy, t, cond_img=lr) | |
| loss = F.mse_loss(pred, noise) | |
| self.opt.zero_grad(set_to_none=True) | |
| loss.backward() | |
| nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) | |
| self.opt.step() | |
| losses.append(loss.item()) | |
| pbar.set_postfix(loss=loss.item()) | |
| avg_loss = sum(losses) / len(losses) | |
| print(f"Epoch {epoch+1}/{self.args.epochs} | Avg Loss: {avg_loss:.4f}") | |
| # ---- checkpoint + sample ---- | |
| if (epoch + 1) % self.args.sample_every == 0: | |
| self.save_ckpt(epoch + 1) | |
| self.model.eval() | |
| with torch.no_grad(): | |
| sr = self.super_resolve(self.sample_imgs, steps=self.args.steps) | |
| grid = vutils.make_grid(sr, nrow=4) | |
| save_path = os.path.join(self.args.outdir, f"sr_epoch{epoch+1}.png") | |
| vutils.save_image(grid, save_path) | |
| print(f"[sample] saved {save_path}") | |
| self.model.train() | |
| def super_resolve(self, lr_imgs, steps=1000): | |
| self.model.eval() | |
| sr_imgs = self.diff.sample((lr_imgs.size(0), 3, self.args.hr_size, self.args.hr_size), | |
| steps=steps, cond=lr_imgs) | |
| return (sr_imgs.clamp(-1,1)+1)*0.5 | |
| # ------------------ CLI ------------------ | |
| def build_parser(): | |
| p = argparse.ArgumentParser(description="Super-Resolution Diffusion (CIFAR 32→HR)") | |
| p.add_argument("--train", action="store_true") | |
| p.add_argument("--sample", action="store_true") | |
| p.add_argument("--data_root", type=str, default="./data") | |
| p.add_argument("--epochs", type=int, default=20) | |
| p.add_argument("--bs", type=int, default=128) | |
| p.add_argument("--lr", type=float, default=1e-4) | |
| p.add_argument("--T", type=int, default=1000) | |
| p.add_argument("--schedule", type=str, default="cosine", choices=["cosine", "linear"]) | |
| p.add_argument("--base", type=int, default=64) | |
| p.add_argument("--cond_dim", type=int, default=256) | |
| p.add_argument("--workers", type=int, default=4) | |
| p.add_argument("--seed", type=int, default=0) | |
| p.add_argument("--outdir", type=str, default="outputs_sr") | |
| p.add_argument("--hr_size", type=int, default=128) | |
| p.add_argument("--steps", type=int, default=50) | |
| p.add_argument("--sample_every", type=int, default=10, | |
| help="Save checkpoint + preview samples every N epochs") | |
| p.add_argument("--resume", action="store_true", | |
| help="Auto-resume from the latest checkpoint in outdir") | |
| p.add_argument("--prompts", nargs="+", type=str, default=None, | |
| help="Text prompts for text→image→super-res pipeline") | |
| return p | |
| def main(): | |
| args = build_parser().parse_args() | |
| trainer = SRTrainer(args) | |
| if args.train: | |
| trainer.train() | |
| if args.sample: | |
| # ---- load checkpoint before sampling ---- | |
| ckpts = sorted(glob.glob(os.path.join(args.outdir, "ckpt_epoch*.pt"))) | |
| if len(ckpts) == 0: | |
| print("[sample] No checkpoints found in outdir, cannot sample") | |
| return | |
| last_ckpt = ckpts[-1] | |
| trainer.load_ckpt(last_ckpt) | |
| if args.prompts: | |
| # -------- Stage-1: text-to-image (32x32) -------- | |
| from t2i_text import Trainer as T2ITrainer, build_parser as t2i_parser | |
| # load stage-1 checkpoint | |
| t2i_args = t2i_parser().parse_args([]) | |
| t2i_trainer = T2ITrainer(t2i_args) | |
| t2i_trainer.load_ckpt("outputs/ckpt_latest.pt") # adjust path if needed | |
| print(f"[stage-1] generating low-res 32x32 from prompts: {args.prompts}") | |
| lr_imgs = t2i_trainer.sample(prompts=args.prompts, steps=50, n_per_prompt=1) | |
| lr_imgs = lr_imgs.to(trainer.device) * 2 - 1 # normalize to [-1,1] | |
| else: | |
| # -------- fallback: use CIFAR test images -------- | |
| ds = datasets.CIFAR10(args.data_root, train=False, | |
| transform=transforms.ToTensor(), download=True) | |
| lr_imgs = torch.stack([ds[i][0] for i in range(16)]) | |
| lr_imgs = (lr_imgs*2-1).to(trainer.device) | |
| # -------- Stage-2: super-resolution (32→128) -------- | |
| sr = trainer.super_resolve(lr_imgs, steps=args.steps) | |
| grid = vutils.make_grid(sr, nrow=4) | |
| os.makedirs(args.outdir, exist_ok=True) | |
| if args.prompts: | |
| save_path = os.path.join(args.outdir, "sr_from_prompt.png") | |
| else: | |
| save_path = os.path.join(args.outdir, "sr_samples.png") | |
| vutils.save_image(grid, save_path) | |
| print(f"[sample] saved {save_path}") | |
| if __name__ == "__main__": | |
| main() | |