from __future__ import annotations import argparse import copy import glob import json import math import os import time import numpy as np import torch import torch.nn.functional as F from diffusers import AutoencoderKL from transformers import AutoModel, CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast from dit_v6 import MMDiT IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) def build_cache(shards_dir, cache): shards = sorted(glob.glob(f"{shards_dir}/shard_*.npz")) if not shards: raise SystemExit(f"no shards in {shards_dir}") lat, t5i, t5m, ci = [], [], [], [] for i, s in enumerate(shards): z = np.load(s) lat.append(z["latents"]); t5i.append(z["t5_ids"]); t5m.append(z["t5_mask"]); ci.append(z["clip_ids"]) if (i + 1) % 50 == 0: print(f" loaded {i+1}/{len(shards)} shards", flush=True) lat = np.concatenate(lat); t5i = np.concatenate(t5i); t5m = np.concatenate(t5m); ci = np.concatenate(ci) np.save(f"{cache}_lat.npy", lat) np.save(f"{cache}_t5ids.npy", t5i) np.save(f"{cache}_t5mask.npy", t5m) np.save(f"{cache}_clipids.npy", ci) return lat, t5i, t5m, ci def repa_weight_at(step, total, warm_frac=0.40, zero_frac=0.70, peak=0.5): p = step / total if p < warm_frac: return peak if p < zero_frac: return peak * (1 - (p - warm_frac) / (zero_frac - warm_frac)) return 0.0 def main(): ap = argparse.ArgumentParser() ap.add_argument("--shards", default="/root/v6cache/shards") ap.add_argument("--cache", default="/root/v6cache/cache") ap.add_argument("--out", default="/root/runs/pm6") ap.add_argument("--steps", type=int, default=150000) ap.add_argument("--batch", type=int, default=192) ap.add_argument("--lr", type=float, default=2e-4) ap.add_argument("--warmup", type=int, default=1500) ap.add_argument("--dim", type=int, default=512) ap.add_argument("--depth", type=int, default=16) ap.add_argument("--heads", type=int, default=8) ap.add_argument("--mlp-hidden", type=int, default=1408) ap.add_argument("--t5-len", type=int, default=32) ap.add_argument("--repa-layer", type=int, default=8) ap.add_argument("--repa-peak", type=float, default=0.5) ap.add_argument("--repa-batch", type=int, default=48) ap.add_argument("--cfg-dropout", type=float, default=0.1) ap.add_argument("--ema", type=float, default=0.9999) ap.add_argument("--val-size", type=int, default=4096) ap.add_argument("--val-every", type=int, default=2000) ap.add_argument("--log-every", type=int, default=200) ap.add_argument("--save-every", type=int, default=5000) ap.add_argument("--vae", default="madebyollin/sdxl-vae-fp16-fix") ap.add_argument("--clip", default="openai/clip-vit-base-patch32") ap.add_argument("--t5", default="google/flan-t5-base") ap.add_argument("--dinov2", default="facebook/dinov2-small") ap.add_argument("--resume", default="") ap.add_argument("--seed", type=int, default=0) ap.add_argument("--grad-ckpt", action="store_true") args = ap.parse_args() dev = "cuda" os.makedirs(args.out, exist_ok=True) torch.manual_seed(args.seed) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True if os.path.exists(f"{args.cache}_lat.npy"): lat = np.load(f"{args.cache}_lat.npy", mmap_mode="r") t5i = np.load(f"{args.cache}_t5ids.npy") t5m = np.load(f"{args.cache}_t5mask.npy") ci = np.load(f"{args.cache}_clipids.npy") else: lat, t5i, t5m, ci = build_cache(args.shards, args.cache) N = len(lat) perm = np.random.RandomState(args.seed).permutation(N) val_i = np.sort(perm[:args.val_size]) tr_i = perm[args.val_size:] print(f"[data] {N} pairs, {len(tr_i)} train, {len(val_i)} val", flush=True) vae = AutoencoderKL.from_pretrained(args.vae).to(dev).half().eval() vae_scale = vae.config.scaling_factor for p in vae.parameters(): p.requires_grad_(False) clip_txt = CLIPTextModel.from_pretrained(args.clip).to(dev).eval() for p in clip_txt.parameters(): p.requires_grad_(False) t5 = T5EncoderModel.from_pretrained(args.t5).to(dev).eval() for p in t5.parameters(): p.requires_grad_(False) dinov2 = AutoModel.from_pretrained(args.dinov2).to(dev).eval() for p in dinov2.parameters(): p.requires_grad_(False) imagenet_mean = IMAGENET_MEAN.to(dev) imagenet_std = IMAGENET_STD.to(dev) @torch.no_grad() def encode_text(t5_ids, t5_mask, clip_ids): t5_out = t5(input_ids=t5_ids, attention_mask=t5_mask).last_hidden_state.float() clip_pool = clip_txt(input_ids=clip_ids).pooler_output.float() return t5_out, clip_pool @torch.no_grad() def dino_features(x1_latent): px = vae.decode((x1_latent / vae_scale).to(vae.dtype)).sample.float() px = (px.clamp(-1, 1) + 1) / 2 px = F.interpolate(px, size=(224, 224), mode="bilinear", align_corners=False) px = (px - imagenet_mean) / imagenet_std out = dinov2(pixel_values=px.to(dinov2.dtype)).last_hidden_state return out[:, 1:, :].float() t5_tok = T5TokenizerFast.from_pretrained(args.t5) clip_tok = CLIPTokenizer.from_pretrained(args.clip) null_t5_enc = t5_tok([""], padding="max_length", max_length=args.t5_len, truncation=True, return_tensors="pt") null_t5_ids = null_t5_enc["input_ids"].to(dev) null_t5_mask = null_t5_enc["attention_mask"].to(dev) null_clip_ids = clip_tok([""], padding="max_length", max_length=ci.shape[1], truncation=True, return_tensors="pt")["input_ids"].to(dev) null_t5_seq, null_clip_pool = encode_text(null_t5_ids, null_t5_mask, null_clip_ids) t5i_t = torch.from_numpy(t5i.astype(np.int64)) t5m_t = torch.from_numpy(t5m.astype(np.int64)) ci_t = torch.from_numpy(ci.astype(np.int64)) vlat = torch.from_numpy(np.asarray(lat[val_i])).float() vg = torch.Generator(device=dev).manual_seed(1234) vx1 = vlat.to(dev) vx0 = torch.randn(vx1.shape, device=dev, generator=vg) vt = torch.sigmoid(torch.randn(vx1.shape[0], device=dev, generator=vg)) with torch.no_grad(): vseq, vmask, vpool = [], [], [] for i in range(0, len(val_i), 256): s, p = encode_text(t5i_t[val_i[i:i+256]].to(dev), t5m_t[val_i[i:i+256]].to(dev), ci_t[val_i[i:i+256]].to(dev)) vseq.append(s); vmask.append(t5m_t[val_i[i:i+256]].to(dev)); vpool.append(p) vseq = torch.cat(vseq); vmask = torch.cat(vmask); vpool = torch.cat(vpool) model = MMDiT(dim=args.dim, depth=args.depth, heads=args.heads, mlp_hidden=args.mlp_hidden, t5_len=args.t5_len).to(dev) ema = copy.deepcopy(model).eval() for p in ema.parameters(): p.requires_grad_(False) opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.99), weight_decay=0.0) print(f"[model] {model.num_params()/1e6:.2f}M trainable, {model.num_backbone_params()/1e6:.2f}M backbone", flush=True) start, best = 0, float("inf") if args.resume and os.path.exists(args.resume): ck = torch.load(args.resume, map_location=dev) model.load_state_dict(ck["model"]); ema.load_state_dict(ck["ema"]) opt.load_state_dict(ck["opt"]); start = ck["step"] + 1; best = ck.get("best", best) print(f"[resume] from step {start}", flush=True) def lr_at(s): if s < args.warmup: return args.lr * (s + 1) / args.warmup p = (s - args.warmup) / max(1, args.steps - args.warmup) return args.lr * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1.0, p)))) @torch.no_grad() def val_loss(): tot, n = 0.0, 0 for j in range(0, vx1.shape[0], args.batch): sl = slice(j, j + args.batch) m = vx1[sl].shape[0] tb = vt[sl].view(-1, 1, 1, 1) xt = (1 - tb) * vx0[sl] + tb * vx1[sl] with torch.autocast("cuda", dtype=torch.bfloat16): v = ema(xt, vt[sl], vseq[sl], vmask[sl], vpool[sl]) tot += F.mse_loss(v.float(), vx1[sl] - vx0[sl]).item() * m n += m return tot / n logf = open(f"{args.out}/log.jsonl", "a") gen = torch.Generator(device=dev).manual_seed(args.seed) run, run_diff, run_repa = 0.0, 0.0, 0.0 t0 = time.time() for step in range(start, args.steps): i = tr_i[np.random.randint(0, len(tr_i), args.batch)] i_sorted = np.sort(i) x1 = torch.from_numpy(np.asarray(lat[i_sorted])).to(dev).float() t5_ids_b = t5i_t[i_sorted].to(dev) t5_mask_b = t5m_t[i_sorted].to(dev) clip_ids_b = ci_t[i_sorted].to(dev) seq, pool = encode_text(t5_ids_b, t5_mask_b, clip_ids_b) mask = t5_mask_b drop = torch.rand(x1.shape[0], device=dev, generator=gen) < args.cfg_dropout seq = torch.where(drop[:, None, None], null_t5_seq, seq) mask = torch.where(drop[:, None], null_t5_mask, mask) pool = torch.where(drop[:, None], null_clip_pool, pool) x0 = torch.randn(x1.shape, device=dev, generator=gen) t = torch.sigmoid(torch.randn(x1.shape[0], device=dev, generator=gen)) tb = t.view(-1, 1, 1, 1) xt = (1 - tb) * x0 + tb * x1 target = x1 - x0 rw = repa_weight_at(step, args.steps, peak=args.repa_peak) for g in opt.param_groups: g["lr"] = lr_at(step) with torch.autocast("cuda", dtype=torch.bfloat16): if rw > 0: v, repa_pred = model(xt, t, seq, mask, pool, return_repa=True, use_checkpoint=args.grad_ckpt) else: v = model(xt, t, seq, mask, pool, use_checkpoint=args.grad_ckpt) loss_diff = F.mse_loss(v.float(), target) if rw > 0: rb = min(args.repa_batch, x1.shape[0]) with torch.no_grad(): dino_tgt = dino_features(x1[:rb]) loss_repa = 1.0 - F.cosine_similarity(repa_pred[:rb].float(), dino_tgt, dim=-1).mean() loss = loss_diff + rw * loss_repa else: loss_repa = torch.zeros((), device=dev) loss = loss_diff opt.zero_grad(set_to_none=True) loss.backward() gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() d = args.ema if step > args.warmup else 0.0 with torch.no_grad(): for pe, pm in zip(ema.parameters(), model.parameters()): pe.mul_(d).add_(pm.detach(), alpha=1 - d) for be, bm in zip(ema.buffers(), model.buffers()): be.copy_(bm) run += loss.item(); run_diff += loss_diff.item(); run_repa += loss_repa.item() if (step + 1) % args.log_every == 0: el = time.time() - t0 sps = args.log_every / el print(f"[s{step+1:06d}] loss={run/args.log_every:.4f} diff={run_diff/args.log_every:.4f} " f"repa={run_repa/args.log_every:.4f} rw={rw:.3f} lr={lr_at(step):.2e} gnorm={gn:.2f} " f"{sps:.2f} steps/s eta={(args.steps-step-1)/sps/3600:.1f}h", flush=True) logf.write(json.dumps({"step": step + 1, "loss": run/args.log_every, "loss_diff": run_diff/args.log_every, "loss_repa": run_repa/args.log_every, "repa_weight": rw, "steps_per_s": sps}) + "\n"); logf.flush() run, run_diff, run_repa, t0 = 0.0, 0.0, 0.0, time.time() if (step + 1) % args.val_every == 0 or step + 1 == args.steps: vl = val_loss() tag = "" if vl < best: best = vl torch.save({"ema": ema.state_dict(), "step": step, "val": vl, "cfg": {"dim": args.dim, "depth": args.depth, "heads": args.heads, "mlp_hidden": args.mlp_hidden, "t5_len": args.t5_len}}, f"{args.out}/best.pt") tag = " *best*" print(f"[s{step+1:06d}] val_loss={vl:.5f}{tag}", flush=True) logf.write(json.dumps({"step": step + 1, "val_loss": vl}) + "\n"); logf.flush() t0 = time.time() if (step + 1) % args.save_every == 0 or step + 1 == args.steps: torch.save({"model": model.state_dict(), "ema": ema.state_dict(), "opt": opt.state_dict(), "step": step, "best": best, "cfg": {"dim": args.dim, "depth": args.depth, "heads": args.heads, "mlp_hidden": args.mlp_hidden, "t5_len": args.t5_len}}, f"{args.out}/latest.pt") print("TRAINDONE best_val", best, flush=True) if __name__ == "__main__": main()