#!/usr/bin/env python3 """ Anima Reflow distillation (InstaFlow / rfpp 流派) ================================================ 理論的根拠: - Anima は既に rectified flow なので「Reflow」(1 round) を当てるのが自然 - rfpp (NeurIPS 2024 "Improving the Training of Rectified Flows") の知見: * U-shape t sampling * LPIPS-Huber (latent Huber + 間欠的に pixel LPIPS) で 1-round で十分 * adversarial 不要 (LADD と独立カテゴリ) - 出力: 4-step (or 1-step) で動く student LoRA データ: 事前に (noise, x0, emb) triplet を precompute (--save-noise) しておく前提。 /dataset/teacher_x0_cache_with_noise/{noise,x0,emb}/{idx}.pt metadata.json に noise_path も書かれていること。 訓練 1 step: 1. (noise_i, x0_i, emb_i) を batch load 2. t ~ U-shape on (0, 1) 3. x_t = (1 - t) * x0 + t * noise 4. v_target = noise - x0 5. v_pred = student.dit_forward(x_t, t, emb) 6. L_huber = huber(v_pred, v_target, delta=0.03) 7. 任意 (lpips-every step): x0_pred = x_t - t*v_pred、VAE decode → LPIPS vs x0 8. backward → student LoRA 更新 """ 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 class ReflowPairDataset(Dataset): """precompute された (noise, x0, emb) triplet を読む。 metadata.json の各 entry に noise_path / x0_path / emb_path が必要。""" 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")) if not self.meta or "noise_path" not in self.meta[0]: raise RuntimeError( f"{self.cache_dir}/metadata.json has no 'noise_path' — " "rerun precompute with --save-noise." ) 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) noise = torch.load(m["noise_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, "noise": noise, "emb": emb} def reflow_collate(batch): return { "x0": torch.stack([b["x0"] for b in batch]), "noise": torch.stack([b["noise"] for b in batch]), "emb": torch.stack([b["emb"] for b in batch]), } def u_shape_t(B: int, device, dtype=torch.float32) -> torch.Tensor: """rfpp 流の U-shape sampler。両端 (t≈0, t≈1) に分布が偏る。 実装: u ~ U(0,1) → t = 0.5 * (1 + sign(u-0.5) * |2u-1|^0.5) (sqrt によって両端密度上昇)""" u = torch.rand(B, device=device, dtype=dtype) centered = 2 * u - 1 t = 0.5 * (1.0 + torch.sign(centered) * centered.abs().sqrt()) return t.clamp(1e-3, 1.0 - 1e-3) def main(): ap = argparse.ArgumentParser() ap.add_argument("--cache-dir", required=True, type=str, help="--save-noise 付きで precompute した cache") 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=4) ap.add_argument("--grad-accum", type=int, default=2) ap.add_argument("--resolution", type=int, default=768) ap.add_argument("--lr", type=float, default=1e-4) 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("--huber-delta", type=float, default=0.03) ap.add_argument("--lpips-weight", type=float, default=0.1) ap.add_argument("--lpips-every", type=int, default=4, help="N step ごとに LPIPS reg を加算 (cost 削減)") ap.add_argument("--t-sampler", default="u_shape", choices=["u_shape", "uniform"]) 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 (teacher は捨てる、cache が teacher 役) ----- print("[load] Anima bundle (student only — teacher cached on disk)") bundle = build_anima(device=device, dtype=dtype) # ----- 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: load_warm_lora(student_transformer, args.warm_lora) # LPIPS (任意) lpips_fn = None if args.lpips_weight > 0: import lpips as _lpips lpips_fn = _lpips.LPIPS(net="alex").to(device).eval() for p in lpips_fn.parameters(): p.requires_grad = False print("[setup] LPIPS(alex) loaded") # ----- optimizer ----- opt = torch.optim.AdamW(student_params, lr=args.lr, betas=(0.9, 0.999), weight_decay=args.weight_decay, eps=1e-8) # ----- dataset ----- print(f"[data] loading cache {args.cache_dir}") dataset = ReflowPairDataset(args.cache_dir) print(f" {len(dataset)} (noise, x0, emb) triplets") loader = DataLoader( dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, collate_fn=reflow_collate, drop_last=True, pin_memory=True, ) # ----- training loop ----- print(f"[train] steps={args.total_steps} bs={args.batch_size} accum={args.grad_accum} " f"lpips_w={args.lpips_weight} every={args.lpips_every}") log_path = out_dir / "reflow_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) noise = batch["noise"].to(device=device, dtype=dtype) emb = batch["emb"].to(device=device, dtype=dtype) B = x0.size(0) # ----- t sampling ----- if args.t_sampler == "u_shape": t = u_shape_t(B, device, dtype=torch.float32).to(dtype=dtype) else: t = torch.rand(B, device=device, dtype=dtype).clamp(1e-3, 1.0 - 1e-3) t_ = t.view(-1, *([1] * (x0.dim() - 1))) x_t = (1 - t_) * x0 + t_ * noise v_target = noise - x0 # ----- student forward (1 forward only) ----- v_pred = AnimaBundle.dit_forward(student_transformer, x_t, t, emb) l_huber = F.huber_loss(v_pred.float(), v_target.float(), delta=args.huber_delta, reduction="mean") # ----- LPIPS reg (every N steps) ----- l_lpips = torch.zeros((), device=device) if lpips_fn is not None and (step % args.lpips_every) == 0: x0_pred = x_t - t_ * v_pred vae_dtype = next(bundle.vae.model.parameters()).dtype img_p = bundle.vae.model.decode(x0_pred.to(dtype=vae_dtype), bundle.vae_scale).squeeze(2) with torch.no_grad(): img_t = bundle.vae.model.decode(x0.to(dtype=vae_dtype), bundle.vae_scale).squeeze(2) l_lpips = lpips_fn(img_p.float(), img_t.float()).mean() loss = (l_huber + args.lpips_weight * l_lpips) / args.grad_accum loss.backward() metrics = { "l_huber": float(l_huber.detach()), "l_lpips": float(l_lpips.detach()), "loss": float((l_huber + args.lpips_weight * l_lpips).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"reflow_step{step:05d}") print(f"[save] reflow_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, "reflow_final") log_f.close() if __name__ == "__main__": main()