train_v6.py

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()