scripts/distill/train_ladd.py

#!/usr/bin/env python3
"""
Anima LADD (Latent Adversarial Diffusion Distillation) — AMD Nitro-1 移植

R3GAN 失敗との対比:
  - R3GAN: CNN-on-latent を一から学習 → Anima 16ch latent prior なし → 崩壊
  - LADD:  D の backbone = teacher MiniTrainDIT (frozen)、head だけ trainable
           + Smooth-L1 recon anchor で mean collapse の引力を断つ
           → R3GAN の "stable zone too narrow" を構造的に回避

訓練ループ (1 step per phase, alternate G:D = 1:1):
  G phase:
    1. student で 1-step rollout (t=1 → t=0)、x0_hat 取得
    2. x0_hat を t_D ∈ [0, 0.75] で re-noise
    3. D で logits_fake、adv_loss = BCE(logits_fake, 1)
    4. recon_loss = smooth_L1(x0_hat, x0_teacher_cached)
    5. G_loss = adv_loss + recon_lambda * recon_loss

  D phase:
    1. student で 1-step rollout (no_grad)、x0_hat
    2. teacher x0 (cached) と x0_hat をそれぞれ別の t_D で re-noise
    3. D で logits_real / logits_fake、BCE(real, 1) + BCE(fake, 0)

precompute 必須:
  modal run modal_app.py::precompute_teacher_x0_cache
"""
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 convert_comfy_to_peft_lora, save_lora_state
from distill.dmd2_official_loss import renoise_rf, x0_from_velocity_rf
from distill.anima_ladd_disc import (
    AnimaLADDDiscriminator, ladd_d_loss, ladd_g_adv_loss, ladd_g_recon_loss,
)


# ----- precomputed cache dataset -------------------------------------------
class PrecomputedCacheDataset(Dataset):
    """teacher_x0_cache から (caption_emb, teacher_x0) を読む。"""
    def __init__(self, cache_dir: str | Path):
        self.cache_dir = Path(cache_dir)
        meta_path = self.cache_dir / "metadata.json"
        self.meta = json.loads(meta_path.read_text(encoding="utf-8"))
        if len(self.meta) == 0:
            raise RuntimeError(f"Empty metadata at {meta_path}")

    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)
        emb = torch.load(m["emb_path"], map_location="cpu", weights_only=True)
        # x0 was saved as (1, 16, 1, H, W), emb as (1, 512, 1024) → squeeze batch dim
        return {"x0": x0.squeeze(0), "emb": emb.squeeze(0), "caption": m["caption"]}


def ladd_collate(batch):
    x0 = torch.stack([b["x0"] for b in batch])
    emb = torch.stack([b["emb"] for b in batch])
    captions = [b["caption"] for b in batch]
    return {"x0": x0, "emb": emb, "captions": captions}


# ----- student 1-step rollout ----------------------------------------------
def student_x0_hat(
    student_v_fn, noise: torch.Tensor, cond_pos: torch.Tensor,
) -> torch.Tensor:
    """1-step distill: t=1 noise から 1 step Euler で t=0 へ。"""
    B = noise.size(0)
    device = noise.device
    dtype = noise.dtype
    t_init = torch.ones(B, device=device, dtype=dtype)  # t=1
    v = student_v_fn(noise, t_init, cond_pos)
    # x0 = x_t - t * v = noise - 1 * v
    return noise - v


