| """Rectified-flow training. |
| |
| Convention: t=0 is data, t=1 is noise. |
| x_t = (1 - t) * x0 + t * eps, target velocity v = dx_t/dt = eps - x0, loss = ||v_theta(x_t,t,c) - v||^2 |
| t ~ logit-normal(0,1) (sigmoid of a standard normal) by default, concentrating on mid t. |
| """ |
| import argparse, copy, json, math, os, time |
| import numpy as np, torch, torch.nn.functional as F |
| from PIL import Image |
| from torch.utils.data import Dataset, DataLoader |
| from conditioning import tokenize |
| from model import build_model |
|
|
|
|
| class FaceDataset(Dataset): |
| def __init__(self, data_dir, resolution=None): |
| self.dir = data_dir |
| with open(os.path.join(data_dir, "metadata.jsonl")) as f: |
| self.rows = [json.loads(l) for l in f if l.strip()] |
| self.res = resolution |
|
|
| def __len__(self): return len(self.rows) |
|
|
| def __getitem__(self, i): |
| r = self.rows[i] |
| img = Image.open(os.path.join(self.dir, r["image_path"])).convert("RGB") |
| if self.res and img.size != (self.res, self.res): |
| img = img.resize((self.res, self.res), Image.BICUBIC) |
| x = torch.from_numpy(np.asarray(img).copy()).permute(2, 0, 1).float() / 127.5 - 1 |
| ids, mask = tokenize([r["caption"]]) |
| return x, ids[0], mask[0] |
|
|
|
|
| def sample_t(n, mode, device): |
| if mode == "logit_normal": |
| return torch.sigmoid(torch.randn(n, device=device)) |
| return torch.rand(n, device=device) |
|
|
|
|
| @torch.no_grad() |
| def ema_update(ema, model, decay): |
| for pe, p in zip(ema.parameters(), model.parameters()): |
| pe.lerp_(p.detach(), 1 - decay) |
|
|
|
|
| def lr_at(step, total, warmup, base): |
| if step < warmup: return base * (step + 1) / warmup |
| prog = min(1.0, (step - warmup) / max(1, total - warmup)) |
| return base * (0.05 + 0.95 * 0.5 * (1 + math.cos(math.pi * prog))) |
|
|
|
|
| def save(path, model, ema, opt, scaler, step, epoch, args): |
| tmp = path + ".tmp" |
| torch.save({"model": model.state_dict(), "ema": ema.state_dict(), "opt": opt.state_dict(), |
| "scaler": scaler.state_dict() if scaler else None, "step": step, "epoch": epoch, |
| "args": vars(args)}, tmp) |
| os.replace(tmp, path) |
|
|
|
|
| @torch.no_grad() |
| def save_sample_grid(ema, prompts, res, steps, path, device): |
| """One row per prompt: 4 EMA samples (CFG 5) with fixed seeds, for tracking progress.""" |
| from sample import sample, parse_prompt |
| rows = [] |
| for j, p in enumerate(prompts): |
| x = sample(ema, parse_prompt(p)[3], res, steps, 5.0, 1000 + j, device, n=4) |
| rows.append(torch.cat(list(x), dim=2)) |
| g = ((torch.cat(rows, dim=1) + 1) * 127.5).round().byte().permute(1, 2, 0).cpu().numpy() |
| Image.fromarray(g).save(path) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--data_dir", default="data") |
| ap.add_argument("--checkpoint_dir", default="checkpoints") |
| ap.add_argument("--batch_size", type=int, default=32) |
| ap.add_argument("--lr", type=float, default=2e-4) |
| ap.add_argument("--epochs", type=int, default=50) |
| ap.add_argument("--max_steps", type=int, default=0, help="stop after N total steps (0 = no limit)") |
| ap.add_argument("--warmup_steps", type=int, default=1000) |
| ap.add_argument("--weight_decay", type=float, default=0.01) |
| ap.add_argument("--ema_decay", type=float, default=0.9999) |
| ap.add_argument("--t_sampling", choices=["logit_normal", "uniform"], default="logit_normal") |
| ap.add_argument("--resolution", type=int, default=0, help="resize images (0 = as stored)") |
| ap.add_argument("--save_every", type=int, default=2000) |
| ap.add_argument("--log_every", type=int, default=50) |
| ap.add_argument("--num_workers", type=int, default=4) |
| ap.add_argument("--precision", choices=["auto", "bf16", "fp16", "fp32"], default="auto") |
| ap.add_argument("--grad_clip", type=float, default=1.0) |
| ap.add_argument("--resume", default="auto", help="'auto' (latest.pt if present), a path, or 'none'") |
| ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--sample_every", type=int, default=0, help="save EMA sample grid every N steps (0 = off)") |
| ap.add_argument("--sample_dir", default="", help="where sample grids go (default: <checkpoint_dir>/samples)") |
| ap.add_argument("--sample_steps", type=int, default=30) |
| ap.add_argument("--sample_prompts", nargs="+", default=[ |
| "a red smiley face with 'HELLO' next to it", "a blue sad face with 'NEWS' next to it", |
| "a green meh face with 'OK' next to it", "a yellow smiley face with 'XQ7' next to it"]) |
| args = ap.parse_args() |
|
|
| torch.manual_seed(args.seed) |
| dev = torch.device(args.device) |
| os.makedirs(args.checkpoint_dir, exist_ok=True) |
| prec = args.precision |
| if prec == "auto": |
| prec = ("bf16" if torch.cuda.is_bf16_supported() else "fp16") if dev.type == "cuda" else "fp32" |
| amp_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": None}[prec] |
| scaler = torch.amp.GradScaler("cuda") if (prec == "fp16" and dev.type == "cuda") else None |
| if dev.type == "cuda": torch.backends.cuda.matmul.allow_tf32 = True; torch.backends.cudnn.allow_tf32 = True |
|
|
| ds = FaceDataset(args.data_dir, args.resolution or None) |
| dl = DataLoader(ds, batch_size=args.batch_size, shuffle=True, drop_last=len(ds) >= args.batch_size, |
| num_workers=args.num_workers, pin_memory=dev.type == "cuda", |
| persistent_workers=args.num_workers > 0) |
| model = build_model().to(dev) |
| ema = copy.deepcopy(model).eval().requires_grad_(False) |
| decay_params = [p for n, p in model.named_parameters() if p.ndim >= 2 and "null_seq" not in n] |
| other = [p for n, p in model.named_parameters() if not (p.ndim >= 2 and "null_seq" not in n)] |
| opt = torch.optim.AdamW([{"params": decay_params, "weight_decay": args.weight_decay}, |
| {"params": other, "weight_decay": 0.0}], lr=args.lr, betas=(0.9, 0.99)) |
| step, start_epoch = 0, 0 |
| rp = os.path.join(args.checkpoint_dir, "latest.pt") if args.resume == "auto" else args.resume |
| if rp != "none" and os.path.exists(rp): |
| ck = torch.load(rp, map_location=dev, weights_only=False) |
| model.load_state_dict(ck["model"]); ema.load_state_dict(ck["ema"]); opt.load_state_dict(ck["opt"]) |
| if scaler and ck.get("scaler"): scaler.load_state_dict(ck["scaler"]) |
| step, start_epoch = ck["step"], ck["epoch"] |
| print(f"resumed from {rp} at step {step}, epoch {start_epoch}") |
| total = args.epochs * len(dl) |
| if args.max_steps: total = min(total, args.max_steps) |
| print(f"device={dev} precision={prec} samples={len(ds)} steps/epoch={len(dl)} total_steps={total}") |
|
|
| t0, done, s0 = time.time(), False, step |
| for epoch in range(start_epoch, args.epochs): |
| model.train() |
| for x0, ids, mask in dl: |
| x0, ids, mask = x0.to(dev, non_blocking=True), ids.to(dev), mask.to(dev) |
| for g in opt.param_groups: g["lr"] = lr_at(step, total, args.warmup_steps, args.lr) |
| eps = torch.randn_like(x0) |
| t = sample_t(x0.shape[0], args.t_sampling, dev) |
| tb = t[:, None, None, None] |
| xt = (1 - tb) * x0 + tb * eps |
| target = eps - x0 |
| with torch.autocast(dev.type, dtype=amp_dtype, enabled=amp_dtype is not None): |
| pred = model(xt, t, ids, mask) |
| loss = F.mse_loss(pred.float(), target) |
| opt.zero_grad(set_to_none=True) |
| if scaler: |
| scaler.scale(loss).backward(); scaler.unscale_(opt) |
| torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) |
| scaler.step(opt); scaler.update() |
| else: |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) |
| opt.step() |
| |
| ema_update(ema, model, min(args.ema_decay, (1 + step) / (10 + step))) |
| step += 1 |
| if step % args.log_every == 0 or step == 1: |
| print(f"ep {epoch} step {step}/{total} loss {loss.item():.4f} lr {opt.param_groups[0]['lr']:.2e} " |
| f"{(time.time()-t0)/max(1, step-s0):.2f}s/it", flush=True) |
| if step % args.save_every == 0: |
| save(os.path.join(args.checkpoint_dir, "latest.pt"), model, ema, opt, scaler, step, epoch, args) |
| print(f"saved checkpoint at step {step}") |
| if args.sample_every and step % args.sample_every == 0: |
| sd = args.sample_dir or os.path.join(args.checkpoint_dir, "samples"); os.makedirs(sd, exist_ok=True) |
| save_sample_grid(ema, args.sample_prompts, x0.shape[-1], args.sample_steps, |
| os.path.join(sd, f"step_{step:07d}.png"), dev) |
| print(f"saved sample grid at step {step}", flush=True) |
| if step >= total: done = True; break |
| save(os.path.join(args.checkpoint_dir, "latest.pt"), model, ema, opt, scaler, step, |
| epoch if done else epoch + 1, args) |
| if done: break |
| torch.save({"ema": ema.state_dict(), "step": step}, os.path.join(args.checkpoint_dir, "ema_final.pt")) |
| print(f"done at step {step}; checkpoints in {args.checkpoint_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
