#!/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()