scripts/distill/train_reflow.py

#!/usr/bin/env python3
"""
Anima Reflow distillation (InstaFlow / rfpp 流派)
================================================

理論的根拠:
  - Anima は既に rectified flow なので「Reflow」(1 round) を当てるのが自然
  - rfpp (NeurIPS 2024 "Improving the Training of Rectified Flows") の知見:
      * U-shape t sampling
      * LPIPS-Huber (latent Huber + 間欠的に pixel LPIPS) で 1-round で十分
      * adversarial 不要 (LADD と独立カテゴリ)
  - 出力: 4-step (or 1-step) で動く student LoRA

データ:
  事前に (noise, x0, emb) triplet を precompute (--save-noise) しておく前提。
  /dataset/teacher_x0_cache_with_noise/{noise,x0,emb}/{idx}.pt
  metadata.json に noise_path も書かれていること。

訓練 1 step:
  1. (noise_i, x0_i, emb_i) を batch load
  2. t ~ U-shape on (0, 1)
  3. x_t = (1 - t) * x0 + t * noise
  4. v_target = noise - x0
  5. v_pred   = student.dit_forward(x_t, t, emb)
  6. L_huber = huber(v_pred, v_target, delta=0.03)
  7. 任意 (lpips-every step): x0_pred = x_t - t*v_pred、VAE decode → LPIPS vs x0
  8. backward → student LoRA 更新
"""
from __future__ import annotations
import argparse
import copy
import json
import os
import sys
import time
from pathlib import Path

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from safetensors.torch import save_file, load_file

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from distill.anima_loader import AnimaPaths, build_anima, AnimaBundle
from distill.dmd2_trainer import attach_wide_lora
from distill.train_traj import load_warm_lora, save_lora_state


class ReflowPairDataset(Dataset):
    """precompute された (noise, x0, emb) triplet を読む。
       metadata.json の各 entry に noise_path / x0_path / emb_path が必要。"""
    def __init__(self, cache_dir: str | Path):
        self.cache_dir = Path(cache_dir)
        self.meta = json.loads((self.cache_dir / "metadata.json").read_text(encoding="utf-8"))
        if not self.meta or "noise_path" not in self.meta[0]:
            raise RuntimeError(
                f"{self.cache_dir}/metadata.json has no 'noise_path' — "
                "rerun precompute with --save-noise."
            )

    def __len__(self):
        return len(self.meta)

    def __getitem__(self, idx):
        m = self.meta[idx]
        x0 = torch.load(m["x0_path"], map_location="cpu", weights_only=True).squeeze(0)
        noise = torch.load(m["noise_path"], map_location="cpu", weights_only=True).squeeze(0)
        emb = torch.load(m["emb_path"], map_location="cpu", weights_only=True).squeeze(0)
        return {"x0": x0, "noise": noise, "emb": emb}


def reflow_collate(batch):
    return {
        "x0": torch.stack([b["x0"] for b in batch]),
        "noise": torch.stack([b["noise"] for b in batch]),
        "emb": torch.stack([b["emb"] for b in batch]),
    }


