#!/usr/bin/env python3 """ Anima PCM (Phased Consistency Model) distillation ================================================= Reference: G-U-N/Phased-Consistency-Model code/text_to_image_sd3/train_pcm_lora_sd3.py SD3 PCM は FlowMatch + v-pred + LoRA で、Anima の rectified flow と math 完全一致。 ε↔v 変換不要、scheduler/loss を直接移植可能。 Algorithm: - num_euler_timesteps=N (e.g. 50) で grid を作り、K phase に等分割 - 各 step ランダムに index in [0, N) を選び、現 phase の終端まで Euler 1 step - student の predicted x_phase と、teacher 1 step 前進 → student EMA-target の x_phase を MSE - pseudo-Huber loss + 任意 adv loss (skip in v1) memory: 1 grad-through student fwd + 3 no_grad fwd (teacher cond, teacher uncond, student-as-EMA) → ~60-80 GB on B200 (LoRA-only, batch=1, 768²) データ: LADD precompute cache の emb を流用 (prompt 再 encode 不要) x0 は teacher rollout で online 生成するので不要 (ただし noise / cap 起点として 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 load_warm_lora, save_lora_state from distill.pcm_scheduler import make_pcm_solver, euler_multiphase, pseudo_huber # ----- dataset: precomputed cache の emb と任意で x0 を使う ---------------- class PCMCacheDataset(Dataset): """LADD/Reflow cache の (caption, emb, x0) を読む。x0 は initial noise の source。""" def __init__(self, cache_dir: str | Path): self.cache_dir = Path(cache_dir) self.meta = json.loads((self.cache_dir / "metadata.json").read_text(encoding="utf-8")) 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).squeeze(0) emb = torch.load(m["emb_path"], map_location="cpu", weights_only=True).squeeze(0) return {"x0": x0, "emb": emb} def pcm_collate(batch): return { "x0": torch.stack([b["x0"] for b in batch]), "emb": torch.stack([b["emb"] for b in batch]), } def main(): ap = argparse.ArgumentParser() ap.add_argument("--cache-dir", required=True, type=str, help="LADD cache (emb と x0 を使用)") ap.add_argument("--out", required=True, type=str) ap.add_argument("--warm-lora", default="", type=str) ap.add_argument("--total-steps", type=int, default=8000) ap.add_argument("--batch-size", type=int, default=1) ap.add_argument("--grad-accum", type=int, default=4) ap.add_argument("--resolution", type=int, default=768) ap.add_argument("--num-euler-timesteps", type=int, default=50) ap.add_argument("--num-phases", type=int, default=4) ap.add_argument("--sigma-shift", type=float, default=3.0) ap.add_argument("--w-min", type=float, default=4.0, help="CFG-aug min") ap.add_argument("--w-max", type=float, default=5.0, help="CFG-aug max") ap.add_argument("--w-fixed", type=float, default=-1.0, help=">=0 で範囲指定無視して固定") ap.add_argument("--huber-c", type=float, default=1e-3) ap.add_argument("--lr", type=float, default=5e-6) ap.add_argument("--weight-decay", type=float, default=0.01) ap.add_argument("--grad-clip", type=float, default=1.0) ap.add_argument("--lora-rank", type=int, default=32) 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) # ----- load Anima ----- print("[load] Anima bundle") bundle = build_anima(device=device, dtype=dtype) # teacher = frozen 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 # 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 if args.warm_lora: load_warm_lora(student_transformer, args.warm_lora) # solver solver = make_pcm_solver( num_euler_timesteps=args.num_euler_timesteps, num_phases=args.num_phases, sigma_shift=args.sigma_shift, device=device, dtype=torch.float32, ) print(f"[schedule] N={solver.num_steps} K={len(solver.phase_ends)} " f"phase_ends={solver.phase_ends.tolist()}") # cond_neg with torch.no_grad(): cond_neg = bundle.text_encode([args.neg_prompt or ""]) # optimizer opt = torch.optim.AdamW(student_params, lr=args.lr, betas=(0.9, 0.999), weight_decay=args.weight_decay, eps=1e-8) # dataset (emb は cache から、x0 は initial noise の source として使う — # 実際の noise は毎 step randn) print(f"[data] {args.cache_dir}") dataset = PCMCacheDataset(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=pcm_collate, drop_last=True, pin_memory=True, ) print(f"[train] steps={args.total_steps} bs={args.batch_size} accum={args.grad_accum} " f"N={args.num_euler_timesteps} K={args.num_phases}") log_path = out_dir / "pcm_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): student_transformer.train() opt.zero_grad() metrics = {} for _ in range(args.grad_accum): batch = _next() x0 = batch["x0"].to(device=device, dtype=dtype) emb = batch["emb"].to(device=device, dtype=dtype) B = x0.size(0) cond_neg_b = cond_neg.expand(B, -1, -1).contiguous() if B > 1 else cond_neg # ----- index sampling ----- index = torch.randint(0, solver.num_steps, (B,), device=device) sigma_cur = solver.sigmas[index].to(dtype=dtype) sigma_prev = solver.sigmas[index + 1].to(dtype=dtype) # ----- x_t (forward process) ----- noise = torch.randn_like(x0) sigma_cur_ = sigma_cur.view(-1, *([1] * (x0.dim() - 1))) x_t = sigma_cur_ * noise + (1.0 - sigma_cur_) * x0 # ----- student forward (grad) ----- v_student = AnimaBundle.dit_forward(student_transformer, x_t, sigma_cur, emb) # phase-snap: student の "予測" を当該 phase 終端まで進める pred_phase, sigma_phase = euler_multiphase(x_t, v_student, index, solver) # ----- target side (no_grad) ----- with torch.no_grad(): # CFG-aug w sample if args.w_fixed >= 0: w = torch.full((B,), args.w_fixed, device=device, dtype=dtype) else: w = torch.empty(B, device=device, dtype=dtype).uniform_(args.w_min, args.w_max) v_t_cond = AnimaBundle.dit_forward(teacher_transformer, x_t, sigma_cur, emb) v_t_uncond = AnimaBundle.dit_forward(teacher_transformer, x_t, sigma_cur, cond_neg_b) w_ = w.view(-1, *([1] * (x_t.dim() - 1))) v_cfg = v_t_uncond + w_ * (v_t_cond - v_t_uncond) # 1 teacher step forward (x_t → x_prev) dt = (sigma_prev - sigma_cur).view(-1, *([1] * (x_t.dim() - 1))) x_prev = x_t + dt * v_cfg # student-as-EMA target prediction at x_prev v_tgt = AnimaBundle.dit_forward(student_transformer, x_prev, sigma_prev, emb) target_phase, _ = euler_multiphase(x_prev, v_tgt, (index + 1).clamp(max=solver.num_steps - 1), solver) # ----- loss ----- diff = pred_phase - target_phase.detach() loss = pseudo_huber(diff, c=args.huber_c).mean() / args.grad_accum loss.backward() metrics = { "loss": float((loss * args.grad_accum).detach()), "pred_abs": float(pred_phase.detach().abs().mean()), "target_abs": float(target_phase.detach().abs().mean()), "w_mean": float(w.mean().detach()), "sigma_cur_mean": float(sigma_cur.mean().detach()), } torch.nn.utils.clip_grad_norm_(student_params, args.grad_clip) opt.step() if step % args.log_every == 0: metrics["step"] = step metrics["elapsed"] = time.time() - t0 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 != "step") print(f"[step {step}/{args.total_steps}] {msg}", flush=True) if step > 0 and step % args.sample_every == 0: save_lora_state(student_transformer, out_dir, f"pcm_step{step:05d}") print(f"[save] pcm_step{step:05d}.safetensors", flush=True) try: import modal modal.Volume.from_name("anima-outputs").commit() except Exception: pass print("[done] saving final") save_lora_state(student_transformer, out_dir, "pcm_final") log_f.close() if __name__ == "__main__": main()