# ---------------------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cache-dir", required=True, type=str)
    ap.add_argument("--out", required=True, type=str)
    ap.add_argument("--warm-lora", default="", type=str)
    ap.add_argument("--total-steps", type=int, default=5000)
    ap.add_argument("--batch-size", type=int, default=4)
    ap.add_argument("--grad-accum", type=int, default=4, help="effective bs = batch_size * grad_accum")
    ap.add_argument("--resolution", type=int, default=768)
    ap.add_argument("--recon-lambda", type=float, default=1.0)
    ap.add_argument("--lr-g", type=float, default=1e-6)
    ap.add_argument("--lr-d", type=float, default=1e-6)
    ap.add_argument("--t-d-max", type=float, default=0.75, help="re-noise t upper bound")
    ap.add_argument("--lora-rank", type=int, default=32)
    ap.add_argument("--grad-clip", type=float, default=1.0)
    ap.add_argument("--block-ids", type=str, default="2,8,14,20,26",
                    help="teacher block indices for D hooks (Anima 28 blocks)")
    ap.add_argument("--head-hidden", type=int, default=512)
    ap.add_argument("--misaligned-pairs-d", action="store_true", default=True,
                    help="text alignment trick: D も misaligned fake (caption roll) を学習")
    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 -----
    print("[load] Anima bundle")
    bundle = build_anima(device=device, dtype=dtype)

    # ----- teacher = deepcopy frozen (D の backbone と shared instance) -----
    print("[setup] teacher = frozen deepcopy")
    teacher_transformer = copy.deepcopy(bundle.transformer).to(device=device, dtype=dtype).eval()
    for p in teacher_transformer.parameters():
        p.requires_grad = False

    # ----- student = wide LoRA -----
    print("[setup] student = wide LoRA on bundle.transformer")
    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:
        from distill.train_traj import load_warm_lora
        load_warm_lora(student_transformer, args.warm_lora)

    # ----- discriminator -----
    print("[setup] LADD discriminator (teacher backbone + multi-scale heads)")
    block_ids = [int(x) for x in args.block_ids.split(",")]
    disc = AnimaLADDDiscriminator(
        teacher_transformer=teacher_transformer,
        block_ids=block_ids, head_hidden=args.head_hidden,
    )
    # heads を lazy init するために dummy forward
    H_lat = args.resolution // 8
    W_lat = args.resolution // 8
    dummy_x = torch.randn(1, 16, 1, H_lat, W_lat, device=device, dtype=dtype)
    dummy_t = torch.tensor([0.5], device=device, dtype=dtype)
    with torch.no_grad():
        dummy_cond = bundle.text_encode([""])
    disc.lazy_init_heads(dummy_x, dummy_t, dummy_cond)
    disc.to(device=device, dtype=torch.float32)  # heads は fp32
    disc_params = disc.trainable_parameters()
    print(f"[setup] D heads trainable: {sum(p.numel() for p in disc_params)/1e6:.1f}M")

    # ----- optimizers -----
    opt_g = torch.optim.AdamW(student_params, lr=args.lr_g, betas=(0.0, 0.999), eps=1e-8)
    opt_d = torch.optim.AdamW(disc_params, lr=args.lr_d, betas=(0.0, 0.999), eps=1e-8)

    # ----- velocity functions -----
    def student_v(x, t, cond):
        return AnimaBundle.dit_forward(student_transformer, x, t, cond)

    # ----- dataset (precomputed cache) -----
    print(f"[data] loading cache {args.cache_dir}")
    dataset = PrecomputedCacheDataset(args.cache_dir)
    print(f"   {len(dataset)} cached samples")
    loader = DataLoader(
        dataset, batch_size=args.batch_size, shuffle=True,
        num_workers=args.num_workers, collate_fn=ladd_collate, drop_last=True,
        pin_memory=True,
    )

    # ----- training loop -----
    print(f"[train] total={args.total_steps} bs={args.batch_size} accum={args.grad_accum} "
          f"recon_lambda={args.recon_lambda}")
    log_path = out_dir / "ladd_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):
        # ---- G phase ----
        student_transformer.train()
        opt_g.zero_grad()
        g_metrics = {}
        for _ in range(args.grad_accum):
            batch = _next()
            x0_teacher = batch["x0"].to(device=device, dtype=dtype)
            cond_pos = batch["emb"].to(device=device, dtype=dtype)
            B = x0_teacher.size(0)
            noise = torch.randn_like(x0_teacher)
            # student forward
            x0_hat = student_x0_hat(student_v, noise, cond_pos)
            # re-noise & D forward (fake side)
            t_D = torch.rand(B, device=device, dtype=dtype) * args.t_d_max
            D_eps = torch.randn_like(x0_hat)
            x_t_fake = renoise_rf(x0_hat, t_D, D_eps)
            # G phase: adv loss が student を訓練するために gradient_to_input=True
            logits_fake = disc(x_t_fake, t_D, cond_pos, gradient_to_input=True)
            l_adv, m_adv = ladd_g_adv_loss(logits_fake)
            l_rec, m_rec = ladd_g_recon_loss(x0_hat, x0_teacher)
            g_loss = (l_adv + args.recon_lambda * l_rec) / args.grad_accum
            g_loss.backward()
            g_metrics = {**m_adv, **m_rec, "l_g_total": (l_adv + args.recon_lambda * l_rec).detach()}
        torch.nn.utils.clip_grad_norm_(student_params, args.grad_clip)
        opt_g.step()

        # ---- D phase ----
        disc.train()
        opt_d.zero_grad()
        d_metrics = {}
        for _ in range(args.grad_accum):
            batch = _next()
            x0_teacher = batch["x0"].to(device=device, dtype=dtype)
            cond_pos = batch["emb"].to(device=device, dtype=dtype)
            B = x0_teacher.size(0)
            with torch.no_grad():
                noise = torch.randn_like(x0_teacher)
                x0_hat = student_x0_hat(student_v, noise, cond_pos)
            t_D_real = torch.rand(B, device=device, dtype=dtype) * args.t_d_max
            t_D_fake = torch.rand(B, device=device, dtype=dtype) * args.t_d_max
            x_t_real = renoise_rf(x0_teacher, t_D_real, torch.randn_like(x0_teacher))
            x_t_fake = renoise_rf(x0_hat, t_D_fake, torch.randn_like(x0_hat))
            cond_for_real = cond_pos
            cond_for_fake = cond_pos
            if args.misaligned_pairs_d:
                # text-alignment: misaligned fake (caption を 1 個 roll) も追加
                cond_misaligned = torch.roll(cond_pos, 1, dims=0)
                t_D_mis = torch.rand(B, device=device, dtype=dtype) * args.t_d_max
                x_t_mis = renoise_rf(x0_teacher, t_D_mis, torch.randn_like(x0_teacher))
                x_t_fake = torch.cat([x_t_fake, x_t_mis], dim=0)
                t_D_fake = torch.cat([t_D_fake, t_D_mis], dim=0)
                cond_for_fake = torch.cat([cond_pos, cond_misaligned], dim=0)
            # D phase: heads だけ訓練、x への gradient 不要 (gradient_to_input=False、省メモリ)
            logits_real = disc(x_t_real, t_D_real, cond_for_real, gradient_to_input=False)
            logits_fake = disc(x_t_fake, t_D_fake, cond_for_fake, gradient_to_input=False)
            d_loss, m_d = ladd_d_loss(logits_real, logits_fake)
            d_loss = d_loss / args.grad_accum
            d_loss.backward()
            d_metrics = {k: float(v) for k, v in m_d.items()}
        torch.nn.utils.clip_grad_norm_(disc_params, args.grad_clip)
        opt_d.step()

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

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

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


if __name__ == "__main__":
    main()