def u_shape_t(B: int, device, dtype=torch.float32) -> torch.Tensor:
    """rfpp 流の U-shape sampler。両端 (t≈0, t≈1) に分布が偏る。
       実装: u ~ U(0,1) → t = 0.5 * (1 + sign(u-0.5) * |2u-1|^0.5)
              (sqrt によって両端密度上昇)"""
    u = torch.rand(B, device=device, dtype=dtype)
    centered = 2 * u - 1
    t = 0.5 * (1.0 + torch.sign(centered) * centered.abs().sqrt())
    return t.clamp(1e-3, 1.0 - 1e-3)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cache-dir", required=True, type=str,
                    help="--save-noise 付きで precompute した cache")
    ap.add_argument("--out", required=True, type=str)
    ap.add_argument("--warm-lora", default="", type=str)
    ap.add_argument("--total-steps", type=int, default=8000)
    ap.add_argument("--batch-size", type=int, default=4)
    ap.add_argument("--grad-accum", type=int, default=2)
    ap.add_argument("--resolution", type=int, default=768)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--weight-decay", type=float, default=0.01)
    ap.add_argument("--grad-clip", type=float, default=1.0)
    ap.add_argument("--lora-rank", type=int, default=32)
    ap.add_argument("--huber-delta", type=float, default=0.03)
    ap.add_argument("--lpips-weight", type=float, default=0.1)
    ap.add_argument("--lpips-every", type=int, default=4,
                    help="N step ごとに LPIPS reg を加算 (cost 削減)")
    ap.add_argument("--t-sampler", default="u_shape", choices=["u_shape", "uniform"])
    ap.add_argument("--log-every", type=int, default=10)
    ap.add_argument("--sample-every", type=int, default=500)
    ap.add_argument("--num-workers", type=int, default=2)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    torch.manual_seed(args.seed)
    device = torch.device("cuda")
    dtype = torch.bfloat16
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    # ----- load Anima base (teacher は捨てる、cache が teacher 役) -----
    print("[load] Anima bundle (student only — teacher cached on disk)")
    bundle = build_anima(device=device, dtype=dtype)

    # ----- student = wide LoRA -----
    student_transformer = attach_wide_lora(bundle.transformer, rank=args.lora_rank)
    student_transformer.to(device=device, dtype=dtype)
    for n, p in student_transformer.named_parameters():
        p.requires_grad = ("lora_" in n)
    student_params = [p for p in student_transformer.parameters() if p.requires_grad]
    print(f"[setup] student trainable: {sum(p.numel() for p in student_params)/1e6:.1f}M")
    bundle.transformer = student_transformer

    # warm-start
    if args.warm_lora:
        load_warm_lora(student_transformer, args.warm_lora)

    # LPIPS (任意)
    lpips_fn = None
    if args.lpips_weight > 0:
        import lpips as _lpips
        lpips_fn = _lpips.LPIPS(net="alex").to(device).eval()
        for p in lpips_fn.parameters():
            p.requires_grad = False
        print("[setup] LPIPS(alex) loaded")

    # ----- optimizer -----
    opt = torch.optim.AdamW(student_params, lr=args.lr, betas=(0.9, 0.999),
                            weight_decay=args.weight_decay, eps=1e-8)

    # ----- dataset -----
    print(f"[data] loading cache {args.cache_dir}")
    dataset = ReflowPairDataset(args.cache_dir)
    print(f"   {len(dataset)} (noise, x0, emb) triplets")
    loader = DataLoader(
        dataset, batch_size=args.batch_size, shuffle=True,
        num_workers=args.num_workers, collate_fn=reflow_collate,
        drop_last=True, pin_memory=True,
    )

    # ----- training loop -----
    print(f"[train] steps={args.total_steps} bs={args.batch_size} accum={args.grad_accum} "
          f"lpips_w={args.lpips_weight} every={args.lpips_every}")
    log_path = out_dir / "reflow_log.jsonl"
    log_f = open(log_path, "a", buffering=1)
    t0 = time.time()
    data_iter = iter(loader)

    def _next():
        nonlocal data_iter
        try:
            return next(data_iter)
        except StopIteration:
            data_iter = iter(loader)
            return next(data_iter)

    for step in range(args.total_steps):
        student_transformer.train()
        opt.zero_grad()
        metrics = {}
        for _ in range(args.grad_accum):
            batch = _next()
            x0 = batch["x0"].to(device=device, dtype=dtype)
            noise = batch["noise"].to(device=device, dtype=dtype)
            emb = batch["emb"].to(device=device, dtype=dtype)
            B = x0.size(0)

            # ----- t sampling -----
            if args.t_sampler == "u_shape":
                t = u_shape_t(B, device, dtype=torch.float32).to(dtype=dtype)
            else:
                t = torch.rand(B, device=device, dtype=dtype).clamp(1e-3, 1.0 - 1e-3)

            t_ = t.view(-1, *([1] * (x0.dim() - 1)))
            x_t = (1 - t_) * x0 + t_ * noise
            v_target = noise - x0

            # ----- student forward (1 forward only) -----
            v_pred = AnimaBundle.dit_forward(student_transformer, x_t, t, emb)
            l_huber = F.huber_loss(v_pred.float(), v_target.float(),
                                   delta=args.huber_delta, reduction="mean")

            # ----- LPIPS reg (every N steps) -----
            l_lpips = torch.zeros((), device=device)
            if lpips_fn is not None and (step % args.lpips_every) == 0:
                x0_pred = x_t - t_ * v_pred
                vae_dtype = next(bundle.vae.model.parameters()).dtype
                img_p = bundle.vae.model.decode(x0_pred.to(dtype=vae_dtype), bundle.vae_scale).squeeze(2)
                with torch.no_grad():
                    img_t = bundle.vae.model.decode(x0.to(dtype=vae_dtype), bundle.vae_scale).squeeze(2)
                l_lpips = lpips_fn(img_p.float(), img_t.float()).mean()

            loss = (l_huber + args.lpips_weight * l_lpips) / args.grad_accum
            loss.backward()
            metrics = {
                "l_huber": float(l_huber.detach()),
                "l_lpips": float(l_lpips.detach()),
                "loss": float((l_huber + args.lpips_weight * l_lpips).detach()),
            }
        torch.nn.utils.clip_grad_norm_(student_params, args.grad_clip)
        opt.step()

        if step % args.log_every == 0:
            metrics["step"] = step
            metrics["elapsed"] = time.time() - t0
            log_f.write(json.dumps(metrics) + "\n")
            msg = " ".join(f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}"
                           for k, v in metrics.items() if k != "step")
            print(f"[step {step}/{args.total_steps}] {msg}", flush=True)

        if step > 0 and step % args.sample_every == 0:
            save_lora_state(student_transformer, out_dir, f"reflow_step{step:05d}")
            print(f"[save] reflow_step{step:05d}.safetensors", flush=True)
            try:
                import modal
                modal.Volume.from_name("anima-outputs").commit()
            except Exception:
                pass

    print("[done] saving final")
    save_lora_state(student_transformer, out_dir, "reflow_final")
    log_f.close()


if __name__ == "__main__":
    main()