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