#!/usr/bin/env python3 """ Anima DMDX (ADM-only) 蒸留 — arxiv 2507.18569v1 移植 ===================================================== Pipeline: - student: wide LoRA on Anima base (DMD2 と同じ wide target、~72M trainable) - discriminator: LADD-style (teacher MiniTrainDIT frozen backbone + spectral norm heads) - real: teacher_x0_cache の x0 を流用 - ADM loss: hinge GAN at t-Δt after teacher 1-step evolution - time schedule: cubic high-noise bias (paper default) - alternation: N_critic × disc → 1 × generator (DMD2 風) DMD2 との違い: - DMD2: dual adapter (student + fake_score)、reverse-KL grad trick - DMDX: single student LoRA + 別 D heads、hinge GAN、TVD 最小化 オプション: - --recon-weight > 0 で LADD 流 Smooth-L1 anchor 追加 (mean collapse 防止) - --misaligned-pairs-d : LADD 流の text-alignment trick Modal CLI 例: modal run modal_app.py::train_dmdx_distill \\ --warm-lora /models/loras/anima_turbo.safetensors \\ --total-outer-steps 5000 """ 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 AnimaPaths, build_anima, AnimaBundle from distill.dmd2_trainer import attach_wide_lora from distill.train_traj import save_lora_state from distill.train_ladd import PrecomputedCacheDataset, ladd_collate from distill.anima_ladd_disc import AnimaLADDDiscriminator from distill.dmdx_loss import adm_generator_loss, adm_discriminator_loss def main(): ap = argparse.ArgumentParser() ap.add_argument("--cache-dir", required=True, type=str, help="teacher_x0_cache (real x0 source)") 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=5000, help="1 outer = N_critic + 1 generator") ap.add_argument("--n-critic-per-gen", type=int, default=2, help="DMDX paper: hinge GAN は 1:1〜2:1 が一般的 (DMD2 の 5:1 より少ない)") ap.add_argument("--n-student-steps", type=int, default=4, help="few-step rollout step 数") ap.add_argument("--batch-size", type=int, default=1) ap.add_argument("--resolution", type=int, default=768) ap.add_argument("--teacher-cfg", type=float, default=4.5, help="Anima 公式推奨 CFG") ap.add_argument("--student-cfg", type=float, default=1.0) ap.add_argument("--dt-ratio", type=float, default=1.0 / 64, help="paper default Δt = T/64") ap.add_argument("--recon-weight", type=float, default=0.0, help="LADD 流 Smooth-L1 anchor (>0 で有効化、stability boost)") ap.add_argument("--lora-rank", type=int, default=32) ap.add_argument("--lr-gen", type=float, default=5e-6) ap.add_argument("--lr-disc", type=float, default=1e-5) ap.add_argument("--weight-decay", type=float, default=0.01) 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") ap.add_argument("--head-hidden", type=int, default=512) 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("--neg-prompt", default="") 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 ----- 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 on bundle.transformer ----- print("[setup] 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: from distill.train_traj import load_warm_lora load_warm_lora(student_transformer, args.warm_lora) # ----- discriminator (LADD-style) ----- print("[setup] DMDX discriminator (teacher backbone + spectral 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, ) 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) disc_params = disc.trainable_parameters() print(f"[setup] D heads trainable: {sum(p.numel() for p in disc_params)/1e6:.1f}M") # ----- optimizers ----- opt_gen = torch.optim.AdamW(student_params, lr=args.lr_gen, betas=(0.0, 0.999), weight_decay=args.weight_decay, eps=1e-8) opt_disc = torch.optim.AdamW(disc_params, lr=args.lr_disc, betas=(0.0, 0.999), weight_decay=args.weight_decay, eps=1e-8) # ----- velocity functions ----- def student_v(x, t, cond): return AnimaBundle.dit_forward(student_transformer, x, t, cond) def teacher_v(x, t, cond): return AnimaBundle.dit_forward(teacher_transformer, x, t, cond) # ----- conditioning (neg prompt cached once) ----- with torch.no_grad(): cond_neg = bundle.text_encode([args.neg_prompt or ""]) # ----- dataset (precompute cache provides real x0 + caption emb) ----- 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] outer={args.total_outer_steps} " f"cycle=({args.n_critic_per_gen} disc + 1 gen) " f"n_student_steps={args.n_student_steps} dt={args.dt_ratio:.4f}") log_path = out_dir / "dmdx_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 _prep_batch(): 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) cond_neg_b = cond_neg.expand(B, -1, -1).contiguous() if B > 1 else cond_neg noise = torch.randn_like(x0_teacher) return noise, x0_teacher, cond_pos, cond_neg_b for outer in range(args.total_outer_steps): # ---- N disc updates ---- disc_metrics = {} for _ in range(args.n_critic_per_gen): noise, x0_teacher, cond_pos, cond_neg_b = _prep_batch() disc.train() student_transformer.eval() opt_disc.zero_grad() loss_d, m_d = adm_discriminator_loss( student_v, teacher_v, disc, noise, x0_teacher, cond_pos, cond_neg_b, teacher_cfg=args.teacher_cfg, student_cfg=args.student_cfg, n_student_steps=args.n_student_steps, dt_ratio=args.dt_ratio, ) loss_d.backward() torch.nn.utils.clip_grad_norm_(disc_params, args.grad_clip) opt_disc.step() disc_metrics = {k: float(v) for k, v in m_d.items()} # ---- 1 generator update ---- noise, x0_teacher, cond_pos, cond_neg_b = _prep_batch() student_transformer.train() disc.eval() opt_gen.zero_grad() loss_g, m_g = adm_generator_loss( student_v, teacher_v, disc, noise, cond_pos, cond_neg_b, real_x0=x0_teacher, teacher_cfg=args.teacher_cfg, student_cfg=args.student_cfg, n_student_steps=args.n_student_steps, dt_ratio=args.dt_ratio, recon_weight=args.recon_weight, ) loss_g.backward() torch.nn.utils.clip_grad_norm_(student_params, args.grad_clip) opt_gen.step() gen_metrics = {k: float(v) for k, v in m_g.items()} # ---- log ---- if outer % args.log_every == 0: metrics = {"outer": outer, "elapsed": time.time() - t0, **disc_metrics, **gen_metrics} log_f.write(json.dumps(metrics) + "\n") msg = " ".join(f"{k}={v:.4f}" for k, v in metrics.items() if k != "outer") print(f"[outer {outer}/{args.total_outer_steps}] {msg}", flush=True) # ---- ckpt (student LoRA only) ---- if outer > 0 and outer % args.sample_every == 0: save_lora_state(student_transformer, out_dir, f"dmdx_student_step{outer:05d}") print(f"[save] dmdx_student_step{outer:05d}.safetensors", flush=True) try: import modal modal.Volume.from_name("anima-outputs").commit() except Exception as e: print(f"[save] volume commit failed: {e}", flush=True) # final print("[done] saving final") save_lora_state(student_transformer, out_dir, "dmdx_student_final") log_f.close() if __name__ == "__main__": main()