scripts/distill/train_sid.py

#!/usr/bin/env python3
"""
Anima SiD2 / SiD-DiT distillation
=================================

DMD2 と同じ「2 LoRA adapter on shared base」パターンで実装。違いは:
  - D も EMA も不要 (D は SiD identity formula に置換)
  - data-free: caption のみで OK (画像は使わない)
  - ψ (fake_score) を 200 step warmup してから generator を活性化
  - ratio: ψ:θ = 1:1 (DMD2 の 5:1 と異なる)

使い方:
  modal run modal_app.py::train_sid_distill \\
      --total-outer-steps 8000 --warm-lora /models/loras/anima_turbo.safetensors

注意:
  cache-dir は emb のみ使うので LADD cache でも reflow cache でも何でも良い。
  画像が無い caption だけのテキストファイル群でも、本質的には動く
  (今回は既存 cache の emb を流用してコスト 0)。
"""
from __future__ import annotations
import argparse
import copy
import json
import os
import sys
import time
from pathlib import Path

import torch
from torch.utils.data import DataLoader

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from distill.anima_loader import build_anima, AnimaBundle
from distill.dmd2_trainer import attach_wide_lora
from distill.train_dmd2_official import (
    attach_dual_lora, make_velocity_fn, make_teacher_velocity_fn,
)
from distill.train_traj import save_lora_state, convert_comfy_to_peft_lora
from distill.sid_loss import sid_generator_loss, sid_score_helper_loss
from distill.train_ladd import PrecomputedCacheDataset, ladd_collate
from safetensors.torch import load_file


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cache-dir", required=True, type=str,
                    help="emb のみ使用 (画像は不要)。LADD/Reflow cache 流用 OK")
    ap.add_argument("--out", required=True, type=str)
    ap.add_argument("--warm-lora", default="", type=str)
    ap.add_argument("--total-outer-steps", type=int, default=8000)
    ap.add_argument("--psi-warmup-steps", type=int, default=200,
                    help="ψ を先に warmup し、それ以降 θ を活性化")
    ap.add_argument("--n-student-steps", type=int, default=4)
    ap.add_argument("--batch-size", type=int, default=2)
    ap.add_argument("--grad-accum", type=int, default=2)
    ap.add_argument("--resolution", type=int, default=768)
    ap.add_argument("--teacher-cfg", type=float, default=4.5)
    ap.add_argument("--student-cfg", type=float, default=1.0)
    ap.add_argument("--alpha", type=float, default=1.2, help="SiD2 identity weight")
    ap.add_argument("--mu-t", type=float, default=0.6931, help="ln 2")
    ap.add_argument("--sigma-t", type=float, default=1.6)
    ap.add_argument("--lora-rank", type=int, default=32)
    ap.add_argument("--lr-gen", type=float, default=1e-5)
    ap.add_argument("--lr-psi", type=float, default=2e-5)
    ap.add_argument("--weight-decay", type=float, default=0.01)
    ap.add_argument("--grad-clip", type=float, default=1.0)
    ap.add_argument("--neg-prompt", default="")
    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)

    print("[load] Anima bundle")
    bundle = build_anima(device=device, dtype=dtype)

    # teacher = deepcopy
    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

    # dual-adapter (student + fake_score / SiD-DiT は ψ と呼ぶ)
    print("[setup] dual-adapter (student + psi)")
    peft_model = attach_dual_lora(bundle.transformer, rank=args.lora_rank,
                                  adapter_names=("student", "psi"))
    peft_model.to(device=device, dtype=dtype)
    for n, p in peft_model.named_parameters():
        p.requires_grad = ("lora_" in n)
    bundle.transformer = peft_model

    student_params = [p for n, p in peft_model.named_parameters()
                      if p.requires_grad and ".student." in n]
    psi_params = [p for n, p in peft_model.named_parameters()
                  if p.requires_grad and ".psi." in n]
    print(f"[setup] student: {sum(p.numel() for p in student_params)/1e6:.1f}M")
    print(f"[setup] psi:     {sum(p.numel() for p in psi_params)/1e6:.1f}M")

    # warm-start (student に注入)
    if args.warm_lora:
        print(f"[warm] loading {args.warm_lora} → student adapter")
        sd_comfy = load_file(args.warm_lora)
        sd_peft_default = convert_comfy_to_peft_lora(sd_comfy)
        sd_student = {}
        for k, v in sd_peft_default.items():
            nk = k.replace(".lora_A.default.weight", ".lora_A.student.weight")
            nk = nk.replace(".lora_B.default.weight", ".lora_B.student.weight")
            sd_student[nk] = v
        model_sd = peft_model.state_dict()
        to_load = {k: v.to(dtype=model_sd[k].dtype) for k, v in sd_student.items()
                   if k in model_sd and model_sd[k].shape == v.shape}
        print(f"[warm] matched {len(to_load)}/{len(sd_student)} keys")
        peft_model.load_state_dict(to_load, strict=False)

    # optimizers
    opt_gen = torch.optim.AdamW(student_params, lr=args.lr_gen,
                                betas=(0.9, 0.999), weight_decay=args.weight_decay)
    opt_psi = torch.optim.AdamW(psi_params, lr=args.lr_psi,
                                betas=(0.9, 0.999), weight_decay=args.weight_decay)

    # velocity functions
    student_v = make_velocity_fn(peft_model, "student")
    psi_v = make_velocity_fn(peft_model, "psi")
    teacher_v = make_teacher_velocity_fn(teacher_transformer)

    # dataset (emb のみ使う、x0 は initial noise の source として扱う)
    print(f"[data] {args.cache_dir}")
    dataset = PrecomputedCacheDataset(args.cache_dir)
    print(f"   {len(dataset)} captions")
    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,
    )

    with torch.no_grad():
        cond_neg = bundle.text_encode([args.neg_prompt or ""])

    H_lat = args.resolution // 8
    W_lat = args.resolution // 8

    print(f"[train] outer={args.total_outer_steps} bs={args.batch_size} accum={args.grad_accum} "
          f"psi_warmup={args.psi_warmup_steps} n_steps={args.n_student_steps}")
    log_path = out_dir / "sid_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)

    def _sample_inputs():
        batch = _next()
        cond_pos = batch["emb"].to(device=device, dtype=dtype)
        B = cond_pos.size(0)
        cond_neg_b = cond_neg.expand(B, -1, -1).contiguous() if B > 1 else cond_neg
        noise = torch.randn(B, 16, 1, H_lat, W_lat, device=device, dtype=dtype)
        return noise, cond_pos, cond_neg_b

    for outer in range(args.total_outer_steps):
        # ---- ψ (score helper) update (every step, including warmup) ----
        peft_model.train()
        opt_psi.zero_grad()
        psi_metrics = {}
        for _ in range(args.grad_accum):
            noise, cond_pos, cond_neg_b = _sample_inputs()
            L_psi, m_psi = sid_score_helper_loss(
                student_v, psi_v, noise, cond_pos, cond_neg_b,
                student_cfg=args.student_cfg, n_steps=args.n_student_steps,
                mu_t=args.mu_t, sigma_t=args.sigma_t,
            )
            (L_psi / args.grad_accum).backward()
            psi_metrics = {k: float(v) for k, v in m_psi.items()}
        torch.nn.utils.clip_grad_norm_(psi_params, args.grad_clip)
        opt_psi.step()

        # ---- θ (generator) update (warmup 後のみ) ----
        gen_metrics = {}
        if outer >= args.psi_warmup_steps:
            opt_gen.zero_grad()
            for _ in range(args.grad_accum):
                noise, cond_pos, cond_neg_b = _sample_inputs()
                L_theta, m_theta = sid_generator_loss(
                    student_v, teacher_v, psi_v, noise, cond_pos, cond_neg_b,
                    teacher_cfg=args.teacher_cfg, student_cfg=args.student_cfg,
                    n_steps=args.n_student_steps, alpha=args.alpha,
                    mu_t=args.mu_t, sigma_t=args.sigma_t,
                )
                (L_theta / args.grad_accum).backward()
                gen_metrics = {k: float(v) for k, v in m_theta.items()}
            torch.nn.utils.clip_grad_norm_(student_params, args.grad_clip)
            opt_gen.step()

        if outer % args.log_every == 0:
            metrics = {"outer": outer, "elapsed": time.time() - t0,
                       "phase": "warmup" if outer < args.psi_warmup_steps else "joint",
                       **psi_metrics, **gen_metrics}
            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 not in ("outer",))
            print(f"[outer {outer}/{args.total_outer_steps}] {msg}", flush=True)

        if outer > 0 and outer % args.sample_every == 0:
            sd = {k: v.detach().cpu() for k, v in peft_model.state_dict().items()
                  if "lora_" in k and ".student." in k}
            from safetensors.torch import save_file as _sf
            _sf(sd, str(out_dir / f"sid_student_step{outer:05d}.safetensors"))
            print(f"[save] sid_student_step{outer:05d}.safetensors", flush=True)
            try:
                import modal
                modal.Volume.from_name("anima-outputs").commit()
            except Exception:
                pass

    print("[done] saving final")
    sd_student_final = {k: v.detach().cpu() for k, v in peft_model.state_dict().items()
                        if "lora_" in k and ".student." in k}
    from safetensors.torch import save_file as _sf
    _sf(sd_student_final, str(out_dir / "sid_student_final.safetensors"))
    log_f.close()


if __name__ == "__main__":
    main()