Initial commit: rapid-anima distillation codebase

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+samples/dmd2_step500.png filter=lfs diff=lfs merge=lfs -text
+samples/dmdx_step500.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,29 @@
+MIT License
+
+Copyright (c) 2026 daraskme
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+---
+
+NOTE: This license covers the **code** in this repository only. Models
+produced or fine-tuned using this code are governed by the upstream model
+licenses (CircleStone Labs Non-Commercial License + NVIDIA Open Model License
+Derivative Model terms for Anima). Such derivative models may only be used
+for non-commercial purposes.

README.md

@@ -0,0 +1,121 @@
+# rapid_anima — Anima 生成速度向上 on Modal
+
+CircleStone Labs の [Anima](https://huggingface.co/circlestone-labs/Anima)
+(2B パラメータの DiT)の **生成速度を改善** するための一式。
+旧名 `darask_anima`、2026-05 にスコープを「fine-tune + 蒸留」から「速度向上特化」に転換。
+
+> ⚠️ **ライセンス**: Anima は CircleStone Labs Non-Commercial License +
+> NVIDIA Open Model License (Derivative Model 条項) で **非商用のみ**。
+> 派生モデルも非商用縛り。詳細は [末尾のライセンス節](#ライセンス) 参照。
+
+## 主な高速化レバー
+
+- **B200 sageattention 実 engage** — sm_100 patched wheel で silent fallback 回避
+  ([darask0/modal_B200_sageattetion_comfyUI](https://huggingface.co/darask0/modal_B200_sageattetion_comfyUI))
+- **batch=8 並列生成** — ComfyUI EmptyLatentImage 経由で 1 submit に 8 枚
+- **蒸留 LoRA** — Civitai 公式 Anima Turbo merge / 自前蒸留 (PCM, Z-Image traj, LADD, Reflow)
+- **ComfyUI 推論最適化** — CapitanFlowMatch scheduler, sageattention, DPM-Solver++ 等
+
+### 実測 (B200, 1024² / 30 step)
+
+| 構成 | per-image | 5000 枚 cost |
+|---|---|---|
+| 旧 batch=1 (no sage) | ~3.3s | ~$44 |
+| **新 batch=8 + sage** | **~3.05s** | **~$26** |
+| 4-step 蒸留 LoRA + sage | ~0.5-1s | (推論時のみ) |
+
+## 配布中の LoRA
+
+[**huggingface.co/darask0/anima-distill-loras**](https://huggingface.co/darask0/anima-distill-loras)
+に Anima v1.0 base 用の蒸留 LoRA を集約 (PEFT + ComfyUI 両形式)。
+
+| サブディレクトリ | 手法 | step / CFG | ステータス | 詳細 |
+|---|---|---|---|---|
+| [`pcm/`](https://huggingface.co/darask0/anima-distill-loras/tree/main/pcm) | Phased Consistency Model | 4-step / CFG=1.0 | ✅ 配布中 | [docs/pcm.md](docs/pcm.md) |
+| `dmd2/` (近日) | DMD2 + TrigFlow | 4-step / CFG=1.0 | 🟢 訓練中 (resume) | [docs/dmd2.md](docs/dmd2.md) |
+| `dmdx/` (近日) | DMDX (ADM) | 4-step / CFG=1.0 | 🟢 訓練中 (新規実装) | [docs/dmdx.md](docs/dmdx.md) |
+
+将来追加検討: LADD / Reflow / Z-Image trajectory / SiD2 など。
+
+**ComfyUI 使い方の最短手順**:
+
+1. `*_comfy.safetensors` を `ComfyUI/models/loras/` に配置
+2. Anima v1.0 base workflow に `LoraLoaderModelOnly` を挿入、`strength_model: 1.0`
+3. KSampler を **steps=4 / cfg=1.0**、推奨 sampler / scheduler を選択
+
+詳細な使い方・訓練詳細は各 LoRA の HF model card / docs/{pcm,dmd2,dmdx}.md を参照。
+
+### サンプル画像 (4-step / CFG=1.0 / er_sde + simple)
+
+DMD2 と DMDX、それぞれ step 500 時点での生成 (warm-start: Civitai Anima Turbo)。
+同 prompt (`2girls, flandre scarlet, remilia scarlet, touhou, ...`)、同 seed=42。
+
+| DMD2 step 500 | DMDX step 500 |
+|---|---|
+| ![dmd2_step500](samples/dmd2_step500.png) | ![dmdx_step500](samples/dmdx_step500.png) |
+
+step 500 ではどちらも warm-start (Civitai Turbo) が支配的で、訓練の差異は微小。
+本番完走 (step 5000) 後に再比較予定。
+
+## ディレクトリ構成 (簡略)
+
+```
+rapid_anima/
+├── modal_app.py                  # Modal アプリ本体 (全 function 定義)
+├── configs/                      # diffusion-pipe 設定 (Phase 1 fine-tune 用)
+├── scripts/
+│   ├── generate_dataset.py       # ComfyUI 経由 self-distill 生成
+│   ├── clean_captions.py         # quality/meta タグ除去
+│   ├── *_workflow.json           # Anima ComfyUI workflow テンプレ
+│   ├── *_prompts.txt             # 各種 prompt セット
+│   └── distill/                  # 蒸留実装 (PCM / DMD2 / DMDX / LADD / Reflow / Z-Image 等)
+├── docs/                         # 詳細ドキュメント (下記参照)
+├── samples/                      # README 用検証サンプル画像
+├── requirements.txt
+└── README.md (このファイル)
+```
+
+> **重要**: `modal_app.py` は `configs/` と `scripts/` をそのままの構造で
+> image に同梱する。ファイルをルート直下にフラット配置すると `add_local_dir`
+> で失敗するので注意。
+
+## ドキュメント
+
+**手法別** (本日 v1.0 base での新規試行):
+
+| ドキュメント | 内容 |
+|---|---|
+| [docs/pcm.md](docs/pcm.md) | PCM 蒸留 (5000 step 完走、HF 配布済、$22) |
+| [docs/dmd2.md](docs/dmd2.md) | DMD2 + TrigFlow 蒸留 (resume 進行中、~$50) |
+| [docs/dmdx.md](docs/dmdx.md) | DMDX (ADM、新規実装、arxiv 2507.18569v1 移植、進行中、$30) |
+
+**インフラ / セットアップ**:
+
+| ドキュメント | 内容 |
+|---|---|
+| [docs/setup.md](docs/setup.md) | Modal セットアップ、HF シークレット、データセット用意 |
+| [docs/workflow.md](docs/workflow.md) | Step 0-5 (モデル DL / dataset / Phase 1 / 検証 / 蒸留)、GPU 選び、コスト概算 |
+| [docs/distillation.md](docs/distillation.md) | Anima 蒸留 — R3GAN 5 連続失敗の deep dive、設計理由、教訓 |
+| [docs/migration_log.md](docs/migration_log.md) | 2026-05 既存実装移植 (Z-Image / DMD2 / LADD / Reflow / PCM)、比較結果、reward hacking 実証 |
+| [docs/operations.md](docs/operations.md) | Modal CLI / diffusion-pipe 罠、実測コスト、トラブルシューティング |
+
+## 参考資料
+
+- 公式モデルカード: https://huggingface.co/circlestone-labs/Anima
+- diffusion-pipe: https://github.com/tdrussell/diffusion-pipe
+- DMD2 論文 (Z-Image): https://arxiv.org/abs/2511.22677
+- LCM 論文: https://arxiv.org/abs/2310.04378
+- Hyper-SD: https://github.com/bytedance/hyper-sd
+- RDBT-Anima (Anima 蒸留先行例): https://civitai.com/models/2364703
+- PCM 元論文: [Wang et al. NeurIPS 2024](https://arxiv.org/abs/2405.18407) / [G-U-N/Phased-Consistency-Model](https://github.com/G-U-N/Phased-Consistency-Model)
+
+## ライセンス
+
+このリポジトリ自体のコードは MIT で公開。生成される派生モデル(蒸留 LoRA や
+fine-tune 結果)は **Anima のライセンス** に従う必要があり、**非商用のみ**:
+
+- CircleStone Labs Non-Commercial License
+- NVIDIA Open Model License(Derivative Model 条項)
+
+詳細・商用利用問い合わせは Anima 公式ページを参照:
+https://huggingface.co/circlestone-labs/Anima

configs/phase1_anima.toml

@@ -0,0 +1,75 @@
+# =============================================================================
+# Anima Phase 1: quality タグ依存性除去 (LoRA)
+# 目的: "masterpiece, best quality, score_9, ..." を書かなくても base が出していた
+#       品質の絵が出るようにする。**審美の方向性は変えない**(artist タグ等は維持)。
+# 性質: 概念追加でなく "prior shift" → rank 小・lr 低・epochs 少で十分。
+#       過剰に回すと base の汎用性を損なう (catastrophic forgetting)。
+# 使用: deepspeed --num_gpus=1 train.py --deepspeed --config phase1_anima.toml
+# =============================================================================
+
+output_dir = '/output/phase1'
+
+# データセット設定 (別ファイル)
+dataset = '/workspace/configs/phase1_dataset.toml'
+
+# 学習設定 ---------------------------------------------------------------------
+# epochs: prior shift だけなら 2 で十分。3 にすると過学習リスク微増。
+# 1 epoch ≈ 3-4h (A100-80GB, 5k枚) → 2 epoch ≈ 6-8h ≈ $15-20
+# 10k 枚使う場合: 1 epoch ≈ 6-7h → 2 epoch ≈ 12-14h ≈ $30-35
+epochs = 2
+micro_batch_size_per_gpu = 1
+pipeline_stages = 1              # 単 GPU は 1
+gradient_accumulation_steps = 4  # 実効 batch = 4 (5k枚以下なら 2 でも可)
+gradient_clipping = 1.0
+warmup_steps = 100
+
+# CPU offload (blocks_to_swap > 0 で VRAM 節約だが速度低下。A100-80GB なら 0 でOK)
+blocks_to_swap = 0
+
+# 評価/保存
+eval_every_n_epochs = 1
+save_every_n_epochs = 1           # 古い ckpt は cleanup_checkpoints で間引く
+checkpoint_every_n_minutes = 60  # 安全保険 (途中落ち対策)
+activation_checkpointing = true
+save_dtype = 'bfloat16'
+
+# モデル ----------------------------------------------------------------------
+[model]
+type = 'anima'
+# Anima preview3-base (公式 ComfyUI 形式)
+transformer_path = '/models/checkpoints/anima-preview3-base.safetensors'
+vae_path = '/models/checkpoints/qwen_image_vae.safetensors'
+llm_path = '/models/checkpoints/qwen_3_06b_base.safetensors'
+dtype = 'bfloat16'
+transformer_dtype = 'bfloat16'
+
+# 公式必須ルール: LLM adapter は学習しない
+# (テキスト埋め込み変換器で、訓練すると壊れやすい)
+llm_adapter_lr = 0
+
+# Anima は flow matching ベース。timestep サンプリングは logit_normal が安定。
+timestep_sample_method = 'logit_normal'
+
+# LoRA --------------------------------------------------------------------
+[adapter]
+type = 'lora'
+rank = 64                # 審美 FT は概念追加でないので 32〜64 で十分
+alpha = 64               # = rank だと scale 1.0
+dropout = 0.0
+dtype = 'bfloat16'
+# 対象モジュール: DiT の attention/MLP のみ (LLM adapter は除外)
+# diffusion-pipe は anima 用にデフォルトで適切な層を選んでくれる
+
+# Optimizer ---------------------------------------------------------------
+[optimizer]
+type = 'AdamW8bitKahan'      # bf16 で安定、VRAM 節約
+lr = 2e-5                    # 公式推奨。広げるなら 1e-5〜5e-5
+betas = [0.9, 0.99]
+weight_decay = 0.01
+eps = 1e-8
+stabilize = false
+
+# Monitoring --------------------------------------------------------------
+[monitoring]
+log_every_n_steps = 10
+enable_wandb = false         # wandb 使うなら true + secret 追加

configs/phase1_dataset.toml

@@ -0,0 +1,34 @@
+# =============================================================================
+# Phase 1 データセット設定
+# 期待構造: /dataset/cleaned/<image>.{png,jpg,webp} と同名の <image>.txt
+# (clean_captions.py 実行後の出力先)
+# =============================================================================
+
+resolutions = [1024]                # Anima 推奨 1MP
+enable_ar_bucket = true             # アスペクト比バケット
+min_ar = 0.5                        # 縦長 1:2 まで
+max_ar = 2.0                        # 横長 2:1 まで
+num_ar_buckets = 9
+frame_buckets = [1]                 # 画像のみ
+
+# キャプション関係 ----------------------------------------------------------
+# danbooru タグ風キャプションを ", " 区切りでシャッフル
+shuffle_tags = true
+# 学習を頑健にするため一部タグをランダムに drop
+# (品質タグはすでに clean_captions.py で除去済みだが、artist 等の頑健化に有効)
+caption_dropout_rate = 0.1
+# プロンプトをまるごと空にする確率 (CFG 学習の暗黙化、5〜10% が定番)
+text_dropout_rate = 0.05
+
+# テキストエンコーダ出力の事前キャッシュ。
+#   効果: 各 step の Qwen3 forward (~0.2s) を省略 → 全体 15-25% 高速化 ≈ $7-13/run
+#   制約: 上の shuffle_tags / *_dropout_rate と相反 (キャッシュは固定文字列で計算)
+# 速度優先する場合: 下記 ↑ の shuffle_tags=false, caption_dropout_rate=0,
+#                  text_dropout_rate=0 にしてこれを true にする。
+cache_text_embeddings = false
+
+# 実データセット --------------------------------------------------------------
+[[directory]]
+path = '/dataset/cleaned'
+# クラス重み付け (複数 dataset を混ぜる時に使う。今は 1.0)
+num_repeats = 1

docs/distillation.md

@@ -0,0 +1,468 @@
+# Anima 蒸留 — 実走で分かった知見
+
+[← README に戻る](../README.md)
+
+> **TL;DR**: 数万円予算で Anima を自前 SOTA 蒸留(DMD2+TSCD+R3GAN)するのは
+> 現実的に困難。**5 回失敗で累計 $165**。実用解は **`merge_turbo_lora`**
+> (公式 Anima Turbo LoRA との合成、$0.5、即動く)。自前で挑む人向けに、
+> 何が起きるか・なぜ難しいかを以下に集約。
+>
+> 2026-05 以降の改善試行 (PCM / LADD / Reflow など既存実装の移植) は
+> [migration_log.md](migration_log.md) を参照。
+
+---
+
+## 1. 概要(結論先出し)
+
+このリポジトリでは Anima(CircleStone Labs, Cosmos-Predict2 派生 2B DiT)を
+**Decoupled DMD2 + TSCD + R3GAN** で **step 削減蒸留** することを目指した。
+B200 GPU(Modal)で 5000 枚 self-distillation データセットを使い、**計 5 回の失敗**:
+
+| 試行 | 構成 | 結果 | コスト |
+|---|---|---|---|
+| Phase A | full DiT, R3GAN gamma=50 | 完全ノイズ(D 爆発) | $70 |
+| Phase B v1 | Q,V LoRA, R3GAN OFF | 緑色化(AdaLN 不触) | $10 |
+| Phase B v2 | wide LoRA, R3GAN OFF | mean collapse | $10 |
+| Phase B v3a | wide LoRA, gamma=0.1 | step 30 で D 爆発 abort | $5 |
+| **Phase B v3b** | wide LoRA, gamma=1.0 | **mean collapse + 後で D 爆発** | $20 |
+
+**最終理解**: R3GAN は「D 爆発」と「無信号 → mean collapse」の間に **安定 zone が
+極めて狭く**、Anima latent(16ch × 128²)で持続的な訓練は現状の構成では困難。
+DMDR / Z-Image の公式実装が **DINOv2 reward に依存** していることからも、
+adv loss 単独の R3GAN だけでは Anima 系の蒸留は厳しい可能性が高い。
+
+### 1.1 結局どうすればいいか
+
+| 用途 | 推奨パス | コスト | 結果 |
+|---|---|---|---|
+| **すぐ 8-12 step 推論したい** | `merge_turbo_lora`(公式 Civitai Turbo LoRA を Anima base に重ねる)| **$0.5** | ✅ 即動く、CFG=1 / 8-12 step |
+| **どうしても自前で蒸留したい** | 本ドキュメントの試行錯誤を踏まえ、**DINOv2 reward 切替** 等を追加実装 | +$50-100 + 開発時間数日 | △ 未検証 |
+| **学習目的・記録目的** | 本ドキュメントを読む | $0 | 知見だけ |
+
+> **強い推奨**: 商用要件がなければ `merge_turbo_lora` で十分。CircleStone Labs
+> 自身が遥かに多い計算資源と試行錯誤の上で作った Turbo LoRA が既に Civitai
+> にある(https://civitai.com/models/2560840)、これに勝つのは数万円では困難。
+
+---
+
+## 2. 結局成功した実用パス: `merge_turbo_lora`
+
+```bash
+# 1. Anima Turbo LoRA を Civitai からダウンロード
+modal run modal_app.py::download_civitai_lora  # version_id=2877687 がデフォルト
+
+# 2. (任意) 自前 Phase 1 LoRA とマージ(なくても可)
+modal run modal_app.py::merge_turbo_lora
+
+# 3. ComfyUI で base + LoRA(strength=1.0)を組み、8 step CFG=1 で生成
+```
+
+実装: `modal_app.py::download_civitai_lora`, `merge_turbo_lora`, `stage_lora_to_models`。
+ComfyUI workflow は `scripts/anima_workflow_turbo.json` 参照。
+
+実測の Turbo LoRA 出力品質: **base 30 step CFG=4.5 と同等のキャラクター品質を
+8 step CFG=1 で達成**(検証画像は `/dataset/compare/turbo_8step/` に保存)。
+
+---
+
+## 3. 自前 SOTA 蒸留(本リポジトリで実装、未完成)
+
+`scripts/distill/` に Decoupled DMD2 + TSCD + R3GAN を実装。**コード自体は動く** が、
+ハイパーパラメータの調整で 5 回失敗。以下「我々が試した推奨設定」だが、これでも
+2000 step で失敗 — **これから挑む人は次節 §4 §5 を読んでから**。
+
+### 3.1 試した(が安定しなかった)設定
+
+```text
+=== モデル / データ ===
+base model:       Anima preview3-base (Cosmos-Predict2 派生)
+text encoder:     Qwen3-0.6B + LLM Adapter (T5 空間ブリッジ)
+VAE:              WanVAE (16-channel latent, scale=[mean, 1/std])
+dataset:          self-distillation 5000 枚(base が 30 step CFG 4.5 で生成)
+                  解像度 1024px 周辺の 7 種 aspect ratio をランダム
+
+=== 学習 (LoRA-only) ===
+trainable:        gen + guidance ともに LoRA、base は frozen
+gen LoRA:         **wide** (all-linear except llm_adapter), rank=32
+                  → AdaLN modulation を含む全 Linear に attach
+guidance LoRA:    Q,V LoRA, rank=32 (DMDR 互換、軽量)
+optimizer:        gen / guidance: AdamW lr=2e-5 β=(0.9, 0.999) wd=0.01
+                  D: Adam lr=2e-4 β=(0, 0.99) (StyleGAN 流)
+grad clip:        1.0
+
+=== Loss 構成 ===
+L_dmd:            DMDR gradient trick (p_real - p_fake) / norm
+L_consist:        TSCD MSE vs EMA generator (decay=0.999)
+L_adv (gen):      R3GAN relativistic RpGAN, weight=0.05
+L_disc:           R3GAN softplus + γ/2 · (R1 + R2)
+                  γ = 0.05〜0.5 (grid search 必須)
+
+=== 更新比率 / スケジュール ===
+guidance update:  5 回 / outer step
+generator update: 1 回 / outer step  
+discriminator:    1 回 / outer step
+CFG Augmentation: real_score を cfg=2.0 で評価 (cold_start 100 step は cfg=0)
+dynamic LoRA r:   real-side scale を cosine で 0 に decay (decay_steps=2000)
+
+=== Phase 構成 ===
+Phase B: base 直接 + wide LoRA, num_inference_steps=4, 2000 step
+Phase C: Phase B から resume または base 直接, wide LoRA, num_steps=2, 2000 step
+※ Full DiT 訓練 (Phase A) は不要 — LoRA-only で十分
+
+=== Sanity / Guard ===
+sanity test:      200 step + 中間生成 1 枚で degeneration 検出
+abort guards:     d_r1 > 1000 / d_r2 > 1000 / loss > 1e6
+checkpoint:       500 step 毎 + 都度 volume.commit()
+PYTHONUNBUFFERED=1 + cwd=/workspace/diffusion-pipe
+
+=== Modal infra ===
+GPU:              B200 (192GB) — LoRA-only で余裕、batch=4
+image:            CUDA 12.4 + Python 3.11 + diffusion-pipe + torchvision
+volumes:          anima-models / anima-dataset / anima-outputs
+secrets:          hf_token, civitai_api_key
+
+=== ComfyUI 出力形式 ===
+LoRA key:         diffusion_model.<module>.lora_A.weight / .lora_B.weight
+                  (PEFT の base_model.model.<...>.lora_A.default.weight から変換)
+基底 ckpt:        anima ckpt 系は net. prefix 必須 (ComfyUI が arch 検出に使う)
+```
+
+実装は `scripts/distill/` + `modal_app.py::train_sota_distill` を参照。
+
+> ⚠️ **重要**: 上の設定で 2000 step training を 5 回試したが、いずれも失敗。
+> Anima latent では R3GAN の安定 zone が極めて狭く、長期持続が困難だった。
+> **これから挑む人は §4(全失敗の詳細)+ §5(残された改善方向)を必ず読むこと**。
+
+### 3.2.1 なぜ LoRA-only(full DiT 訓練ではなく)
+
+| | full DiT(Phase A) | LoRA-only(Phase B/C) |
+|---|---|---|
+| trainable params | 2B | 30-150M |
+| step 時間 (B200) | 14 秒 | **2.8-3.0 秒**(5x 速) |
+| 出力 | 4 GiB ckpt | **30-150 MiB LoRA** |
+| 配布性 | 他 FT と排他 | **他 LoRA とスタック可** |
+| optimizer state | 8 GB | ~200 MB |
+| failure 局所化 | base 全体破壊リスク | LoRA だけ |
+
+Anima のように既に高品質な base を蒸留する場合、LoRA-only がコスト・配布性・
+安全性のすべてで優位。Phase A 流の full DiT 訓練は不要だった。
+
+### 3.2.2 なぜ「wide LoRA」(Q,V LoRA ではなく)
+
+Anima/Cosmos の DiT は **AdaLN modulation で timestep 情報を処理**する架構:
+
+```
+block(x, t):
+    t_embed = t_embedder(t)
+    # AdaLN modulation = block_input をどう変調するかを timestep から決める
+    scale, shift, gate = adaln_modulation(t_embed)  ← ここが timestep の入り口
+    x = LayerNorm(x) * (1 + scale) + shift
+    x = x + gate * attention(x, ...)
+```
+
+蒸留とは「**timestep の解釈を変える**」こと(少ない step でも denoise しきる)。
+それには adaln_modulation の Linear を学習対象に含める必要がある。Q,V LoRA だと
+attention pattern だけ変えられて timestep 入り口を触れない → 学習が degenerate
+な「色 / スタイル shift」に逃げる(Phase B v1 の緑色化はこれ)。
+
+公式 Anima Turbo LoRA も同じ理由で adaln_modulation を target(1016 keys、
+我々の wide LoRA 980 keys とほぼ同等)。
+
+### 3.2.3 なぜ R3GAN が必要(adv loss なしではダメ)
+
+DMD2 単体は **mean collapse** に陥る(Phase B v2 で実証)。理由:
+- DMD gradient `(real_score - fake_score)` は「real distribution の方向」を示す
+- LoRA capacity が大きいほど、「real の **mean** を出力するのが最も安全」という解に収束しやすい
+- TSCD(EMA self との consistency)は collapse を加速こそすれ防げない
+
+R3GAN(または同等の adv / reward 信号)があると:
+- D が「real は sharp、fake は blurry」と学習
+- gen は「blurry な mean を出すと D に見抜かれる」と分かり、sharp な多様性を保つ
+- これが mean collapse の唯一の防御策
+
+実際、DMD2 / Hyper-SD / DMDR 系の全公式実装で adv loss(or reward loss)が必須。
+
+### 3.2.4 なぜ R3GAN gamma の grid search が必須
+
+R3GAN の R1+R2 gradient penalty 重み `gamma` は **入力空間のスケールに比例**:
+
+| ドメイン | 入力 channel | 解像度 | gamma |
+|---|---|---|---|
+| CIFAR-10 | 3 | 32x32 | **0.05** |
+| FFHQ-64 | 3 | 64x64 | 2 |
+| FFHQ-256 | 3 | 256x256 | 150 |
+| **Anima latent** | **16** | **128x128** | **~0.05-1**(grid search 必須) |
+
+Anima latent は per-sample 262k 要素で、典型 pixel 空間より gradient norm が
+大きい → gamma は数値が小さくないと penalty が爆発。Phase A は gamma=50 で **500x 過大**
+にしたら discriminator が崩壊した(詳細は §4 Phase A 参照)。
+
+### 3.2.5 なぜ guard rail(`d_r1/r2 > 1000` で abort)
+
+Phase A は **loss は全部健全に推移**(`l_dmd` 0.005, `l_adv_g` 0.69 で一定)しつつ
+裏で discriminator gradient が爆発し、completed ckpt が完全ノイズ出力という
+**ステルス失敗**。後付けの分析で「step 1500 時点 d_r1 ~700 / step 2300 で 500k」と
+わかったため、しきい値 1000 で abort すれば即停止できる。
+
+しきい値 1000 の根拠: 正常上限 ~100 の 10x、致命傷 500k の 1/500 で、確実に
+ステルス damage が始まる前に止まる。
+
+### 3.2.6 なぜ Decoupled DMD2 + TSCD(LCM / Hyper-SD ではなく)
+
+| 手法 | 性質 | 我々の選択理由 |
+|---|---|---|
+| LCM | consistency 単体、シンプル | low-step 品質に限界 |
+| Hyper-SD | TSCD + adv、segment-wise progressive | gamma calibration 必要、参照実装 SDXL only |
+| **DMD2 / DMDR** | **分布マッチング + adv、Anima のような flow matching DiT に最適** | Z-Image(Cosmos 派生)で実証済、適応容易 |
+
+DMD2 系は **flow matching ベース(Anima の rectified flow と同じ)** で、math
+変換不要なのが決定打。TSCD は Hyper-SD から借りた多 step 一貫性で、純粋 DMD2
+の単 step 偏りを補正。
+
+### 3.2.7 なぜ 5000 枚 self-distillation データセット
+
+| データソース | コスト | 利点 | 欠点 |
+|---|---|---|---|
+| Anima 自己生成 | $44 (10 B200 並列 40min) | base の出力分布に最適化 / ライセンス問題なし | base の癖をそのまま継承 |
+| Danbooru / Pixiv | $0 + キュレーション時間 | 多様性 | ライセンス・タグ品質ばらつき |
+| 公開データセット(LAION 等) | $0 + DL コスト | 多様性最大 | アニメ系少ない |
+
+蒸留の本質は「**base と同じ出力を低 step で**」なので、self-distillation で
+分布を base に一致させるのが理論的に正しい(distillation の定義)。
+
+### 3.2.8 なぜ B200(H100 / A100 ではなく)
+
+| GPU | Modal $/h | LoRA-only step 時間 | 192GB VRAM の意味 |
+|---|---|---|---|
+| A100-80GB | $2.50 | 7-9 秒 | OOM 危険、batch=1 |
+| H100-80GB | $3.95 | 4-6 秒 | OOM 危険、batch=1-2 |
+| **B200** | **$6.25** | **2.8-3.0 秒** | **3 model 同時 + batch=4-8 余裕** |
+
+B200 の per-hour 単価は高いが、**3 model(real/fake/gen)同時+batch 増**で
+**実質的な per-experiment コストは最安**。R3GAN を入れても OOM 不安なし。
+
+---
+
+## 4. 実際の失敗(全 5 件、詳細)
+
+| # | 試行 | 構成 | 結果 | 累計 |
+|---|---|---|---|---|
+| 1 | **Phase A** | full DiT + Q,V LoRA, R3GAN γ=**50**, 3000 step | 完全ノイズ(D 爆発) | $70 |
+| 2 | **Phase B v1** | Q,V LoRA, R3GAN **OFF**, 2000 step | 緑色化(AdaLN 不触) | $80 |
+| 3 | **Phase B v2** | wide LoRA, R3GAN **OFF**, 2000 step | mean collapse(灰緑平面)| $90 |
+| 4 | **Phase B v3a** | wide LoRA, R3GAN γ=**0.1**, 2000 step | step 30 で d_r1=1648 → abort | $95 |
+| 5 | **Phase B v3b** | wide LoRA, R3GAN γ=**1.0**, 2000 step | 早期 mean collapse → 後半 D 爆発 | $115 |
+
+各失敗が **異なる failure mode** で起きた点が重要。「1 つ直しても他で詰まる」
+直交性が、SOTA 蒸留を Anima 系で安定させる難しさの本質。
+
+### 4.1 Phase A の詳細 postmortem(5 層因果分析)
+
+**ステルス失敗**: 訓練ログ上は全 loss 健全(`l_dmd` 0.003-0.008, `l_adv_g` 0.69 一定)
+だったのに、完走 ckpt が完全ノイズ出力。原因は 5 層の連鎖:
+
+**Layer 1 — gamma calibration ミス(根本)**
+- Anima latent(16ch × 128x128)に対し gamma=50 は 50-500 倍過大
+- Agent リサーチが「latent norm 用 grid search 必須」と明示警告していたが、
+  手抜きで推測値を採用したのが発端
+
+**Layer 2 — R3GAN gradient が gen に流入(見えない damage)**
+```
+l_adv_g = softplus(-(d_fake - d_real)) ≈ log 2 = 0.6914  ← 中立点
+```
+`l_adv_g` が一定で「adv 信号効いてない」と誤読。**しかし gradient norm は別物**:
+```
+∂L_adv_g/∂gen_output = (sigmoid 微分 ~0.5) · ∂(D の全層 Jacobian)
+```
+D 重みが gamma penalty 最小化で異常拡大 → D の Jacobian 爆発 → gen に流れる
+gradient が正常時の **数百〜数千倍**。adv_weight=0.05 で抑えても scale 負け。
+
+> 教訓: **adv_loss の絶対値だけ見て安全と判断するな**。gradient norm を別途モニタ。
+
+**Layer 3 — D 自体の破滅スパイラル**
+```
+最初:    l_disc = 0.69 + 25·0       = 0.69
+膨張後:  l_disc = 0.69 + 25·600,000 = 15,000,000
+```
+gamma が大きすぎて R1+R2 penalty 項が adv 項を支配 → D は adv 学習する余地なく、
+ひたすら penalty 最小化に動く → ある瞬間 scale 崩れて d_r1/r2 が指数的に膨張。
+
+**Layer 4 — Sanity test の見落とし**
+- 50 step では destabilization が緩やかすぎて d_r1/r2 が正常範囲
+- 本番 step 500-1000 から急速悪化
+- 後から見れば step 1000 で d_r1/r2 が 100 超え始めていた($50 救えた)
+
+> 教訓: **Sanity は 200-500 step に延長**。緩慢な発散も検出可能に。
+
+**Layer 5 — 監視/アラート不足**
+- d_r1/r2 はログ出していたが **しきい値アラート未実装**
+- NaN/Inf check はあったが「大きすぎる値」検出なし
+- gen �� weight L2 norm を時系列で見れば damage 早期検出可能だった
+- 中間 ckpt 毎の自動生成テストも未実装
+
+修正後の guard rail(実装済):
+```python
+# scripts/distill/dmd2_trainer.py
+if d_r1 > 1000.0 or d_r2 > 1000.0:
+    raise RuntimeError(f"R3GAN penalty exploded (r1={d_r1:.1f}, r2={d_r2:.1f}). "
+                       f"Reduce r3gan_gamma. Aborting.")
+```
+
+**Phase A から学んだ悪設計のチェックリスト**
+
+| 問題 | 影響 |
+|---|---|
+| gamma を grid search せず推測 | 50-500 倍ズレ |
+| adv_weight=0.05 で「安全装置」になると誤解 | gradient norm は別軸、効かなかった |
+| sanity 50 step だけで本番投入 | 緩慢な発散見逃し |
+| `d_r1/r2` にアラートしない | 早期 abort 機会逸失 |
+| 中間 ckpt で部分検証しなかった | 完走してから初めて発覚($70 ロス) |
+
+### 4.2 Phase B v1 — Q,V LoRA だと AdaLN を触れない
+
+R3GAN を外して安全側に振った構成。loss は 0.02-0.03 で安定したが、生成すると
+**緑色化した anime キャラ**。LoRA strength を 1.0 → 0.3 に下げると base に戻る
+(= LoRA は「均一に緑方向に押している」)。
+
+原因: Q,V LoRA は cross-attention / self-attention の query / value だけを変更。
+**timestep modulation の入り口である AdaLN modulation Linear を触れない**。
+DiT が「少ない step でちゃんと denoise する」ために必要な timestep 解釈の変更が
+できず、勾配が attention の color / style shift に逃げて degenerate に収束。
+
+→ 解決: LoRA target を **「AdaLN + attention + MLP の全 Linear」(wide LoRA)** に拡張。
+公式 Anima Turbo LoRA も同じ範囲を target にしている(検証で 1016 keys 確認)。
+
+### 4.3 Phase B v2 — wide LoRA + R3GAN OFF で mean collapse
+
+wide LoRA(980 keys)に拡張したら、今度は **完全に灰緑色の平面**を出力。loss は
+v1 より遥かに小さく(`l_dmd` 0.001-0.007)、**訓練ログ上は最も健全に見える**
+失敗。LoRA strength を変えても改善せず。
+
+原因: DMD2 単体は「real distribution の方向」を示すだけ。LoRA capacity が大きい
+ほど **「data distribution の mean を出力すれば DMD gradient はゼロになる」**
+という degenerate な解に収束しやすい。R3GAN なしでは「sharp なものを出せ」と
+gen に教える信号がないので、訓練が進むほど mean に潰れていく。
+
+→ 解決: **R3GAN(または同等の reward / classifier loss)を必ず投入**。DMD2 系の
+全公式実装(NVIDIA DMD2, Z-Image DMDR, Hyper-SD)で adv 系 loss が必須なのは
+このため。
+
+### 4.4 Phase B v3a — gamma=0.1 で D が即爆発
+
+grid search の 100-200 step では gamma=0.1 が安定に見えていた:
+```
+[step 60/100] d_r1=0.020  d_r2=0.015  l_disc=0.017  ← 健全
+```
+
+しかし本番(2000 step)で開始 30 step 以内に爆発:
+```
+[step 10/2000] d_r1=67.5   d_r2=74.8   l_disc=7.1   ← 既に異常
+[step 30/2000] d_r1=1648.0 d_r2=1821.6 ← guard 発火 abort
+```
+
+原因推測:
+- 短時間 grid search は **ランダム初期化の運**で偶然安定 zone に居ただけ
+- 本番の異なる初期化 + 長時間 → 不安定境界を越えた
+- gamma=0.1 は Anima latent では **安定境界の崖際**
+
+### 4.5 Phase B v3b — gamma=1.0 で長期 mean collapse + 最終 D 爆発
+
+gamma=1.0 で training は確かに **1500 step まで全 metric 健全**:
+```
+[step 1000/2000] l_dmd=0.001 l_consist=0.005 d_r1=0.005 l_disc=0.004  ← 全部小さい
+```
+
+**しかし step 500/1000/1500 の中間 ckpt で生成テストすると全部 mean collapse**:
+- step 500: 白い平面
+- step 1000: 白い平面
+- step 1500: 青+白の blob
+
+しかも step 1380 から D が「目覚め」始め、step ~1600 で d_r1=2656 で abort:
+```
+[step 1380/2000] d_adv=3.73  d_r1=0.36   d_r2=0.22  ← D が動き始める
+[step 1510/2000] d_adv=0.00  d_r1=0.67   d_r2=0.45  ← じわじわ拡大
+[step ~1600]    d_r1=2656.0  d_r2=3043.2 ← abort
+```
+
+**核心の理解**:
+- gamma=1.0 は D が **静か → 突然動く → 爆発** のパターン
+- 静かな間(d_r1 ~ 0.005)は **D が adv 信号を gen に渡せず**、gen は mean に collapse
+- D が動き始めた頃には gen は既に degenerate state
+- ckpt が「健全に見える」のは loss だけ。実態は破綻
+
+これが **「R3GAN の安定 zone が存在しない」** という最終診断の根拠。
+
+### 4.6 失敗パターンの直交性
+
+| | Phase A | B v1 | B v2 | B v3a | B v3b |
+|---|---|---|---|---|---|
+| failure mode | gradient damage(ノイズ)| degenerate(緑) | mean collapse | D 爆発 abort | D 静止 → mean collapse → D 爆発 |
+| 根本原因 | R3GAN γ 過大 | LoRA target 狭 | adv なし | γ 小すぎて D 暴走 | γ 大すぎて D が adv 信号出さず |
+| 検出 | step 1000+ d_r1 | 生成テスト即 | 生成テスト即 | step 30 で guard | 中間 ckpt 生成テスト |
+| 対策 | γ grid search + guard | wide LoRA | R3GAN を入れる | γ 上げる | γ 下げる(でも v3a に戻る)|
+
+**5 つの failure mode は互いに直交かつ循環**(γ 下げると v3a、上げると v3b)。
+Anima latent で R3GAN を安定させる γ は、**存在するとしても極めて狭く**、
+我々の grid 解像度では発見できなかった。
+
+### 4.7 周辺で踏んだ Modal / diffusion-pipe / ComfyUI 罠
+
+蒸留の本体とは別に、インフラ側で踏んだバグ:
+
+**diffusion-pipe を programmatic に使う際の 8 件**
+
+| 罠 | 症状 | 対処 |
+|---|---|---|
+| `models/`, `utils/` が namespace package(`__init__.py` なし)| site-packages 側の `utils` に shadow されて `from utils.common import` 失敗 | `importlib.util.spec_from_file_location` で強制ロード(`scripts/distill/anima_loader.py`)|
+| `CosmosPredict2Pipeline(cfg)` の dtype が **torch.dtype 必須** | `tensor() argument 'dtype' must be torch.dtype, not str` | `cfg["model"]["dtype"] = torch.bfloat16`(文字列禁止) |
+| `load_text_encoder()` / `load_vae()` メソッドが **存在しない** | AttributeError | `load_diffusion_model()` だけ呼ぶ、残りは `__init__` で完了 |
+| `MiniTrainDIT.forward` は `padding_mask` 必須(None 不可) | `transforms.functional.resize(None, ...)` TypeError | `torch.zeros(B, 1, H, W)` を渡す |
+| VAE input が bfloat16 必須 | `Input type (float) and bias type (BFloat16) should be the same` | `vae_encode` 内で auto-cast |
+| DiT 入力 `(noisy, t, cond)` の dtype 統一必須 | `expected mat1 and mat2 same dtype` | dit_forward helper で全部 weight dtype に揃える |
+| `torchvision` import するが requirements に無い | `ModuleNotFoundError` | image に明示インストール |
+| save の filter で integer buffer を落とすと壊れる | ロード後ノイズ | float 限定 filter は廃止、全 tensor 保存 |
+
+**ComfyUI LoRA 形式**
+
+Anima/Cosmos 用 LoRA は `lora_A.weight` / `lora_B.weight`(NOT `lora_down/up`)、
+`diffusion_model.` prefix 必須。PEFT 出力からの変換は `modal_app.py::_convert_peft_to_comfy_lora`。
+SDXL 流 `lora_down/up` に変換すると ComfyUI が key 認識できずランダム scaling で
+適用 → 出力が緑化(これも Phase B v1 のもう 1 つの寄与因子だった可能性)。
+
+**Modal 運用**
+
+| 罠 | 対処 |
+|---|---|
+| `modal run --detach` の local CLI が早く戻る | 完了確認は `modal app list` / `modal app logs <id>` |
+| `print()` の block buffering で logs 滞留 | `PYTHONUNBUFFERED=1` を env に追加 |
+| Volume commit は subprocess 終了時のみ | 訓練ループ内で `modal.Volume.from_name(...).commit()` を定期発行 |
+| `modal volume cp` は cross-volume 不可 | 両 volume mount した Modal function で `shutil.copy` |
+| `run_in_background=true` + TaskStop で app が "detached_disconnected" に残る | 必ず `--detach` を一緒に付ける(明示 detach)|
+| Anima ckpt は `net.` prefix 前提で ComfyUI が arch 検出 | 保存時に `{"net." + k: v for k, v in sd.items()}` |
+
+### 4.8 コスト・時間の実測値
+
+| タスク | 時間 | コスト |
+|---|---|---|
+| dataset 生成(B200 × 10 並列) | 40 分 | $44 |
+| Path A baseline 確認(Turbo LoRA DL + 10 枚生成) | 5 分 | ~$1 |
+| **Phase A: full DiT 訓練(失敗)** | **11.6 h** | **~$70**(全損)|
+| Phase B v1: Q,V LoRA(失敗) | 1.6 h | ~$10 |
+| Phase B v2: wide LoRA, R3GAN OFF(失敗) | 1.6 h | ~$10 |
+| Phase B v3a: wide LoRA, γ=0.1(早期 abort) | ~5 分 | ~$5 |
+| Phase B v3b: wide LoRA, γ=1.0(失敗) | 4 h | ~$20 |
+| R3GAN gamma grid search | 30 分 | ~$3 |
+| 比較生成 30 枚 + その他 | | ~$2 |
+| Modal volume(50 GB × 月)| 常時 | ~$7.5/月 |
+| **累計** | | **~$165 ≈ 25,000円** |
+
+R3GAN ON で 1 step 約 +20-30%(D forward + R1/R2 backward 追加分)。
+
+---
+
+## 5. 続き
+
+§4 の 5 回失敗を受け、「自前で書き下す」を諦めて **既存の動く実装を最小改造で
+移植**した試行が 2026-05 から始まる。各手法 (Z-Image traj / LADD / Reflow / DMD2 /
+SiD2 / Shortcut / PCM / DRaFT+) の移植ログ・コスト・結果は
+[migration_log.md](migration_log.md) を参照。

docs/dmd2.md

@@ -0,0 +1,115 @@
+# DMD2 (NVIDIA cosmos-predict2.5 流派) 蒸留
+
+[← README に戻る](../README.md) — 関連: [PCM](pcm.md) / [DMDX](dmdx.md) / [蒸留全般](distillation.md)
+
+## 概要
+
+NVIDIA cosmos-predict2.5 公式の **DMD2 + TrigFlow** 流派を Anima 移植。同 base に **2 つの独立 PEFT adapter** (student / fake_score) を attach し `set_adapter()` で runtime 切替、TTUR (alt 5 critic : 1 generator)。
+
+**実装**: `scripts/distill/{dmd2_official_loss,train_dmd2_official}.py`
+**Modal**: `modal_app.py::train_dmd2_official_distill`
+
+## 設計の鍵
+
+- **2 adapter / single base** → memory 効率 (teacher は別 deepcopy)
+- **TrigFlow time sampling**: `s ∈ shifted_uniform(shift=3-5) → t = arctan(s/(1-s))`
+- **DMD2 gradient trick** (逆 KL): `grad = (x0_fake - x0_teacher) / |x_hat - x0_teacher|.mean()`
+  → `loss = ((x_hat - (x_hat - grad).detach())**2).mean()`
+- **Critic phase**: weighted denoise with `1/sin(t)**2` factor
+- **few-step rollout** (n=1..4)、grad は最終 step のみ (memory efficient)
+
+## 訓練設定 (実走、2026-05-18 - 19)
+
+| 項目 | 値 |
+|---|---|
+| Base | `anima-base-v1.0.safetensors` |
+| Method | DMD2 + TrigFlow (cosmos-predict2.5) |
+| LoRA target | wide (980 keys) × 2 (student + fake_score) |
+| LoRA rank | 32 |
+| Total outer | 5000 (1 outer = 5 critic + 1 generator) |
+| n_student_steps | 4 |
+| Teacher CFG | **4.5** (Anima 公式推奨に整合、過去 preview3 試行の 3.0 から変更) |
+| Student CFG | 1.0 (CFG embedded) |
+| Shift | **3.0** (Anima sigma_shift と一致、過去の 5.0 から変更) |
+| LR (gen / critic) | 5e-6 / 1e-5 (TTUR) |
+| Warm-start | **`/models/loras/anima_turbo.safetensors`** (Civitai 公式 Turbo、968/1016 keys 一致) |
+| GPU | B200 |
+| 訓練時間 | step 660 まで ~40 min、resume + 4500 outer ~7h |
+| コスト | 累計 ~$48-55 想定 (rate 6.2s/outer parallel 時) |
+
+## 試行履歴 (2026-05-18 - 19)
+
+### 1st run (`/output/dmd2_v1/`)
+- 設定: warm-start (Anima Turbo) + 5000 outer
+- 経過: step 660/5000 で `KeyboardInterrupt`
+- 原因: **Modal workspace billing cycle spend limit reached** (DMDX 並列起動で予算限度超過)
+- 保存: step 500 ckpt のみ無事 (volume commit 済)
+
+### Resume (`/output/dmd2_v1_resume/`)
+- 設定: step 500 ckpt を warm-lora、残り 4500 outer
+- warm-load: **980/980 keys 完全一致** (自前 LoRA 同 architecture)
+- 進行中、step 500 (= 元 1000 相当) / step 1000 (= 元 1500) で ckpt 保存
+
+## Loss 推移 (順調)
+
+| outer | l_dmd_critic | l_dmd_gen | 評価 |
+|---|---|---|---|
+| 0 | 0.0364 | 1.5780 | warm-start 直後 |
+| 100 | 0.07 | 1.18 | gen 急速降下中 |
+| 500 | 0.06 | 1.30 | spike あるが healthy |
+| 660 (中断) | 0.0471 | 0.16 | 健全な状態で中断 |
+| resume 0 | 0.0364 | 1.5771 | warm-load から再開、原 outer 0 と類似 |
+| resume 500 | 0.058 | 1.29 | 期待通り |
+| resume 1000 | 0.054 | 0.89 | 順調収束 |
+
+magnitude (x_hat / teacher / fake): ~0.3-0.5 範囲で揃って distill 方向収束。
+
+## 検証結果 (step 500、4-step / er_sde+simple)
+
+サンプル画像: `/dataset/ckpt_health_check/dmd2_v1_step00500_step4_cfg1.0/p0000_s000.png`
+
+- 2 キャラ (Flandre / Remilia) 正確分離
+- Civitai Turbo に近い品質 (warm-start 支配的、step 500 で大幅変化なし)
+- 生成時間 10.95s (4-step、CFG 1.0、B200 sageattention)
+
+ローカル DL 済: `~/Downloads/anima_loras/dmd2/dmd2_v1_step00500.safetensors`
+
+## 配布
+
+完走後 HF `darask0/anima-distill-loras/dmd2/` に upload 予定。
+
+## ComfyUI 推奨設定
+
+```
+LoraLoaderModelOnly:
+  lora_name: dmd2_v1_step00500.safetensors
+  strength_model: 1.0
+
+ModelSamplingAuraFlow:
+  shift: 3.0
+
+KSampler:
+  steps: 4
+  cfg: 1.0
+  sampler_name: er_sde
+  scheduler: simple
+```
+
+## 過去の preview3 試行との対比
+
+[migration_log.md §5.7](migration_log.md#57-実画像比較-2026-05同-prompt--seed) の preview3 DMD2 では:
+- 8-step: 絵画調シフト、洗色傾向
+- 4-step: シルエットのみ、ボケすぎ
+
+→ 改善要因 (v1.0 試行):
+- ✅ **warm-start (Civitai Turbo)** 必須化
+- ✅ **teacher_cfg 3.0 → 4.5** (Anima 公式推奨)
+- ✅ **shift 5.0 → 3.0** (Anima sigma_shift 整合、sampler 互換性向上)
+- ✅ v1.0 base 直接 (preview3 ベース LoRA の style drift 回避)
+
+## 知見
+
+- **2 adapter peft set_adapter() で必須**: 毎 forward で明示切替しないと片方の重みが leak ([operations.md](operations.md) 参照)
+- **`No module named 'modal'` の volume commit 失敗**: subprocess 内で modal lib 未参照、harmless warning (function 終了時に Modal 自動 commit)
+- **billing limit による KeyboardInterrupt**: Modal の workspace spend limit に達すると **cloud function に SIGINT** が送られて停止する。ローカル CLI の interrupt とは別経路、Modal dashboard で billing を確認・上げる必要
+- **parallel rate 低下**: 並列 (DMDX と) で 3.66s/outer → 6.2s/outer (~1.7x 遅い)、別 B200 アロケート時の��雑度に依存

docs/dmdx.md

@@ -0,0 +1,141 @@
+# DMDX (Adversarial Distribution Matching) 蒸留
+
+[← README に戻る](../README.md) — 関連: [PCM](pcm.md) / [DMD2](dmd2.md) / [蒸留全般](distillation.md)
+
+## 論文
+
+**"Adversarial Distribution Matching for Diffusion Distillation Towards Efficient Image and Video Synthesis"**
+- arxiv: [2507.18569v1](https://arxiv.org/html/2507.18569v1)
+- 著者: Yanzuo Lu et al. (Sun Yat-Sen University + ByteDance Seed Vision)
+- 公開コード: なし (論文のみ)
+
+## 核心アイデア
+
+DMD2 の `(real_score - fake_score)` の **逆 KL 発散** はモード崩壊しがち (Phase B v2 と同症状)。DMDX は **学習可能 discriminator による hinge GAN (TVD = Total Variation Distance、対称) 最小化** に置換し、根本解決を狙う。
+
+### vs DMD2
+
+| 観点 | DMD2 | DMDX (ADM) |
+|---|---|---|
+| 発散尺度 | 逆 KL (非対称、mean collapse 寄り) | **TVD (対称)** |
+| 損失実装 | grad trick で score difference 直接最小化 | **学習可能 discriminator** で hinge loss 敵対学習 |
+| GAN 役割 | optional 正則化 | **本体損失** |
+| 時刻情報 | t での score 比較のみ | `t → t-Δt` の ODE 進化中間点も考慮 |
+| Initialization | MSE pretrain | 論文では ADP (= adversarial pre-training)、本実装は warm-start で代替 |
+
+## 本実装の範囲
+
+論文の **ADM のみ移植**、**ADP は省略** (SAM 依存 pixel-space discriminator が重い、コスト効果が低い)。
+
+代替: **Civitai Anima Turbo を warm-start** で「分布レベルの良い初期化」を確保。
+
+**実装**: `scripts/distill/{dmdx_loss,train_dmdx}.py`
+**Modal**: `modal_app.py::train_dmdx_distill`
+
+### ADM の構造
+
+- **Discriminator backbone** = teacher MiniTrainDIT (frozen) — **LADD の `AnimaLADDDiscriminator` を流用**
+- **Trainable heads** × 5 (blocks 2/8/14/20/26 を hook、各 spectral norm + BatchNormLocal)
+- **Cubic time schedule**: `u ~ U(0,1), t = 1 - u**3` (high-noise バイアス)
+- **Δt evolution**: teacher で 1-step Euler で `t → t-Δt` 進化、Δt = T/64
+- **Hinge loss**: G = `mean(-D(fake))`、D = `relu(1-D(real)) + relu(1+D(fake))`
+
+## 訓練設定 (実走、2026-05-19)
+
+| 項目 | 値 |
+|---|---|
+| Base | `anima-base-v1.0.safetensors` |
+| Method | DMDX ADM-only |
+| Student LoRA | wide (980 keys、single adapter) |
+| LoRA rank | 32 |
+| Disc heads | 5 (block 2/8/14/20/26)、各 ~3.4M = 17.1M total |
+| Total outer | 5000 (1 outer = 2 disc + 1 gen) |
+| n_critic_per_gen | **2** (DMD2 の 5 より少なめ、hinge GAN 系は 1:1〜2:1 が一般的) |
+| n_student_steps | 4 |
+| Teacher CFG | 4.5 |
+| Student CFG | 1.0 |
+| dt_ratio | 1/64 ≈ 0.0156 |
+| recon_weight | **0.0** (pure ADM、LADD の Smooth-L1 anchor なし) |
+| LR (gen / disc) | 5e-6 / 1e-5 |
+| Warm-start | **`/models/loras/anima_turbo.safetensors`** (Civitai 公式 Turbo) |
+| GPU | B200 |
+| 訓練時間 | smoke ~35s (10 outer)、本番 ~5h 想定 (3.6s/outer) |
+| コスト | smoke ~$0.5、本番 ~$30 想定 |
+
+## Smoke 結果 (10 outer、2026-05-19)
+
+| outer | l_d_real | l_d_fake | l_d_total | l_g_adv | 評価 |
+|---|---|---|---|---|---|
+| 0 | 0.717 | 1.295 | 2.012 | -0.292 | 初期 hinge equilibrium |
+| 5 | 0.724 | 1.277 | 2.001 | -0.276 | 安定継続 |
+| 9 | 0.740 | 1.268 | 2.008 | -0.270 | smoke 完了 |
+
+OOM / NaN / divergence なし。setup・warm-start (968/1016)・D init・ADM loss 全部正常。
+
+## 本番訓練 (進行中、2026-05-19)
+
+### Step 500 検証
+
+| outer | l_d_real | l_d_fake | l_d_total | l_g_adv |
+|---|---|---|---|---|
+| 0 | 0.717 | 1.295 | 2.012 | -0.292 |
+| 90 | 0.774 | 1.227 | 2.001 | -0.218 |
+| 340 | 0.938 | 1.034 | 1.972 | +0.048 (D 優勢) |
+| **500** | **0.822** | **0.976** | **1.798** | **-0.146** |
+| 1000 | 0.933 | 0.971 | 1.904 | -0.093 |
+
+GAN dynamics:
+- `l_d_total` 2.012 → 1.798 → 1.904 (D 学習、後半は equilibrium 周辺で振動)
+- `l_g_adv` -0.292 → +0.05 (一時 D 優勢) → -0.15 → -0.09 (G 持ち直し)
+- D と G の拮抗状態に到達、healthy GAN training
+
+### 生成検証 (step 500、4-step / er_sde+simple)
+
+サンプル: `/dataset/ckpt_health_check/dmdx_v1_step00500_step4_cfg1.0/p0000_s000.png`
+
+- 2 キャラ正確分離、Civitai Turbo に近い品質 (warm-start 支配)
+- 生成時間 **5.77s** (PCM 8.8s / DMD2 11s より速い、B200 contention 軽い時間帯)
+
+## 配布
+
+完走後 HF `darask0/anima-distill-loras/dmdx/` に upload 予定。
+
+## ComfyUI 推奨設定
+
+```
+LoraLoaderModelOnly:
+  lora_name: dmdx_student_final_comfy.safetensors
+  strength_model: 1.0
+
+ModelSamplingAuraFlow:
+  shift: 3.0
+
+KSampler:
+  steps: 4
+  cfg: 1.0
+  sampler_name: er_sde
+  scheduler: simple
+```
+
+## 期待品質 (paper 報告ベース)
+
+SDXL 1-step での DMDX vs DMD2:
+
+| metric | DMD2 | DMDX | 差 |
+|---|---|---|---|
+| CLIP Score | 35.22 | 35.26 | +0.04 |
+| PickScore | 22.10 | 22.27 | +0.17 |
+| HPSv2 | 27.45 | 27.70 | +0.25 |
+| MPS | 10.69 | 11.20 | +0.51 |
+| **多様性 LPIPS** | 0.6715 | **0.7156** | **+0.044** |
+
+→ **多様性 +6.5% が最大の利点**、画像品質は誤差レベル。Video (CogVideoX) では 8-step で 100-step 比 92-96% 加速の大きな成果。本実装は image only。
+
+## 知見
+
+- **LADD discriminator の再利用**: spectral norm + BatchNormLocal + multi-block hook の構造はそのまま流用可能、DMDX 専用 D を新規開発する必要なし
+- **cubic time schedule**: shifted uniform より extreme な high-noise バイアス、`t_mean=0.0259-0.99` の wide range (cubic の variance 大)
+- **n_critic_per_gen 削減**: DMD2 の 5 から **2** に。hinge GAN は D vs G の balance が critical で、disc 過多だと G が学習機会失う
+- **`recon_weight=0` で開始**: pure ADM の純粋性 verify。不安定なら LADD 流 Smooth-L1 anchor (`>0`) で stability boost 可能
+- **rate 3.6s/outer**: DMD2 の 3.7s と同等以下 (n_critic 削減効果)、parallel でも slowdown 小さい
+- **billing limit 再発防止**: 並列 launch 前に Modal spend cycle limit を確認 (DMD2 v1 が KeyboardInterrupt'd した先例あり)

docs/migration_log.md

@@ -0,0 +1,583 @@
+# 既存実装の移植 — 2026-05 試行ログ
+
+[← README に戻る](../README.md) — 前章: [distillation.md](distillation.md)
+
+> ⚠️ 本ドキュメントは **2026-05-17 時点のスナップショット**を含む。Modal volume の
+> 状態 (LoRA path、cache 内容など) は時間とともに変動するので、最新状態は
+> `modal volume ls` で確認すること。配布済 LoRA は HF の
+> [darask0/anima-distill-loras](https://huggingface.co/darask0/anima-distill-loras) に集約。
+
+[§4 失敗分析](distillation.md#4-実際の失敗全-5-件詳細) の 5 回失敗を受け、「自前で書き下す」のを諦めて **既存の動く実装を最小改造で
+Anima に移植する** 方針に転換。survey と smoke test を経て **5 手法** を実装、
+本リポジトリに収録 (`scripts/distill/`)。
+
+## 5.0 全体マップ
+
+```
+                         ┌── Z-Image trajectory imitation (DiffSynth-Studio 由来)
+                         ├── DMD2 + TrigFlow      (NVIDIA cosmos-predict2.5 公式)
+        蒸留 5 候補 ────┼── LADD                 (AMD Nitro-1 移植、PixArt 由来)
+                         ├── Reflow / InstaFlow   (rfpp NeurIPS 2024 知見適用)
+                         └── PCM                  (G-U-N/Phased-Consistency-Model、SD3 変種)
+        
+        無料補完  ────── ComfyUI-CapitanFlowMatch / ZiT-Scheduler (推論側 scheduler)
+        
+        survey で却下 ── DMDR / Hyper-SD / NitroFusion / ADD (訓練コード非公開)
+```
+
+## 5.1 採用 5 手法の概要
+
+| # | 手法 | カテゴリ | 中核ファイル | 設計の鍵 |
+|---|---|---|---|---|
+| ① | Z-Image trajectory imitation | trajectory matching | `scripts/distill/{traj_scheduler,traj_loss,train_traj}.py` | critic なし、teacher 50-step CFG=2 trajectory を student 8-step に align |
+| ② | DMD2 + TrigFlow (cosmos-predict2.5) | score matching | `scripts/distill/{dmd2_official_loss,train_dmd2_official}.py` | **同じ base に 2 LoRA adapter** (student + fake_score)、`set_adapter()` 切替、alt 5:1 |
+| ③ | LADD (AMD Nitro-1) | adversarial + recon | `scripts/distill/{anima_ladd_disc,train_ladd}.py` + `precompute_teacher_x0.py` | D backbone = teacher MiniTrainDIT (frozen)、5 spectral-norm head、Smooth-L1 recon anchor |
+| ④ | Reflow / InstaFlow | flow alignment | `scripts/distill/train_reflow.py` (cache は `--save-noise` 必須) | Anima は元から RF、U-shape t + Huber + 間欠 LPIPS、1 grad-through で最安 |
+| ⑤ | PCM (Phased Consistency Model) | consistency | `scripts/distill/{pcm_scheduler,train_pcm}.py` | SD3-PCM 流派 (FlowMatch + v-pred)、Anima RF と math 一致、N=50 を K=4 phase に分割 |
+
+全手法に共通の前提:
+- **wide LoRA** (`attach_wide_lora()` で AdaLN + attn + MLP 全 Linear、454 modules、~72M trainable)
+- **warm-start = Civitai 公式 Anima Turbo LoRA** (matched 968/1016 keys、x_embedder/t_embedder の 12 keys のみ skip)
+- **B200 GPU** + dataset は `/dataset/raw` (5000 caption、品質タグは未掃除でも蒸留には影響なし)
+
+## 5.2 R3GAN との決定的な違い (なぜ LADD は動く可能性が高いか)
+
+| 失敗パターン | R3GAN | LADD |
+|---|---|---|
+| D が無信号 → mean collapse | CNN を一から学習、16ch latent prior なし | D backbone = **teacher MiniTrainDIT (frozen)**、Anima latent を既に理解 |
+| D 爆発 | γ (gradient penalty) calibration 困難 | **spectral norm + BatchNormLocal** で D の Lipschitz 制約 |
+| anchor 不在で発散 | DMD2 単独だと mean に潰れる | **Smooth-L1 recon loss** (vs teacher x0) で「絶対に teacher から離れない」床を作る |
+
+LADD smoke (5 step、`l_g_adv ≈ ln(2) ≈ 0.69`, `l_d_total ≈ 2·ln(2) ≈ 1.39`) — R3GAN 試行で
+1 度も達成できなかった「**stable initialization**」が初手で出る。
+
+## 5.3 LoRA-only + warm-start + precompute cache の三位一体
+
+3 つを組み合わせて実コストを大幅圧縮:
+
+1. **LoRA-only (~72M trainable / 2B base)** — gen と D backbone は frozen、optimizer state も小さい
+2. **warm-start (Civitai Turbo LoRA)** — cold-start より convergence が早く、品質下限が保証される
+3. **precompute teacher x0 cache** — teacher の 20-step CFG=4.5 rollout を 1 度だけ実行 (~$11 / 5000 サンプル)、
+   その後の蒸留で **複数手法が同じ cache を共有** (LADD / PCM が利用)
+
+```
+                  precompute_teacher_x0_cache  (1回, $11)
+                          │
+                          ├── LADD train     (cache の x0 を real として使用)
+                          └── PCM train      (cache の emb を流用、x0 は teacher rollout source)
+        
+        precompute_for_reflow  (--save-noise 付き、別 cache, $11)
+                          │
+                          └── Reflow train   (noise + x0 + emb triplet)
+```
+
+## 5.4 採用しなかった手法 (survey 結果)
+
+| 手法 | 却下理��� |
+|---|---|
+| **DMDR** (vvvvvjdy/dmdr) | ImageNet/SiT のクラス条件 toy のみ公開、T2I (Z-Image, SD3.5) パイプラインは未公開 |
+| **Hyper-SD / Hyper-FLUX** (ByteDance) | 訓練コード非公開 (HF model card のみ)、コミュニティ port なし |
+| **NitroFusion / NitroSD-Vibrant** (ChenDarYen) | 訓練コード "Coming Soon" のまま停滞、リポジトリは 6 commit で凍結 |
+| **ADD / LADD-DINOv2** (StabilityAI) | 訓練コード非公開 (issues wontfix)、AMD 版で代替 |
+
+## 5.5 inference 側補完: AYS / TeaCache / GGUF / sageattention の適用可否
+
+| 手法 | Anima 適用可否 | 理由 |
+|---|---|---|
+| **AYS** (NVIDIA Align Your Steps) | ❌ | ComfyUI 公式実装は **SD1/SDXL/SVD の sigma schedule を hardcode** で rectified flow 未対応 |
+| **TeaCache** (timestep 残差 cache) | ⚠️ **base 高 step 用のみ** | 4-8 step 蒸留 LoRA とは**構造的に non-compatible** (連続 step の差分が大きすぎてスキップ不可)。base 28-50 step での **2x 速** には使える、porting ~250 LoC |
+| **ComfyUI-GGUF** Q8/Q4 量子化 | ⚠️ Anima 未対応 | FLUX/SD3/Hunyuan 系のみ。Anima loader 追加 + GGUF 変換が必要 |
+| **CapitanFlowMatch / CapitanZiT-Scheduler** | ✅ | rectified flow / Cosmos-Predict2 兄弟向け scheduler、ComfyUI ノード追加だけ |
+| **sageattention** | ✅ | torch SDPA 置換で 1.2-1.5x 推論速、`pip install sageattention` で `_base_image` に同梱済 |
+| **DPM-Solver++ 高次** / **UniPC** | ✅ | sampler ノード切替のみ、Anima 標準 `er_sde` の代替 |
+
+**実用推奨**: 蒸留 LoRA + sageattention + CapitanFlowMatch scheduler のスタックが最も効率的。
+TeaCache は「LoRA 使わず base 高品質モード」用の補完として将来 porting する価値あり (~250 LoC, ~$10 calibration)。
+
+## 5.6 これから挑む人へ
+
+- **survey をサボらない**: 4 つの「有名手法」が実は訓練コード非公開で詰む。GitHub の README を見るだけでなく `train.py` / `run.sh` を実際に開いて確認
+- **`scripts/distill/anima_loader.py` は再利用** — `build_anima()` + `AnimaBundle` で diffusion-pipe の罠 (namespace shadow、`padding_mask` 必須、dtype object vs str など) を全部吸収済
+- **smoke test を必ず 1 step 走らせる** — sign / dtype / 配線ミスを **$1-2 で発見**、本番 $30-50 をドブに捨てるリスク回避
+- **`MSYS_NO_PATHCONV=1 PYTHONIOENCODING=utf-8 PYTHONUTF8=1` を modal CLI に毎回つける** (Windows JP locale 必須)
+
+## 5.7 実画像比較 (2026-05、同 prompt × seed)
+
+5 prompt × 2 seed × {8 step, 4 step} × {base CFG=4.5, Civitai 公式 Turbo CFG=1, ① Z-Image
+trajectory imitation CFG=1, ② DMD2 CFG=1} = **80 枚** を Modal `compare_distill_loras` で
+生成 + v1.0 base 上で **13 条件 1 prompt 1 seed** を `verify_completed_loras` で生成。
+同 noise seed なので 1:1 で比較可能 (実装: [modal_app.py](../modal_app.py))。
+
+**観察された明確な順位** (8-step CFG=1、preview3 base、5 prompt × 2 seed = 80 枚):
+
+| Rank | 手法 | 評価 |
+|---|---|---|
+| 🥇 | **Civitai 公式 Anima Turbo** | シャープ、彩度高、盾の紋章/松明の光まで完備、キャラ識別正確 |
+| 🥈 | ① Z-Image traj imitation (本 repo) | 鮮明・安定、ディテール OK |
+| 🥉 | ② DMD2 (本 repo) | 絵画調シフト、やや洗色傾向 |
+| ❌ | base 30step→8step | 暗め・シンプル、品質劣化大 |
+
+**4-step での頑健性順位**:
+
+| Rank | 手法 | 評価 |
+|---|---|---|
+| 🥇 | **Civitai 公式 Anima Turbo** | 8-step とほぼ同等品質を維持 |
+| 🥈 | ① Z-Image | エッジ甘くなる (8-step 訓練だったため) |
+| ❌ | ② DMD2 | シルエットのみ、ボケすぎ |
+| ❌ | base | ほぼ崩壊、使用不可 |
+
+**v1.0 base での生成時間 (sageattention 有効、1024×1024、er_sde/simple/shift=3)**:
+
+| step / 条件 | 時間 | 備考 |
+|---|---|---|
+| 4-step LoRA (全 6 LoRA) | **3.5-4.0s** | 公式 Turbo / ① Z-Image / ② DMD2 / ④ Reflow / ⑩ DRaFT+ 系すべて |
+| 8-step LoRA (全 6 LoRA) | **4.0-4.7s** | 同上 |
+| 30-step base CFG=4.5 | **10.7s** | LoRA 無効、品質基準 |
+
+→ **4-step LoRA = base 30 step の 3 倍速**、sageattention 効果込み。Modal B200 のコスト換算で
+1 画像 ~$0.006 (4-step LoRA) vs ~$0.018 (base 30 step)。
+
+**正直な結論**: **CircleStone Labs 公式 Turbo (Civitai) に予算 ~$300 で勝つのは現状困難**。
+公式は遥かに多い GPU 時間と専門知見を投入した結果。本 repo の自前蒸留は:
+
+- base からの **品質向上は明白** (8-step CFG=1 が実用レベル)
+- 公式 Turbo に **代替/補完** できる位置 (ライセンス遵守、自前カスタマイズ可)
+- **学習目的・特定スタイル特化・蒸留パイプライン理解** には十分価値
+
+### 5.7.1 v1.0 base 切替で観察された style drift
+
+2026-05-17 に base を `anima-preview3-base.safetensors` → `anima-base-v1.0.safetensors` に
+切替。**既存 LoRA はすべて preview3 で訓練済だが v1.0 でも動作する** (アーキ互換)。ただし:
+
+- **出力 style がデフォルメ寄りに変化** — verify_v1_sage の 13 条件すべてで chibi/SD 風が強く
+  出た。preview3 base での同じ LoRA の出力 (compare_z_vs_dmd2) は比較的写実寄り
+- 原因仮説: v1.0 base が「young girl appearance」prompt を v1.0 の aesthetic prior でより
+  強く解釈する。preview3 とは prior 分布が異なる
+- LoRA を再訓練するなら v1.0 base 上で行うのが筋。preview3 ベース LoRA を v1.0 で使う場合は
+  style ドリフトを受け入れる前提
+
+## 5.8 まだ試せる方向 (諦める理由はない)
+
+公式 Turbo に勝てないと判明しても、以下は試す価値がある:
+
+1. **訓練 step を増やす** — ① Z-Image を 2000 → 8000 step に延長 (+$54)、loss はまだ
+   plateau の手前で振動していた
+2. ~~**蒸留 LoRA に DRaFT+ HPSv2 reward を追加学習**~~ — **検証済、推奨せず** (§5.10 参照)。
+   HPSv2 score は +13 上昇したが視覚的には **reward hacking で劣化** (anime → 西洋イラスト
+   風にドリフト、キャラ固有要素曖昧化、prompt adherence 落下)
+3. **より高解像度 teacher rollout** — 768 → 1024 で precompute (+$22 / cache、+$20 / train)
+4. **スタイル/キャラ特化蒸留** — 公式 Turbo は汎用向け、特定ドメイン (例: 単一キャラ、
+   特定 artist style) なら自前のほうが優位の可能性
+5. **LoRA rank を上げる** — 32 → 64 / 128 (`--lora-rank 64`)、より大きな容量で
+   teacher の細部を学べる可能性
+6. **既存実装 (LADD/Reflow/PCM/SiD2/Shortcut) の本番結果を待つ** — 進行中、Z-Image を
+   超える手法が出る可能性は残る
+7. **8 手法の **アンサンブル** — 複数 LoRA を runtime で blend (`ComfyUI` で
+   `LoraLoaderModelOnly` を 2 個直列に挟むだけ)、相補的に効くケースあり
+8. **v1.0 base で再訓練** — 既存 LoRA は preview3 起点、v1.0 で訓練し直すと aesthetic
+   一致でさらに伸びる可能性 (§5.7.1 参照)
+
+「公式に勝つ」のではなく「**公式とは違う強みを持つ自前 LoRA**」を作る方が現実解。
+本 repo の 7 実装は **どれもライセンス上自由に派生可能** で、warm-start の柔軟性が
+あるのが公式 Turbo にはない利点。
+
+### 5.8.1 ファインチューニング dataset 生成 (進行中、2026-05-17)
+
+**目的**: hakushiMixAnima_v02 base + anima-highres-aesthetic-boost LoRA を組み合わせて、
+**特定 4 artist のスタイル mix LoRA** を fine-tune するための合成 dataset を生成。
+
+**Artist tags (固定 4 名、本 repo では秘匿 ♥)**:
+
+```
+@artist_1, @artist_2, @artist_3, @artist_4
+```
+
+- **本ドキュメントではセキュリティ / プライバシーの観点から具体的 artist 名を伏せる**
+- 実 artist tag は `scripts/finetune_prompts*.txt` (Modal volume の private prompts ファイル) と
+  user の脳内にのみ存在する
+- **Claude (AI assistant) が処理上 artist tag の実値が必要になった場合は、必ず user に質問**
+  して聞き出すこと。推測や前回会話履歴からの引用で済まそうとしない (privacy 保護)
+- 4 artist は毎 prompt 必ず含まれる (artist drop なし、style mix を強固に学習させる狙い)
+
+**Dataset 構成**:
+
+| subdir | safety | themes | seeds | 想定枚数 | 内容 |
+|---|---|---|---|---|---|
+| `/dataset/finetune_raw/` | `safe` | 60 | 16 | 960 | SFW v1: 単体キャラ + couple + group + 環境 |
+| `/dataset/finetune_raw_v2/` | `safe` | 60 | 16 | 960 | SFW v2: animal/mytho/sports/cyberpunk/etc |
+| `/dataset/finetune_raw_v3/` | `safe` | 60 | 16 | 960 | SFW v3: 職業/世界観 (idol/detective/scholar/etc) |
+| `/dataset/finetune_raw_nsfw/` | `nsfw` | 30 | 16 | 480 | suggestive (bath/lingerie/bedroom/etc) |
+| `/dataset/finetune_raw_explicit/` | `explicit` | 30 | 16 | 480 | explicit sex acts #1-30 |
+| `/dataset/finetune_raw_explicit_v2/` | `explicit` | 60 | 16 | 960 | explicit sex acts #31-90、consensual framing 8 件修正 |
+| **合計** | | **300 themes** | | **~4800** | |
+
+訓練時は AnimaImageCaptionDataset の rglob で全 subdir 自動取り込み (親 `/dataset/` を指定)。
+
+### 5.8.2 ファインチューニング訓練計画
+
+完成後 `train_finetune_4variants` で 4 hyperparam 並列実行 (各 B200):
+
+| variant | rank | lr | epochs | drop_artist | 狙い |
+|---|---|---|---|---|---|
+| A baseline | 32 | 2e-5 | 3 | 0.0 | Anima 公式推奨そのまま |
+| B higher capacity | 64 | 2e-5 | 3 | 0.0 | rank 倍で細部表現↑ |
+| C safer LR | 32 | 1e-5 | 4 | 0.0 | 過学習回避 + step ↑ |
+| D artist focus | 32 | 3e-5 | 3 | 0.0 | lr 強めで style 強化 |
+
+**`llm_adapter_lr = 0` 必須** (Anima 公式既知識保護、本 repo `configs/phase1_anima.toml` 既設定)。
+
+## 5.9 実走中の状態スナップショット (2026-05-17 時点)
+
+> ⚠️ このセクションは **2026-05-17 時点のスナップショット**。実際の Modal volume 上の
+> path / 中身は時間と共に変わる。最新は `modal volume ls anima-outputs` で確認。
+> 現在配布済の LoRA は HF [darask0/anima-distill-loras](https://huggingface.co/darask0/anima-distill-loras) に集約。
+
+このリポジトリで現在 Modal 上に保存されている / 訓練中の LoRA とその想定用途。
+中断/再開時の参照用:
+
+| LoRA | path on `anima-outputs` volume | source method | 状態 (当時) | コスト |
+|---|---|---|---|---|
+| **traj_final** | `/output/traj_full/traj_final.safetensors` | ① Z-Image trajectory imitation | ✅ 完走 (2000 step) | $18 |
+| **traj_extended** | `/output/traj_extended_8000/traj_final.safetensors` | ① Z-Image 延長 (6000 step 追加) | 🟢 進行中 | ~$54 |
+| **traj_rank128** | `/output/traj_rank128/traj_final.safetensors` | ① rank 32→128 cold-start | 🟢 進行中 | ~$50 |
+| **dmd2_student_final** | `/output/dmd2_full/dmd2_student_final.safetensors` | ② DMD2 + TrigFlow | ✅ 完走 (3000 outer) | $19 |
+| **ladd_v2_fast** | `/output/ladd_v2_fast/ladd_student_final.safetensors` | ③ LADD-v2 (bug 修正版) | 🟢 進行中 | ~$17 |
+| **reflow_final** | `/output/reflow_full/reflow_final.safetensors` | ④ Reflow (rfpp) | 🟢 進行中 (~7000 step) | ~$8 |
+| **pcm_final** | `/output/pcm_full/pcm_final.safetensors` | ⑤ PCM (SD3 流派) | 🟢 進行中 | ~$36 |
+| **sid_student_final** | `/output/sid_full/sid_student_final.safetensors` | ⑥ SiD2 (data-free) | 🟢 進行中 | ~$18 |
+| **shortcut_final** | `/output/shortcut_full/shortcut_final.safetensors` | ⑦ Shortcut Models (d-head) | 🟢 進行中 | ~$25 |
+| **draftp_final** (Z-Image warm) | `/output/draftp_full/draftp_final.safetensors` | ⑧ DRaFT+ HPSv2 on ① | ✅ 完走 | $3 |
+| **draftp_on_turbo** | `/output/draftp_on_turbo/draftp_final.safetensors` | ⑧ DRaFT+ HPSv2 on Civitai Turbo | ✅ 完走 | $3 |
+
+中間 cache:
+- `/dataset/teacher_x0_cache/` — 5000 サンプル、teacher 20-step CFG=4.5 の x0 + Qwen3 emb
+- `/dataset/reflow_cache/` — 同上 + `noise` も保存 (Reflow / Shortcut 用)
+- `/dataset/teacher_x0_smoke/` — 50 サンプル smoke 用
+- `/models/loras/anima_turbo.safetensors` — Civitai 公式 (warm-start 用)
+- `/models/hpsv2/HPS_v2_compressed.pt` — DRaFT+ 用 reward model
+
+比較生成: `compare_distill_loras` (4 条件) と `compare_all_methods` (全 9-11 条件、final LoRA
+存在を自動判定) の 2 種を `modal_app.py` に実装済。後者は 5 残り手法の本番完了後に走らせる。
+
+実コスト累計 (2026-05-17 時点): **~$200** (smoke + 失敗 + 完走 + 進行中の使用済分込み)、
+予算 $300 の 67%。
+
+### 5.9.2 Modal volume 構成 (4 volumes)
+
+本リポジトリは **4 つの Modal volume** を使い分けて状態を持つ。volume mount は
+`modal_app.py::VOLUMES` で定義。
+
+#### A. `anima-models` (~30 GB) — base モデル + 全 LoRA + reward weights
+
+mount: `/models`、create_if_missing=True
+
+```
+checkpoints/
+  anima-base-v1.0.safetensors          ← 現用 base (4.18 GB)
+  hakushiMixAnima_v02.safetensors      ← (symlink, comfyui-anima-models から)
+  qwen_3_06b_base.safetensors          ← text encoder
+  qwen_image_vae.safetensors           ← VAE
+  phase_a_distilled.safetensors        ← 旧 Phase A 失敗 ckpt (~$70 で失敗、残ってる)
+  phase_a_step1000.safetensors         ← 旧 Phase A 中間
+  anima-preview3-base.safetensors      ← ❌ 削除済 (v1.0 切替時)
+
+loras/
+  anima_turbo.safetensors               ← Civitai 公式 (warm-start のデファクト)
+  z_image_traj_final.safetensors        ← ① Z-Image 蒸留 (preview3 起点、完走 LoRA)
+  dmd2_student_final.safetensors        ← ② DMD2 (preview3 起点、完走 LoRA)
+  reflow_final.safetensors              ← ④ Reflow (preview3 起点、完走 LoRA)
+  draftp_on_zimage.safetensors          ← ⑩ DRaFT+ on Z-Image (品質追加学習)
+  draftp_on_turbo.safetensors           ← B: DRaFT+ on Civitai Turbo (reward hacking 検証済、品質劣化)
+  ladd_v2_step500.safetensors           ← ③ LADD v2 部分 (停止前 step 500)
+  pcm_step500.safetensors               ← ⑤ PCM 部分 (停止前 step 500)
+  anima-highres-aesthetic-boost.safetensors ← (symlink, comfyui-anima-models から)
+  phase_b/d/f_4step_lora.safetensors    ← 旧 R3GAN 試行の残骸 LoRA
+
+hpsv2/
+  HPS_v2_compressed.pt                  ← DRaFT+ 用 reward model (1.97 GB)
+
+torchhub/                                ← torch.hub cache
+cache/                                   ← misc cache (sageattention 等)
+```
+
+#### B. `anima-dataset` (~40 GB) — 全データセット + 比較生成出力
+
+mount: `/dataset`、create_if_missing=True
+
+**A. 元データ / self-distillation**:
+- `raw/` — Anima base 自己生成 5000 (.png + .txt)、初期 distill 用
+- `samples/`, `turbo_test/`, `compare/` — ��去の動作確認生成
+
+**B. 蒸留用 precompute cache**:
+- `teacher_x0_smoke/` — 50 サンプル smoke 用
+- `teacher_x0_cache/` — 5000 サンプル × (x0 + emb)、LADD/PCM/SiD2 用
+- `reflow_cache/` — 同上 + noise 保存、Reflow / Shortcut 用
+
+**C. 比較 / 検証生成**:
+- `compare_z_vs_dmd2/` — 80 PNG (base / turbo / ① / ② × 8/4 step)
+- `verify_completed/` — 旧 verify (workflow bug で 0 images、失敗)
+- `verify_v1_sage/` — v1.0 + sageattention で 13 条件 verify (1 prompt × 13)
+- `ckpt_health_check/` — quick_check_ckpt 出力 (PCM step 500 等)
+
+**D. ファインチューニング dataset (現在進行中、~4800 枚予定)**:
+- `finetune_raw/` — SFW v1 (60 themes × 16 seeds = 960)
+- `finetune_raw_v2/` — SFW v2 (animal/mytho/sports 60×16 = 960)
+- `finetune_raw_v3/` — SFW v3 (職業/世界観 60×16 = 960)
+- `finetune_raw_nsfw/` — NSFW suggestive (30×16 = 480)
+- `finetune_raw_explicit/` — EXPLICIT v1 (30×16 = 480)
+- `finetune_raw_explicit_v2/` — EXPLICIT v2 #31-90 (60×16 = 960)
+
+#### C. `anima-outputs` (~30 GB) — 訓練 LoRA / ckpt の保存先
+
+mount: `/output`、create_if_missing=True
+
+**完走 (preview3 起点)**:
+- `traj_full/` — ① Z-Image final ($18)
+- `dmd2_full/` — ② DMD2 final ($19)
+- `reflow_full/` — ④ Reflow final ($13)
+- `draftp_full/` — ⑩ DRaFT+ on Z final ($3)
+- `draftp_on_turbo/` — B: DRaFT+ on Turbo final ($3)
+
+**中断 / 部分 (preview3、~$80-100 浪費)**:
+- `ladd_v2_fast/` — ③ LADD v2 step 500 ckpt まで
+- `pcm_full/` — ⑤ PCM step 500 ckpt まで
+- `sid_full/` — ⑦ SiD2 outer ~4200 まで (停止時)
+- `traj_extended_8000/` — A: Z-Image extend (停止)
+- `traj_rank128/` — E: rank 128 (停止)
+- `shortcut_full/` — ⑦ Shortcut 異常遅延 step 500 まで ($63 損)
+
+**旧 R3GAN 試行 (5 連続失敗の残骸)**:
+- `distill/`, `distill_e/`, `distill_f/`, `distill_f_full/` — Phase A/B/C 試行
+
+**smoke 各種**:
+- `*_smoke/` — 各手法の smoke test 出力 (1-5 step、debug 用)
+
+#### D. `comfyui-anima-models` (~5 GB) — user アップロード外部モデル
+
+mount: `/comfyui_anima_models`、create_if_missing=**False** (user が事前アップロード)
+
+```
+diffusion_models/
+  hakushiMixAnima_v02.safetensors      ← ファインチューニング用 base
+loras/
+  anima-highres-aesthetic-boost.safetensors ← ファインチューニング用 quality boost LoRA
+text_encoders/, vae/                    ← 空 (将来用 placeholder)
+```
+
+ファインチューニング用 dataset 生成時に `/models/{checkpoints,loras}` へ symlink される
+(`generate_finetune_chunk` 内で自動)。
+
+#### 累計サイズ + コスト
+
+| volume | 概算サイズ | 月額保管料 (Modal $0.15/GB/月) |
+|---|---|---|
+| anima-models | ~30 GB | ~$4.5 |
+| anima-dataset | ~40 GB | ~$6 |
+| anima-outputs | ~30 GB | ~$4.5 |
+| comfyui-anima-models | ~5 GB | ~$0.75 |
+| **合計** | **~105 GB** | **~$16/月** |
+
+不要 dir は `cleanup_checkpoints` Modal function 等で適時間引きが推奨 (古い phase_a/b/c 系、
+shortcut_full、放置 smoke 各種、要らない finetune subdir は数十 GB 占有してる)。
+
+### 5.9.1 Anima 公式推奨 prompt / 推論設定 / ファインチューニング設定
+
+[CircleStone Labs/Anima HF README](https://huggingface.co/circlestone-labs/Anima) から
+抜粋した推奨設定。本リポジトリの全 workflow / generate_dataset / fine-tune script は
+この推奨に準拠している。
+
+**Prompt format**:
+
+```
+[quality/meta/year/safety] [1girl/1boy/etc] [character] [series] [@artist] [general]
+```
+
+- Quality prefix (positive 先頭): `masterpiece, best quality, score_7, safe`
+- Negative: `worst quality, low quality, score_1, score_2, score_3, artist name`
+- tag は **lowercase + space 区切り** (underscore は score タグだけ: `score_7`)
+- Danbooru タグより **Gelbooru タグを優先**
+- artist tag は **`@` プレフィックス必須** (なしだと効果激減)、`[artist]` 位置に置く
+  - 多語 artist は半角 space (例: `@nnn yryr`)
+  - 別解釈 disambiguator がある場合は paren escape (例: `@artist_name \(circle_name\)`)
+  - 複数 artist を `,` 区切りで列挙可
+
+**Recommended sampler / scheduler / CFG / steps**:
+
+| 推奨 | 値 / 説明 |
+|---|---|
+| sampler | **`er_sde`** (中立、フラット彩色、シャープライン)、`euler_a` (柔らかめ、2.5D 傾向)、`dpmpp_2m_sde_gpu` (バラエティ豊か、暴走しがち) |
+| scheduler | **`simple`** (デフォルト)、`beta57` (RES4LYF custom node、リアル系・絵画調用) |
+| sigma_shift | **`3.0`** (`ModelSamplingAuraFlow` で必須、本リポジトリ全 workflow で設定済) |
+| CFG | **4-5** |
+| step | **30-50** |
+| resolution | **512² から 1536²** までサポート (1024² 中心が無難) |
+| VAE | `qwen_image_vae.safetensors` |
+| Text encoder | `qwen_3_06b_base.safetensors` |
+| Diffusion model | `anima-base-v1.0.safetensors` |
+
+**Fine-tuning best practices** (Anima ��式抜粋):
+
+- **`llm_adapter_lr = 0` 必須** — adapter は影響力が強く既知識を多く持つため訓練すると壊れる
+  (本 repo `configs/phase1_anima.toml` で設定済)
+- 低 LR start: rank 32 LoRA で **`lr = 2e-5`** から
+- light touch training — モデルが既に多様な visual concept を持っているので無理に押し込まない
+- captioning は **Danbooru tag + 自然言語の併用** OK
+- 例: 公式 [style LoRA on Civitai](https://civitai.com/models/2536147) (dataset + config 公開)
+
+**Content policy / restrictions**:
+
+- safety tag: `safe / sensitive / nsfw / explicit` を positive / negative 両方で使い分ける
+- text rendering は弱い (長文を画像に書かせない)
+- 複数キャラ: character 名 → 基本属性の順で書くと混同回避
+- prompt weighting は SDXL より強めに必要 (例: `(chibi:2)`)
+- 自然言語 prompt は **最低 2 文以上** が安定 (極端に短いと予測不能)
+- **非商用ライセンス縛り**、商用問い合わせ `tdrussell@circlestone.ai`
+
+## 5.10 品質評価の方法論 (reward hacking 実証あり)
+
+**「ディテール多い・彩度高い」≠「良い」**。蒸留 LoRA や reward fine-tuning の出力を評価
+する際は、必ず以下の軸で分解して判断する:
+
+| 評価軸 | 確認内容 |
+|---|---|
+| (a) **target style 維持** | anime 系を期待しているなら anime のままか? 西洋イラスト / リアル系にドリフトしていないか? |
+| (b) **キャラクター固有要素の精度** | 例: Touhou Flandre Scarlet なら金髪+赤目+クリスタル翼、Remilia なら青髪+赤目+コウモリ翼 が正確か? |
+| (c) **prompt adherence** | 指示したポーズ (例: hands raised pointing toward viewer with open palms)、構図、背景要素を反映しているか? |
+| (d) **破綻ポイント** | 手指の本数・関節・方向、複数キャラの顔混じり、関節破綻、bad anatomy |
+| (e) **生成時間** | per-image gen_time (`_summary.json` から)、step 数 × sampler × LoRA 構成で比較 |
+
+**DRaFT+ HPSv2 reward fine-tuning の reward hacking 実証 (2026-05-17)**:
+
+- 仮説: 公式 Anima Turbo に HPSv2 reward fine-tuning (DRaFT+ K=1 LV、kl_coeff=0.2、1500 step、
+  $3) を後付けすれば品質が公式超えできる?
+- 結果: HPSv2 score は **25.78 → 38.99 (+13.21)** と大幅向上 (loss 順調、KL 0.2-0.3 安定)
+- **しかし視覚的には劣化** (ユーザー評価): 西洋イラスト風にドリフト、Flandre/Remilia の固有
+  翼が曖昧化、両キャラの衣装がほぼ同じに、prompt の手のポーズが崩れる、装飾過多 (シャンデリア
+  + ランプ + 飾り) でゴチャゴチャ
+- 原因: HPSv2 は「人間が好きそうな generic な見栄え」preference を学習しており、Touhou anime
+  style や character accuracy を直接最適化していない。reward が指す方向と用途が乖離した
+
+教訓:
+- **reward score 上昇 ≠ 品質向上**。reward 系は必ず視覚検証
+- HPSv2 / PickScore / ImageReward 等は「平均的に綺麗」を上げるだけ、特定 niche style には不向き
+- 蒸留 LoRA の評価では、必ず **「複数の評価軸を分解」して判定**。1 軸だけ (例: ディテール量)
+  で判断しない
+- 「score が上がった」「ディテールが増えた」を理由に勝者宣言しない、視覚的に必ず確認する
+
+**verify 用 prompt の設計指針** ([scripts/verify_prompts.txt](../scripts/verify_prompts.txt)):
+
+単純な single-character prompt では LoRA の差が出ない。**「壊れにくいもの」より「壊れやすい
+もの」をテストする**:
+
+- 複数キャラクター (face mix / 属性入れ替えが起きやすい)
+- キャラ毎に異なる属性 (髪色、目色、翼種別)
+- 手・指の特定ポーズ (open palms、pointing)
+- 詳細背景 + foreground キャラ両立
+- target style の制約 (anime style 維持を強制)
+
+これで初めて LoRA 間の優劣が顕在化する。
+
+## 5.11 次セッション引き継ぎ (2026-05-17 23:xx 中断時点)
+
+> ⚠️ 2026-05-17 時点のスナップショット。**現在は v1.0 base で PCM が cold-start 完走済**
+> ([darask0/anima-distill-loras/pcm](https://huggingface.co/darask0/anima-distill-loras/tree/main/pcm))。
+> 本セクションは当時の状態を歴史記録として保持。
+
+### 1. 完走済 (Anima preview3 base 起点、すべて使用可)
+
+| LoRA | 出力 path | コスト |
+|---|---|---|
+| ① Z-Image traj imitation | `/output/traj_full/traj_final.safetensors` | $18 |
+| ② DMD2 + TrigFlow | `/output/dmd2_full/dmd2_student_final.safetensors` | $19 |
+| ④ Reflow rfpp | `/output/reflow_full/reflow_final.safetensors` | $13 |
+| ⑩ DRaFT+ on ① Z-Image | `/output/draftp_full/draftp_final.safetensors` | $3 |
+| B: DRaFT+ on Civitai Turbo | `/output/draftp_on_turbo/draftp_final.safetensors` | $3 ⚠️ reward hacking 検証で **品質劣化** 確認 |
+
+### 2. 中断 (preview3 起点、user が「古いから止めて」と判断し全停止 2026-05-17)
+
+| 手法 | 部分 ckpt | 中断時 step | 浪費コスト |
+|---|---|---|---|
+| ③ LADD v2 | `/output/ladd_v2_fast/ladd_student_step00500.safetensors` | ~500/5000 | ~$5 |
+| ⑤ PCM | `/output/pcm_full/pcm_step00500.safetensors` | ~760/5000 | ~$30 |
+| ⑦ SiD2 | `/output/sid_full/` (sample_every=500 で step ~4000 までの ckpt) | ~4200/5000 | ~$15 |
+| ⑧ Shortcut | `/output/shortcut_full/shortcut_step00500.safetensors` | ~560/4000 (64s/step 異常遅延で abort) | $63 |
+| A: Z-Image extend | `/output/traj_extended_8000/` (step ~1360 までの ckpt) | ~1360/6000 | ~$13 |
+| E: rank 128 | `/output/traj_rank128/` (step ~1440 までの ckpt) | ~1440/2000 | ~$13 |
+
+中断合計浪費: **~$139**。**ファイルは残してある**、user 判断で再開 or 削除。
+
+### 3. ファインチューニング dataset (新規、hakushi base、ほぼ完成)
+
+| subdir | safety | 完了状況 |
+|---|---|---|
+| `/dataset/finetune_raw/` | safe | ✅ 960 |
+| `/dataset/finetune_raw_v2/` | safe | ✅ 960 |
+| `/dataset/finetune_raw_v3/` | safe | ✅ 960 |
+| `/dataset/finetune_raw_nsfw/` | nsfw | ✅ 480 |
+| `/dataset/finetune_raw_explicit/` | explicit | ✅ 480 |
+| `/dataset/finetune_raw_explicit_v2/` | explicit | ✅ 960 |
+| **合計** | | ✅ **4800 / 4800 完成** |
+
+訓練 prompt: `scripts/finetune_prompts*.txt`、artist tag は **秘密 ♥** (§5.8.1 参照)。
+
+### 4. 次セッションで user 判断待ちのこと (最重要)
+
+ファインチューニング dataset 完成後、**どちらに進むか user 判断**:
+
+**Path A (推奨): ファインチューニング集中**
+- `train_finetune_4variants` で 4 hyperparam 並列 (A: rank 32 lr 2e-5 / B: rank 64 lr 2e-5 / C: rank 32 lr 1e-5 epochs 4 / D: rank 32 lr 3e-5)
+- 各 ~3-4h、計 ~$120
+- 蒸留は公式 Anima Turbo に任せる (compare_z_vs_dmd2 で Turbo 圧勝が実証済)
+- 自前 fine-tune LoRA × 公式 Turbo の **スタック** が現実解
+
+**Path B: v1.0 base で蒸留 LoRA を再訓練**
+- 既存の `/dataset/raw/` (preview3 self-distill) は base 違いで不適
+- v1.0 base で self-distillation dataset 新規生成必要 (~$25)
+- その後 5 蒸留手法を v1.0 で再訓練 (~$80-120)
+- 計 ~$120-145
+- 公式 Turbo に勝てる見込み低 (前回比較で実証)
+
+**Path C: 両方やる**
+- Path A 先行で完成見て、余力あれば Path B
+- 計 ~$250-280
+
+> **2026-05-18 update**: PCM を Path B 流に v1.0 base で cold-start 完走、
+> HF 配布まで完了 ([詳細は HF model card](https://huggingface.co/darask0/anima-distill-loras/tree/main/pcm))。
+
+### 5. 次 Claude が必ず守るルール
+
+- **Artist tag は秘密** (§5.8.1)。Claude が prompt 生成で必要な場合、user に直接質問する。
+  推測 / 履歴引用しない。`scripts/finetune_prompts*.txt` には実 tag が入ってるが README / chat
+  出力では具体名を引用しない
+- **Modal CLI on Windows 必須 prefix**: `MSYS_NO_PATHCONV=1 PYTHONIOENCODING=utf-8 PYTHONUTF8=1`
+- **画像品質評価**: 「ディテール多い ≠ 良い」、必ず user judgment を仰ぐ (§5.10 reward hacking 実証)
+- **workflow JSON node ID**: positive=5, latent=7, KSampler=8 (`generate_dataset.py::patch_workflow` が hardcode)
+- **base 切替の影響**: preview3 LoRA を v1.0 で使うと style drift する (§5.7.1)
+- **rank 違いの LoRA は warm-start 不可** (`load_warm_lora` で 0 keys matched エラー)、cold-start 必要
+- **LADD discriminator bug fix**: `gradient_to_input` フラグ追加済、G phase は True、D phase は False
+- **Modal function param は lowercase**: 大文字 `K: int` は Modal CLI で `k=1` に変換され抹消、`k_grad` などに
+- **sageattention 有効**: `_base_image` に同梱、`generate_dataset.py` で auto-enable
+
+### 6. 現コスト累計 + 予算
+
+- 完走 LoRA + smoke + 失敗 retry + dataset: **~$200**
+- 中断蒸留訓練の浪費: **~$139**
+- 当初予算 $300 → **~$339 で +$39 超過**
+- 推定 ファインチューニング 4 並列: +$120 → 完成時総額 ~$460
+- user が「金は OK」と承認 (実コストは全部 Modal dashboard で監視可能)
+
+### 7. 中断時 Modal で active な task
+
+セッション終了時に bash background は全部止めてよい (Modal cloud は無影響)。
+ただし Modal app は **app_stop コマンドで明示停止が必要** (止めないと課金続行):
+
+```bash
+modal app list                      # 全 app 確認
+modal app stop <ap-xxxxx> --yes     # 個別停止
+```
+
+次セッション開始時は `modal app list` で再確認して、無駄に残った app が無いか確認推奨。
+本セッション中断時には:
+- 6 訓練 (LADD v2 / PCM / SiD2 / A / E / Shortcut) はすべて停止済
+- dataset 生成: ✅ **全 6 subdir 4800 枚完成** (SFW v3 も 2026-05-17 終わり際で完了)
+- 残 active task なし、Modal 課金は volume 保管料 ($16/月) のみ

docs/operations.md

@@ -0,0 +1,158 @@
+# 実装 / 運用 Tips + トラブルシューティング
+
+[← README に戻る](../README.md)
+
+[distillation.md §4.7](distillation.md#47-周辺で踏んだ-modal--diffusion-pipe--comfyui-罠) の diffusion-pipe / ComfyUI 罠に加え、
+2026-05 試行 (5 手法移植) で踏んだ追加の罠と対処、よくあるトラブルの解決策。
+
+## Modal CLI on Windows + JP (cp932) locale
+
+毎回必須:
+```bash
+MSYS_NO_PATHCONV=1 PYTHONIOENCODING=utf-8 PYTHONUTF8=1 modal run ...
+```
+
+| 罠 | 症状 | 対処 |
+|---|---|---|
+| cp932 codec error | Modal CLI 出力の Unicode ✓ や ● を encode できず即死 | `PYTHONIOENCODING=utf-8 PYTHONUTF8=1` |
+| Git Bash MSYS path mangling | `--dataset-path /dataset/raw` が `C:/Program Files/Git/dataset/raw` に化ける | `MSYS_NO_PATHCONV=1` |
+| Modal token が silent revoke | `.modal.toml` は valid、`modal config show` で見える、しかし API 呼び出しで "Token not found" | `MODAL_LOGLEVEL=DEBUG modal app list` でサーバー側のエラーを確認 → ユーザーが `modal token new` で再発行 |
+| `modal app stop` が対話確認で停止 | `Are you sure...? [y/N]:` で固まる | `--yes` を必ず付ける |
+
+## diffusion-pipe + Anima 移植時の罠
+
+| 罠 | 症状 | 対処 |
+|---|---|---|
+| `_find_blocks` ヘルパーが `llm_adapter.blocks` を選ぶ | Anima DiT は本体 28 blocks + llm_adapter 6 blocks。名前 heuristic で adapter 側を拾うと **5 head すべて学習されない** | `transformer.blocks` を direct 参照、または `"llm_adapter" not in name` で filter (`anima_ladd_disc.py::_find_blocks` 参照) |
+| LPIPS reg が OOM | `compute_regularization` 内で **student rollout を grad-through で 2 回目** 走らせると 16 forward が graph に乗り 178 GiB 食い尽くす | (1) `--lpips-weight 0` で disable (2) `--lpips-every N` で間欠的に有効化 (Reflow パターン) (3) 解像度 1024 → 768 に下げる |
+| PEFT LoRA 2 adapter 切替 | DMD2 で同じ base に student + fake_score を attach する場合、**毎 forward で `set_adapter(name)` を呼ばないと 片方の重みが leak する** | `make_velocity_fn(peft_model, adapter_name)` のような callable wrapper で隠蔽 (`train_dmd2_official.py` 参照) |
+| Civitai Turbo LoRA → PEFT 形式変換 | warm-start 時、ComfyUI 形式 `diffusion_model.<…>.lora_A.weight` を PEFT 形式 `base_model.model.<…>.lora_A.default.weight` に変換必要 | `train_traj.py::convert_comfy_to_peft_lora` (`_convert_peft_to_comfy_lora` の逆向き) |
+| warm-start key match | wide LoRA target が `x_embedder / t_embedder` を含まないため、Anima Turbo の 1016 keys のうち **968 keys のみ match** (12 keys skip、48 keys missing) | 致命的ではない、warm-start として十分機能 |
+| subprocess 内 `import modal` 失敗 | 訓練 script が `modal.Volume.from_name(...).commit()` を呼ぶと `ModuleNotFoundError: No module named 'modal'` (subprocess の Python path に Modal がない) | `try/except` で囲んで warning 出力のみ。Modal は subprocess 終了時に自動で volume commit する |
+| **LADD discriminator の `torch.no_grad()` + `feat.detach()` で adv 勾配が student に届かない** | G phase で adv loss を `.backward()` しても勾配が disc の heads までで止まる。実質「Smooth L1 recon distill only」になり LADD のメリット消失 | `AnimaLADDDiscriminator.forward(...)` に `gradient_to_input` フラグを足す。G phase は `True` (teacher は frozen なので weight 更新されないが activations が残り x_input → student に勾配が流れる)、D phase は `False` (heads だけ訓練、メモリ節約) |
+| **LADD bs=4 × accum=4 × misaligned_pairs_d で 170 s/step (smoke の 200x 遅い)** | bs=4 で D forward の cost が線形増、accum=4 で 4 倍 micro-step、misaligned で D batch を 2 倍に → 20+ full-DiT forward/step、B200 上で 170 s | **bs=2 accum=1 + misaligned 無効** に下げる → ~2 s/step (85x 高速化)。LADD は GAN なので大 batch 不要 (small noise の方が D の安定性に貢献)。`train_ladd.py` 引数: `--batch-size 2 --grad-accum 1` (`--misaligned-pairs-d` フラグを付けない) |
+| **異なる rank の LoRA を warm-start に使えない** | `--lora-rank 128` で `--warm-lora /models/loras/anima_turbo.safetensors` (rank 32) を指定 → `load_warm_lora` 内で shape mismatch、`No LoRA keys matched` で abort | rank 変更時は **cold-start** で訓練 (warm-start を省略)。または rank padding helper を書く |
+| **modal CLI の uppercase arg は失敗** | function param `K: int` を CLI から `--K 1` で渡しても Modal が `k=1` に lowercase 変換、`unexpected keyword argument 'k'` で abort | function param は **必ず lowercase** にする (例 `k_grad: int`)、CLI 側 `--k-grad 1` を渡す |
+| **ComfyUI workflow JSON の `_comment` キーが node 扱いされて crash** | workflow JSON を iterate して `node.get("class_type")` する処理で、`_comment` (string value) が来ると `'str' object has no attribute 'get'` | `for node_id, node in wf.items(): if not isinstance(node, dict): continue` で skip |
+| **workflow JSON の node ID が `patch_workflow` の hardcode と一致しないと全 generation が HTTP 400 で fail** | `scripts/generate_dataset.py::patch_workflow` は positive prompt を node "5"、latent を "7"、KSampler を "8" に hardcode。ユーザー由来 workflow の node ID が違うと `$PROMPT/$SEED` プレースホルダが置換されず ComfyUI が prompt graph reject | 新 workflow を作る時は **必ず node ID を 5=positive, 7=latent, 8=KSampler** で構築。`anima_workflow.json` 系を雛形にする |
+| **reward fine-tuning (DRaFT+) で reward score 上昇 ≠ 視覚品質向上** | HPSv2 score +13.21 上昇したが視覚的には reward hacking が起き、anime → 西洋イラスト ドリフト + キャラ固有要素曖昧化 + prompt adherence 落下 | 必ず視覚検証する、score だけで「勝った」と判断しない。kl_coeff 上げる (0.2 → 0.5) や reward 種類変更 (HPSv2 → ImageReward) で軽減可能性、ただし根本解決ではない ([migration_log §5.10](migration_log.md#510-品質評価の方法論-reward-hacking-実証あり) 参照) |
+| **base model 切替で既存 LoRA の出力 style が drift する** | preview3 で訓練した LoRA を v1.0 base で使うと chibi/SD 寄りに style が変わる (LoRA は動作する、出力分布が違う) | LoRA は train 時の base と一致させて使うのが筋。base 切替時は LoRA を再訓練するか style drift を許容する |
+| **sageattention は .so 不要 (Triton JIT)** | `cache_sageattention_to_volume` で `/usr/local/lib/python3.11/site-packages/sageattention/` を volume copy しても `.so` が 0 個。一見「コンパイル失敗」に見える | 仕様。sageattention は Triton で runtime JIT、import 時に kernel が compile される。動作確認は CUDA kernel smoke test (`sageattn(q,k,v)` を呼ぶ) で行う |
+| **ComfyUI で sageattention を効かせる** | デフォルトは torch SDPA、sageattention 入っていても自動使用されない | ComfyUI 起動引数に `--use-sage-attention` を追加。`generate_dataset.py` は `import sageattention` 成功時に自動で付与するよう改修済 |
+| **Modal volume get で Japanese 文字パスに download すると後から Bash で見つけられない** | `modal volume get anima-dataset path/ "/c/Users/micro/Downloads/新しいフォルダー/dest/"` で download は成功するが、後続の `cd` / `find` / Python `os.walk` で日本語パスが broken or 空に見える | 一次 download 先は ASCII path に (例: `./verify_v1_local/`)、必要なら後で copy で日本語フォルダに移す |
+| **Modal volume `rm` は subprocess 内では使えない** | training script が訓練中に preview3 を消そうとしてもファイルハンドルが mmap で開いている (PyTorch safetensors loader) | volume cleanup は別の Modal function で実行、training 中は触らない (実害は少ないが orchestration 注意) |
+
+## コスト / 時間の実測値 (本リポジトリ 5 手法移植時)
+
+すべて B200 ($6.25/h Modal)、batch=1-4、768-1024 解像度:
+
+| 操作 | 時間 | コスト |
+|---|---|---|
+| Image rebuild (lpips 追加分) | ~5 分 (キャッシュ後 0 分) | ~$0.5 |
+| B200 cold start | ~3 分 | ~$0.30 |
+| Smoke test (1 step) | ~5-10 分 | ~$0.5-2 |
+| Z-Image train 2000 step | ~3 h (5.2s/step) | **~$18** (見積もり $42 の半額) |
+| DMD2 train 3000 outer (5 critic + 1 gen) | ~5.8 h (3.6s/step) | **~$19** (見積もり $36 の半額) |
+| precompute_teacher_x0 5000 サンプル | ~1.5 h (0.69/s) | ~$11 |
+| LADD train 5000 step bs=4 accum=4 | ~4.8 h (0.86s/step) | ~$30 |
+| Reflow train 8000 step | ~3-4 h | ~$25 |
+| PCM train 5000 step (1 grad + 3 no_grad) | ~6.5 h (4.7s/step) | ~$41 |
+| **5 手法 + precompute × 2 トータル** | wall-clock ~8 h (並列) | **~$170** |
+
+実測 step 時間は **predicted の概ね 半分** — Modal の B200 は H100 比 2-3x 速いことを織り込まないと過大見積もりになる。
+
+**推論側 (verify 段階) の実測 (v1.0 base + sageattention 有効、B200、1024×1024)**:
+
+| 条件 | per-image | 内訳 |
+|---|---|---|
+| LoRA 4-step CFG=1 | **3.5-4.0s** | sage 効果あり、ほぼ全 LoRA 同等 (Z-Image / DMD2 / Turbo / Reflow / DRaFT+ 系) |
+| LoRA 8-step CFG=1 | **4.0-4.7s** | sage 効果あり、step 数比例 |
+| base 30-step CFG=4.5 (LoRA strength=0) | **10.7s** | 品質基準ライン |
+
+→ 蒸留 LoRA 4-step は base 30-step の **3 倍速 = $0.006/枚 vs $0.018/枚** (B200 $6.25/h 換算)。
+
+## トラブルシューティング
+
+### OOM (CUDA out of memory)
+- `phase1_anima.toml` の `activation_checkpointing = true` を確認
+- `[adapter].rank` を 32 まで下げる
+- 解像度を 768 に落とす(`phase1_dataset.toml` の `resolutions`)
+- それでもダメなら H100 80GB に上げる
+
+### 学習が進まない
+- データセット数枚で `epochs = 1` 試走 → エラーログを見る
+- `modal app logs rapid-anima` でリアルタイム監視
+- `llm_adapter_lr = 0` が効いているか(消えていると壊れる)
+
+### 出力の絵が地味
+- これは正常。Anima base は素なので、データセットを `score≥200` 等で
+  絞って色彩豊かな絵を集めると改善
+- artist タグを残して `@favorite_artist1, @favorite_artist2` を
+  caption に入れたままにすると、それらの平均が出やすくなる
+
+### diffusion-pipe が image build で失敗
+- `requirements.txt` 内のバージョン競合が原因のことが多い
+- `modal_app.py` の `.pip_install` 部分でバージョン固定済み
+- それでも詰まる時は GitHub Issues 参照: https://github.com/tdrussell/diffusion-pipe/issues
+
+### attention をもっと速くしたい
+デフォルトは torch SDPA(flash-attention 2 が built-in で使われる)。
+さらに 1.2-1.5x 速くしたい場合は、`modal_app.py` の image 定義の最後に追加:
+
+```python
+.run_commands(
+    # diffusion-pipe install 後にビルドして torch ABI を合わせる
+    "pip install --no-build-isolation sageattention",
+    # または flash-attn (source build, 30-40分):
+    # "pip install --no-build-isolation flash-attn",
+)
+```
+
+依存解決の都合で順序が重要(diffusion-pipe より後ろ)。
+
+### `modal volume put` で大きいデータが上がらない
+- `--force` を付けて差分上書き
+- 5GB 超えるなら一旦 tar.gz → put → 中で展開する関数を作る方が安定
+
+### `Secret 'hf_token' not found`
+Modal 上に `hf_token` という名前で secret が無い場合に出る。
+[setup.md](setup.md) の Hugging Face シークレット節を参照して作成、または
+`modal_app.py` の `from_name("hf_token", ...)` を既存 secret 名に書き換え。
+
+### HF upload で `401 Unauthorized`
+`hf_token` secret が read-only token の場合。write 権限が必要 (repo create / upload):
+
+```bash
+modal secret create HF_TOKEN_WRITE HF_TOKEN_WRITE=hf_xxxxx
+```
+
+`modal_app.py::upload_lora_to_hf` が `HF_TOKEN_WRITE` secret を参照する。
+詳細は [setup.md](setup.md#hf-hub-への-upload-lora-配布など)。
+
+### Modal `workspace billing cycle spend limit reached`
+新規 app 起動時:
+```
+App creation failed: workspace billing cycle spend limit reached
+```
+
+**現在進行中の cloud function にも SIGINT が送られて停止する**(KeyboardInterrupt として cloud log に出る、ローカル CLI の Ctrl-C と区別不可)。先に並列起動した training が **両方** とも preempt される事例あり (2026-05-18 DMD2 v1 と DMDX v1 並列起動時)。
+
+対処:
+1. https://modal.com/settings/billing で **Spend limits** タブを確認
+2. current cycle の limit を引き上げる、または cycle reset 待ち
+3. 引き上げ後は新規 app 起動も既存 detach app も復旧
+
+予防:
+- 並列で高負荷 training を起動する前に Modal dashboard で remaining budget 確認
+- DMD2 + DMDX 並列は B200 1 つあたり ~$30 × 2 = $60+ 高速消費するので、limit に余裕が必要
+
+### 並列訓練時の rate 低下
+本日実測 (2026-05-19、DMD2 resume + DMDX 並列):
+- DMD2 単独: 3.66s/outer
+- DMD2 並列 (with DMDX): **6.2s/outer** (1.7x 遅い)
+- DMDX 単独 (smoke): 3.45s/outer
+- DMDX 並列: 3.6s/outer (slowdown 小さい)
+
+原因推定: Modal が異なる B200 instance を割当てるが、混雑度や近隣 workload で per-step rate に差。DMDX は forward 数が少ない (2 disc + 1 gen = 3 stage vs DMD2 の 5+1=6 stage) ので絶対的に slowdown 耐性が高い。
+
+並列予算は単独の 1.5-2x で見積もる (e.g., DMD2 単独 $30 → 並列時 $50-60)。

docs/pcm.md

@@ -0,0 +1,98 @@
+# PCM (Phased Consistency Model) 蒸留
+
+[← README に戻る](../README.md) — 関連: [DMD2](dmd2.md) / [DMDX](dmdx.md) / [蒸留全般](distillation.md)
+
+## 概要
+
+[Wang et al. NeurIPS 2024](https://arxiv.org/abs/2405.18407)、G-U-N/Phased-Consistency-Model SD3 流派を Anima v1.0 base に移植。**FlowMatch + v-pred** で Anima の rectified flow と math 一致、ε↔v 変換不要。
+
+**実装**: `scripts/distill/{pcm_scheduler,train_pcm}.py`
+**Modal**: `modal_app.py::train_pcm_distill`
+
+## 設計の鍵
+
+- num_euler_timesteps **N=50** を K=4 phase に均等分割 → 4-step inference に最適化
+- ランダム index sampling → phase 終端まで Euler 1 step 進める
+- pseudo-Huber loss (`c=1e-3`) で stable convergence
+- 1 grad-through + 3 no_grad forward / step → ~60-80 GB on B200 (LoRA-only)
+- CFG-augmentation: `w ∈ [4.0, 5.0]` を訓練に embed → 推論 CFG=1.0 で teacher CFG=4.5 効果
+
+## 訓練設定 (実走、2026-05-18)
+
+| 項目 | 値 |
+|---|---|
+| Base | `anima-base-v1.0.safetensors` |
+| Method | Phased Consistency Model |
+| LoRA target | wide (AdaLN + attn + MLP の全 Linear、980 keys) |
+| LoRA rank | 32 |
+| Total steps | 5000 (sample_every=500) |
+| Batch / grad accum | 1 / 4 |
+| Resolution | 768 |
+| Euler N / Phases K | 50 / 4 |
+| sigma_shift | 3.0 |
+| CFG-aug w range | [4.0, 5.0] |
+| Huber c | 1e-3 |
+| LR | 5e-6 |
+| Optimizer | AdamW (wd 0.01) |
+| Warm-start | **無し** (cold-start on v1.0 base) |
+| GPU | B200 |
+| 訓練時間 | **~3.4h** (1.72s/step、想定 6.5h の半分) |
+| コスト | **~$22** (B200 $6.25/h) |
+
+## Loss 推移
+
+- 早期 (step 0-100): 0.001-0.14 範囲、warm-start 無しでも安定
+- 中盤 (step 1000-3500): 0.001-0.007 で healthy plateau
+- 後半 (step 4000-5000): 同水準維持、divergence なし
+
+```
+step 20  loss=0.0013  (初期 random spike 後すぐ低下)
+step 130 loss=0.4468  (single spike、shifted sampling の hard spot)
+step 500-4500 loss=0.001-0.007 範囲で平均化
+```
+
+mean collapse・NaN・OOM の発生 **なし**。
+
+## 検証結果 (4-step 生成)
+
+同 prompt × seed=42、Anima v1.0 base、er_sde / simple / sigma_shift=3.0:
+
+| ckpt | sampler / scheduler | 時間 | 評価 |
+|---|---|---|---|
+| step 3500 | er_sde + simple | 8.8s | ✅ 2 キャラ正確分離、softer 構図 |
+| step 4000 | res_multistep + beta | 11.2s | ✅ 同等品質、anime tone 強め |
+| final (step 5000) | — | — | HF model card のサンプル参照 |
+
+両 sampler とも 4-step CFG=1.0 で破綻なし。`er_sde + simple` を default 推奨 (Anima 公式準拠)。
+
+## 配布
+
+[**HF: darask0/anima-distill-loras/pcm/**](https://huggingface.co/darask0/anima-distill-loras/tree/main/pcm)
+
+- `pcm_final_peft.safetensors` (diffusers / peft 用)
+- `pcm_final_comfy.safetensors` (ComfyUI 用、980 keys)
+- README.md (使い方 + 訓練詳細)
+- samples/ (step 3500 / 4000 検証画像)
+
+## ComfyUI 推奨設定
+
+```
+LoraLoaderModelOnly:
+  lora_name: pcm_final_comfy.safetensors
+  strength_model: 1.0
+
+ModelSamplingAuraFlow:
+  shift: 3.0
+
+KSampler:
+  steps: 4
+  cfg: 1.0
+  sampler_name: er_sde  (or res_multistep)
+  scheduler: simple     (or beta with res_multistep)
+```
+
+## 知見
+
+- **cold-start でも安定** — 他手法 (DMD2/LADD) は warm-start 必須だが PCM は phase consistency が anchor になる
+- **rate 想定より速い** — README 試算 4.7s/step → 実測 1.72s/step (B200 sageattention 効果含む)
+- **single grad-through で memory 余裕** — DMD2 の dual adapter 構成より軽量

docs/setup.md

@@ -0,0 +1,61 @@
+# 事前準備
+
+[← README に戻る](../README.md)
+
+## 1. Modal セットアップ
+
+```bash
+pip install modal
+modal token new
+```
+
+## 2. Hugging Face シークレット(必須)
+
+`modal_app.py` は Modal 上の secret 名 `hf_token`(キー: `HF_TOKEN`)を参照する。
+既に作成済みであれば追加作業は不要。新規の場合:
+
+```bash
+modal secret create hf_token HF_TOKEN=hf_xxxxxxxxxxxxx
+```
+
+トークンは https://huggingface.co/settings/tokens で発行。Anima は public モデルなので
+レート制限回避のために推奨。secret 名を変えたい場合は `modal_app.py` の
+`from_name("hf_token", ...)` を書き換える。
+
+### HF Hub への upload (LoRA 配布など)
+
+`hf_token` は read 専用のことが多い。配布 / repo 作成には別途 write 権限のトークンを
+`HF_TOKEN_WRITE` secret で持たせる:
+
+```bash
+modal secret create HF_TOKEN_WRITE HF_TOKEN_WRITE=hf_xxxxxx_write_xxxxx
+```
+
+`modal_app.py` の `upload_lora_to_hf` がこの secret を参照する。
+
+## 3. データセットの用意
+
+Modal に上げる前に手元で:
+
+```
+my_images/
+├── 0001.png
+├── 0001.txt   # 画像と同名 .txt にタグ/キャプション
+├── 0002.png
+├── 0002.txt
+└── ...
+```
+
+タグは Danbooru 風 (", " 区切り) を推奨。WD-EVA02-Tagger v3 などで
+自動付与してから手で軽くクリーンするのが現実的。**この時点では
+品質タグ・年タグ・メタタグを書いていても OK**(後で
+`clean_captions.py` が全部抜く)。
+
+**推奨枚数**: 3,000〜8,000 枚で十分。
+このプロジェクトの Phase 1 は「**品質タグ依存性を消すだけ**」がゴールであり、
+審美の方向性は変えない (= base モデルが本来出していた品質を tag なしで再現)。
+そのため:
+- 自分の "好み" でキュレーションする必要はない
+- Danbooru `score≥150` など **品質フィルタだけ** かけてランダムサンプル
+- artist タグも残す (`drop_artist_prob=0` デフォルト)
+- 多様性 > 美的偏向(モデルの汎用性を保つため)

docs/workflow.md

@@ -0,0 +1,210 @@
+# 実行手順 + GPU + コスト概算
+
+[← README に戻る](../README.md) — 前提: [setup.md](setup.md)
+
+## Step 0: モデルダウンロード(初回のみ)
+
+```bash
+cd anima_modal
+modal run modal_app.py::download_models
+```
+
+`anima-preview3-base.safetensors` (4.18GB) + Qwen3 + VAE が
+`anima-models` Volume に保存される。
+
+## Step 1: データセットを用意
+
+### 1a. self-distillation で自動生成(おすすめ、約 11h / $28)
+
+Anima base 自身に生成させて quality タグ依存性を消す pair を作る。
+[scripts/gen_prompts.txt](../scripts/gen_prompts.txt) の 100 プロンプト × 50 seed で
+5,000 枚生成。クライアント側でデータ用意不要。
+
+```bash
+# 試走: ComfyUI 起動 + 10 枚生成 (~10分, $0.10)
+modal run modal_app.py::generate_dataset --max-images 10
+
+# 本番(シリアル A100): 5,000 枚 (12.5h, $31)
+modal run --detach modal_app.py::generate_dataset
+
+# 本番(B200 × 10 並列): 5,000 枚 (~40分, $44)  ← 速度優先
+modal run --detach modal_app.py::generate_dataset_parallel
+
+# パラメータ調整可
+modal run --detach modal_app.py::generate_dataset \
+    --seeds-per-prompt 80 --max-images 8000
+
+# 途中から再開 (既存 .png はスキップ)
+modal run --detach modal_app.py::generate_dataset --start-from 50
+```
+
+ワークフローは [scripts/anima_workflow.json](../scripts/anima_workflow.json)、
+**Anima 公式 README の推奨設定**:
+- Sampler: `er_sde` / Scheduler: `simple` / Steps: 30 / CFG: 4.5
+- ModelSamplingAuraFlow shift: 3.0 (公式 `anima_comparison.json` 由来)
+- Positive prefix: `masterpiece, best quality, score_7, safe`
+- Negative: `worst quality, low quality, score_1, score_2, score_3, artist name`
+- タグ順: `[quality/meta/safety] [1girl/1boy] [character] [series] [artist] [general]`
+
+### 1b. ローカル画像を持ち込み
+
+```bash
+modal volume put anima-dataset ./my_images /raw
+```
+
+`/dataset/raw/` 以下に展開される。
+
+## Step 2: キャプションを掃除
+
+このプロジェクトの Phase 1 は審美シフトではないので **artist タグは残す**:
+
+```bash
+modal run modal_app.py::clean_captions
+```
+
+artist タグを drop するオプションは「Anima 平均スタイル化」用で、
+今回の目的(タグ依存性除去のみ)では基本不要:
+
+```bash
+# 必要なら確率 drop / 完全 drop も可能
+modal run modal_app.py::clean_captions --drop-artist-prob 0.5
+modal run modal_app.py::clean_captions --keep-artist false
+```
+
+`/dataset/cleaned/` に処理後の画像と .txt が出る。中を確認:
+
+```bash
+modal volume get anima-dataset cleaned/0001.txt
+```
+
+`masterpiece, best quality, score_9, year 2025, newest, highres, safe`
+が綺麗に消えていればOK。
+
+## Step 3: Phase 1 学習(quality タグ依存性除去)
+
+```bash
+modal run --detach modal_app.py::train_phase1
+```
+
+`--detach` で**バックグラウンド実行**(数時間続くため)。
+Modal ダッシュボード(modal.com/apps)で進捗確認できる。
+
+**目安時間 (A100-80GB)**:
+| データセット | 1 epoch | 2 epoch | コスト |
+|---|---|---|---|
+| 3,000 枚 | 2 h | 4 h | **$10** |
+| 5,000 枚 | 3 h | 6-7 h | **$15-18** |
+| 10,000 枚 | 6-7 h | 12-14 h | **$30-35** |
+
+デフォルトは `epochs = 2`(prior shift には十分)。H100 にすれば 1.4-1.6x 速い。
+
+**設定変更したい時**: `configs/phase1_anima.toml` を編集して
+`modal run` を再実行(Modal が image を rebuild する)。
+
+主な調整点:
+- `epochs`: 試走 1、デフォルト 2、念のため 3 まで
+  (それ以上は base 性能を損なうリスク)
+- `[optimizer].lr`: 1e-5 → 効きが弱ければ 2e-5
+- `[adapter].rank`: 32-64 で十分(概念追加でないので大きくしても無駄)
+- `gradient_accumulation_steps`: 5k 枚以下なら 2 で十分
+
+## Step 4: Phase 1 結果の検証
+
+```bash
+modal run modal_app.py::generate_samples
+```
+
+`scripts/eval_prompts.txt` の短い「品質タグなし」プロンプトで生成して
+`/output/phase1_samples/` に保存。
+※ 現状は ComfyUI ベースで生成するスケルトンなので、本番では
+ComfyUI を Modal の別 function で立てて API 経由で呼ぶ構成を別途追加することを想定。
+
+## Step 5: Phase 2 蒸留
+
+**数万円予算なら 5a (merge) を強く推奨**。5b/5c は数十時間の試行錯誤前提。
+
+### 5a. Turbo LoRA との weight merge (即動く、$0.5、10分)
+
+```bash
+# Civitai の Turbo LoRA URL を指定して 1 ファイルに合成
+modal run modal_app.py::merge_turbo_lora \
+    --turbo-url "https://civitai.com/api/download/models/<id>"
+
+# alpha 調整 (デフォルトはどちらも 1.0)
+modal run modal_app.py::merge_turbo_lora \
+    --turbo-url "..." --alpha-phase1 1.0 --alpha-turbo 0.8
+```
+
+候補の Turbo LoRA:
+- 公式 Anima Turbo: https://civitai.com/models/2560840 (CFG=1, 8-12 step)
+- RDBT-Anima (DMD2 蒸留先行例): https://civitai.com/models/2364703
+
+出力は `/output/merged/anima_phase1_plus_turbo.safetensors` (LoRA だけの
+ファイル)。ComfyUI で base + この LoRA を組めば即推論可能。
+
+### 5b. SOTA 自前蒸留(Decoupled DMD2 + TSCD + R3GAN、B200 — **未完成、要追加開発**)
+
+`scripts/distill/` に実装あり。**ただし 5 回試行して全て失敗**。詳細は
+[distillation.md](distillation.md) を参照。これから取り組む人は、まず
+そのドキュメントの §1.1 / §4 / §5 を読むことを強く推奨。
+
+```bash
+# (動くが収束しない) 我々が試したコマンド:
+modal run --detach modal_app.py::generate_dataset_parallel  # データ生成 $44
+modal run --detach modal_app.py::train_sota_distill --phase b --total-steps 2000  # 失敗
+```
+
+> **強い推奨**: 商用要件なしなら 5a で十分。Anima の Turbo LoRA は **CircleStone
+> Labs 公式** が **遥かに多い計算資源と試行錯誤** で作っており、数万円の予算で
+> これを超える自前蒸留は現実的に困難。
+
+## GPU の選び方
+
+| GPU | Modal 料金 | 速度 | 推奨用途 |
+|---|---|---|---|
+| A100-40GB | $2.10/hr | 遅 | OOM のリスク。非推奨 |
+| **A100-80GB** | **$2.50/hr** | 標準 | **Phase 1 デフォルト** |
+| L40S 48GB | $1.95/hr | やや遅 | Phase 1 試走に使える |
+| H100 80GB | $3.95/hr | 1.4-1.6x 速 | 時間最優先 / Phase 2 DMD2 |
+| H200 | H100 と同額 | さらに速 | Modal が自動で割り当てる場合あり |
+
+`modal_app.py` の `gpu="A100-80GB"` を書き換えるだけで切替可能。
+
+## コスト概算 (1USD ≒ 150円換算)
+
+### プラン A: ミニマム検証 (約 5,000円 / $33)
+
+| 項目 | 想定 | コスト |
+|---|---|---|
+| Modal Volume | 10GB × 1ヶ月 | $0.90 |
+| モデル DL (CPU 並列) | 5-10 分 | $0.01 |
+| キャプション掃除 (CPU) | 5 分 | ~$0.00 |
+| Phase 1 学習 (A100-80GB, 2,000枚 × 3 epoch) | 8 時間 | $20 |
+| 検証生成 (L40S) | 10 分 | $0.30 |
+| Phase 2 merge | 10 分 | $0.50 |
+| 余裕枠 (失敗 retry 1 回ぶん) | | ~$10 |
+| **合計** | | **~$32** |
+
+### プラン B: スタンダード (約 1〜1.2万円 / $70-80)
+
+| 項目 | 想定 | コスト |
+|---|---|---|
+| Modal Volume | 20GB × 1ヶ月 | $1.80 |
+| **Phase 1 (A100-80GB, 10,000枚 × 3 epoch)** | 18-21 時間 | **$45-53** |
+| 検証 + merge + 試行 retry | | ~$5 |
+| 余裕 | | ~$15 |
+| **合計** | | **~$70** |
+
+### プラン C: 余裕プラン / 蒸留も自前 (約 3〜4万円 / $200-270)
+
+| 項目 | 想定 | コスト |
+|---|---|---|
+| Phase 1 (A100-80GB, 10k × 5 epoch) | 30-35 時間 | $75-88 |
+| Phase 2 LCM 自前蒸留 (A100, 試行2回) | 20 時間 | $50 |
+| ↑ DMD2 にする場合 (H100, 試行1回) | 15 時間 | $60 |
+| Volume + 検証 + 余裕 | | ~$30 |
+| **合計** | | **$200-270** |
+
+> ⚠️ `--detach` で背景実行中も課金は続く。`modal app stop rapid-anima` で
+> 即停止できるので、想定時間を超えたら必ず確認。`cleanup_checkpoints` で
+> Volume 課金 ($0.15/GB/月) も適時間引き。

requirements.txt

@@ -0,0 +1,4 @@
+# Modal Image 側でインストール済み。ローカル開発時用。
+modal>=0.66
+huggingface_hub>=0.27
+Pillow>=10.0

scripts/anima_finetune_dataset_workflow.json

@@ -0,0 +1,90 @@
+{
+  "_comment": "Fine-tune 用データセット生成 workflow。base=hakushiMixAnima_v02 + LoRA=anima-highres-aesthetic-boost strength 1.0、Anima 公式推奨 30 step CFG=4.5 er_sde simple、ModelSamplingAuraFlow shift=3.0 込み。node ID は scripts/generate_dataset.py の patch_workflow hardcode に合わせて positive=5 latent=7 KSampler=8。",
+  "1": {
+    "class_type": "UNETLoader",
+    "inputs": {
+      "unet_name": "hakushiMixAnima_v02.safetensors",
+      "weight_dtype": "default"
+    }
+  },
+  "2": {
+    "class_type": "ModelSamplingAuraFlow",
+    "inputs": {
+      "model": ["1", 0],
+      "shift": 3.0
+    }
+  },
+  "2b": {
+    "class_type": "LoraLoaderModelOnly",
+    "inputs": {
+      "model": ["2", 0],
+      "lora_name": "anima-highres-aesthetic-boost.safetensors",
+      "strength_model": 1.0
+    }
+  },
+  "3": {
+    "class_type": "CLIPLoader",
+    "inputs": {
+      "clip_name": "qwen_3_06b_base.safetensors",
+      "type": "stable_diffusion",
+      "device": "default"
+    }
+  },
+  "4": {
+    "class_type": "CLIPTextEncode",
+    "inputs": {
+      "clip": ["3", 0],
+      "text": "worst quality, low quality, score_1, score_2, score_3, score_4, blurry, deformed, bad anatomy, extra limbs, fused fingers, artist name, text, watermark, lowres, jpeg artifacts, censored"
+    }
+  },
+  "5": {
+    "class_type": "CLIPTextEncode",
+    "inputs": {
+      "clip": ["3", 0],
+      "text": "$PROMPT"
+    }
+  },
+  "6": {
+    "class_type": "VAELoader",
+    "inputs": {
+      "vae_name": "qwen_image_vae.safetensors"
+    }
+  },
+  "7": {
+    "class_type": "EmptyLatentImage",
+    "inputs": {
+      "width": "$WIDTH",
+      "height": "$HEIGHT",
+      "batch_size": 1
+    }
+  },
+  "8": {
+    "class_type": "KSampler",
+    "inputs": {
+      "model": ["2b", 0],
+      "positive": ["5", 0],
+      "negative": ["4", 0],
+      "latent_image": ["7", 0],
+      "seed": "$SEED",
+      "steps": 30,
+      "cfg": 4.5,
+      "sampler_name": "er_sde",
+      "scheduler": "simple",
+      "denoise": 1.0
+    }
+  },
+  "9": {
+    "class_type": "VAEDecode",
+    "inputs": {
+      "samples": ["8", 0],
+      "vae": ["6", 0]
+    }
+  },
+  "10": {
+    "class_type": "SaveImage",
+    "inputs": {
+      "filename_prefix": "finetune_data",
+      "images": ["9", 0]
+    }
+  }
+}

scripts/anima_verify_workflow.json

@@ -0,0 +1,90 @@
+{
+  "_comment": "ユーザー指定 Anima_simple.json + Anima 公式必須の ModelSamplingAuraFlow shift=3.0 + LoraLoaderModelOnly 追加版。node ID は scripts/generate_dataset.py の patch_workflow が hardcode (positive=5, latent=7, sampler=8) するため、それに合わせて配置。LoRA を使わない場合は LoraLoaderModelOnly の strength_model を 0.0 にする。",
+  "1": {
+    "class_type": "UNETLoader",
+    "inputs": {
+      "unet_name": "anima-base-v1.0.safetensors",
+      "weight_dtype": "default"
+    }
+  },
+  "2": {
+    "class_type": "ModelSamplingAuraFlow",
+    "inputs": {
+      "model": ["1", 0],
+      "shift": 3.0
+    }
+  },
+  "2b": {
+    "class_type": "LoraLoaderModelOnly",
+    "inputs": {
+      "model": ["2", 0],
+      "lora_name": "anima_turbo.safetensors",
+      "strength_model": 1.0
+    }
+  },
+  "3": {
+    "class_type": "CLIPLoader",
+    "inputs": {
+      "clip_name": "qwen_3_06b_base.safetensors",
+      "type": "stable_diffusion",
+      "device": "default"
+    }
+  },
+  "4": {
+    "class_type": "CLIPTextEncode",
+    "inputs": {
+      "clip": ["3", 0],
+      "text": "worst quality, low quality, score_1, score_2, score_3, score_4, blurry, deformed, bad anatomy, extra limbs, fused fingers, artist name, text, watermark, lowres, jpeg artifacts, censored"
+    }
+  },
+  "5": {
+    "class_type": "CLIPTextEncode",
+    "inputs": {
+      "clip": ["3", 0],
+      "text": "$PROMPT"
+    }
+  },
+  "6": {
+    "class_type": "VAELoader",
+    "inputs": {
+      "vae_name": "qwen_image_vae.safetensors"
+    }
+  },
+  "7": {
+    "class_type": "EmptyLatentImage",
+    "inputs": {
+      "width": "$WIDTH",
+      "height": "$HEIGHT",
+      "batch_size": 1
+    }
+  },
+  "8": {
+    "class_type": "KSampler",
+    "inputs": {
+      "model": ["2b", 0],
+      "positive": ["5", 0],
+      "negative": ["4", 0],
+      "latent_image": ["7", 0],
+      "seed": "$SEED",
+      "steps": 30,
+      "cfg": 4.5,
+      "sampler_name": "er_sde",
+      "scheduler": "simple",
+      "denoise": 1.0
+    }
+  },
+  "9": {
+    "class_type": "VAEDecode",
+    "inputs": {
+      "samples": ["8", 0],
+      "vae": ["6", 0]
+    }
+  },
+  "10": {
+    "class_type": "SaveImage",
+    "inputs": {
+      "filename_prefix": "anima_verify",
+      "images": ["9", 0]
+    }
+  }
+}

scripts/anima_workflow.json

@@ -0,0 +1,82 @@
+{
+  "_comment": "Anima ComfyUI workflow (API format). Used by generate_dataset.py. Placeholders: $PROMPT, $WIDTH, $HEIGHT, $SEED are substituted at submit time.",
+  "1": {
+    "class_type": "UNETLoader",
+    "inputs": {
+      "unet_name": "anima-base-v1.0.safetensors",
+      "weight_dtype": "default"
+    }
+  },
+  "2": {
+    "class_type": "ModelSamplingAuraFlow",
+    "inputs": {
+      "model": ["1", 0],
+      "shift": 3.0
+    }
+  },
+  "3": {
+    "class_type": "CLIPLoader",
+    "inputs": {
+      "clip_name": "qwen_3_06b_base.safetensors",
+      "type": "stable_diffusion",
+      "device": "default"
+    }
+  },
+  "4": {
+    "class_type": "VAELoader",
+    "inputs": {
+      "vae_name": "qwen_image_vae.safetensors"
+    }
+  },
+  "5": {
+    "class_type": "CLIPTextEncode",
+    "inputs": {
+      "text": "$PROMPT",
+      "clip": ["3", 0]
+    }
+  },
+  "6": {
+    "class_type": "CLIPTextEncode",
+    "inputs": {
+      "text": "worst quality, low quality, score_1, score_2, score_3, artist name",
+      "clip": ["3", 0]
+    }
+  },
+  "7": {
+    "class_type": "EmptyLatentImage",
+    "inputs": {
+      "width": "$WIDTH",
+      "height": "$HEIGHT",
+      "batch_size": 1
+    }
+  },
+  "8": {
+    "class_type": "KSampler",
+    "inputs": {
+      "model": ["2", 0],
+      "positive": ["5", 0],
+      "negative": ["6", 0],
+      "latent_image": ["7", 0],
+      "seed": "$SEED",
+      "steps": 30,
+      "cfg": 4.5,
+      "sampler_name": "er_sde",
+      "scheduler": "simple",
+      "denoise": 1.0
+    }
+  },
+  "9": {
+    "class_type": "VAEDecode",
+    "inputs": {
+      "samples": ["8", 0],
+      "vae": ["4", 0]
+    }
+  },
+  "10": {
+    "class_type": "SaveImage",
+    "inputs": {
+      "images": ["9", 0],
+      "filename_prefix": "anima_gen"
+    }
+  }
+}

scripts/anima_workflow_phase_a.json

@@ -0,0 +1,82 @@
+{
+  "_comment": "Phase A 蒸留出力 (gen_final.safetensors) を base として使う workflow。LoRA 不要、CFG=1、step 数は --override-steps で切替可。",
+  "1": {
+    "class_type": "UNETLoader",
+    "inputs": {
+      "unet_name": "phase_a_distilled.safetensors",
+      "weight_dtype": "default"
+    }
+  },
+  "2": {
+    "class_type": "ModelSamplingAuraFlow",
+    "inputs": {
+      "model": ["1", 0],
+      "shift": 3.0
+    }
+  },
+  "3": {
+    "class_type": "CLIPLoader",
+    "inputs": {
+      "clip_name": "qwen_3_06b_base.safetensors",
+      "type": "stable_diffusion",
+      "device": "default"
+    }
+  },
+  "4": {
+    "class_type": "VAELoader",
+    "inputs": {
+      "vae_name": "qwen_image_vae.safetensors"
+    }
+  },
+  "5": {
+    "class_type": "CLIPTextEncode",
+    "inputs": {
+      "text": "$PROMPT",
+      "clip": ["3", 0]
+    }
+  },
+  "6": {
+    "class_type": "CLIPTextEncode",
+    "inputs": {
+      "text": "",
+      "clip": ["3", 0]
+    }
+  },
+  "7": {
+    "class_type": "EmptyLatentImage",
+    "inputs": {
+      "width": "$WIDTH",
+      "height": "$HEIGHT",
+      "batch_size": 1
+    }
+  },
+  "8": {
+    "class_type": "KSampler",
+    "inputs": {
+      "model": ["2", 0],
+      "positive": ["5", 0],
+      "negative": ["6", 0],
+      "latent_image": ["7", 0],
+      "seed": "$SEED",
+      "steps": 8,
+      "cfg": 1.0,
+      "sampler_name": "er_sde",
+      "scheduler": "simple",
+      "denoise": 1.0
+    }
+  },
+  "9": {
+    "class_type": "VAEDecode",
+    "inputs": {
+      "samples": ["8", 0],
+      "vae": ["4", 0]
+    }
+  },
+  "10": {
+    "class_type": "SaveImage",
+    "inputs": {
+      "images": ["9", 0],
+      "filename_prefix": "anima_phase_a"
+    }
+  }
+}

scripts/anima_workflow_phase_a_step1000.json

@@ -0,0 +1,32 @@
+{
+  "_comment": "Phase A step 1000 checkpoint で動作確認用",
+  "1": {
+    "class_type": "UNETLoader",
+    "inputs": {
+      "unet_name": "phase_a_step1000.safetensors",
+      "weight_dtype": "default"
+    }
+  },
+  "2": {
+    "class_type": "ModelSamplingAuraFlow",
+    "inputs": {"model": ["1", 0], "shift": 3.0}
+  },
+  "3": {
+    "class_type": "CLIPLoader",
+    "inputs": {"clip_name": "qwen_3_06b_base.safetensors", "type": "stable_diffusion", "device": "default"}
+  },
+  "4": {"class_type": "VAELoader", "inputs": {"vae_name": "qwen_image_vae.safetensors"}},
+  "5": {"class_type": "CLIPTextEncode", "inputs": {"text": "$PROMPT", "clip": ["3", 0]}},
+  "6": {"class_type": "CLIPTextEncode", "inputs": {"text": "", "clip": ["3", 0]}},
+  "7": {"class_type": "EmptyLatentImage", "inputs": {"width": "$WIDTH", "height": "$HEIGHT", "batch_size": 1}},
+  "8": {
+    "class_type": "KSampler",
+    "inputs": {
+      "model": ["2", 0], "positive": ["5", 0], "negative": ["6", 0], "latent_image": ["7", 0],
+      "seed": "$SEED", "steps": 8, "cfg": 1.0,
+      "sampler_name": "er_sde", "scheduler": "simple", "denoise": 1.0
+    }
+  },
+  "9": {"class_type": "VAEDecode", "inputs": {"samples": ["8", 0], "vae": ["4", 0]}},
+  "10": {"class_type": "SaveImage", "inputs": {"images": ["9", 0], "filename_prefix": "anima_step1000"}}
+}

scripts/anima_workflow_phase_b.json

@@ -0,0 +1,37 @@
+{
+  "_comment": "Phase B 4-step distillation LoRA + Anima base. CFG=1, 4 step (Phase C で 2 step も試す予定)",
+  "1": {
+    "class_type": "UNETLoader",
+    "inputs": {"unet_name": "anima-base-v1.0.safetensors", "weight_dtype": "default"}
+  },
+  "2": {
+    "class_type": "ModelSamplingAuraFlow",
+    "inputs": {"model": ["1", 0], "shift": 3.0}
+  },
+  "2b": {
+    "class_type": "LoraLoaderModelOnly",
+    "inputs": {
+      "model": ["2", 0],
+      "lora_name": "phase_b_4step_lora.safetensors",
+      "strength_model": 1.0
+    }
+  },
+  "3": {
+    "class_type": "CLIPLoader",
+    "inputs": {"clip_name": "qwen_3_06b_base.safetensors", "type": "stable_diffusion", "device": "default"}
+  },
+  "4": {"class_type": "VAELoader", "inputs": {"vae_name": "qwen_image_vae.safetensors"}},
+  "5": {"class_type": "CLIPTextEncode", "inputs": {"text": "$PROMPT", "clip": ["3", 0]}},
+  "6": {"class_type": "CLIPTextEncode", "inputs": {"text": "", "clip": ["3", 0]}},
+  "7": {"class_type": "EmptyLatentImage", "inputs": {"width": "$WIDTH", "height": "$HEIGHT", "batch_size": 1}},
+  "8": {
+    "class_type": "KSampler",
+    "inputs": {
+      "model": ["2b", 0], "positive": ["5", 0], "negative": ["6", 0], "latent_image": ["7", 0],
+      "seed": "$SEED", "steps": 4, "cfg": 1.0,
+      "sampler_name": "er_sde", "scheduler": "simple", "denoise": 1.0
+    }
+  },
+  "9": {"class_type": "VAEDecode", "inputs": {"samples": ["8", 0], "vae": ["4", 0]}},
+  "10": {"class_type": "SaveImage", "inputs": {"images": ["9", 0], "filename_prefix": "anima_phase_b"}}
+}

scripts/anima_workflow_turbo.json

@@ -0,0 +1,90 @@
+{
+  "_comment": "Anima + Turbo LoRA workflow. 8 step, CFG 1 (Civitai 公式推奨)。$PROMPT/$WIDTH/$HEIGHT/$SEED は generate_dataset.py が置換。",
+  "1": {
+    "class_type": "UNETLoader",
+    "inputs": {
+      "unet_name": "anima-base-v1.0.safetensors",
+      "weight_dtype": "default"
+    }
+  },
+  "2": {
+    "class_type": "ModelSamplingAuraFlow",
+    "inputs": {
+      "model": ["1", 0],
+      "shift": 3.0
+    }
+  },
+  "2b": {
+    "class_type": "LoraLoaderModelOnly",
+    "inputs": {
+      "model": ["2", 0],
+      "lora_name": "anima_turbo.safetensors",
+      "strength_model": 1.0
+    }
+  },
+  "3": {
+    "class_type": "CLIPLoader",
+    "inputs": {
+      "clip_name": "qwen_3_06b_base.safetensors",
+      "type": "stable_diffusion",
+      "device": "default"
+    }
+  },
+  "4": {
+    "class_type": "VAELoader",
+    "inputs": {
+      "vae_name": "qwen_image_vae.safetensors"
+    }
+  },
+  "5": {
+    "class_type": "CLIPTextEncode",
+    "inputs": {
+      "text": "$PROMPT",
+      "clip": ["3", 0]
+    }
+  },
+  "6": {
+    "class_type": "CLIPTextEncode",
+    "inputs": {
+      "text": "",
+      "clip": ["3", 0]
+    }
+  },
+  "7": {
+    "class_type": "EmptyLatentImage",
+    "inputs": {
+      "width": "$WIDTH",
+      "height": "$HEIGHT",
+      "batch_size": 1
+    }
+  },
+  "8": {
+    "class_type": "KSampler",
+    "inputs": {
+      "model": ["2b", 0],
+      "positive": ["5", 0],
+      "negative": ["6", 0],
+      "latent_image": ["7", 0],
+      "seed": "$SEED",
+      "steps": 8,
+      "cfg": 1.0,
+      "sampler_name": "er_sde",
+      "scheduler": "simple",
+      "denoise": 1.0
+    }
+  },
+  "9": {
+    "class_type": "VAEDecode",
+    "inputs": {
+      "samples": ["8", 0],
+      "vae": ["4", 0]
+    }
+  },
+  "10": {
+    "class_type": "SaveImage",
+    "inputs": {
+      "images": ["9", 0],
+      "filename_prefix": "anima_turbo"
+    }
+  }
+}

scripts/clean_captions.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python3
+"""
+キャプション掃除スクリプト
+==========================
+
+Anima 公式が推奨する以下のタグを ".txt" キャプションから除去:
+  - Quality (Human): masterpiece, best quality, good quality, normal quality, low quality, worst quality
+  - Quality (PonyV7 score): score_9, score_8, ..., score_1
+  - Time period: year 2025, year 2024, ..., newest, recent, mid, early, old
+  - Meta: highres, absurdres, anime screenshot, jpeg artifacts, official art,
+          lowres, blurry, watermark, signature, web address, artist name,
+          censored, uncensored 等
+
+これらをデータから抜くことで、学習後のモデルは「タグなし = 高品質」を学ぶ。
+
+使い方:
+  python clean_captions.py --input ./raw --output ./cleaned \
+      [--drop-artist-prob 0.0] [--drop-all-artists]
+"""
+import argparse
+import re
+import shutil
+import random
+from pathlib import Path
+
+
+# 削除対象タグ (大文字小文字無視、前後空白許容)
+QUALITY_TAGS_HUMAN = {
+    "masterpiece", "best quality", "good quality",
+    "normal quality", "low quality", "worst quality",
+}
+
+QUALITY_TAGS_SCORE = {f"score_{i}" for i in range(1, 10)}
+
+# 公式 README より: 期間タグ
+PERIOD_TAGS = {"newest", "recent", "mid", "early", "old"}
+
+# Year タグは正規表現で吸収
+YEAR_RE = re.compile(r"^year\s+(19|20)\d{2}$", re.IGNORECASE)
+
+META_TAGS = {
+    "highres", "absurdres", "lowres",
+    "anime screenshot", "official art",
+    "jpeg artifacts", "blurry", "bad quality",
+    "watermark", "signature", "web address", "twitter username",
+    "artist name", "logo",
+    "censored", "uncensored", "mosaic censoring", "bar censor",
+    "safe", "sensitive", "nsfw", "explicit",  # safety タグも除去 (要なら別フラグに)
+}
+
+FIXED_DROP = (
+    {t.lower() for t in QUALITY_TAGS_HUMAN}
+    | QUALITY_TAGS_SCORE
+    | PERIOD_TAGS
+    | META_TAGS
+)
+
+
+def should_drop(tag: str) -> bool:
+    t = tag.strip().lower()
+    if not t:
+        return True
+    if t in FIXED_DROP:
+        return True
+    if YEAR_RE.match(t):
+        return True
+    return False
+
+
+def clean_caption(
+    text: str,
+    drop_artist_prob: float = 0.0,
+    drop_all_artists: bool = False,
+    rng: random.Random | None = None,
+) -> str:
+    rng = rng or random.Random()
+    # ", " 区切りを基本に、Danbooru の "_" 区切りも空白へ統一されている前提
+    parts = [p.strip() for p in text.split(",")]
+    out = []
+    for p in parts:
+        if should_drop(p):
+            continue
+        # artist タグは "@xxx" 形式 (Anima 仕様)
+        if p.startswith("@"):
+            if drop_all_artists:
+                continue
+            if drop_artist_prob > 0 and rng.random() < drop_artist_prob:
+                continue
+        out.append(p)
+    return ", ".join(out)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--input", required=True, type=Path)
+    ap.add_argument("--output", required=True, type=Path)
+    ap.add_argument(
+        "--drop-artist-prob",
+        type=float,
+        default=0.0,
+        help="artist タグ(@xxx)をランダムに drop する確率 (0.0〜1.0)",
+    )
+    ap.add_argument(
+        "--drop-all-artists",
+        action="store_true",
+        help="artist タグを全部除去 (Anima 平均スタイル化)",
+    )
+    ap.add_argument(
+        "--exts",
+        default="png,jpg,jpeg,webp",
+        help="画像拡張子(カンマ区切り)",
+    )
+    args = ap.parse_args()
+
+    args.output.mkdir(parents=True, exist_ok=True)
+    exts = tuple("." + e.lower() for e in args.exts.split(","))
+    rng = random.Random(42)
+
+    n_img, n_cap, n_missing = 0, 0, 0
+    for img in args.input.rglob("*"):
+        if not img.is_file() or not img.suffix.lower() in exts:
+            continue
+        n_img += 1
+        rel = img.relative_to(args.input)
+        out_img = args.output / rel
+        out_img.parent.mkdir(parents=True, exist_ok=True)
+        if not out_img.exists():
+            shutil.copy(img, out_img)
+
+        cap_in = img.with_suffix(".txt")
+        cap_out = out_img.with_suffix(".txt")
+        if not cap_in.exists():
+            n_missing += 1
+            cap_out.write_text("", encoding="utf-8")
+            continue
+
+        text = cap_in.read_text(encoding="utf-8", errors="ignore")
+        cleaned = clean_caption(
+            text,
+            drop_artist_prob=args.drop_artist_prob,
+            drop_all_artists=args.drop_all_artists,
+            rng=rng,
+        )
+        cap_out.write_text(cleaned, encoding="utf-8")
+        n_cap += 1
+
+    print(f"[clean_captions] images={n_img}  captions={n_cap}  missing_txt={n_missing}")
+    print(f"[clean_captions] output -> {args.output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/compare_prompts.txt

@@ -0,0 +1,6 @@
+# Phase A 完了後の比較用 5 prompt (sample_prompts.txt から抜粋、絶対 index は本 file 内 0..4)
+1boy, short blonde hair, green eyes, knight armor, castle hallway, holding shield and sword, torchlight, noble pose
+2girls, friends, one with long brown hair other with short blonde, casual clothes, shopping together, mall interior, laughing, holding shopping bags
+a fluffy white cat, curled up on windowsill, sunlight, lace curtains, peaceful sleep, soft focus background
+fantasy landscape, floating islands with waterfalls, sky bridges, sunset, magical atmosphere, distant mountains
+magical spellcasting scene, glowing rune circle, swirling energy, dramatic shadows, fantasy atmosphere, ethereal

scripts/distill/anima_ladd_disc.py

@@ -0,0 +1,271 @@
+"""
+Anima LADD discriminator (AMD Nitro-1 流派、PixArt → Anima 移植)
+
+設計の根拠 (R3GAN 失敗との対比):
+  - R3GAN は CNN-on-latent を一から学習 → Anima 16ch latent の prior を持たず崩壊
+  - LADD は teacher の frozen MiniTrainDIT を **feature extractor** として再利用
+    teacher は Anima latent を理解済 → head は判別だけ学習すればよい
+  - 加えて spectral norm + BatchNormLocal で D 自体の安定性も確保
+
+Architecture:
+  - D backbone = teacher MiniTrainDIT (frozen) を deepcopy
+  - forward 時に N 個の block 出力を hook で集める (multi-scale)
+  - 各 hook 出力に DiscHead を 1 個ずつ通して logits を得る
+  - 最終: list[Tensor] of logits を返し、loss 側で BCE
+"""
+from __future__ import annotations
+from typing import List, Sequence
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import spectral_norm
+
+from .anima_loader import AnimaBundle
+
+
+# ----- BatchNormLocal -------------------------------------------------------
+class BatchNormLocal(nn.Module):
+    """各 sample 内で per-channel normalize。BatchNorm のような分布依存性なし。
+       AMD Nitro-1 で D の安定化に効いた知見。"""
+    def __init__(self, channels: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(channels))
+        self.bias = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, C, L) or (B, C)
+        if x.dim() == 2:
+            mean = x.mean(dim=1, keepdim=True)
+            var = x.var(dim=1, keepdim=True, unbiased=False)
+        else:
+            mean = x.mean(dim=2, keepdim=True)
+            var = x.var(dim=2, keepdim=True, unbiased=False)
+        x = (x - mean) / torch.sqrt(var + self.eps)
+        if x.dim() == 2:
+            return x * self.weight + self.bias
+        return x * self.weight.view(1, -1, 1) + self.bias.view(1, -1, 1)
+
+
+# ----- DiscHead -------------------------------------------------------------
+class DiscHead(nn.Module):
+    """1 つの hooked feature を scalar logit に。
+       SpectralConv1d → BNL → LeakyReLU → residual block → Linear projection
+    """
+    def __init__(self, in_dim: int, hidden: int = 512):
+        super().__init__()
+        # in: (B, L, D) → permute → (B, D, L)
+        self.conv1 = spectral_norm(nn.Conv1d(in_dim, hidden, kernel_size=1))
+        self.bnl1 = BatchNormLocal(hidden)
+        self.act1 = nn.LeakyReLU(0.2, inplace=True)
+
+        self.conv2 = spectral_norm(nn.Conv1d(hidden, hidden, kernel_size=9, padding=4, padding_mode="circular"))
+        self.bnl2 = BatchNormLocal(hidden)
+        self.act2 = nn.LeakyReLU(0.2, inplace=True)
+
+        self.proj = spectral_norm(nn.Conv1d(hidden, 1, kernel_size=1))
+
+    def forward(self, feat: torch.Tensor) -> torch.Tensor:
+        """
+        feat: (B, L, D)  ※L = num tokens、D = channel dim
+        return: (B,) per-sample logit (mean over L)
+        """
+        x = feat.transpose(1, 2)  # (B, D, L)
+        h = self.act1(self.bnl1(self.conv1(x)))
+        h = h + self.act2(self.bnl2(self.conv2(h)))
+        logits = self.proj(h)  # (B, 1, L)
+        return logits.mean(dim=2).squeeze(1)  # (B,)
+
+
+# ----- AnimaLADDDiscriminator ----------------------------------------------
+class AnimaLADDDiscriminator(nn.Module):
+    """teacher MiniTrainDIT を frozen で抱え、forward 時に N ブロックの出力を
+    hook で集めて N 個の DiscHead に通す。head の出力 (logits) を list で返す。
+
+    teacher が Anima latent (16ch × 128² → 1024 token) を理解しているので、
+    head は判別だけを学習。スクラッチ CNN より遥かに stable。
+    """
+    def __init__(
+        self,
+        teacher_transformer: nn.Module,
+        block_ids: Sequence[int] = (2, 8, 14, 20, 26),
+        head_hidden: int = 512,
+    ):
+        super().__init__()
+        # teacher 本体は frozen で保持 (forward に使う)
+        # 重要: deepcopy はしない。teacher_transformer はすでに caller 側で
+        # frozen な状態で与えられる前提 (gen の teacher と共有 OK、forward は no_grad)
+        self.teacher = teacher_transformer
+        for p in self.teacher.parameters():
+            p.requires_grad = False
+
+        # block_ids が teacher の block 数を超えていないか
+        blocks = self._find_blocks()
+        n_blocks = len(blocks)
+        self.block_ids = [bid for bid in block_ids if bid < n_blocks]
+        print(f"[ladd-D] teacher has {n_blocks} blocks; hooking {self.block_ids}")
+
+        # 各 hook 出力の channel 数を 1 度 dummy forward で取って head を作成
+        self._hook_outputs: dict[int, torch.Tensor] = {}
+        self._handles = []
+        for bid in self.block_ids:
+            h = blocks[bid].register_forward_hook(self._make_hook(bid))
+            self._handles.append(h)
+
+        # heads は dummy forward 後に lazy init
+        self.heads = nn.ModuleDict()
+        self.head_hidden = head_hidden
+
+    @torch.no_grad()
+    def lazy_init_heads(self, sample_x: torch.Tensor, sample_t: torch.Tensor, sample_cond: torch.Tensor):
+        """heads を 1 度 dummy forward して channel 数を確定させてから build。"""
+        self._hook_outputs.clear()
+        _ = self._teacher_forward(sample_x, sample_t, sample_cond)
+        # 各 hook の出力 shape を見て head を作る
+        for bid in self.block_ids:
+            out = self._hook_outputs[bid]
+            # out: (B, L, D) を想定 (MiniTrainDIT の block 出力)
+            if out.dim() == 4:  # (B, T, L, D) など video 系
+                out = out.flatten(1, -2)  # (B, T*L, D)
+            B, L, D = out.shape
+            print(f"[ladd-D] block {bid} feature shape: B={B} L={L} D={D}")
+            head = DiscHead(in_dim=D, hidden=self.head_hidden)
+            self.heads[str(bid)] = head
+        self._hook_outputs.clear()
+
+    def _find_blocks(self) -> list[nn.Module]:
+        """MiniTrainDIT 内部の **DiT 本体 block list** を見つける。
+        Anima/Cosmos の DiT blocks は ~28 個、llm_adapter.blocks は 6 個程度なので
+        区別する必要がある。優先順位:
+          1. transformer.blocks (top-level、direct attribute) を最優先
+          2. それ以外で、llm_adapter を含まない blocks ModuleList で最大のもの
+        """
+        # 1) 最優先: top-level の `.blocks`
+        if hasattr(self.teacher, "blocks") and isinstance(self.teacher.blocks, nn.ModuleList):
+            n = len(self.teacher.blocks)
+            print(f"[ladd-D] found top-level teacher.blocks ({n} blocks)")
+            return list(self.teacher.blocks)
+
+        # 2) llm_adapter を含まない候補から最大のものを採用
+        candidates: list[tuple[str, list[nn.Module]]] = []
+        for name, module in self.teacher.named_modules():
+            if not isinstance(module, nn.ModuleList):
+                continue
+            if "block" not in name.lower():
+                continue
+            if "llm_adapter" in name:
+                continue  # text encoder の bridge ブロックを除外
+            if len(module) >= 4:
+                candidates.append((name, list(module)))
+
+        if not candidates:
+            raise RuntimeError(
+                "Could not find DiT blocks ModuleList. Inspected teacher module tree but "
+                "found no non-llm_adapter `blocks`-named list with >=4 children."
+            )
+        # 最大数の block list を採用 (DiT 本体 28+ blocks > llm_adapter の 6 blocks)
+        candidates.sort(key=lambda x: -len(x[1]))
+        name, children = candidates[0]
+        print(f"[ladd-D] selected block list at: {name} ({len(children)} blocks)")
+        return children
+
+    def _make_hook(self, bid: int):
+        def _h(module, inputs, output):
+            # output の正規化: tuple なら最初の要素、5D なら 3D に flatten
+            if isinstance(output, tuple):
+                out = output[0]
+            else:
+                out = output
+            if out.dim() == 5:
+                # (B, T, H, W, D) → (B, T*H*W, D)
+                B = out.size(0)
+                D = out.size(-1)
+                out = out.reshape(B, -1, D)
+            elif out.dim() == 4:
+                # (B, H, W, D) → (B, H*W, D)
+                B = out.size(0)
+                D = out.size(-1)
+                out = out.reshape(B, -1, D)
+            self._hook_outputs[bid] = out
+        return _h
+
+    def _teacher_forward(self, x: torch.Tensor, t: torch.Tensor, cond: torch.Tensor):
+        """teacher を forward して hook を発火させる。返り値は使わない (hook が中身)。"""
+        return AnimaBundle.dit_forward(self.teacher, x, t, cond)
+
+    def forward(
+        self, x: torch.Tensor, t: torch.Tensor, cond: torch.Tensor,
+        gradient_to_input: bool = False,
+    ) -> List[torch.Tensor]:
+        """各 head の logit を list で返す。
+
+        gradient_to_input=True (G phase で adv loss が student を訓練するために必要):
+          teacher は frozen だが、forward を grad-on で走らせ activation を残す。
+          そうすると logits の grad が teacher backbone を通って x → student LoRA に届く。
+        gradient_to_input=False (D phase: heads だけ訓練、x は固定):
+          teacher を no_grad で走らせ、features を detach (省メモリ + 安全)。
+        """
+        self._hook_outputs.clear()
+        if gradient_to_input:
+            # teacher params は frozen なので weights は更新されないが、
+            # activations が残り gradient は input → student に流れる
+            _ = self._teacher_forward(x, t, cond)
+        else:
+            with torch.no_grad():
+                _ = self._teacher_forward(x, t, cond)
+        logits = []
+        for bid in self.block_ids:
+            feat = self._hook_outputs[bid]
+            if not gradient_to_input:
+                feat = feat.detach().requires_grad_(False)
+            # feat の dtype を head と揃える (head は fp32、bf16 混在を避ける)
+            head = self.heads[str(bid)]
+            head_dtype = next(head.parameters()).dtype
+            logits.append(head(feat.to(dtype=head_dtype)))
+        self._hook_outputs.clear()
+        return logits
+
+    def trainable_parameters(self):
+        return list(self.heads.parameters())
+
+
+# ----- losses ---------------------------------------------------------------
+def ladd_d_loss(
+    logits_real: List[torch.Tensor], logits_fake: List[torch.Tensor],
+) -> tuple[torch.Tensor, dict]:
+    """BCE on each head, sum across heads."""
+    loss_real, loss_fake = 0.0, 0.0
+    for lr in logits_real:
+        loss_real = loss_real + F.binary_cross_entropy_with_logits(
+            lr, torch.ones_like(lr)
+        )
+    for lf in logits_fake:
+        loss_fake = loss_fake + F.binary_cross_entropy_with_logits(
+            lf, torch.zeros_like(lf)
+        )
+    loss = (loss_real + loss_fake) / max(len(logits_real), 1)
+    return loss, {
+        "l_d_real": loss_real.detach() / max(len(logits_real), 1),
+        "l_d_fake": loss_fake.detach() / max(len(logits_fake), 1),
+        "l_d_total": loss.detach(),
+    }
+
+
+def ladd_g_adv_loss(logits_fake: List[torch.Tensor]) -> tuple[torch.Tensor, dict]:
+    """non-saturating: BCE with label=1 (D に real と思わせる)"""
+    loss = 0.0
+    for lf in logits_fake:
+        loss = loss + F.binary_cross_entropy_with_logits(
+            lf, torch.ones_like(lf)
+        )
+    loss = loss / max(len(logits_fake), 1)
+    return loss, {"l_g_adv": loss.detach()}
+
+
+def ladd_g_recon_loss(
+    x0_hat: torch.Tensor, x0_teacher: torch.Tensor,
+) -> tuple[torch.Tensor, dict]:
+    """smooth L1 = Huber loss、mean collapse 防止の anchor。"""
+    loss = F.smooth_l1_loss(x0_hat.float(), x0_teacher.detach().float())
+    return loss, {"l_g_recon": loss.detach()}

scripts/distill/anima_loader.py

@@ -0,0 +1,275 @@
+"""
+Anima loader: diffusion-pipe の cosmos_predict2 Pipeline を活用して
+DiT / VAE / Qwen3 text encoder / LLM adapter をロードする薄いラッパ。
+
+diffusion-pipe 本体には sampling/inference 関数が無いので、ここで
+- text_encode():  prompts -> crossattn_emb (LLM adapter 通過後)
+- vae_encode():   pixels -> latents
+- vae_decode():   latents -> pixels
+- velocity():     DiT forward (rectified flow velocity prediction)
+- add_noise():    noisy = (1-t)*latents + t*noise  (rectified flow forward)
+- euler_step():   inference 用 1-step 前進
+までを実装する。
+"""
+from __future__ import annotations
+import os
+import sys
+from pathlib import Path
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+
+# diffusion-pipe を import path に強制配置。
+# 重要: site-packages に top-level `utils` を持つ別パッケージがある場合、
+# それが namespace package 化して shadow するため、(1) /workspace/diffusion-pipe を
+# 先頭に固定 (2) cached `utils` モジュールを sys.modules から除外 する。
+_DPIPE = Path("/workspace/diffusion-pipe")
+if _DPIPE.exists():
+    _dp = str(_DPIPE)
+    sys.path = [_dp] + [p for p in sys.path if p != _dp]
+    # 既にロード済みの top-level utils / models をクリア(他パッケージの shadow を排除)
+    for _mod in list(sys.modules.keys()):
+        if _mod == "utils" or _mod.startswith("utils.") \
+           or _mod == "models" or _mod.startswith("models."):
+            del sys.modules[_mod]
+
+
+@dataclass
+class AnimaPaths:
+    transformer: str = "/models/checkpoints/anima-base-v1.0.safetensors"
+    vae: str = "/models/checkpoints/qwen_image_vae.safetensors"
+    llm: str = "/models/checkpoints/qwen_3_06b_base.safetensors"
+
+
+class AnimaBundle:
+    """Anima のモデル一式を保持。.transformer / .vae / .text_encoder / .tokenizer 等を直接触る。"""
+
+    def __init__(self, pipeline, device: str | torch.device = "cuda"):
+        self.pipeline = pipeline  # diffusion-pipe の CosmosPredict2Pipeline
+        self.device = torch.device(device)
+
+        self.transformer = pipeline.transformer
+        self.vae = pipeline.vae
+        self.text_encoder = pipeline.text_encoder  # Qwen3Model (inner)
+        self.qwen_tokenizer = pipeline.tokenizer
+        # LLM adapter は transformer 内部にある (Anima のみ)
+        self.llm_adapter = getattr(pipeline.transformer, "llm_adapter", None)
+        self.t5_tokenizer = getattr(pipeline, "t5_tokenizer", None)
+        self.is_generic_llm = getattr(pipeline, "is_generic_llm", False)
+        self.vae_scale = pipeline.vae.scale
+
+    # ---- text encoding (Qwen3 → LLM adapter) ------------------------------
+    @torch.no_grad()
+    def text_encode(self, prompts: list[str]) -> torch.Tensor:
+        """prompts -> crossattn_emb (B, 512, 1024) (DiT 用、LLM adapter 通過後)"""
+        # diffusion-pipe の _tokenize を再現
+        def _tok(tokenizer, prompts):
+            return tokenizer(prompts, return_tensors="pt", truncation=True,
+                              padding="max_length", max_length=512)
+
+        qwen_enc = _tok(self.qwen_tokenizer, prompts)
+        input_ids = qwen_enc.input_ids.to(self.device)
+        attn_mask = qwen_enc.attention_mask.to(self.device)
+
+        outputs = self.text_encoder(input_ids=input_ids, attention_mask=attn_mask)
+        encoded = outputs.last_hidden_state
+        encoded = encoded.masked_fill(~attn_mask.bool().unsqueeze(-1), 0.0)
+
+        if self.llm_adapter is None or self.t5_tokenizer is None:
+            return encoded
+
+        t5_enc = _tok(self.t5_tokenizer, prompts)
+        t5_ids = t5_enc.input_ids.to(self.device)
+        t5_mask = t5_enc.attention_mask.to(self.device)
+
+        crossattn = self.llm_adapter(
+            source_hidden_states=encoded,
+            target_input_ids=t5_ids,
+            target_attention_mask=t5_mask,
+            source_attention_mask=attn_mask,
+        )
+        crossattn = crossattn.masked_fill(~t5_mask.bool().unsqueeze(-1), 0.0)
+        return crossattn
+
+    # ---- VAE --------------------------------------------------------------
+    @torch.no_grad()
+    def vae_encode(self, pixels: torch.Tensor) -> torch.Tensor:
+        """pixels (B, 3, H, W) in [-1, 1] -> latents (B, 16, 1, H/8, W/8)
+        ※ Anima は静止画でも T=1 の 3D latent を扱う"""
+        if pixels.dim() == 4:
+            pixels = pixels.unsqueeze(2)  # (B, 3, 1, H, W)
+        # VAE の weights dtype に合わせる
+        vae_dtype = next(self.vae.model.parameters()).dtype
+        pixels = pixels.to(device=self.device, dtype=vae_dtype)
+        return self.vae.model.encode(pixels, self.vae_scale)
+
+    @torch.no_grad()
+    def vae_decode(self, latents: torch.Tensor) -> torch.Tensor:
+        """latents (B, 16, 1, H_lat, W_lat) -> pixels (B, 3, 1, H, W) in [-1, 1]"""
+        vae_dtype = next(self.vae.model.parameters()).dtype
+        latents = latents.to(device=self.device, dtype=vae_dtype)
+        return self.vae.model.decode(latents, self.vae_scale)
+
+    # ---- DiT (rectified flow velocity prediction) -------------------------
+    def velocity(
+        self,
+        latents: torch.Tensor,           # (B, 16, T, H_lat, W_lat) - noisy
+        timesteps: torch.Tensor,         # (B,) in [0, 1]
+        crossattn_emb: torch.Tensor,     # (B, 512, 1024)
+        padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Anima DiT forward. 返り値は velocity (B, 16, T, H_lat, W_lat)。"""
+        if padding_mask is None:
+            padding_mask = self.zero_padding_mask(latents)
+        return self.transformer(
+            x_B_C_T_H_W=latents,
+            timesteps_B_T=timesteps,
+            crossattn_emb=crossattn_emb,
+            padding_mask=padding_mask,
+        )
+
+    @staticmethod
+    def zero_padding_mask(latents: torch.Tensor) -> torch.Tensor:
+        """MiniTrainDIT.concat_padding_mask=True なので必須。shape=(B, 1, H_lat, W_lat)。"""
+        B, _, _, H, W = latents.shape
+        return torch.zeros(B, 1, H, W, dtype=latents.dtype, device=latents.device)
+
+    @staticmethod
+    def dit_forward(
+        transformer: "torch.nn.Module",
+        latents: torch.Tensor,
+        t: torch.Tensor,
+        crossattn_emb: torch.Tensor,
+    ) -> torch.Tensor:
+        """trainer から任意の transformer (gen / guidance) を呼ぶための薄い wrapper。
+        全入力を transformer の weights dtype に揃え、padding_mask も自動生成。"""
+        # weight dtype に揃える (t は float なので broadcast で他がアップキャストされやすい)
+        w_dtype = next(transformer.parameters()).dtype
+        latents = latents.to(dtype=w_dtype)
+        t = t.to(dtype=w_dtype)
+        crossattn_emb = crossattn_emb.to(dtype=w_dtype)
+        padding_mask = AnimaBundle.zero_padding_mask(latents)
+        return transformer(
+            x_B_C_T_H_W=latents,
+            timesteps_B_T=t,
+            crossattn_emb=crossattn_emb,
+            padding_mask=padding_mask,
+        )
+
+    # ---- rectified-flow utilities ----------------------------------------
+    @staticmethod
+    def add_noise(
+        latents: torch.Tensor, noise: torch.Tensor, t: torch.Tensor
+    ) -> torch.Tensor:
+        """forward process: x_t = (1-t)*x_0 + t*noise"""
+        t_ = t.view(-1, *([1] * (latents.dim() - 1)))
+        return (1 - t_) * latents + t_ * noise
+
+    @staticmethod
+    def velocity_target(latents: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
+        """target velocity = noise - latents (rectified flow の training target)"""
+        return noise - latents
+
+    @staticmethod
+    def x0_from_velocity(
+        latents_t: torch.Tensor, v_pred: torch.Tensor, t: torch.Tensor
+    ) -> torch.Tensor:
+        """x_0 estimate = x_t - t * v_pred (linear FM のため)"""
+        t_ = t.view(-1, *([1] * (latents_t.dim() - 1)))
+        return latents_t - t_ * v_pred
+
+    @staticmethod
+    def euler_step(
+        latents_t: torch.Tensor,
+        v_pred: torch.Tensor,
+        t: torch.Tensor,
+        t_next: torch.Tensor,
+    ) -> torch.Tensor:
+        """Euler step: x_{t_next} = x_t + (t_next - t) * v_pred"""
+        dt = (t_next - t).view(-1, *([1] * (latents_t.dim() - 1)))
+        return latents_t + dt * v_pred
+
+
+def build_anima(
+    paths: AnimaPaths | None = None,
+    device: str | torch.device = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+) -> AnimaBundle:
+    """diffusion-pipe の CosmosPredict2Pipeline を初期化して AnimaBundle で返す。
+
+    内部的には toml config を最小限作って Pipeline に渡す。
+    """
+    import importlib
+    import importlib.util
+    _dp = str(_DPIPE)
+
+    # diffusion-pipe の utils/ と models/ に空 __init__.py を作成して regular package 化
+    for sub in ("utils", "models", "optimizers"):
+        sub_dir = os.path.join(_dp, sub)
+        init_file = os.path.join(sub_dir, "__init__.py")
+        if os.path.isdir(sub_dir) and not os.path.exists(init_file):
+            with open(init_file, "w"):
+                pass
+
+    # sys.path 強制 + cache クリア
+    sys.path = [_dp] + [p for p in sys.path if p != _dp]
+    for _mod in list(sys.modules.keys()):
+        if _mod in ("utils", "models", "optimizers") \
+           or _mod.startswith(("utils.", "models.", "optimizers.")):
+            del sys.modules[_mod]
+
+    # importlib で **明示的に** 各 package を /workspace/diffusion-pipe 配下から
+    # 読み込み、sys.modules に固定して shadow を完全排除する。
+    def _load_pkg(name: str):
+        pkg_dir = os.path.join(_dp, name)
+        init_file = os.path.join(pkg_dir, "__init__.py")
+        spec = importlib.util.spec_from_file_location(
+            name, init_file,
+            submodule_search_locations=[pkg_dir],
+        )
+        mod = importlib.util.module_from_spec(spec)
+        sys.modules[name] = mod
+        spec.loader.exec_module(mod)
+        print(f"[setup] forced load {name} from {pkg_dir}")
+
+    for _pkg in ("utils", "models", "optimizers"):
+        _load_pkg(_pkg)
+
+    # 動作確認: utils.common が解決できることを試す
+    import utils.common  # noqa: F401
+    print("[setup] utils.common import OK")
+
+    from models import cosmos_predict2  # diffusion-pipe 側
+
+    paths = paths or AnimaPaths()
+
+    # CosmosPredict2Pipeline は config dict から **torch.dtype object** を期待する
+    # (TOML 読み込み時に diffusion-pipe 側で str→dtype 変換するが、ここでは直接渡す)
+    cfg = {
+        "model": {
+            "type": "anima",
+            "transformer_path": str(paths.transformer),
+            "vae_path": str(paths.vae),
+            "llm_path": str(paths.llm),
+            "dtype": dtype,
+            "transformer_dtype": dtype,
+            "timestep_sample_method": "logit_normal",
+        },
+        # adapter / optimizer 等は trainer 側で attach するので空でOK
+        "adapter": {"type": "lora", "rank": 1, "alpha": 1, "dropout": 0.0, "dtype": dtype},
+    }
+
+    # __init__ で VAE + text encoder + tokenizers が CPU 上に load 済み
+    pipeline = cosmos_predict2.CosmosPredict2Pipeline(cfg)
+    # DiT をロード (transformer.llm_adapter も内部で attach される)
+    pipeline.load_diffusion_model()
+
+    # device / dtype に配置
+    pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
+    pipeline.text_encoder = pipeline.text_encoder.to(device=device, dtype=dtype)
+    pipeline.vae.model = pipeline.vae.model.to(device=device, dtype=dtype)
+    # mean/std は __init__ で 'cuda' に moved 済 (cosmos_predict2.py:200-201)
+
+    return AnimaBundle(pipeline, device=device)

scripts/distill/dataset.py

@@ -0,0 +1,84 @@
+"""
+キャプション付き画像データセット。/dataset/raw 以下の .png + .txt ペアを返す。
+DMD2 では「real image を見る」のは fake-score 更新時 (の v-pred MSE) のみで
+あって、generator/real-score 側は noise から直接サンプル → image は不要。
+従ってここでは pixels + caption を返すシンプルな実装で十分。
+"""
+from __future__ import annotations
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from PIL import Image
+
+
+def _resize_short_side(img: Image.Image, target: int) -> Image.Image:
+    """短辺を target に合わせる aspect-preserving resize"""
+    w, h = img.size
+    if min(w, h) == target:
+        return img
+    if w < h:
+        new_w = target
+        new_h = int(round(h * target / w))
+    else:
+        new_h = target
+        new_w = int(round(w * target / h))
+    return img.resize((new_w, new_h), Image.BICUBIC)
+
+
+def _center_crop(img: Image.Image, size: int) -> Image.Image:
+    w, h = img.size
+    left = (w - size) // 2
+    top = (h - size) // 2
+    return img.crop((left, top, left + size, top + size))
+
+
+def _to_tensor_normalize(img: Image.Image) -> torch.Tensor:
+    """PIL RGB -> torch (3, H, W) in [-1, 1]"""
+    arr = np.asarray(img, dtype=np.float32) / 127.5 - 1.0  # (H, W, 3) in [-1, 1]
+    return torch.from_numpy(arr).permute(2, 0, 1).contiguous()
+
+
+class AnimaImageCaptionDataset(Dataset):
+    """
+    Args:
+        root: 画像/キャプションのルート。再帰的に *.png + 同名.txt を拾う。
+        resolution: 単一解像度に統一する出力サイズ(短辺合わせ → center crop)。
+    """
+
+    def __init__(
+        self,
+        root: str | Path,
+        resolution: int = 1024,
+        exts: tuple[str, ...] = (".png", ".jpg", ".jpeg", ".webp"),
+    ):
+        self.root = Path(root)
+        self.resolution = resolution
+        self.items: list[tuple[Path, Path]] = []
+        for img in sorted(self.root.rglob("*")):
+            if not img.is_file() or img.suffix.lower() not in exts:
+                continue
+            cap = img.with_suffix(".txt")
+            if cap.exists():
+                self.items.append((img, cap))
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __getitem__(self, idx: int) -> dict:
+        img_path, cap_path = self.items[idx]
+        img = Image.open(img_path).convert("RGB")
+        img = _resize_short_side(img, self.resolution)
+        img = _center_crop(img, self.resolution)
+        pixels = _to_tensor_normalize(img)  # (3, H, W) in [-1, 1]
+        caption = cap_path.read_text(encoding="utf-8").strip()
+        return {"pixels": pixels, "caption": caption, "path": str(img_path)}
+
+
+def collate_fn(batch: list[dict]) -> dict:
+    """default collate と同じだが captions は list で保持。"""
+    pixels = torch.stack([b["pixels"] for b in batch])
+    captions = [b["caption"] for b in batch]
+    paths = [b["path"] for b in batch]
+    return {"pixels": pixels, "captions": captions, "paths": paths}

scripts/distill/dmd2_official_loss.py

@@ -0,0 +1,244 @@
+"""
+Anima 用 DMD2 + TrigFlow distillation (NVIDIA cosmos-predict2.5 official 流派)
+
+参照: cosmos_predict2/_src/predict2/distill/models/video2world_model_distill_dmd2.py
+       (lines 165-310 の step_generator / step_critic を直接移植)
+
+要点:
+  - TrigFlow time sampling: s ∈ shifted_uniform(shift=5) → t = arctan(s/(1-s))
+  - DMD2 gradient trick: loss = ((x_hat - (x_hat - grad).detach())**2).mean()
+        where grad = (x0_fake - x0_teacher_cfg) / |x_hat - x0_teacher_cfg|.mean()
+  - Critic phase weighted denoise: loss = ((x_hat - x0_fake)**2 / sin(t)**2).mean()
+  - Alternation: critic × N (=5), generator × 1
+  - Few-step student rollout (n=1..4), **grad は最終 step のみ** (memory efficient)
+
+Anima の差分 (vs upstream):
+  - upstream は video2world、condition.is_video 経路あり → image only に折り畳む (T=1)
+  - LoRA-only: student/fake_score それぞれに別 PEFT adapter
+  - rectified flow (sigma_data=1.0 想定)、本 file の sCM wrapper 計算は upstream と同じ
+"""
+from __future__ import annotations
+import math
+from typing import Callable
+
+import torch
+import torch.nn.functional as F
+
+from .anima_loader import AnimaBundle
+
+
+# ----- TrigFlow / sCM coefficient helpers (sigma_data=1.0, rectified flow) ---
+
+def trigflow_t_to_sigma(t: torch.Tensor) -> torch.Tensor:
+    """trigflow t ∈ (0, π/2) を rectified-flow sigma に変換。
+       sigma = tan(t)、sigma_data=1.0 で normalize。"""
+    return torch.tan(t)
+
+
+def rf_s_to_trigflow_t(s: torch.Tensor) -> torch.Tensor:
+    """rectified-flow time s ∈ [0, 1] を trigflow t = arctan(s / (1 - s)) に変換。"""
+    s = s.clamp(1e-6, 1.0 - 1e-6)
+    return torch.atan(s / (1.0 - s))
+
+
+def sample_shifted_uniform_s(B: int, shift: float = 5.0, device=None, dtype=torch.float32) -> torch.Tensor:
+    """shifted uniform on [0,1]: u ~ U(0,1) → s = shift*u / (1 + (shift-1)*u)。
+       upstream sample_shifted_uniform 相当。shift>1 で low-s (high-noise) 側にシフト。"""
+    u = torch.rand(B, device=device, dtype=dtype)
+    return shift * u / (1.0 + (shift - 1.0) * u)
+
+
+def scm_coefficients(t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """sigma_data=1.0 の RectifiedFlow_sCMWrapper:
+       c_skip = cos(t)、c_out = -sin(t)、c_in = 1 / sqrt(sigma**2 + 1) = cos(t)、c_noise = t
+
+       使い方:
+         x_t  = c_in.recip() * x_t_raw  (DiT 入力前の正規化、cos(t) スケール)
+         x_t_raw = c_skip * x_t + c_out * v_pred  (x0 復元)
+       簡単化のため Anima では x_t を直接渡し、v_pred を直接消費。
+    """
+    c_skip = torch.cos(t)
+    c_out  = -torch.sin(t)
+    c_in   = torch.cos(t)  # 1/sqrt(tan(t)**2+1) = cos(t)
+    c_noise = t
+    return c_skip, c_out, c_in, c_noise
+
+
+# ----- noise / re-noise (rectified flow native, sigma_data=1.0) -------------
+
+def renoise_rf(x0: torch.Tensor, t_rf: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
+    """rectified flow forward: x_t = (1-t)*x0 + t*noise。t_rf は ∈ [0,1]。"""
+    t_ = t_rf.view(-1, *([1] * (x0.dim() - 1)))
+    return (1.0 - t_) * x0 + t_ * noise
+
+
+def x0_from_velocity_rf(x_t: torch.Tensor, v: torch.Tensor, t_rf: torch.Tensor) -> torch.Tensor:
+    """v = noise - x0 から x0 を復元: x0 = x_t - t * v"""
+    t_ = t_rf.view(-1, *([1] * (x_t.dim() - 1)))
+    return x_t - t_ * v
+
+
+# ----- few-step student rollout --------------------------------------------
+
+def backward_simulation(
+    student_v_fn: Callable[..., torch.Tensor],
+    noise: torch.Tensor,
+    n_steps: int,
+    cond_pos: torch.Tensor,
+    cond_neg: torch.Tensor | None,
+    cfg_scale: float = 1.0,
+    grad_last_only: bool = True,
+) -> torch.Tensor:
+    """t=1 (pure noise) から t=0 (clean) へ n_steps の Euler でロールアウト。
+    grad_last_only=True なら最終 1 step だけ grad を通し、それ以前は no_grad。
+    DMD2 paper 流派、memory efficient。
+
+    Returns: x_hat at t=0 (clean estimate)
+    """
+    from .traj_loss import cfg_guided, _broadcast_t
+
+    B = noise.size(0)
+    device = noise.device
+    dtype = noise.dtype
+
+    # n_steps Euler points (t=1 → t=0)、両端含めて n+1 個
+    ts = torch.linspace(1.0, 0.0, n_steps + 1, device=device, dtype=torch.float32)
+
+    x = noise
+    for i in range(n_steps):
+        t_cur = ts[i]
+        t_next = ts[i + 1]
+        is_last = (i == n_steps - 1)
+        ctx = torch.enable_grad() if (is_last and grad_last_only) else \
+              (torch.no_grad() if grad_last_only else torch.enable_grad())
+        t_in = _broadcast_t(t_cur, B, device, dtype)
+        with ctx:
+            v = cfg_guided(student_v_fn, x, t_in, cond_pos, cond_neg, cfg_scale)
+            dt = (t_next - t_cur).to(device=device, dtype=dtype)
+            x = x + dt * v
+    return x
+
+
+# ----- generator phase ------------------------------------------------------
+
+def dmd2_generator_loss(
+    student_v_fn: Callable[..., torch.Tensor],   # PEFT student (LoRA active)
+    teacher_v_fn: Callable[..., torch.Tensor],   # frozen base, no LoRA
+    fake_score_v_fn: Callable[..., torch.Tensor],  # PEFT fake_score (LoRA active, no grad here)
+    init_noise: torch.Tensor,
+    cond_pos: torch.Tensor,
+    cond_neg: torch.Tensor | None,
+    teacher_cfg: float = 3.0,
+    student_cfg: float = 1.0,
+    n_steps: int = 4,
+    shift: float = 5.0,
+) -> tuple[torch.Tensor, dict]:
+    """DMD2 generator step。
+       - student で few-step rollout、x_hat 推定 (grad は最終 step のみ)
+       - x_hat を t_D で re-noise
+       - teacher で CFG'd x0 推定
+       - fake_score で x0_fake 推定 (no_grad)
+       - DMD2 gradient trick で loss 計算
+    """
+    from .traj_loss import _broadcast_t
+
+    B = init_noise.size(0)
+    device = init_noise.device
+    dtype = init_noise.dtype
+
+    # 1) few-step student rollout (grad only last)
+    x_hat = backward_simulation(
+        student_v_fn, init_noise, n_steps, cond_pos, cond_neg,
+        cfg_scale=student_cfg, grad_last_only=True,
+    )
+
+    # 2) t_D for re-noise: shifted uniform on [0,1]
+    s = sample_shifted_uniform_s(B, shift=shift, device=device, dtype=torch.float32)
+    t_D = s  # rectified flow time domain
+    D_eps = torch.randn_like(x_hat)
+    x_t = renoise_rf(x_hat, t_D, D_eps)
+
+    # 3) teacher CFG'd velocity → x0 estimate
+    with torch.no_grad():
+        t_in = t_D.to(dtype=dtype)
+        v_cond = teacher_v_fn(x_t, t_in, cond_pos)
+        if teacher_cfg > 1.0 and cond_neg is not None:
+            v_uncond = teacher_v_fn(x_t, t_in, cond_neg)
+            v_teacher = v_uncond + teacher_cfg * (v_cond - v_uncond)
+        else:
+            v_teacher = v_cond
+        x0_teacher = x0_from_velocity_rf(x_t, v_teacher, t_D)
+
+        # 4) fake_score velocity → x0_fake estimate (no_grad in gen phase)
+        v_fake = fake_score_v_fn(x_t, t_in, cond_pos)
+        x0_fake = x0_from_velocity_rf(x_t, v_fake, t_D)
+
+    # 5) DMD2 gradient trick
+    diff = x_hat - x0_teacher
+    # per-sample L1 normalization (sCM2 weight)
+    w = diff.abs().mean(dim=list(range(1, diff.dim())), keepdim=True).clamp(min=1e-5)
+    grad = (x0_fake.detach() - x0_teacher.detach()) / w
+    loss = ((x_hat - (x_hat - grad).detach()).float() ** 2).mean()
+
+    metrics = {
+        "l_dmd_gen": loss.detach(),
+        "t_D_mean": t_D.mean().detach(),
+        "x_hat_abs_mean": x_hat.detach().abs().mean(),
+        "teacher_abs_mean": x0_teacher.detach().abs().mean(),
+        "fake_abs_mean": x0_fake.detach().abs().mean(),
+        "w_mean": w.detach().mean(),
+    }
+    return loss, metrics
+
+
+# ----- critic (fake_score) phase -------------------------------------------
+
+def dmd2_critic_loss(
+    student_v_fn: Callable[..., torch.Tensor],
+    fake_score_v_fn: Callable[..., torch.Tensor],
+    init_noise: torch.Tensor,
+    cond_pos: torch.Tensor,
+    cond_neg: torch.Tensor | None,
+    student_cfg: float = 1.0,
+    n_steps: int = 4,
+    shift: float = 5.0,
+) -> tuple[torch.Tensor, dict]:
+    """fake_score を student 分布の denoiser として学習。
+       weighted denoising loss with `1 / sin(t)**2` factor (上流の trigflow weighting に等価)。"""
+    from .traj_loss import _broadcast_t
+
+    B = init_noise.size(0)
+    device = init_noise.device
+    dtype = init_noise.dtype
+
+    # 1) student rollout (全部 no_grad: critic phase では student は固定扱い)
+    with torch.no_grad():
+        x_hat = backward_simulation(
+            student_v_fn, init_noise, n_steps, cond_pos, cond_neg,
+            cfg_scale=student_cfg, grad_last_only=False,
+        )
+
+    # 2) t_D で re-noise
+    s = sample_shifted_uniform_s(B, shift=shift, device=device, dtype=torch.float32)
+    t_D = s
+    D_eps = torch.randn_like(x_hat)
+    x_t = renoise_rf(x_hat, t_D, D_eps)
+
+    # 3) fake_score forward (grad on)
+    t_in = t_D.to(dtype=dtype)
+    v_fake = fake_score_v_fn(x_t, t_in, cond_pos)
+    x0_fake = x0_from_velocity_rf(x_t, v_fake, t_D)
+
+    # 4) weighted denoising loss: target = x_hat (= student output)
+    #    weight = 1 / sin(arctan(s/(1-s)))**2、small s → big weight (high-noise emphasized)
+    t_trig = rf_s_to_trigflow_t(t_D)
+    w = 1.0 / torch.sin(t_trig).clamp(min=1e-3) ** 2
+    w = w.view(-1, *([1] * (x_hat.dim() - 1)))
+    loss = (w * (x_hat.detach() - x0_fake).float() ** 2).mean()
+
+    metrics = {
+        "l_dmd_critic": loss.detach(),
+        "t_D_mean": t_D.mean().detach(),
+        "w_mean": w.detach().mean(),
+    }
+    return loss, metrics

scripts/distill/dmd2_trainer.py

@@ -0,0 +1,430 @@
+"""
+DMD2 trainer for Anima (DMDR pattern + R3GAN + TSCD)
+====================================================
+
+DMDR paper の "2 モデル + LoRA on/off で real/fake 切替" 方式を踏襲。
+ここに R3GAN discriminator と TSCD consistency loss を統合する。
+
+Models:
+  gen_model       : Anima DiT (full trainable, no LoRA)
+  guidance_model  : Anima DiT (Q,V LoRA, scale 切替で real/fake)
+  discriminator   : R3GAN projection D (text-conditional)
+  ema_gen         : gen の EMA (TSCD target)
+
+Update schedule per outer step:
+  - guidance × N_GUIDANCE   (default 5)
+  - generator × 1
+  - discriminator × 1
+"""
+from __future__ import annotations
+import copy
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import peft
+
+from .anima_loader import AnimaBundle
+from .r3gan_disc import R3GANDiscriminator, d_loss_r3gan, g_loss_r3gan
+
+
+@dataclass
+class DMD2Args:
+    # Training
+    batch_size: int = 4
+    lr_gen: float = 2e-5
+    lr_guidance: float = 2e-5
+    lr_disc: float = 2e-4
+    grad_clip: float = 1.0
+
+    # DMD2 / DMDR
+    guidance_updates: int = 5     # 1 outer step あたりの guidance update 回数
+    lora_rank: int = 32
+    lora_scale_f: float = 2.0     # fake-score (LoRA ON)
+    lora_scale_r: float = 0.75    # real-score start (cosine decay → 0)
+    cfg_r: float = 2.0            # real-score 側 CFG (Augmentation の Spear)
+    dynamic_decay_steps: int = 2000  # lora_scale_r を cosine で 0 にする step
+
+    # Timestep sampling
+    gui_alpha: float = 4.0        # Beta(α, β) for continuous t
+    gui_beta: float = 1.5
+    num_inference_steps: int = 8  # gen の discrete grid (Phase A=8, B=4, C=2)
+
+    # TSCD
+    tscd_weight: float = 0.5
+    ema_decay: float = 0.999
+
+    # R3GAN
+    adv_weight: float = 0.1
+    r3gan_gamma: float = 50.0
+
+    # CFG-Aug schedule
+    cfg_aug_cold_steps: int = 200  # 最初は cfg_r=0 で wait
+
+
+def attach_qv_lora(transformer: nn.Module, rank: int = 32) -> nn.Module:
+    """attention Q/V projection に LoRA を attach。MiniTrainDIT の実 module 名を
+    inspection で見つける。"""
+    Q_KEYS = ("to_q", "q_proj", "wq", "q")
+    V_KEYS = ("to_v", "v_proj", "wv", "v")
+    target_modules = []
+    for name, module in transformer.named_modules():
+        if not isinstance(module, nn.Linear):
+            continue
+        leaf = name.rsplit(".", 1)[-1]
+        if leaf in Q_KEYS or leaf in V_KEYS:
+            target_modules.append(name)
+
+    if not target_modules:
+        # Fallback 1: 名前に 'attn' が入る Linear すべて
+        for name, module in transformer.named_modules():
+            if isinstance(module, nn.Linear) and (
+                "attn" in name.lower() or "attention" in name.lower()
+            ):
+                target_modules.append(name)
+
+    if not target_modules:
+        print("[lora] WARN: no attention Q/V found, falling back to all-linear")
+        target_modules = "all-linear"
+    else:
+        print(f"[lora] target_modules={len(target_modules)} layers (Q/V or attn)")
+
+    cfg = peft.LoraConfig(
+        r=rank,
+        lora_alpha=rank,
+        lora_dropout=0.0,
+        bias="none",
+        target_modules=target_modules,
+    )
+    return peft.get_peft_model(transformer, cfg)
+
+
+def attach_wide_lora(transformer: nn.Module, rank: int = 32) -> nn.Module:
+    """全 nn.Linear に LoRA を attach (LLM adapter 内部は除外)。
+    Anima/Cosmos の AdaLN modulation・attention・MLP すべてを学習対象にする。
+    timestep 解釈を含む蒸留タスクには Q,V LoRA では不足、これが必要。"""
+    target_modules = []
+    for name, module in transformer.named_modules():
+        if not isinstance(module, nn.Linear):
+            continue
+        # LLM adapter (Qwen3→T5 bridge) は触らない (壊れやすい)
+        if "llm_adapter" in name:
+            continue
+        target_modules.append(name)
+
+    if not target_modules:
+        raise RuntimeError("No nn.Linear found in transformer for wide LoRA")
+
+    # 概要
+    cats = {"attn": 0, "adaln": 0, "mlp": 0, "other": 0}
+    for n in target_modules:
+        nl = n.lower()
+        if "attn" in nl or "attention" in nl:
+            cats["attn"] += 1
+        elif "adaln" in nl or "modulation" in nl:
+            cats["adaln"] += 1
+        elif "mlp" in nl or "ffn" in nl or "feed" in nl:
+            cats["mlp"] += 1
+        else:
+            cats["other"] += 1
+    print(f"[lora] wide target: total={len(target_modules)} "
+          f"(attn={cats['attn']}, adaln={cats['adaln']}, mlp={cats['mlp']}, other={cats['other']})")
+
+    cfg = peft.LoraConfig(
+        r=rank,
+        lora_alpha=rank,
+        lora_dropout=0.0,
+        bias="none",
+        target_modules=target_modules,
+    )
+    return peft.get_peft_model(transformer, cfg)
+
+
+def set_lora_scale(model: nn.Module, scale: float) -> None:
+    """PEFT LoRA の scaling を runtime で書き換え。real/fake 切替用。"""
+    for module in model.modules():
+        if hasattr(module, "scaling") and isinstance(module.scaling, dict):
+            for key in module.scaling:
+                # alpha/r = 1.0 を基準に scale 倍する
+                module.scaling[key] = scale
+
+
+def sample_continuous_t(
+    batch_size: int, alpha: float, beta: float, device: torch.device
+) -> torch.Tensor:
+    """logit-normal Beta(α, β) サンプリング。返り値 t ∈ (0, 1)"""
+    # Beta -> logit-normal 風の重み (DMDR の sample_continue 簡略版)
+    u = torch.rand(batch_size, device=device)
+    # Beta(α, β) inverse CDF を近似:正規 logit に変換するシンプル版
+    # ここでは log-normal で代用 (実装簡易のため)
+    t = torch.distributions.Beta(alpha, beta).sample((batch_size,)).to(device)
+    return t.clamp(1e-3, 1 - 1e-3)
+
+
+def sample_discrete_t(
+    batch_size: int, num_steps: int, device: torch.device
+) -> torch.Tensor:
+    """generator 側: 離散 grid {1/N, 2/N, ..., (N-1)/N} から uniform サンプル。"""
+    grid = torch.linspace(1.0 / num_steps, 1.0 - 1.0 / num_steps, num_steps - 1, device=device)
+    idx = torch.randint(0, num_steps - 1, (batch_size,), device=device)
+    return grid[idx]
+
+
+def cosine_decay(step: int, total: int) -> float:
+    if step >= total:
+        return 0.0
+    return 0.5 * (1 + math.cos(math.pi * step / total))
+
+
+class EMA:
+    """生成器の EMA を別 model copy として保持。TSCD target として直接 forward 可能。
+
+    only_trainable=True (default) なら trainable param のみ EMA 更新する。
+    LoRA-only モードでは base は frozen なので大半が skip され、EMA cost が激減。"""
+
+    def __init__(self, model: nn.Module, decay: float = 0.999,
+                 only_trainable: bool = True):
+        self.decay = decay
+        self.only_trainable = only_trainable
+        # 同 device/dtype の eval-only copy
+        self.ema_model = copy.deepcopy(model).eval()
+        for p in self.ema_model.parameters():
+            p.requires_grad = False
+
+    @torch.no_grad()
+    def update(self, model: nn.Module) -> None:
+        for ep, p in zip(self.ema_model.parameters(), model.parameters()):
+            if self.only_trainable and not p.requires_grad:
+                continue
+            ep.mul_(self.decay).add_(p.detach(), alpha=1 - self.decay)
+        # buffers (RoPE 等) は最新を採用
+        for eb, b in zip(self.ema_model.buffers(), model.buffers()):
+            eb.copy_(b.detach())
+
+    def __call__(self, **kwargs):
+        return self.ema_model(**kwargs)
+
+
+# ============================================================================
+# Trainer
+# ============================================================================
+
+class DMD2Trainer:
+    def __init__(
+        self,
+        bundle: AnimaBundle,
+        gen_transformer: nn.Module,
+        guidance_transformer: nn.Module,  # PEFT-wrapped (LoRA attached)
+        discriminator: R3GANDiscriminator,
+        args: DMD2Args,
+    ):
+        self.bundle = bundle
+        self.gen = gen_transformer
+        self.guidance = guidance_transformer
+        self.disc = discriminator
+        self.args = args
+        self.device = bundle.device
+        self.step = 0
+
+        # Optimizers (gen も frozen 部分があるなら filter)
+        self.opt_gen = torch.optim.AdamW(
+            [p for p in self.gen.parameters() if p.requires_grad],
+            lr=args.lr_gen, betas=(0.9, 0.999),
+            weight_decay=0.01, eps=1e-8,
+        )
+        self.opt_guidance = torch.optim.AdamW(
+            [p for p in self.guidance.parameters() if p.requires_grad],
+            lr=args.lr_guidance, betas=(0.9, 0.999), weight_decay=0.0, eps=1e-8,
+        )
+        self.opt_disc = torch.optim.Adam(
+            self.disc.parameters(), lr=args.lr_disc, betas=(0.0, 0.99), eps=1e-8,
+        )
+
+        # EMA for TSCD
+        self.ema = EMA(self.gen, decay=args.ema_decay)
+
+    # -- helpers ------------------------------------------------------------
+    def _set_train(self, gen: bool, guidance: bool, disc: bool):
+        self.gen.train(gen); self.guidance.train(guidance); self.disc.train(disc)
+
+    def _t_continuous(self, B: int) -> torch.Tensor:
+        return sample_continuous_t(B, self.args.gui_alpha, self.args.gui_beta, self.device)
+
+    def _t_discrete(self, B: int) -> torch.Tensor:
+        return sample_discrete_t(B, self.args.num_inference_steps, self.device)
+
+    def _cfg_r_current(self) -> float:
+        """cold start 中は cfg_r=0、それ以降は args.cfg_r。"""
+        return 0.0 if self.step < self.args.cfg_aug_cold_steps else self.args.cfg_r
+
+    def _lora_scale_r_current(self) -> float:
+        decay = cosine_decay(self.step, self.args.dynamic_decay_steps)
+        return self.args.lora_scale_r * decay
+
+    # -- guidance update ----------------------------------------------------
+    def step_guidance(self, real_latents: torch.Tensor, cond: torch.Tensor) -> dict:
+        """fake-score (LoRA ON) を v-prediction MSE で更新"""
+        self._set_train(gen=False, guidance=True, disc=False)
+        set_lora_scale(self.guidance, self.args.lora_scale_f)
+
+        B = real_latents.size(0)
+        noise = torch.randn_like(real_latents)
+        t = self._t_continuous(B)
+        noisy = AnimaBundle.add_noise(real_latents, noise, t)
+        v_target = AnimaBundle.velocity_target(real_latents, noise)
+
+        v_pred = AnimaBundle.dit_forward(self.guidance, noisy, t, cond)
+        loss = F.mse_loss(v_pred, v_target)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(
+            [p for p in self.guidance.parameters() if p.requires_grad],
+            self.args.grad_clip)
+        self.opt_guidance.step(); self.opt_guidance.zero_grad()
+        return {"l_guidance": loss.detach()}
+
+    # -- CFG forward (real-score side) -------------------------------------
+    def _real_score_forward(
+        self, noisy: torch.Tensor, t: torch.Tensor, cond: torch.Tensor, cfg: float,
+    ) -> torch.Tensor:
+        """LoRA OFF / decay 状態の guidance_model で real-side velocity を取得。
+        CFG > 1 のときは null condition と blend。"""
+        set_lora_scale(self.guidance, self._lora_scale_r_current())
+        v_cond = AnimaBundle.dit_forward(self.guidance, noisy, t, cond)
+        if cfg > 1.0:
+            null_cond = torch.zeros_like(cond)
+            v_uncond = AnimaBundle.dit_forward(self.guidance, noisy, t, null_cond)
+            v = v_uncond + cfg * (v_cond - v_uncond)
+        else:
+            v = v_cond
+        return v
+
+    # -- TSCD consistency loss ---------------------------------------------
+    def _tscd_loss(
+        self, x0_pred: torch.Tensor, cond: torch.Tensor, t: torch.Tensor
+    ) -> torch.Tensor:
+        """segment-wise consistency vs EMA gen。
+        TSCD 簡易版: 同じ x0 ground 周辺で noisy 化 → EMA で x0 推定 → MSE。"""
+        with torch.no_grad():
+            noise = torch.randn_like(x0_pred)
+            noisy = AnimaBundle.add_noise(x0_pred.detach(), noise, t)
+            v_ema = AnimaBundle.dit_forward(self.ema.ema_model, noisy, t, cond)
+            x0_ema = AnimaBundle.x0_from_velocity(noisy, v_ema, t)
+        return F.mse_loss(x0_pred, x0_ema.detach())
+
+    # -- generator update ---------------------------------------------------
+    def step_generator(self, real_latents: torch.Tensor, cond: torch.Tensor) -> dict:
+        self._set_train(gen=True, guidance=False, disc=False)
+        B = real_latents.size(0)
+        noise = torch.randn_like(real_latents)
+        t = self._t_discrete(B)
+        noisy = AnimaBundle.add_noise(real_latents, noise, t)
+
+        # Generator forward
+        v_gen = AnimaBundle.dit_forward(self.gen, noisy, t, cond)
+        x0_pred = AnimaBundle.x0_from_velocity(noisy, v_gen, t)
+
+        # === DMD gradient ===
+        with torch.no_grad():
+            cfg = self._cfg_r_current()
+            v_real = self._real_score_forward(noisy, t, cond, cfg=cfg)
+            set_lora_scale(self.guidance, self.args.lora_scale_f)
+            v_fake = AnimaBundle.dit_forward(self.guidance, noisy, t, cond)
+        x0_real = AnimaBundle.x0_from_velocity(noisy, v_real, t)
+        x0_fake = AnimaBundle.x0_from_velocity(noisy, v_fake, t)
+        p_real = x0_pred - x0_real
+        p_fake = x0_pred - x0_fake
+        denom = p_real.abs().mean(dim=list(range(1, p_real.dim())), keepdim=True) + 1e-8
+        grad = (p_real - p_fake) / denom
+        dmd_loss = 0.5 * F.mse_loss(x0_pred.float(), (x0_pred - grad).detach().float())
+
+        # === TSCD consistency ===
+        consist_loss = self._tscd_loss(x0_pred, cond, t)
+
+        # === R3GAN adversarial (generator side) — adv_weight>0 のときだけ ===
+        if self.args.adv_weight > 0:
+            d_fake = self.disc(x0_pred, cond)
+            with torch.no_grad():
+                d_real_for_g = self.disc(real_latents, cond)
+            adv_loss = g_loss_r3gan(d_real_for_g, d_fake)
+        else:
+            adv_loss = torch.zeros((), device=self.device)
+
+        loss = (
+            dmd_loss
+            + self.args.tscd_weight * consist_loss
+            + self.args.adv_weight * adv_loss
+        )
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.gen.parameters(), self.args.grad_clip)
+        self.opt_gen.step(); self.opt_gen.zero_grad()
+
+        self.ema.update(self.gen)
+        return {
+            "l_dmd": dmd_loss.detach(),
+            "l_consist": consist_loss.detach(),
+            "l_adv_g": adv_loss.detach(),
+            "loss_gen_total": loss.detach(),
+        }
+
+    # -- discriminator update ----------------------------------------------
+    def step_discriminator(self, real_latents: torch.Tensor, cond: torch.Tensor) -> dict:
+        self._set_train(gen=False, guidance=False, disc=True)
+        B = real_latents.size(0)
+
+        # 生成器で fake samples を作成 (no_grad)
+        with torch.no_grad():
+            noise = torch.randn_like(real_latents)
+            t_init = torch.ones(B, device=self.device)
+            v = AnimaBundle.dit_forward(self.gen, noise, t_init, cond)
+            fake_latents = AnimaBundle.x0_from_velocity(noise, v, t_init)
+
+        real_samples = real_latents.detach().clone().requires_grad_(True)
+        fake_samples = fake_latents.detach().clone().requires_grad_(True)
+
+        d_real = self.disc(real_samples, cond)
+        d_fake = self.disc(fake_samples, cond)
+        loss, metrics = d_loss_r3gan(
+            d_real, d_fake, real_samples, fake_samples, gamma=self.args.r3gan_gamma,
+        )
+
+        # GUARD: R1/R2 が異常に大きくなったら abort
+        # latent 空間 (16ch × 128×128 = 262k) では正常時でも R1/R2 ~ 10-100 になる。
+        # Phase A は 700 -> 700k で爆発。1000 を超えたら abort。
+        r1_val = float(metrics["d_r1"])
+        r2_val = float(metrics["d_r2"])
+        if r1_val > 1000.0 or r2_val > 1000.0:
+            raise RuntimeError(
+                f"R3GAN penalty exploded (r1={r1_val:.1f}, r2={r2_val:.1f}). "
+                f"Reduce r3gan_gamma (current={self.args.r3gan_gamma}). "
+                f"Aborting to prevent gen damage."
+            )
+
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.disc.parameters(), self.args.grad_clip)
+        self.opt_disc.step(); self.opt_disc.zero_grad()
+        metrics["l_disc"] = loss.detach()
+        return metrics
+
+    # -- main step ----------------------------------------------------------
+    def train_step(self, batch: dict) -> dict:
+        # 1) text encode + VAE encode (no grad)
+        with torch.no_grad():
+            cond = self.bundle.text_encode(batch["captions"])
+            real_latents = self.bundle.vae_encode(batch["pixels"].to(self.device))
+
+        # 2) guidance × N
+        log = {}
+        for _ in range(self.args.guidance_updates):
+            log.update(self.step_guidance(real_latents, cond))
+
+        # 3) generator × 1
+        log.update(self.step_generator(real_latents, cond))
+
+        # 4) discriminator × 1 (adv_weight=0 のときは skip して R3GAN を完全 disable)
+        if self.args.adv_weight > 0:
+            log.update(self.step_discriminator(real_latents, cond))
+
+        self.step += 1
+        return log

scripts/distill/dmdx_loss.py

@@ -0,0 +1,195 @@
+"""
+Anima DMDX (ADM = Adversarial Distribution Matching) distillation
+=================================================================
+Reference: arxiv 2507.18569v1
+  "Adversarial Distribution Matching for Diffusion Distillation"
+
+DMD2 の grad trick (reverse KL) を、**学習可能 discriminator による hinge GAN
+(TVD 最小化)** に置換した手法。本実装は image 用 ADM 単体 (ADP 省略)。
+
+要点:
+  - student rollout は DMD2 と同じ (few-step、grad-last-only)
+  - sample t (cubic schedule、high-noise バイアス)
+  - x_hat を t に re-noise → x_t
+  - **teacher 1 step で t → t-Δt 進化** (DMDX の核、時刻情報を D に与える)
+  - D(x_{t-Δt}, t-Δt) で hinge generator loss
+  - real path は teacher_x0_cache の x0 を使用 (LADD と同じ流派)
+  - D backbone = teacher MiniTrainDIT (frozen) + spectral norm heads (LADD と共有)
+
+vs DMD2:
+  - DMD2: real_score - fake_score (逆 KL gradient trick)
+  - DMDX: D(x_{t-Δt}_fake) ↑ / D(x_{t-Δt}_real) ↓ (hinge、対称 TVD)
+"""
+from __future__ import annotations
+from typing import Callable, List, Optional
+
+import torch
+import torch.nn.functional as F
+
+from .dmd2_official_loss import backward_simulation, renoise_rf
+
+
+# ----- cubic time sampling -------------------------------------------------
+
+def sample_cubic_t(B: int, device, dtype=torch.float32,
+                   bias_high_noise: bool = True) -> torch.Tensor:
+    """Cubic time schedule sampling.
+       bias_high_noise=True: u ~ U(0,1), t = 1 - u^3 → 多くは t > 0.5 (high noise)
+       paper の "高ノイズレベルへのバイアス" に対応。
+    """
+    u = torch.rand(B, device=device, dtype=dtype)
+    if bias_high_noise:
+        return 1.0 - u ** 3
+    return u ** 3
+
+
+# ----- teacher Δt evolution -------------------------------------------------
+
+def teacher_evolve_dt(
+    x_t: torch.Tensor, t: torch.Tensor, dt: torch.Tensor,
+    teacher_v_fn: Callable, cond_pos: torch.Tensor, cond_neg: torch.Tensor,
+    teacher_cfg: float = 4.5,
+) -> torch.Tensor:
+    """teacher CFG'd 1-step Euler で x_t (at time t) → x_{t-dt} (at time t-dt) へ進化。
+    rectified flow: x_{t-dt} = x_t - dt * v_teacher(x_t, t, cond)
+    """
+    dtype = x_t.dtype
+    B = x_t.size(0)
+    with torch.no_grad():
+        t_in = t.to(dtype=dtype)
+        v_cond = teacher_v_fn(x_t, t_in, cond_pos)
+        if teacher_cfg > 1.0 and cond_neg is not None:
+            v_uncond = teacher_v_fn(x_t, t_in, cond_neg)
+            v = v_uncond + teacher_cfg * (v_cond - v_uncond)
+        else:
+            v = v_cond
+        dt_b = dt.view(-1, *([1] * (x_t.dim() - 1))).to(dtype=dtype)
+        x_next = x_t - dt_b * v
+    return x_next
+
+
+# ----- ADM generator loss ---------------------------------------------------
+
+def adm_generator_loss(
+    student_v_fn: Callable, teacher_v_fn: Callable, disc,
+    init_noise: torch.Tensor, cond_pos: torch.Tensor, cond_neg: torch.Tensor,
+    real_x0: Optional[torch.Tensor] = None,
+    teacher_cfg: float = 4.5, student_cfg: float = 1.0,
+    n_student_steps: int = 4, dt_ratio: float = 1.0 / 64,
+    recon_weight: float = 0.0,
+) -> tuple[torch.Tensor, dict]:
+    """ADM generator loss (hinge -D + optional recon anchor)。
+    1. student backward simulation で x_hat 推定 (grad-last-only)
+    2. cubic t sampling, t > dt_ratio
+    3. x_hat を t に re-noise → x_t_fake
+    4. teacher 1 step で x_{t-dt}_fake へ進化
+    5. D(x_{t-dt}_fake, t-dt) hinge GAN gen loss (mean(-D))
+    6. optional: recon_weight > 0 なら Smooth-L1 anchor (vs real_x0)
+    """
+    B = init_noise.size(0)
+    device, dtype = init_noise.device, init_noise.dtype
+
+    # 1) student rollout (grad on last step only)
+    x_hat = backward_simulation(
+        student_v_fn, init_noise, n_student_steps, cond_pos, cond_neg,
+        cfg_scale=student_cfg, grad_last_only=True,
+    )
+
+    # 2) sample t (cubic, biased high-noise)
+    t = sample_cubic_t(B, device=device, dtype=torch.float32)
+    t = t.clamp(min=dt_ratio + 1e-3, max=1.0 - 1e-3)
+
+    # 3) re-noise x_hat to t
+    eps = torch.randn_like(x_hat)
+    x_t_fake = renoise_rf(x_hat, t, eps)
+
+    # 4) teacher Δt evolution
+    dt_vec = torch.full_like(t, dt_ratio)
+    t_next = t - dt_vec
+    x_next_fake = teacher_evolve_dt(
+        x_t_fake, t, dt_vec, teacher_v_fn, cond_pos, cond_neg, teacher_cfg=teacher_cfg,
+    )
+
+    # 5) discriminator at t-dt; gradient_to_input=True で student に grad が届く
+    t_next_in = t_next.to(dtype=dtype)
+    logits_fake = disc(x_next_fake, t_next_in, cond_pos, gradient_to_input=True)
+    # hinge generator: maximize D(fake) → minimize -D(fake)
+    loss_g = sum(-l.mean() for l in logits_fake) / max(len(logits_fake), 1)
+
+    metrics = {
+        "l_g_adv": loss_g.detach(),
+        "t_mean": t.mean().detach(),
+        "x_hat_abs_mean": x_hat.detach().abs().mean(),
+        "x_next_fake_abs_mean": x_next_fake.detach().abs().mean(),
+    }
+
+    # 6) optional recon anchor (LADD 流の mean collapse 防��)
+    if recon_weight > 0.0 and real_x0 is not None:
+        recon = F.smooth_l1_loss(x_hat.float(), real_x0.detach().float())
+        loss_g = loss_g + recon_weight * recon
+        metrics["l_g_recon"] = recon.detach()
+
+    return loss_g, metrics
+
+
+# ----- ADM discriminator loss -----------------------------------------------
+
+def adm_discriminator_loss(
+    student_v_fn: Callable, teacher_v_fn: Callable, disc,
+    init_noise: torch.Tensor, real_x0: torch.Tensor,
+    cond_pos: torch.Tensor, cond_neg: torch.Tensor,
+    teacher_cfg: float = 4.5, student_cfg: float = 1.0,
+    n_student_steps: int = 4, dt_ratio: float = 1.0 / 64,
+) -> tuple[torch.Tensor, dict]:
+    """ADM discriminator hinge loss。
+    Fake path: student rollout (no grad) → re-noise t → teacher t-dt
+    Real path: real_x0 → re-noise t → teacher t-dt
+    Hinge: relu(1 - D(real)) + relu(1 + D(fake))
+    """
+    B = init_noise.size(0)
+    device, dtype = init_noise.device, init_noise.dtype
+
+    # Fake path (no student grad in D phase)
+    with torch.no_grad():
+        x_hat = backward_simulation(
+            student_v_fn, init_noise, n_student_steps, cond_pos, cond_neg,
+            cfg_scale=student_cfg, grad_last_only=False,
+        )
+
+    # sample t (cubic)
+    t = sample_cubic_t(B, device=device, dtype=torch.float32)
+    t = t.clamp(min=dt_ratio + 1e-3, max=1.0 - 1e-3)
+
+    # re-noise (fake / real)
+    eps_f = torch.randn_like(x_hat)
+    eps_r = torch.randn_like(real_x0)
+    x_t_fake = renoise_rf(x_hat, t, eps_f)
+    x_t_real = renoise_rf(real_x0, t, eps_r)
+
+    # Δt evolution (both)
+    dt_vec = torch.full_like(t, dt_ratio)
+    t_next = t - dt_vec
+    x_next_fake = teacher_evolve_dt(
+        x_t_fake, t, dt_vec, teacher_v_fn, cond_pos, cond_neg, teacher_cfg=teacher_cfg,
+    )
+    x_next_real = teacher_evolve_dt(
+        x_t_real, t, dt_vec, teacher_v_fn, cond_pos, cond_neg, teacher_cfg=teacher_cfg,
+    )
+
+    # D at t-dt (heads only, x detached、メモリ節約)
+    t_next_in = t_next.to(dtype=dtype)
+    logits_fake = disc(x_next_fake.detach(), t_next_in, cond_pos, gradient_to_input=False)
+    logits_real = disc(x_next_real.detach(), t_next_in, cond_pos, gradient_to_input=False)
+
+    # hinge: real → +1、fake → -1
+    loss_real = sum(F.relu(1.0 - l).mean() for l in logits_real) / max(len(logits_real), 1)
+    loss_fake = sum(F.relu(1.0 + l).mean() for l in logits_fake) / max(len(logits_fake), 1)
+    loss_d = loss_real + loss_fake
+
+    metrics = {
+        "l_d_real": loss_real.detach(),
+        "l_d_fake": loss_fake.detach(),
+        "l_d_total": loss_d.detach(),
+        "t_mean": t.mean().detach(),
+    }
+    return loss_d, metrics

scripts/distill/hps_reward.py

@@ -0,0 +1,85 @@
+"""
+HPSv2 (Human Preference Score v2) を differentiable reward として使うラッパ。
+
+公式 hpsv2 package は score を file path 経由でしか返さないので、
+内部の OpenCLIP H/14 + custom checkpoint を直接ロードして grad-through 化する。
+
+Repo: https://github.com/tgxs002/HPSv2
+Weights: HPS_v2_compressed.pt (1.97 GB、HF Space xswu/HPSv2)
+"""
+from __future__ import annotations
+import os
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class HPSv2Reward(nn.Module):
+    """HPSv2 differentiable reward (higher is better)。
+
+    forward(images, captions) → (B,) reward
+      images: (B, 3, H, W) in [-1, 1] (CLIP preprocessing は内部で実施)
+      captions: list[str]
+    """
+    def __init__(self, weights_path: str, device="cuda", dtype=torch.float32):
+        super().__init__()
+        import open_clip
+        self.device = device
+        self.dtype = dtype
+        # OpenCLIP ViT-H/14 を laion2b プリトレで初期化
+        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
+            "ViT-H-14", pretrained="laion2b_s32b_b79k", device=device,
+        )
+        self.tokenizer = open_clip.get_tokenizer("ViT-H-14")
+        # HPSv2 の preference-finetuned weights を上書き
+        if weights_path and os.path.exists(weights_path):
+            sd = torch.load(weights_path, map_location=device, weights_only=False)
+            if "state_dict" in sd:
+                sd = sd["state_dict"]
+            # remove 'module.' prefix if present
+            sd = {k.replace("module.", ""): v for k, v in sd.items()}
+            missing, unexpected = self.model.load_state_dict(sd, strict=False)
+            print(f"[hpsv2] loaded weights: missing={len(missing)} unexpected={len(unexpected)}")
+        else:
+            print(f"[hpsv2] WARNING: weights_path '{weights_path}' not found, using OpenCLIP baseline only")
+
+        # freeze
+        for p in self.model.parameters():
+            p.requires_grad = False
+        self.model.eval()
+
+        # CLIP normalization 定数
+        self.register_buffer("mean", torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1))
+        self.register_buffer("std", torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1))
+        # device / dtype を明示的に移動 (register_buffer は default CPU なので必須)
+        self.to(device=device, dtype=dtype)
+
+    def preprocess_images(self, images: torch.Tensor) -> torch.Tensor:
+        """images (B, 3, H, W) in [-1, 1] → 224 bicubic + CLIP normalize、grad-through。"""
+        # [-1, 1] → [0, 1]
+        x = (images + 1.0) / 2.0
+        x = x.clamp(0.0, 1.0)
+        # resize to 224
+        x = F.interpolate(x, size=(224, 224), mode="bicubic", align_corners=False, antialias=True)
+        # normalize
+        x = (x - self.mean) / self.std
+        return x
+
+    def score(self, images: torch.Tensor, captions: list[str]) -> torch.Tensor:
+        """grad-through HPSv2 score。"""
+        img = self.preprocess_images(images.to(self.device, dtype=self.dtype))
+        with torch.no_grad():
+            tokens = self.tokenizer(captions).to(self.device)
+        # text features (no grad)
+        with torch.no_grad():
+            text_feats = self.model.encode_text(tokens)
+            text_feats = text_feats / text_feats.norm(dim=-1, keepdim=True)
+        # image features (grad on)
+        image_feats = self.model.encode_image(img)
+        image_feats = image_feats / image_feats.norm(dim=-1, keepdim=True)
+        # cosine similarity scaled (HPSv2 uses temperature)
+        logit_scale = self.model.logit_scale.exp().detach()
+        scores = (image_feats * text_feats).sum(dim=-1) * logit_scale
+        return scores

scripts/distill/pcm_scheduler.py

@@ -0,0 +1,107 @@
+"""
+PCM (Phased Consistency Model) 用 Euler solver + phase snap
+=========================================================
+
+Reference: G-U-N/Phased-Consistency-Model
+            code/text_to_image_sd3/pcm_fm_deterministic_scheduler.py
+
+Anima は rectified flow なので SD3 用 PCM (FlowMatchEulerDiscrete) と math 完全一致。
+sigma == t、x_t = sigma*noise + (1-sigma)*x0、v = noise - x0、x_prev = x_t + dt*v
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+
+import torch
+import numpy as np
+
+
+@dataclass
+class PCMEulerSolver:
+    """num_euler_timesteps の grid (1.0 → 0.0) と、K phase の境界 index を保持。
+
+    Attributes:
+        sigmas:      (N+1,) tensor、t=1 → 0 (= sigma_max → 0)
+        sigmas_prev: (N,)   tensor、各 sigmas[i+1]
+        phase_ends:  (K,)   int tensor、各 phase の **終端 sigmas index**
+    """
+    sigmas: torch.Tensor       # (N+1,)
+    sigmas_prev: torch.Tensor  # (N,)
+    phase_ends: torch.Tensor   # (K,) ints in [1, N]
+
+    @property
+    def num_steps(self) -> int:
+        return len(self.sigmas_prev)
+
+
+def make_pcm_solver(
+    num_euler_timesteps: int = 50,
+    num_phases: int = 4,
+    sigma_shift: float = 3.0,
+    device=None, dtype=torch.float32,
+) -> PCMEulerSolver:
+    """linear schedule + sigmoid shift。K 個に等分割。
+       phase_ends は各 phase の終端 sigma の **index** (= 次 phase の 開始 sigma index と一致)。"""
+    from .traj_scheduler import sigmoid_shift
+    ts = torch.linspace(1.0, 0.0, num_euler_timesteps + 1, dtype=torch.float64)
+    ts = sigmoid_shift(ts, sigma_shift)
+    sigmas = ts.to(device=device, dtype=dtype)
+    sigmas_prev = sigmas[1:]
+    # phase boundaries: np.linspace(0, N, K, endpoint=False) → starts
+    # 終端は次 phase の開始 - 1。最後の phase の終端は N (clean)。
+    starts = np.linspace(0, num_euler_timesteps, num_phases, endpoint=False, dtype=int)
+    ends = list(starts[1:]) + [num_euler_timesteps]  # K 個
+    phase_ends = torch.tensor(ends, dtype=torch.long, device=device)
+    return PCMEulerSolver(sigmas=sigmas, sigmas_prev=sigmas_prev, phase_ends=phase_ends)
+
+
+def euler_step_pcm(
+    x: torch.Tensor, v: torch.Tensor,
+    sigma_cur: torch.Tensor, sigma_next: torch.Tensor,
+) -> torch.Tensor:
+    """rectified-flow Euler 1 step: x_next = x + (sigma_next - sigma_cur) * v
+       (sigma 減少方向、v は noise - x0)"""
+    dt = (sigma_next - sigma_cur).view(-1, *([1] * (x.dim() - 1)))
+    return x + dt * v
+
+
+def euler_multiphase(
+    x_t: torch.Tensor, v: torch.Tensor, index: torch.Tensor,
+    solver: PCMEulerSolver,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """index で指定した時点 sigma から、その所属 phase の終端 sigma まで 1-step Euler。
+
+    Args:
+        x_t:   (B, ...)
+        v:     (B, ...) velocity prediction
+        index: (B,) long、各サンプルの現在の sigmas grid index
+        solver: PCMEulerSolver
+    Returns:
+        x_at_phase_end: (B, ...)
+        sigma_at_phase_end: (B,)
+    """
+    B = x_t.size(0)
+    device = x_t.device
+    # 各 sample が属する phase の終端 index を引く
+    phase_ends = solver.phase_ends  # (K,)
+    # index → phase index: searchsorted で「index 以上の最初の phase_end」
+    # ただし PCM 流は「index <= phase_end - 1 を満たす最小 phase_end」
+    # ここでは現在 sigmas[index] からその phase の sigmas[phase_end] まで進める
+    # = sigma_cur = sigmas[index], sigma_target = sigmas[phase_end]
+    idx_np = index.detach().cpu().numpy()
+    end_np = phase_ends.detach().cpu().numpy()
+    target_end = []
+    for ix in idx_np:
+        # 最初に phase_end >= ix+1 になる phase を選ぶ (= 現 phase の終端)
+        end = end_np[end_np > ix][0] if (end_np > ix).any() else end_np[-1]
+        target_end.append(int(end))
+    target_end_t = torch.tensor(target_end, dtype=torch.long, device=device)
+
+    sigma_cur = solver.sigmas[index]              # (B,)
+    sigma_target = solver.sigmas[target_end_t]    # (B,)
+    x_next = euler_step_pcm(x_t, v, sigma_cur, sigma_target)
+    return x_next, sigma_target
+
+
+def pseudo_huber(diff: torch.Tensor, c: float = 1e-3) -> torch.Tensor:
+    """sqrt(diff^2 + c^2) - c。PCM 公式の robust loss。"""
+    return torch.sqrt(diff.float() ** 2 + c * c) - c

scripts/distill/precompute_teacher_x0.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+"""
+Precompute teacher x0 cache for AMD LADD training.
+
+For each caption in the dataset:
+  1. Sample fresh noise
+  2. Run teacher (= base Anima, no LoRA) for N steps with CFG=W
+  3. Save final x0 (decoded latent stays in latent space) as .pt
+  4. Save caption embedding too (avoid re-encoding during training)
+
+Output structure:
+  /dataset/teacher_x0_cache/
+    metadata.json    [{caption, x0_path, emb_path}, ...]
+    x0/0001.pt       (B=1, 16, 1, H_lat, W_lat) bf16
+    emb/0001.pt      (B=1, 512, 1024) bf16
+
+LADD train ループはこの cache を読むだけで、teacher を都度 forward しないので
+1 step あたりの時間を半分以下に。
+"""
+from __future__ import annotations
+import argparse
+import json
+import sys
+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.train_traj import TextOnlyDataset, text_collate
+
+
+@torch.no_grad()
+def teacher_rollout(
+    bundle: AnimaBundle,
+    cond_pos: torch.Tensor,
+    cond_neg: torch.Tensor,
+    H_lat: int, W_lat: int,
+    num_steps: int = 20,
+    cfg_scale: float = 4.5,
+    sigma_shift: float = 3.0,
+    device=None, dtype=torch.bfloat16,
+    init_noise: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """teacher (= base Anima、frozen) を N step Euler で回して x0 を返す。
+       sigma_shift は Anima 公式 workflow 推奨 3.0。
+       init_noise を渡せば外で seed 管理可能。返り値は (x0_final, noise_used)。"""
+    from distill.traj_scheduler import make_schedule
+    sched = make_schedule(num_steps, sigma_shift, device=device, dtype=torch.float32)
+    ts = sched.timesteps  # (N+1,) t=1 → 0
+    B = cond_pos.size(0)
+    if init_noise is None:
+        init_noise = torch.randn(B, 16, 1, H_lat, W_lat, device=device, dtype=dtype)
+    x = init_noise.clone()
+    for i in range(num_steps):
+        t_cur = ts[i].expand(B).to(dtype=dtype)
+        v_cond = bundle.velocity(x, t_cur, cond_pos)
+        if cfg_scale > 1.0:
+            v_uncond = bundle.velocity(x, t_cur, cond_neg)
+            v = v_uncond + cfg_scale * (v_cond - v_uncond)
+        else:
+            v = v_cond
+        dt = (ts[i + 1] - ts[i]).to(device=device, dtype=dtype)
+        x = x + dt * v
+    return x, init_noise
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dataset", required=True, type=str, help="caption dir")
+    ap.add_argument("--out", required=True, type=str, help="cache output dir")
+    ap.add_argument("--num-steps", type=int, default=20)
+    ap.add_argument("--cfg-scale", type=float, default=4.5)
+    ap.add_argument("--sigma-shift", type=float, default=3.0)
+    ap.add_argument("--resolution", type=int, default=768)
+    ap.add_argument("--max-samples", type=int, default=-1, help="-1 = all")
+    ap.add_argument("--neg-prompt", default="")
+    ap.add_argument("--start-from", type=int, default=0, help="resume index")
+    ap.add_argument("--save-noise", action="store_true",
+                    help="initial noise も保存 (Reflow 用 (noise, x0) pair に必須)")
+    ap.add_argument("--metadata-out", default="",
+                    help="metadata 出力先 (省略時は <out>/metadata.json)。"
+                         "並列実行時に worker ごと別ファイルにして race を回避するため。")
+    args = ap.parse_args()
+
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+
+    out_dir = Path(args.out)
+    (out_dir / "x0").mkdir(parents=True, exist_ok=True)
+    (out_dir / "emb").mkdir(parents=True, exist_ok=True)
+    if args.save_noise:
+        (out_dir / "noise").mkdir(parents=True, exist_ok=True)
+
+    print("[load] Anima bundle (teacher only)")
+    bundle = build_anima(device=device, dtype=dtype)
+
+    print(f"[data] {args.dataset}")
+    ds = TextOnlyDataset(args.dataset)
+    N = len(ds) if args.max_samples < 0 else min(len(ds), args.max_samples)
+    print(f"   precomputing {N - args.start_from} samples (resolution={args.resolution})")
+
+    H_lat = args.resolution // 8
+    W_lat = args.resolution // 8
+
+    with torch.no_grad():
+        cond_neg = bundle.text_encode([args.neg_prompt or ""])
+
+    metadata = []
+    metadata_path = Path(args.metadata_out) if args.metadata_out else out_dir / "metadata.json"
+    metadata_path.parent.mkdir(parents=True, exist_ok=True)
+    if metadata_path.exists() and args.start_from > 0:
+        metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+
+    import time
+    t_start = time.time()
+    for i in range(args.start_from, N):
+        caption = ds[i]
+        with torch.no_grad():
+            cond_pos = bundle.text_encode([caption])
+            x0, noise = teacher_rollout(
+                bundle, cond_pos, cond_neg, H_lat, W_lat,
+                num_steps=args.num_steps, cfg_scale=args.cfg_scale,
+                sigma_shift=args.sigma_shift, device=device, dtype=dtype,
+            )
+        x0_p = out_dir / "x0" / f"{i:05d}.pt"
+        emb_p = out_dir / "emb" / f"{i:05d}.pt"
+        torch.save(x0.cpu(), x0_p)
+        torch.save(cond_pos.cpu(), emb_p)
+        entry = {"caption": caption, "x0_path": str(x0_p), "emb_path": str(emb_p)}
+        if args.save_noise:
+            noise_p = out_dir / "noise" / f"{i:05d}.pt"
+            torch.save(noise.cpu(), noise_p)
+            entry["noise_path"] = str(noise_p)
+        metadata.append(entry)
+
+        if i % 50 == 0 or i == N - 1:
+            metadata_path.write_text(json.dumps(metadata, ensure_ascii=False), encoding="utf-8")
+            elapsed = time.time() - t_start
+            rate = (i - args.start_from + 1) / max(elapsed, 1e-3)
+            eta = (N - i - 1) / max(rate, 1e-3)
+            print(f"[precompute {i}/{N}] rate={rate:.2f}/s eta={eta/60:.1f}min", flush=True)
+            try:
+                import modal
+                modal.Volume.from_name("anima-dataset").commit()
+            except Exception:
+                pass
+
+    metadata_path.write_text(json.dumps(metadata, ensure_ascii=False), encoding="utf-8")
+    print(f"[done] {len(metadata)} samples cached to {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/r3gan_disc.py

@@ -0,0 +1,200 @@
+"""
+R3GAN (NeurIPS 2024) discriminator + losses (RpGAN + R1 + R2)。
+
+Anima 用に latent-space で動かす:
+  - 入力: VAE latents (B, 16, T, H_lat, W_lat) — T=1 (静止画)
+  - 条件: Qwen3+LLMAdapter の crossattn_emb を平均プールして projection D
+  - 出力: scalar critic value per sample
+
+Normalization 不使用 (R3GAN の "modernized GAN" の要)、MSR init、
+activation gain で variance を解析的に維持。
+"""
+from __future__ import annotations
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def msr_init(weight: torch.Tensor, gain: float = 1.0) -> None:
+    """MSR (He) initialization, fan-in 基準。"""
+    fan_in = weight.shape[1] * weight[0][0].numel()
+    std = gain / math.sqrt(fan_in)
+    with torch.no_grad():
+        weight.normal_(0, std)
+
+
+class BiasedActivation(nn.Module):
+    """Leaky ReLU + 学習可能 bias。R3GAN の活性化代替。"""
+
+    def __init__(self, channels: int, negative_slope: float = 0.2):
+        super().__init__()
+        self.bias = nn.Parameter(torch.zeros(channels))
+        self.negative_slope = negative_slope
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.leaky_relu(x + self.bias.view(1, -1, *([1] * (x.dim() - 2))),
+                            self.negative_slope, inplace=True)
+
+
+class ResidualBlock(nn.Module):
+    """ConvNeXt-style: 1x1 expand -> grouped depthwise -> 1x1 project, no norm."""
+
+    def __init__(self, channels: int, expansion: int = 2, cardinality: int = 1, kernel: int = 3):
+        super().__init__()
+        hidden = channels * expansion
+        groups = max(1, hidden // cardinality)  # depthwise-ish grouping
+
+        self.expand = nn.Conv2d(channels, hidden, kernel_size=1, bias=False)
+        self.dw = nn.Conv2d(hidden, hidden, kernel_size=kernel, padding=kernel // 2,
+                             groups=groups, bias=False)
+        self.project = nn.Conv2d(hidden, channels, kernel_size=1, bias=False)
+
+        self.act = BiasedActivation(hidden)
+
+        # MSR init, project は ActivationGain=0 で残差ゼロ初期化(安定)
+        msr_init(self.expand.weight, gain=1.0)
+        msr_init(self.dw.weight, gain=1.0)
+        nn.init.zeros_(self.project.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.expand(x)
+        h = self.act(h)
+        h = self.dw(h)
+        h = self.act(h)
+        h = self.project(h)
+        return x + h
+
+
+class DownsampleLayer(nn.Module):
+    """[1,2,1] bilinear / anti-aliased downsample by factor 2."""
+
+    def __init__(self):
+        super().__init__()
+        kernel = torch.tensor([1.0, 2.0, 1.0]) / 4.0
+        kernel = kernel[:, None] * kernel[None, :]  # (3, 3)
+        self.register_buffer("kernel", kernel[None, None].repeat(1, 1, 1, 1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = x.shape
+        k = self.kernel.expand(c, 1, 3, 3)
+        return F.conv2d(x, k, stride=2, padding=1, groups=c)
+
+
+class DiscriminatorStage(nn.Module):
+    """[ResidualBlock × N] + Downsample"""
+
+    def __init__(self, channels: int, num_blocks: int = 2, downsample: bool = True):
+        super().__init__()
+        self.blocks = nn.Sequential(*[ResidualBlock(channels) for _ in range(num_blocks)])
+        self.down = DownsampleLayer() if downsample else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down(self.blocks(x))
+
+
+class R3GANDiscriminator(nn.Module):
+    """
+    Anima latent-space text-conditional discriminator。
+
+    Args:
+        latent_channels: VAE latent channels (Anima: 16)
+        widths: stage 毎の channel 数 (粗→細の解像度遷移とは逆順)
+        blocks_per_stage: 各 stage の ResidualBlock 数
+        cond_dim: text embedding 次元(LLMAdapter 出力次元 = 1024)
+        cond_embedding_dim: projection D の中間次元
+    """
+
+    def __init__(
+        self,
+        latent_channels: int = 16,
+        widths: tuple[int, ...] = (128, 256, 512, 768),
+        blocks_per_stage: int = 2,
+        cond_dim: int = 1024,
+        cond_embedding_dim: int = 512,
+    ):
+        super().__init__()
+        # Extraction: latent -> initial channels
+        self.extract = nn.Conv2d(latent_channels, widths[0], kernel_size=1, bias=False)
+        msr_init(self.extract.weight, gain=1.0)
+
+        stages = []
+        for i, w in enumerate(widths):
+            next_w = widths[i + 1] if i + 1 < len(widths) else w
+            stages.append(DiscriminatorStage(w, num_blocks=blocks_per_stage,
+                                              downsample=(i < len(widths) - 1)))
+            if i < len(widths) - 1 and next_w != w:
+                # channel 変換 (1x1)
+                stages.append(nn.Conv2d(w, next_w, kernel_size=1, bias=False))
+                msr_init(stages[-1].weight, gain=1.0)
+        self.stages = nn.Sequential(*stages)
+
+        # Discriminative basis: 空間集約 -> per-sample feature vec
+        # 入力サイズによっては adaptive pool
+        self.spatial_pool = nn.AdaptiveAvgPool2d(4)
+        self.final_conv = nn.Conv2d(widths[-1], cond_embedding_dim, kernel_size=4, bias=False)
+        msr_init(self.final_conv.weight, gain=1.0)
+
+        # Projection conditioning (Miyato & Koyama 2018)
+        self.cond_proj = nn.Linear(cond_dim, cond_embedding_dim, bias=False)
+        msr_init(self.cond_proj.weight, gain=1.0 / math.sqrt(cond_embedding_dim))
+
+    def forward(self, latents: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        """
+        latents: (B, 16, T, H_lat, W_lat) — T=1 想定なので squeeze
+        cond:    (B, N_tokens, cond_dim) — 平均プールして使う
+        returns: (B,) critic value
+        """
+        if latents.dim() == 5:
+            assert latents.shape[2] == 1, "static-image distillation: T must be 1"
+            latents = latents.squeeze(2)  # (B, 16, H, W)
+
+        # weights dtype に揃える(generator 出力が float32 で来る場合あり)
+        w_dtype = next(self.parameters()).dtype
+        latents = latents.to(dtype=w_dtype)
+        cond = cond.to(dtype=w_dtype)
+
+        h = self.extract(latents)
+        h = self.stages(h)
+        h = self.spatial_pool(h)
+        feat = self.final_conv(h).flatten(1)  # (B, cond_embedding_dim)
+
+        # text embedding 平均プール -> projection
+        cond_mean = cond.mean(dim=1)  # (B, cond_dim)
+        cond_emb = self.cond_proj(cond_mean)  # (B, cond_embedding_dim)
+
+        # projection-D logit = <feat, cond_emb>
+        logit = (feat * cond_emb).sum(dim=1)  # (B,)
+        return logit
+
+
+# ---------- losses (RpGAN + R1 + R2) ----------------------------------------
+
+def zero_centered_grad_penalty(samples: torch.Tensor, logits: torch.Tensor) -> torch.Tensor:
+    """∑(∂D/∂x)^2 per sample."""
+    grads, = torch.autograd.grad(
+        outputs=logits.sum(), inputs=samples, create_graph=True,
+    )
+    return grads.flatten(1).square().sum(dim=1)
+
+
+def d_loss_r3gan(
+    d_real_logits: torch.Tensor,
+    d_fake_logits: torch.Tensor,
+    real_samples: torch.Tensor,
+    fake_samples: torch.Tensor,
+    gamma: float = 50.0,
+) -> tuple[torch.Tensor, dict]:
+    """Discriminator loss = RpGAN softplus + (γ/2)(R1 + R2)."""
+    adv = F.softplus(-(d_real_logits - d_fake_logits)).mean()
+    r1 = zero_centered_grad_penalty(real_samples, d_real_logits).mean()
+    r2 = zero_centered_grad_penalty(fake_samples, d_fake_logits).mean()
+    loss = adv + 0.5 * gamma * (r1 + r2)
+    return loss, {"d_adv": adv.detach(), "d_r1": r1.detach(), "d_r2": r2.detach()}
+
+
+def g_loss_r3gan(d_real_logits: torch.Tensor, d_fake_logits: torch.Tensor) -> torch.Tensor:
+    """Generator loss (relativistic): softplus(-(D(fake) - D(real)))."""
+    return F.softplus(-(d_fake_logits - d_real_logits)).mean()

scripts/distill/shortcut_module.py

@@ -0,0 +1,111 @@
+"""
+Shortcut Models 用の d-head injection。
+
+設計:
+  - transformer.d_head: nn.Linear(1 → cond_dim) (zero-init、初期は影響ゼロ)
+  - transformer の t_embedder 出力に hook で d_head(d) を加算
+  - 推論/訓練時に transformer._current_d = d を set してから forward
+"""
+from __future__ import annotations
+import torch
+import torch.nn as nn
+
+
+def find_t_embedder(transformer: nn.Module) -> tuple[str, nn.Module]:
+    """transformer 内部の t_embedder を発見。短い path を優先。"""
+    candidates = []
+    for name, module in transformer.named_modules():
+        n = name.lower()
+        if ("t_embedder" in n or "time_embedder" in n or "timestep_embedder" in n) and \
+           ("llm_adapter" not in name):  # llm_adapter 内の embedder は除外
+            candidates.append((name, module))
+    if not candidates:
+        # fallback: find by "time" in name (more inclusive)
+        for name, module in transformer.named_modules():
+            if "time" in name.lower() and "llm_adapter" not in name and \
+               isinstance(module, (nn.Linear, nn.Sequential, nn.Module)) and \
+               len(list(module.parameters())) > 0:
+                candidates.append((name, module))
+    if not candidates:
+        raise RuntimeError("Could not find t_embedder module")
+    # shortest path
+    candidates.sort(key=lambda x: len(x[0].split(".")))
+    name, module = candidates[0]
+    print(f"[shortcut] found t_embedder at: {name}")
+    return name, module
+
+
+def _detect_t_embed_dim(transformer: nn.Module, t_embedder: nn.Module) -> int:
+    """t_embedder の output dim を dummy forward で推定。"""
+    device = next(t_embedder.parameters()).device
+    dtype = next(t_embedder.parameters()).dtype
+    with torch.no_grad():
+        # 多くの実装は (B,) or (B,1) の timestep を受ける
+        for shape in [(1,), (1, 1)]:
+            try:
+                t_in = torch.tensor([[0.5]] if len(shape) == 2 else [0.5],
+                                    device=device, dtype=dtype)
+                out = t_embedder(t_in)
+                if isinstance(out, tuple):
+                    out = out[0]
+                return out.shape[-1]
+            except Exception:
+                continue
+    raise RuntimeError("Could not detect t_embedder output dim")
+
+
+def attach_shortcut_d_head(transformer: nn.Module) -> nn.Module:
+    """transformer に d_head を attach、t_embedder 出力に hook で加算する。
+    Returns: transformer (mutated)。`transformer._current_d` を set してから forward。"""
+    name, t_emb = find_t_embedder(transformer)
+    cond_dim = _detect_t_embed_dim(transformer, t_emb)
+    print(f"[shortcut] t_embedder output dim: {cond_dim}")
+
+    # d_head: 1 → cond_dim → cond_dim (small MLP, zero-init last layer)
+    device = next(t_emb.parameters()).device
+    dtype = next(t_emb.parameters()).dtype
+    d_head = nn.Sequential(
+        nn.Linear(1, cond_dim),
+        nn.SiLU(),
+        nn.Linear(cond_dim, cond_dim),
+    ).to(device=device, dtype=dtype)
+    # zero-init last linear so initial output is 0
+    nn.init.zeros_(d_head[-1].weight)
+    nn.init.zeros_(d_head[-1].bias)
+
+    transformer.d_head = d_head
+    transformer._current_d = None  # set per-forward
+
+    def _hook(module, input, output):
+        d = getattr(transformer, "_current_d", None)
+        if d is None:
+            return output
+        # output shape: (B, cond_dim) or (B, 1, cond_dim) etc
+        d_in = d.view(-1, 1).to(device=output.device if hasattr(output, 'device') else next(d_head.parameters()).device,
+                                dtype=next(d_head.parameters()).dtype)
+        d_emb = d_head(d_in)  # (B, cond_dim)
+        # broadcast to match output shape
+        if isinstance(output, tuple):
+            base = output[0]
+            d_emb_b = d_emb.view(base.size(0), *([1] * (base.dim() - 2)), base.size(-1)) if base.dim() > 2 else d_emb
+            return (base + d_emb_b.to(dtype=base.dtype),) + output[1:]
+        else:
+            d_emb_b = d_emb.view(output.size(0), *([1] * (output.dim() - 2)), output.size(-1)) if output.dim() > 2 else d_emb
+            return output + d_emb_b.to(dtype=output.dtype)
+
+    handle = t_emb.register_forward_hook(_hook)
+    transformer._d_head_hook_handle = handle
+
+    return transformer
+
+
+def set_shortcut_d(transformer: nn.Module, d: torch.Tensor | None):
+    """forward 前に d を set。d=None で d-head 効果無効。"""
+    transformer._current_d = d
+
+
+def shortcut_d_head_params(transformer: nn.Module) -> list[nn.Parameter]:
+    """d_head の trainable params (LoRA とは別の optimizer に渡す用)。"""
+    if hasattr(transformer, "d_head"):
+        return list(transformer.d_head.parameters())
+    return []

scripts/distill/sid_loss.py

@@ -0,0 +1,209 @@
+"""
+SiD2 / SiD-DiT 流派 (Score Identity Distillation) for Anima
+
+Reference:
+  - arXiv 2404.04057 (original SiD)
+  - arXiv 2509.25127 (SiD-DiT、RF + TrigFlow native 適応)
+  - github.com/mingyuanzhou/SiD-LSG (SD 用 reference)
+
+特徴:
+  - **data-free**: caption だけあれば画像不要
+  - critic 不要 (D 自体が存在しない)
+  - EMA 不要
+  - 2 LoRA pattern (generator + fake_score) を DMD2 と共有
+  - SiD identity formula で **mean collapse をスコア空間で防ぐ**
+  - RF (rectified flow) を直接サポート (teacher の再学習不要)
+
+Generator loss (SiD-DiT, Eq. paper 2509.25127):
+  w(t) = (1 - t)
+  x_g  = student rollout (n-step, last-step grad only)
+  x_t  = (1-t) * x_g + t * eps
+  x0_φ = teacher(x_t, t, c) [+ CFG]
+  x0_ψ = fake_score(x_t, t, c)
+
+  diff_main = x0_φ - x0_ψ                                ← SiD core
+  diff_corr = x0_ψ - x_g                                 ← identity correction
+  L_θ = (1-α) * w(t) * ||x_g - x0_φ||²
+       + w(t) * <diff_main, diff_corr>
+       (per-sample mean over spatial dims)
+
+Score helper loss (flow-matching MSE):
+  L_ψ = ||fake_score(x_t, t, c) - x_g.detach()||²       ← simple distillation of x_g distribution
+"""
+from __future__ import annotations
+from typing import Callable
+
+import torch
+import torch.nn.functional as F
+
+from .anima_loader import AnimaBundle
+
+
+# ----- RF helpers ----------------------------------------------------------
+def renoise_rf(x0: torch.Tensor, t_rf: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
+    """rectified flow forward: x_t = (1-t)*x0 + t*noise"""
+    t_ = t_rf.view(-1, *([1] * (x0.dim() - 1)))
+    return (1.0 - t_) * x0 + t_ * noise
+
+
+def x0_from_velocity_rf(x_t: torch.Tensor, v: torch.Tensor, t_rf: torch.Tensor) -> torch.Tensor:
+    """x0 = x_t - t * v (v = noise - x0)"""
+    t_ = t_rf.view(-1, *([1] * (x_t.dim() - 1)))
+    return x_t - t_ * v
+
+
+# ----- time sampler --------------------------------------------------------
+def sample_logit_normal_t(B: int, mu: float = 0.6931, sigma: float = 1.6,
+                          device=None, dtype=torch.float32) -> torch.Tensor:
+    """LogitNormal(mu=ln 2, sigma=1.6) on t ∈ (0, 1)。
+       SiD-DiT 流の time sampler、mid-noise (t≈0.66) に分布が集中。"""
+    z = torch.randn(B, device=device, dtype=dtype) * sigma + mu
+    t = torch.sigmoid(z)  # (0, 1)
+    return t.clamp(1e-3, 1.0 - 1e-3)
+
+
+# ----- few-step student rollout (DMD2 と同じパターン) ----------------------
+def backward_simulation_grad_last(
+    student_v_fn: Callable[..., torch.Tensor],
+    noise: torch.Tensor,
+    n_steps: int,
+    cond_pos: torch.Tensor,
+    cond_neg: torch.Tensor | None,
+    cfg_scale: float = 1.0,
+) -> torch.Tensor:
+    """t=1 → t=0 へ Euler n_steps、grad は last step のみ。"""
+    from .traj_loss import cfg_guided, _broadcast_t
+    B = noise.size(0)
+    device = noise.device
+    dtype = noise.dtype
+    ts = torch.linspace(1.0, 0.0, n_steps + 1, device=device, dtype=torch.float32)
+    x = noise
+    for i in range(n_steps):
+        t_cur, t_next = ts[i], ts[i + 1]
+        is_last = (i == n_steps - 1)
+        ctx = torch.enable_grad() if is_last else torch.no_grad()
+        t_in = _broadcast_t(t_cur, B, device, dtype)
+        with ctx:
+            v = cfg_guided(student_v_fn, x, t_in, cond_pos, cond_neg, cfg_scale)
+            dt = (t_next - t_cur).to(device=device, dtype=dtype)
+            x = x + dt * v
+    return x
+
+
+# ----- generator loss ------------------------------------------------------
+def sid_generator_loss(
+    student_v_fn: Callable,
+    teacher_v_fn: Callable,
+    fake_score_v_fn: Callable,
+    init_noise: torch.Tensor,
+    cond_pos: torch.Tensor,
+    cond_neg: torch.Tensor | None,
+    teacher_cfg: float = 4.5,
+    student_cfg: float = 1.0,
+    n_steps: int = 4,
+    alpha: float = 1.2,
+    mu_t: float = 0.6931,
+    sigma_t: float = 1.6,
+) -> tuple[torch.Tensor, dict]:
+    """SiD-DiT generator loss。data-free、Anima RF native。"""
+    B = init_noise.size(0)
+    device = init_noise.device
+    dtype = init_noise.dtype
+
+    # 1) student rollout
+    x_g = backward_simulation_grad_last(
+        student_v_fn, init_noise, n_steps, cond_pos, cond_neg,
+        cfg_scale=student_cfg,
+    )
+
+    # 2) re-noise at sampled t
+    t = sample_logit_normal_t(B, mu_t, sigma_t, device, dtype=torch.float32)
+    D_eps = torch.randn_like(x_g)
+    x_t = renoise_rf(x_g, t, D_eps)
+
+    # 3) teacher x0 (CFG) and fake_score x0 (no_grad both)
+    with torch.no_grad():
+        t_in = t.to(dtype=dtype)
+        v_t_cond = teacher_v_fn(x_t, t_in, cond_pos)
+        if teacher_cfg > 1.0 and cond_neg is not None:
+            v_t_uncond = teacher_v_fn(x_t, t_in, cond_neg)
+            v_teacher = v_t_uncond + teacher_cfg * (v_t_cond - v_t_uncond)
+        else:
+            v_teacher = v_t_cond
+        x0_phi = x0_from_velocity_rf(x_t, v_teacher, t)
+
+        v_fake = fake_score_v_fn(x_t, t_in, cond_pos)
+        x0_psi = x0_from_velocity_rf(x_t, v_fake, t)
+
+    # 4) SiD-DiT fused loss
+    # w(t) = (1 - t)、per-sample weight
+    w = (1.0 - t).clamp(min=0.05).view(-1, *([1] * (x_g.dim() - 1)))
+    diff_main = (x0_phi - x0_psi).detach()
+    diff_corr = (x0_psi - x_g).detach()
+    # term 1: baseline match (mean reverse-KL-like)
+    # ||x_g - x0_phi||² の grad は x_g を x0_phi に引き寄せる
+    term_baseline = ((x_g - x0_phi.detach()).float() ** 2)
+    # term 2: identity correction (negative for "push away from fake_score")
+    # <x_g - (x_g - identity_grad)> として書き直すと DMD2 trick と同じ形
+    # ここでは直接 SiD-DiT 表式: w * <diff_main, diff_corr> (定数項なので grad は 0)
+    # 実装上は generator が ψ - x_g 方向に動くよう、x_g に grad を流す:
+    grad_signal = (diff_main * diff_corr).detach()  # 形状 (B, ...) 、generator に伝わる勾配は diff_main 方向
+    term_correction = - (x_g * (diff_main.detach())).float()  # 内積の grad part (∂/∂x_g)
+
+    L_theta = (
+        (1.0 - alpha) * (w * term_baseline).mean()
+        + (w * term_correction).mean()
+    )
+
+    metrics = {
+        "l_sid_gen": L_theta.detach(),
+        "t_mean": t.mean(),
+        "x_g_abs": x_g.detach().abs().mean(),
+        "x0_phi_abs": x0_phi.abs().mean(),
+        "x0_psi_abs": x0_psi.abs().mean(),
+        "diff_main_abs": diff_main.abs().mean(),
+        "w_mean": w.mean(),
+    }
+    return L_theta, metrics
+
+
+# ----- score helper loss (fake_score の更新) -------------------------------
+def sid_score_helper_loss(
+    student_v_fn: Callable,
+    fake_score_v_fn: Callable,
+    init_noise: torch.Tensor,
+    cond_pos: torch.Tensor,
+    cond_neg: torch.Tensor | None,
+    student_cfg: float = 1.0,
+    n_steps: int = 4,
+    mu_t: float = 0.6931,
+    sigma_t: float = 1.6,
+) -> tuple[torch.Tensor, dict]:
+    """fake_score を student 分布の denoiser として学習。
+       student 側は no_grad、fake_score 側に grad。"""
+    B = init_noise.size(0)
+    device = init_noise.device
+    dtype = init_noise.dtype
+
+    # 1) student rollout (no_grad)
+    with torch.no_grad():
+        x_g = backward_simulation_grad_last(
+            # all no_grad: 内部の last-step grad はもとから no_grad コンテキストでマスク
+            student_v_fn, init_noise, n_steps, cond_pos, cond_neg,
+            cfg_scale=student_cfg,
+        )
+
+    # 2) re-noise
+    t = sample_logit_normal_t(B, mu_t, sigma_t, device, dtype=torch.float32)
+    D_eps = torch.randn_like(x_g)
+    x_t = renoise_rf(x_g, t, D_eps)
+
+    # 3) fake_score x0 prediction (grad on)
+    t_in = t.to(dtype=dtype)
+    v_psi = fake_score_v_fn(x_t, t_in, cond_pos)
+    x0_psi = x0_from_velocity_rf(x_t, v_psi, t)
+
+    # 4) plain MSE to x_g (denoising target)
+    L_psi = F.mse_loss(x0_psi.float(), x_g.detach().float())
+
+    return L_psi, {"l_sid_psi": L_psi.detach()}

scripts/distill/train_dmd2_official.py

@@ -0,0 +1,305 @@
+#!/usr/bin/env python3
+"""
+Anima DMD2 distillation (NVIDIA cosmos-predict2.5 official 流派)
+==============================================================
+
+特徴:
+  - 同じ base 2B DiT に **2 つの独立 PEFT adapter** を attach (student / fake_score)
+  - peft.PeftModel.add_adapter() + set_adapter() で runtime 切替
+  - Alternating: critic × N_critic (=5) → generator × 1
+  - few-step student rollout (1-4 step)、grad は最終 step だけ通す (memory efficient)
+  - teacher は frozen base (no LoRA、deepcopy で別 instance)
+
+要点:
+  - student / fake_score の grad は **独立した optimizer** で更新
+  - 各 step 開始時に明示的に set_adapter() で切替
+  - teacher は別 module、active adapter 概念とは無関係
+
+使い方 (Modal):
+  # Smoke (1 outer step、cycle=5 critic+1 gen)
+  modal run modal_app.py::train_dmd2_official_distill \\
+      --total-outer-steps 1 --resolution 768 --n-student-steps 4
+
+  # 本番 (5000 outer step、~$21)
+  modal run --detach modal_app.py::train_dmd2_official_distill \\
+      --total-outer-steps 5000 --resolution 768 --warm-lora /models/loras/anima_turbo.safetensors
+"""
+from __future__ import annotations
+import argparse
+import copy
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+import torch
+import peft
+from torch.utils.data import DataLoader
+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.dmd2_official_loss import (
+    dmd2_generator_loss, dmd2_critic_loss,
+)
+from distill.train_traj import (
+    TextOnlyDataset, text_collate,
+    convert_comfy_to_peft_lora, save_lora_state,
+)
+
+
+# ---------------------------------------------------------------------------
+# 2 adapter 管理ヘルパー
+# ---------------------------------------------------------------------------
+def attach_dual_lora(
+    transformer: torch.nn.Module,
+    rank: int = 32,
+    adapter_names: tuple[str, str] = ("student", "fake_score"),
+):
+    """同じ base DiT に 2 つの独立 LoRA adapter を attach する。
+    PEFT は LoraConfig + adapter_name で複数 adapter を保持できる。
+    set_adapter(name) で active adapter を切替。
+    """
+    # まず 1 つ目の adapter を attach (peft.get_peft_model がベース wrap)
+    target_modules = []
+    for name, module in transformer.named_modules():
+        if not isinstance(module, torch.nn.Linear):
+            continue
+        if "llm_adapter" in name:
+            continue
+        target_modules.append(name)
+    print(f"[lora] wide target: {len(target_modules)} linear modules")
+
+    cfg = peft.LoraConfig(
+        r=rank, lora_alpha=rank, lora_dropout=0.0, bias="none",
+        target_modules=target_modules,
+    )
+    peft_model = peft.get_peft_model(transformer, cfg, adapter_name=adapter_names[0])
+    # 2 つ目を追加
+    peft_model.add_adapter(adapter_names[1], cfg)
+    print(f"[lora] attached adapters: {adapter_names}")
+    return peft_model
+
+
+def make_velocity_fn(peft_model, adapter_name: str):
+    """指定 adapter を active にして dit_forward する callable を返す。"""
+    def _v(x, t, cond):
+        # 毎 forward で adapter を確実に切替 (PEFT は内部で全モジュールに反映)
+        peft_model.set_adapter(adapter_name)
+        return AnimaBundle.dit_forward(peft_model, x, t, cond)
+    return _v
+
+
+def make_teacher_velocity_fn(teacher_transformer):
+    """frozen teacher 用 (LoRA なし、別 module instance)。"""
+    def _v(x, t, cond):
+        return AnimaBundle.dit_forward(teacher_transformer, x, t, cond)
+    return _v
+
+
+# ---------------------------------------------------------------------------
+# main
+# ---------------------------------------------------------------------------
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dataset", required=True, type=str)
+    ap.add_argument("--out", required=True, type=str)
+    ap.add_argument("--warm-lora", default="", type=str,
+                    help="Civitai Anima Turbo LoRA (ComfyUI 形式) を student adapter 初期値に")
+    ap.add_argument("--total-outer-steps", type=int, default=5000,
+                    help="1 outer step = N_critic + 1 generator")
+    ap.add_argument("--n-critic-per-gen", type=int, default=5)
+    ap.add_argument("--n-student-steps", type=int, default=4,
+                    help="few-step rollout の step 数 (1-4)、1=single-step distill")
+    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=3.0)
+    ap.add_argument("--student-cfg", type=float, default=1.0)
+    ap.add_argument("--shift", type=float, default=5.0,
+                    help="shifted uniform sampling shift (high-noise 側に偏らせる)")
+    ap.add_argument("--lora-rank", type=int, default=32)
+    ap.add_argument("--lr-gen", type=float, default=5e-6)
+    ap.add_argument("--lr-critic", 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("--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)
+    ap.add_argument("--neg-prompt", default="")
+    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, no LoRA) -----
+    print("[setup] deepcopy DiT for teacher (frozen, no LoRA)")
+    teacher_transformer = copy.deepcopy(bundle.transformer).to(device=device, dtype=dtype).eval()
+    for p in teacher_transformer.parameters():
+        p.requires_grad = False
+
+    # ----- 元 DiT (= bundle.transformer) に 2 adapter を attach -----
+    print("[setup] dual-adapter on base DiT (student + fake_score)")
+    peft_model = attach_dual_lora(bundle.transformer, rank=args.lora_rank,
+                                  adapter_names=("student", "fake_score"))
+    peft_model.to(device=device, dtype=dtype)
+    # base 凍結、LoRA だけ trainable
+    for n, p in peft_model.named_parameters():
+        p.requires_grad = ("lora_" in n)
+    bundle.transformer = peft_model
+
+    # 各 adapter の trainable params を log
+    student_params = [p for n, p in peft_model.named_parameters() if p.requires_grad and ".student." in n]
+    fake_params = [p for n, p in peft_model.named_parameters() if p.requires_grad and ".fake_score." in n]
+    print(f"[setup] student trainable: {sum(p.numel() for p in student_params)/1e6:.1f}M")
+    print(f"[setup] fake_score trainable: {sum(p.numel() for p in fake_params)/1e6:.1f}M")
+
+    # warm-start: Civitai Anima Turbo を student adapter に注入
+    if args.warm_lora:
+        print(f"[warm] loading {args.warm_lora} → student adapter")
+        # PEFT の 2-adapter 形式: key = base_model.model.<...>.lora_A.<adapter_name>.weight
+        sd_comfy = load_file(args.warm_lora)
+        sd_peft_default = convert_comfy_to_peft_lora(sd_comfy)  # adapter_name="default" 想定
+        # "default" → "student" にリネーム
+        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
+        # filter shape mismatch
+        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 (student / critic 独立) -----
+    opt_gen = torch.optim.AdamW(student_params, lr=args.lr_gen, betas=(0.9, 0.999),
+                                weight_decay=args.weight_decay, eps=1e-8)
+    opt_critic = torch.optim.AdamW(fake_params, lr=args.lr_critic, betas=(0.9, 0.999),
+                                   weight_decay=args.weight_decay, eps=1e-8)
+
+    # ----- velocity functions -----
+    student_v = make_velocity_fn(peft_model, "student")
+    fake_score_v = make_velocity_fn(peft_model, "fake_score")
+    teacher_v = make_teacher_velocity_fn(teacher_transformer)
+
+    # ----- dataset -----
+    print(f"[data] loading {args.dataset}")
+    dataset = TextOnlyDataset(args.dataset)
+    print(f"   {len(dataset)} captions")
+    loader = DataLoader(
+        dataset, batch_size=args.batch_size, shuffle=True,
+        num_workers=args.num_workers, collate_fn=text_collate, drop_last=True,
+    )
+
+    # ----- conditioning preencoded once for negative -----
+    with torch.no_grad():
+        cond_neg = bundle.text_encode([args.neg_prompt or ""])
+
+    H_lat = args.resolution // 8
+    W_lat = args.resolution // 8
+
+    # ----- training loop -----
+    print(f"[train] outer={args.total_outer_steps} cycle=({args.n_critic_per_gen} critic + 1 gen) "
+          f"n_student_steps={args.n_student_steps} res={args.resolution}")
+    log_path = out_dir / "dmd2_log.jsonl"
+    log_f = open(log_path, "a", buffering=1)
+    t0 = time.time()
+    data_iter = iter(loader)
+
+    def _next_batch():
+        nonlocal data_iter
+        try:
+            return next(data_iter)
+        except StopIteration:
+            data_iter = iter(loader)
+            return next(data_iter)
+
+    def _sample_inputs():
+        captions = _next_batch()
+        B = len(captions)
+        with torch.no_grad():
+            cond_pos = bundle.text_encode(captions)
+            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):
+        # ---- N critic updates ----
+        critic_metrics = {}
+        for _ in range(args.n_critic_per_gen):
+            noise, cond_pos, cond_neg_b = _sample_inputs()
+            peft_model.train()
+            loss_c, m_c = dmd2_critic_loss(
+                student_v, fake_score_v, noise, cond_pos, cond_neg_b,
+                student_cfg=args.student_cfg, n_steps=args.n_student_steps,
+                shift=args.shift,
+            )
+            opt_critic.zero_grad()
+            loss_c.backward()
+            torch.nn.utils.clip_grad_norm_(fake_params, args.grad_clip)
+            opt_critic.step()
+            critic_metrics = {k: float(v) for k, v in m_c.items()}
+
+        # ---- 1 generator update ----
+        noise, cond_pos, cond_neg_b = _sample_inputs()
+        peft_model.train()
+        loss_g, m_g = dmd2_generator_loss(
+            student_v, teacher_v, fake_score_v, noise, cond_pos, cond_neg_b,
+            teacher_cfg=args.teacher_cfg, student_cfg=args.student_cfg,
+            n_steps=args.n_student_steps, shift=args.shift,
+        )
+        opt_gen.zero_grad()
+        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,
+                       **critic_metrics, **gen_metrics}
+            log_f.write(json.dumps(metrics) + "\n")
+            msg = " ".join(f"{k}={v:.4f}" for k, v in metrics.items() if k not in ("outer",))
+            print(f"[outer {outer}/{args.total_outer_steps}] {msg}", flush=True)
+
+        # ---- checkpoint (student adapter のみ保存) ----
+        if outer > 0 and outer % args.sample_every == 0:
+            # student adapter のみを抽出
+            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
+            (out_dir).mkdir(parents=True, exist_ok=True)
+            _sf(sd, str(out_dir / f"dmd2_student_step{outer:05d}.safetensors"))
+            print(f"[save] dmd2_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 adapters")
+    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 / "dmd2_student_final.safetensors"))
+    sd_fake_final = {k: v.detach().cpu() for k, v in peft_model.state_dict().items()
+                     if "lora_" in k and ".fake_score." in k}
+    _sf(sd_fake_final, str(out_dir / "dmd2_fake_score_final.safetensors"))
+    log_f.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/train_dmdx.py

@@ -0,0 +1,248 @@
+#!/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()

scripts/distill/train_draftp.py

@@ -0,0 +1,250 @@
+#!/usr/bin/env python3
+"""
+Anima DRaFT+ / AlignProp training with HPSv2 reward
+====================================================
+
+これは **品質向上** (速度向上ではない) 用の蒸留 LoRA fine-tuning。
+既に蒸留された student LoRA を warm-start として、HPSv2 (Human Preference Score)
+を maximize する方向に追加学習する。
+
+Algorithm (DRaFT-K LV with KL regularization):
+  1. caption → cond_pos (no_grad)
+  2. init noise → student で N step rollout
+     - 前 N-K step: no_grad
+     - 後 K step: grad on (K=1 が paper の best)
+  3. final x0 → VAE decode (grad on) → image
+  4. reward = HPSv2(image, prompt)
+  5. KL term: ||v_pred - v_pred_base||² (LoRA disable して frozen base v_pred と比較)
+  6. loss = -reward + kl_coeff * KL
+"""
+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
+
+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_traj import (
+    TextOnlyDataset, text_collate, load_warm_lora, save_lora_state,
+)
+from distill.traj_scheduler import make_schedule
+
+
+def vae_decode_with_grad(bundle: AnimaBundle, latents: torch.Tensor) -> torch.Tensor:
+    """grad を通す VAE decode (anima_loader の vae_decode は @no_grad なので別実装)。"""
+    vae_dtype = next(bundle.vae.model.parameters()).dtype
+    latents = latents.to(dtype=vae_dtype)
+    return bundle.vae.model.decode(latents, bundle.vae_scale)
+
+
+def student_rollout_with_truncation(
+    student_v_fn,
+    base_v_fn_for_kl,
+    init_noise: torch.Tensor,
+    schedule_ts: torch.Tensor,    # (N+1,)
+    cond_pos: torch.Tensor,
+    student_cfg: float,
+    K: int,
+    capture_kl: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """rollout、最後の K step で grad on。KL 項用に最後の v_pred も別途 base で評価。
+
+    Returns: (x0_final, kl_loss)
+    """
+    B = init_noise.size(0)
+    device = init_noise.device
+    dtype = init_noise.dtype
+    N = len(schedule_ts) - 1
+    truncate_idx = N - K
+
+    x = init_noise
+    kl_loss = torch.zeros((), device=device)
+
+    for i in range(N):
+        t_cur = schedule_ts[i].expand(B).to(dtype=dtype)
+        t_next = schedule_ts[i + 1]
+        is_grad_step = (i >= truncate_idx)
+        ctx = torch.enable_grad() if is_grad_step else torch.no_grad()
+        with ctx:
+            v = student_v_fn(x, t_cur, cond_pos)
+            if is_grad_step and capture_kl and i == N - 1:
+                # 最後の step で KL term: student LoRA v vs frozen base v
+                with torch.no_grad():
+                    v_base = base_v_fn_for_kl(x, t_cur, cond_pos)
+                kl_loss = ((v - v_base.detach()).float() ** 2).mean()
+            dt = (t_next - schedule_ts[i]).to(device=device, dtype=dtype)
+            x = x + dt * v
+    return x, kl_loss
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dataset", required=True, type=str,
+                    help="caption-only dir (画像不要)")
+    ap.add_argument("--out", required=True, type=str)
+    ap.add_argument("--warm-lora", required=True, type=str,
+                    help="必須: 既に蒸留された student LoRA (例 ① Z-Image の出力)")
+    ap.add_argument("--hps-weights", default="/models/hpsv2/HPS_v2_compressed.pt",
+                    help="HPSv2 weights path、無ければ OpenCLIP baseline fallback")
+    ap.add_argument("--total-steps", type=int, default=1500)
+    ap.add_argument("--batch-size", type=int, default=2)
+    ap.add_argument("--grad-accum", type=int, default=1)
+    ap.add_argument("--n-student-steps", type=int, default=8)
+    ap.add_argument("--K", type=int, default=1, help="gradient truncation depth, paper best=1")
+    ap.add_argument("--n-lv-samples", type=int, default=2,
+                    help="DRaFT-LV: extra noise samples at last step, averaged")
+    ap.add_argument("--resolution", type=int, default=768)
+    ap.add_argument("--student-cfg", type=float, default=1.0)
+    ap.add_argument("--sigma-shift", type=float, default=3.0)
+    ap.add_argument("--lr", type=float, default=1e-4)
+    ap.add_argument("--kl-coeff", type=float, default=0.2,
+                    help="NeMo DRaFT+ default、reward hacking 防止")
+    ap.add_argument("--lora-rank", type=int, default=32)
+    ap.add_argument("--grad-clip", type=float, default=1.0)
+    ap.add_argument("--log-every", type=int, default=5)
+    ap.add_argument("--sample-every", type=int, default=200)
+    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)
+
+    # base = frozen deepcopy (KL term の reference として使う)
+    print("[setup] base = frozen deepcopy (for KL reference)")
+    base_transformer = copy.deepcopy(bundle.transformer).to(device=device, dtype=dtype).eval()
+    for p in base_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
+
+    # warm-start 必須
+    load_warm_lora(student_transformer, args.warm_lora)
+
+    # HPSv2 reward
+    print(f"[setup] loading HPSv2 from {args.hps_weights}")
+    from distill.hps_reward import HPSv2Reward
+    hps = HPSv2Reward(args.hps_weights, device=device, dtype=torch.float32)
+
+    # schedule (固定 N step)
+    sched = make_schedule(args.n_student_steps, args.sigma_shift,
+                          device=device, dtype=torch.float32)
+    print(f"[schedule] N={sched.num_steps} timesteps={sched.timesteps.tolist()}")
+
+    # dataset (caption-only)
+    print(f"[data] {args.dataset}")
+    dataset = TextOnlyDataset(args.dataset)
+    print(f"   {len(dataset)} captions")
+    loader = DataLoader(
+        dataset, batch_size=args.batch_size, shuffle=True,
+        num_workers=args.num_workers, collate_fn=text_collate,
+        drop_last=True,
+    )
+
+    opt = torch.optim.AdamW(student_params, lr=args.lr, betas=(0.9, 0.999),
+                            weight_decay=0.01, eps=1e-8)
+
+    def student_v_fn(x, t, cond):
+        return AnimaBundle.dit_forward(student_transformer, x, t, cond)
+
+    def base_v_fn(x, t, cond):
+        return AnimaBundle.dit_forward(base_transformer, x, t, cond)
+
+    H_lat = args.resolution // 8
+    W_lat = args.resolution // 8
+
+    print(f"[train] steps={args.total_steps} bs={args.batch_size} N={args.n_student_steps} "
+          f"K={args.K} lv={args.n_lv_samples} kl={args.kl_coeff}")
+    log_path = out_dir / "draftp_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):
+            captions = _next()
+            with torch.no_grad():
+                cond_pos = bundle.text_encode(captions)
+            B = cond_pos.size(0)
+            init_noise = torch.randn(B, 16, 1, H_lat, W_lat, device=device, dtype=dtype)
+
+            # DRaFT-LV: rollout once + n_lv_samples 個の last-step alternative を試して平均
+            x0_final, kl = student_rollout_with_truncation(
+                student_v_fn, base_v_fn, init_noise, sched.timesteps,
+                cond_pos, args.student_cfg, args.K, capture_kl=True,
+            )
+
+            # VAE decode (grad on)
+            img = vae_decode_with_grad(bundle, x0_final).squeeze(2)  # (B,3,H,W) in [-1,1]
+            reward = hps.score(img, captions)  # (B,)
+            r_mean = reward.mean()
+
+            # DRaFT+ loss: -reward + kl_coeff * kl
+            loss = (-r_mean + args.kl_coeff * kl) / args.grad_accum
+            loss.backward()
+            metrics = {
+                "reward_mean": float(r_mean.detach()),
+                "reward_std": float(reward.std().detach()),
+                "kl": float(kl.detach()),
+                "loss": float((-r_mean + args.kl_coeff * kl).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"draftp_step{step:05d}")
+            print(f"[save] draftp_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, "draftp_final")
+    log_f.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/train_ladd.py

@@ -0,0 +1,293 @@
+#!/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()

scripts/distill/train_pcm.py

@@ -0,0 +1,250 @@
+#!/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()

scripts/distill/train_reflow.py

@@ -0,0 +1,246 @@
+#!/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()

scripts/distill/train_shortcut.py

@@ -0,0 +1,226 @@
+#!/usr/bin/env python3
+"""
+Anima Shortcut Models distillation (Frans et al., 2024)
+
+特徴:
+  - 単一 LoRA、d を入力で取って 1/2/4/8/128-step 自在に切替可能
+  - Flow-matching half (d=0) + Bootstrap half (d>0) を 1 step 内で混在
+  - PCM と違い phase 固定なし、d を連続値で扱える
+
+データ:
+  Reflow cache (--save-noise 付き) を流用。(noise, x0, emb) triplet。
+"""
+from __future__ import annotations
+import argparse
+import copy
+import json
+import math
+import os
+import sys
+import time
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+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_traj import load_warm_lora, save_lora_state
+from distill.train_reflow import ReflowPairDataset, reflow_collate
+from distill.shortcut_module import attach_shortcut_d_head, set_shortcut_d, shortcut_d_head_params
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--cache-dir", required=True, type=str,
+                    help="Reflow cache (--save-noise 付き)")
+    ap.add_argument("--out", required=True, type=str)
+    ap.add_argument("--warm-lora", default="", type=str)
+    ap.add_argument("--total-steps", type=int, default=2000)
+    ap.add_argument("--batch-size", type=int, default=4)
+    ap.add_argument("--grad-accum", type=int, default=2)
+    ap.add_argument("--denoise-timesteps", type=int, default=128,
+                    help="discrete grid for sampling t and d (2^k granularity)")
+    ap.add_argument("--bootstrap-every", type=int, default=8,
+                    help="batch 内 bootstrap 比率: 1/N (paper default 8)")
+    ap.add_argument("--resolution", type=int, default=768)
+    ap.add_argument("--lr", type=float, default=2e-5)
+    ap.add_argument("--lr-d-head", type=float, default=5e-4,
+                    help="d_head は zero-init なので高めの lr で立ち上げ")
+    ap.add_argument("--lora-rank", type=int, default=32)
+    ap.add_argument("--grad-clip", type=float, default=1.0)
+    ap.add_argument("--clip-x-bootstrap", type=float, default=4.0,
+                    help="bootstrap 中の x_t clip range")
+    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)
+
+    # student = wide LoRA + d-head
+    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)
+    attach_shortcut_d_head(student_transformer)  # adds d_head + hook + _current_d attr
+    # d_head の params は trainable に
+    d_head_params = shortcut_d_head_params(student_transformer)
+    for p in d_head_params:
+        p.requires_grad = True
+    student_lora_params = [p for n, p in student_transformer.named_parameters()
+                           if p.requires_grad and "lora_" in n]
+    print(f"[setup] student LoRA: {sum(p.numel() for p in student_lora_params)/1e6:.1f}M")
+    print(f"[setup] d_head:       {sum(p.numel() for p in d_head_params)/1e6:.1f}M")
+    bundle.transformer = student_transformer
+
+    if args.warm_lora:
+        load_warm_lora(student_transformer, args.warm_lora)
+
+    # optimizers (LoRA と d_head で別 lr)
+    opt_lora = torch.optim.AdamW(student_lora_params, lr=args.lr,
+                                 betas=(0.9, 0.999), weight_decay=0.01)
+    opt_d_head = torch.optim.AdamW(d_head_params, lr=args.lr_d_head,
+                                   betas=(0.9, 0.999), weight_decay=0.0)
+
+    # dataset
+    print(f"[data] {args.cache_dir}")
+    dataset = ReflowPairDataset(args.cache_dir)
+    print(f"   {len(dataset)} 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,
+    )
+
+    T = args.denoise_timesteps
+    log2_sections = int(math.log2(T))  # 7 for T=128
+
+    def student_v_with_d(x, t, cond, d):
+        set_shortcut_d(student_transformer, d)
+        try:
+            return AnimaBundle.dit_forward(student_transformer, x, t, cond)
+        finally:
+            set_shortcut_d(student_transformer, None)
+
+    print(f"[train] steps={args.total_steps} bs={args.batch_size} accum={args.grad_accum} "
+          f"T={T} bootstrap_every={args.bootstrap_every}")
+    log_path = out_dir / "shortcut_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_lora.zero_grad()
+        opt_d_head.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)
+            B_boot = max(1, B // args.bootstrap_every)
+            B_flow = B - B_boot
+
+            # ----- Flow-matching half (d=0, infinitesimal) -----
+            t_fm_idx = torch.randint(0, T, (B_flow,), device=device)
+            t_fm = t_fm_idx.float() / T
+            t_fm_b = t_fm.view(-1, *([1] * (x0.dim() - 1)))
+            x0_fm = x0[:B_flow]; noise_fm = noise[:B_flow]; emb_fm = emb[:B_flow]
+            x_t_fm = (1 - t_fm_b) * x0_fm + t_fm_b * noise_fm
+            v_tgt_fm = noise_fm - x0_fm
+            d_fm = torch.zeros(B_flow, device=device, dtype=dtype)
+
+            # ----- Bootstrap half (d > 0, self-consistency) -----
+            if B_boot > 0:
+                k = torch.randint(0, log2_sections, (B_boot,), device=device)
+                d_b = (1.0 / (2.0 ** k.float())).to(dtype=dtype)  # d ∈ {1/2^k}
+                # t aligned to d grid: t_b = (random int) / 2^k * 1
+                t_b_max = (2 ** k).float()
+                t_b_idx = (torch.rand(B_boot, device=device) * t_b_max).floor()
+                t_b = t_b_idx / t_b_max.clamp(min=1.0)
+                t_b_b = t_b.view(-1, *([1] * (x0.dim() - 1)))
+                x0_b = x0[B_flow:]; noise_b = noise[B_flow:]; emb_b = emb[B_flow:]
+                x_t_b = (1 - t_b_b) * x0_b + t_b_b * noise_b
+
+                # 2 sub-step bootstrap target (no_grad)
+                with torch.no_grad():
+                    d_half = d_b * 0.5
+                    v1 = student_v_with_d(x_t_b, t_b.to(dtype=dtype), emb_b, d_half)
+                    dt_half = d_half.view(-1, *([1] * (x_t_b.dim() - 1)))
+                    x_t2 = (x_t_b + dt_half * v1).clamp(-args.clip_x_bootstrap, args.clip_x_bootstrap)
+                    t_b_half = (t_b + d_half.float()).clamp(0.0, 1.0)
+                    v2 = student_v_with_d(x_t2, t_b_half.to(dtype=dtype), emb_b, d_half)
+                    v_tgt_b = 0.5 * (v1 + v2)
+
+                # concat batch
+                x_cat = torch.cat([x_t_fm, x_t_b], dim=0)
+                t_cat = torch.cat([t_fm.to(dtype=dtype), t_b.to(dtype=dtype)], dim=0)
+                d_cat = torch.cat([d_fm, d_b], dim=0)
+                emb_cat = torch.cat([emb_fm, emb_b], dim=0)
+                v_tgt_cat = torch.cat([v_tgt_fm, v_tgt_b], dim=0)
+            else:
+                x_cat, t_cat, d_cat, emb_cat, v_tgt_cat = x_t_fm, t_fm.to(dtype=dtype), d_fm, emb_fm, v_tgt_fm
+
+            # 1 trainable forward
+            v_pred = student_v_with_d(x_cat, t_cat, emb_cat, d_cat)
+            loss = F.mse_loss(v_pred.float(), v_tgt_cat.detach().float()) / args.grad_accum
+            loss.backward()
+            metrics = {
+                "loss": float((loss * args.grad_accum).detach()),
+                "B_flow": B_flow, "B_boot": B_boot,
+                "v_pred_abs": float(v_pred.detach().abs().mean()),
+                "v_tgt_abs": float(v_tgt_cat.detach().abs().mean()),
+            }
+        torch.nn.utils.clip_grad_norm_(student_lora_params + d_head_params, args.grad_clip)
+        opt_lora.step()
+        opt_d_head.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:
+            # LoRA + d_head 両方保存
+            save_lora_state(student_transformer, out_dir, f"shortcut_step{step:05d}")
+            torch.save(student_transformer.d_head.state_dict(),
+                       out_dir / f"shortcut_d_head_step{step:05d}.pt")
+            print(f"[save] shortcut_step{step:05d}", 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, "shortcut_final")
+    torch.save(student_transformer.d_head.state_dict(), out_dir / "shortcut_d_head_final.pt")
+    log_f.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/train_sid.py

@@ -0,0 +1,236 @@
+#!/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()

scripts/distill/train_sota.py

@@ -0,0 +1,260 @@
+#!/usr/bin/env python3
+"""
+Anima SOTA distillation main entrypoint
+========================================
+
+3 phase progressive 蒸留:
+  Phase A: 8-step, DMD2 + light TSCD + light adv
+  Phase B: 4-step, balanced
+  Phase C: 2-step, full adv weight
+
+各 phase 完了時に LoRA を保存し、次 phase へ。
+
+使い方 (Modal 経由):
+  modal run --detach modal_app.py::train_sota_distill --phase a
+  modal run --detach modal_app.py::train_sota_distill --phase b --resume /output/distill/phase_a_final
+  modal run --detach modal_app.py::train_sota_distill --phase c --resume /output/distill/phase_b_final
+"""
+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
+from safetensors.torch import save_file
+
+# import-path: 同ディレクトリ前提
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from distill.anima_loader import AnimaPaths, build_anima
+from distill.dataset import AnimaImageCaptionDataset, collate_fn
+from distill.dmd2_trainer import (
+    DMD2Args, DMD2Trainer, attach_qv_lora, attach_wide_lora, set_lora_scale,
+)
+from distill.r3gan_disc import R3GANDiscriminator
+
+
+# ----- phase 設定 -----------------------------------------------------------
+# Phase A postmortem の教訓: R3GAN gamma calibration ミスで gen 破壊。
+# Phase B/C は **R3GAN を完全 disable** (adv_weight=0) で DMD+TSCD のみで蒸留。
+# R3GAN は後付け実験(別途 phase b_r3gan 等)で導入する想定。
+PHASE_CONFIGS = {
+    "a": DMD2Args(
+        num_inference_steps=8,
+        tscd_weight=0.3, adv_weight=0.0, r3gan_gamma=0.0,  # disabled
+        cfg_aug_cold_steps=200, dynamic_decay_steps=3000,
+    ),
+    "b": DMD2Args(
+        num_inference_steps=4,
+        tscd_weight=0.5, adv_weight=0.0, r3gan_gamma=0.0,  # disabled
+        cfg_aug_cold_steps=100, dynamic_decay_steps=2000,
+    ),
+    "c": DMD2Args(
+        num_inference_steps=2,
+        tscd_weight=0.5, adv_weight=0.0, r3gan_gamma=0.0,  # disabled
+        cfg_aug_cold_steps=50, dynamic_decay_steps=1500,
+    ),
+}
+
+
+def save_lora_state(model, path: Path, name: str) -> None:
+    """PEFT-wrapped model から LoRA weight だけ抽出して safetensors で保存。"""
+    path.mkdir(parents=True, exist_ok=True)
+    sd = {k: v.detach().cpu() for k, v in model.state_dict().items() if "lora_" in k}
+    save_file(sd, str(path / f"{name}.safetensors"))
+
+
+def save_full_state(state_dict: dict, path: Path) -> None:
+    """フル状態(gen の全 weights + buffers)を safetensors で保存。
+    重要: float 限定で filter すると RoPE 等 integer buffer が落ちて model 壊れる。"""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    cpu_sd = {k: v.detach().cpu().contiguous() for k, v in state_dict.items()
+              if isinstance(v, torch.Tensor)}
+    save_file(cpu_sd, str(path))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--phase", choices=["a", "b", "c"], required=True)
+    ap.add_argument("--dataset", required=True, type=str,
+                    help="/dataset/raw のような image+caption ディレクトリ")
+    ap.add_argument("--out", required=True, type=str,
+                    help="checkpoint 出力先ディレクトリ")
+    ap.add_argument("--resume", default=None, type=str,
+                    help="前 phase の出力ディレクトリ (phase B/C で必須)")
+    ap.add_argument("--total-steps", type=int, default=3000)
+    ap.add_argument("--sample-every", type=int, default=500)
+    ap.add_argument("--log-every", type=int, default=10)
+    ap.add_argument("--resolution", type=int, default=1024)
+    ap.add_argument("--num-workers", type=int, default=4)
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--bf16", action="store_true", default=True)
+    ap.add_argument("--gen-lora-only", action="store_true", default=False,
+                    help="gen にも Q,V LoRA を attach し base を凍結。Phase B/C 用。"
+                         "出力は LoRA weights のみ (数百MB) になる。")
+    ap.add_argument("--override-adv-weight", type=float, default=-1.0,
+                    help="PHASE_CONFIGS の adv_weight を上書き (−1 = 上書きしない)")
+    ap.add_argument("--override-r3gan-gamma", type=float, default=-1.0,
+                    help="PHASE_CONFIGS の r3gan_gamma を上書き (−1 = 上書きしない)")
+    args = ap.parse_args()
+
+    torch.manual_seed(args.seed)
+    device = torch.device("cuda")
+    dtype = torch.bfloat16 if args.bf16 else torch.float32
+
+    out_dir = Path(args.out)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # ----- load Anima base -----
+    print("[load] Anima bundle (DiT + Qwen3 + WanVAE + LLMAdapter)")
+    bundle = build_anima(device=device, dtype=dtype)
+
+    # ----- gen / guidance を base から複製 -----
+    print("[setup] cloning DiT to gen + guidance")
+    gen_transformer = copy.deepcopy(bundle.transformer).to(device=device, dtype=dtype)
+    guidance_transformer = copy.deepcopy(bundle.transformer).to(device=device, dtype=dtype)
+    # bundle.transformer は推論用途で残してもよいが、メモリ節約のため None 化:
+    bundle.transformer = None
+    torch.cuda.empty_cache()
+
+    # ----- LoRA を guidance に attach -----
+    args_phase = PHASE_CONFIGS[args.phase]
+    print(f"[setup] attaching Q,V LoRA (rank={args_phase.lora_rank}) to guidance")
+    guidance_transformer = attach_qv_lora(guidance_transformer, rank=args_phase.lora_rank)
+    guidance_transformer.to(device=device, dtype=dtype)
+
+    # base 部分の grad を切る (LoRA だけ訓練)
+    for n, p in guidance_transformer.named_parameters():
+        p.requires_grad = ("lora_" in n)
+
+    # ----- resume (前 phase の gen 重みを新 base として引き継ぐ) -----
+    # Phase B/C は Phase A の蒸留済 DiT を base にする想定。LoRA-only 学習時は
+    # この段階で gen_transformer に Phase A の重みをロード → その後 LoRA attach + 凍結。
+    if args.resume:
+        resume_dir = Path(args.resume)
+        gen_ckpt = resume_dir / "gen_final.safetensors"
+        if gen_ckpt.exists():
+            print(f"[resume] loading gen from {gen_ckpt}")
+            from safetensors.torch import load_file
+            sd = load_file(str(gen_ckpt))
+            missing, unexpected = gen_transformer.load_state_dict(sd, strict=False)
+            print(f"   missing={len(missing)} unexpected={len(unexpected)}")
+
+    # ----- CLI override for grid search / experimentation -----
+    if args.override_adv_weight >= 0:
+        args_phase.adv_weight = args.override_adv_weight
+        print(f"[override] adv_weight -> {args_phase.adv_weight}")
+    if args.override_r3gan_gamma >= 0:
+        args_phase.r3gan_gamma = args.override_r3gan_gamma
+        print(f"[override] r3gan_gamma -> {args_phase.r3gan_gamma}")
+
+    # ----- LoRA-only モード: gen に wide LoRA (all-linear) を attach + base 凍結 -----
+    # 注: Q,V LoRA だと AdaLN modulation を触れず、Anima のような timestep 依存
+    #     architecture では蒸留信号が degenerate に収束する (Phase B v1 で失敗)。
+    #     wide にして AdaLN + attention + MLP を全部学習対象に。
+    if args.gen_lora_only:
+        print(f"[setup] gen-lora-only mode: attaching WIDE LoRA (rank={args_phase.lora_rank}) to gen")
+        gen_transformer = attach_wide_lora(gen_transformer, rank=args_phase.lora_rank)
+        gen_transformer.to(device=device, dtype=dtype)
+        # base を凍結、LoRA だけ trainable
+        for n, p in gen_transformer.named_parameters():
+            p.requires_grad = ("lora_" in n)
+        trainable = sum(p.numel() for p in gen_transformer.parameters() if p.requires_grad)
+        total = sum(p.numel() for p in gen_transformer.parameters())
+        print(f"   gen trainable params: {trainable/1e6:.1f}M / {total/1e6:.1f}M total")
+
+    # ----- discriminator -----
+    # adv_weight=0 のときは構築しない (Phase A 教訓: R3GAN 自体が不安定要因)
+    if args_phase.adv_weight > 0:
+        print("[setup] R3GAN discriminator (latent space, text-conditional)")
+        disc = R3GANDiscriminator(latent_channels=16, cond_dim=1024).to(device=device, dtype=dtype)
+    else:
+        print("[setup] R3GAN DISABLED (adv_weight=0) — DMD+TSCD のみで蒸留")
+        # placeholder: trainer は disc を受け取るが forward は呼ばれない
+        disc = R3GANDiscriminator(latent_channels=16, cond_dim=1024).to(device=device, dtype=dtype)
+
+    # ----- dataset -----
+    print(f"[data] loading {args.dataset}")
+    dataset = AnimaImageCaptionDataset(args.dataset, resolution=args.resolution)
+    print(f"   {len(dataset)} samples")
+    loader = DataLoader(
+        dataset,
+        batch_size=args_phase.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        collate_fn=collate_fn,
+        drop_last=True,
+        pin_memory=True,
+    )
+
+    # ----- trainer -----
+    trainer = DMD2Trainer(
+        bundle=bundle,
+        gen_transformer=gen_transformer,
+        guidance_transformer=guidance_transformer,
+        discriminator=disc,
+        args=args_phase,
+    )
+
+    # ----- training loop -----
+    print(f"[train] phase={args.phase} steps={args.total_steps} batch={args_phase.batch_size}")
+    log_path = out_dir / f"phase_{args.phase}_log.jsonl"
+    log_f = open(log_path, "a", buffering=1)
+    t0 = time.time()
+    data_iter = iter(loader)
+
+    for step in range(args.total_steps):
+        try:
+            batch = next(data_iter)
+        except StopIteration:
+            data_iter = iter(loader)
+            batch = next(data_iter)
+
+        metrics = trainer.train_step(batch)
+
+        if step % args.log_every == 0:
+            scalars = {k: float(v) for k, v in metrics.items()}
+            scalars["step"] = step
+            scalars["elapsed"] = time.time() - t0
+            log_f.write(json.dumps(scalars) + "\n")
+            msg = " ".join(f"{k}={v:.4f}" for k, v in scalars.items() if k not in ("step", "elapsed"))
+            print(f"[step {step}/{args.total_steps}] elapsed={scalars['elapsed']:.0f}s | {msg}")
+
+        if step > 0 and step % args.sample_every == 0:
+            if args.gen_lora_only:
+                save_lora_state(trainer.gen, out_dir, f"phase_{args.phase}_step{step:05d}_gen_lora")
+                print(f"[save] phase_{args.phase}_step{step:05d}_gen_lora.safetensors", flush=True)
+            else:
+                ckpt = out_dir / f"phase_{args.phase}_step{step:05d}_gen.safetensors"
+                save_full_state(trainer.gen.state_dict(), ckpt)
+                print(f"[save] {ckpt}", flush=True)
+            # Modal volume の commit を内部で発行(途中クラッシュからの保全)
+            try:
+                import modal
+                vol = modal.Volume.from_name("anima-outputs")
+                vol.commit()
+                print(f"[save] volume committed at step {step}", flush=True)
+            except Exception as e:
+                print(f"[save] volume commit failed: {e}", flush=True)
+
+    # final save
+    print("[done] saving final outputs")
+    if args.gen_lora_only:
+        # LoRA-only: gen の LoRA weights だけ保存(数百 MB)
+        save_lora_state(trainer.gen, out_dir, "gen_lora_final")
+        print(f"  -> gen_lora_final.safetensors (LoRA-only mode)")
+    else:
+        # Full DiT: gen の全 weights を保存(~4 GB)
+        save_full_state(trainer.gen.state_dict(), out_dir / "gen_final.safetensors")
+        print(f"  -> gen_final.safetensors (full DiT)")
+    # guidance LoRA は両モード共通で残しておく(分析・実験用)
+    save_lora_state(trainer.guidance, out_dir, "guidance_lora_final")
+    log_f.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/train_traj.py

@@ -0,0 +1,363 @@
+#!/usr/bin/env python3
+"""
+Anima trajectory imitation distillation
+========================================
+
+DiffSynth-Studio の Z-Image trajectory imitation を Anima に最小移植したもの。
+**critic なし、単一ネットワーク**で前回 5 回失敗の R3GAN 不安定性を完全回避。
+
+アルゴリズム:
+  - teacher = Anima base (frozen, no LoRA) を 50-step CFG=2.0 で rollout
+  - student = Anima base + wide LoRA (rank 32) を 8-step CFG=1.0 で rollout
+  - L_align = MSE(student velocity, teacher segment velocity) on 8 segments
+  - L_reg   = LPIPS(student final image, teacher final image)  (任意、--lpips-weight>0 で有効)
+
+warm-start:
+  --warm-lora /models/loras/anima_turbo.safetensors  で Civitai 公式 Anima Turbo LoRA を
+  student LoRA 初期値に注入できる。format は ComfyUI 形式 (diffusion_model.<...>.lora_A.weight)、
+  内部で PEFT 形式に変換してロード。
+
+使い方 (Modal 経由):
+  # Smoke test (1 step, sanity check)
+  modal run modal_app.py::train_traj_imitation --total-steps 1 --batch-size 1 \\
+      --teacher-steps 12 --student-steps 8 --lpips-weight 0.0
+
+  # 本番 (2000 step, ~$45, B200)
+  modal run --detach modal_app.py::train_traj_imitation \\
+      --total-steps 2000 --batch-size 1 --teacher-steps 50 --student-steps 8 \\
+      --warm-lora /models/loras/anima_turbo.safetensors --lpips-weight 0.1
+"""
+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
+
+# import-path: 同ディレクトリ前提
+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.traj_scheduler import make_schedule, snap_to_targets
+from distill.traj_loss import (
+    fetch_trajectory, align_trajectory, compute_regularization,
+)
+
+
+# ---------------------------------------------------------------------------
+# Text-only dataset: caption だけ使う (trajectory imitation は init=noise なので image 不要)
+# ---------------------------------------------------------------------------
+class TextOnlyDataset(Dataset):
+    def __init__(self, root: str | Path):
+        self.root = Path(root)
+        # *.txt を再帰的に拾う。/dataset/cleaned 配下を想定。
+        self.files = sorted(self.root.rglob("*.txt"))
+        if not self.files:
+            raise RuntimeError(f"No .txt files under {self.root}")
+
+    def __len__(self) -> int:
+        return len(self.files)
+
+    def __getitem__(self, idx: int) -> str:
+        return self.files[idx].read_text(encoding="utf-8").strip()
+
+
+def text_collate(batch: list[str]) -> list[str]:
+    return list(batch)
+
+
+# ---------------------------------------------------------------------------
+# ComfyUI Anima LoRA -> PEFT 形式変換 (warm-start 用)
+# 既存の _convert_peft_to_comfy_lora の逆向き
+# ---------------------------------------------------------------------------
+def convert_comfy_to_peft_lora(sd: dict) -> dict:
+    """ComfyUI(Anima/Cosmos)形式 LoRA を PEFT 形式に変換。
+       Comfy:  'diffusion_model.<module>.lora_A.weight'
+       PEFT:   'base_model.model.<module>.lora_A.default.weight'
+    """
+    out = {}
+    for k, v in sd.items():
+        nk = k
+        if nk.startswith("diffusion_model."):
+            nk = nk[len("diffusion_model."):]
+        # adapter name '.default.' を挿入
+        nk = nk.replace(".lora_A.weight", ".lora_A.default.weight")
+        nk = nk.replace(".lora_B.weight", ".lora_B.default.weight")
+        nk = nk.replace(".lora_A.bias", ".lora_A.default.bias")
+        nk = nk.replace(".lora_B.bias", ".lora_B.default.bias")
+        nk = "base_model.model." + nk
+        out[nk] = v
+    return out
+
+
+def load_warm_lora(student_model, warm_lora_path: str) -> None:
+    """Civitai 形式の Anima Turbo LoRA を student の PEFT 形式に変換して注入。
+       shape / rank が合わない key は skip して警告。"""
+    print(f"[warm] loading {warm_lora_path}")
+    sd_comfy = load_file(warm_lora_path)
+    print(f"[warm] {len(sd_comfy)} comfy keys, sample: {list(sd_comfy.keys())[:3]}")
+    sd_peft = convert_comfy_to_peft_lora(sd_comfy)
+    print(f"[warm] converted to {len(sd_peft)} peft keys, sample: {list(sd_peft.keys())[:3]}")
+
+    # student の現 state_dict と shape を照合して filter
+    model_sd = student_model.state_dict()
+    matched, skipped_shape, missing = 0, 0, 0
+    to_load = {}
+    for k, v in sd_peft.items():
+        if k not in model_sd:
+            missing += 1
+            continue
+        if model_sd[k].shape != v.shape:
+            skipped_shape += 1
+            continue
+        to_load[k] = v.to(dtype=model_sd[k].dtype)
+        matched += 1
+    print(f"[warm] matched={matched} skipped_shape={skipped_shape} missing={missing}")
+    if matched == 0:
+        raise RuntimeError(
+            "No LoRA keys matched. Check format conversion & target_modules of wide LoRA."
+        )
+    missing_keys, _ = student_model.load_state_dict(to_load, strict=False)
+    # PEFT 内部の LoRA キーだけ載れば OK、その他 missing は base 側なので無視
+    lora_missing = [k for k in missing_keys if "lora_" in k]
+    if lora_missing:
+        print(f"[warm] WARN: {len(lora_missing)} lora keys not loaded (e.g. {lora_missing[:3]})")
+
+
+# ---------------------------------------------------------------------------
+# save / restore
+# ---------------------------------------------------------------------------
+def save_lora_state(model, path: Path, name: str) -> None:
+    path.mkdir(parents=True, exist_ok=True)
+    sd = {k: v.detach().cpu() for k, v in model.state_dict().items() if "lora_" in k}
+    save_file(sd, str(path / f"{name}.safetensors"))
+
+
+# ---------------------------------------------------------------------------
+# main
+# ---------------------------------------------------------------------------
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dataset", required=True, type=str,
+                    help="caption (*.txt) を含むディレクトリ")
+    ap.add_argument("--out", required=True, type=str,
+                    help="LoRA 出力ディレクトリ")
+    ap.add_argument("--warm-lora", default="", type=str,
+                    help="Civitai Anima Turbo LoRA path (ComfyUI 形式)。空なら cold-start")
+    ap.add_argument("--total-steps", type=int, default=2000)
+    ap.add_argument("--batch-size", type=int, default=1)
+    ap.add_argument("--teacher-steps", type=int, default=50)
+    ap.add_argument("--student-steps", type=int, default=8)
+    ap.add_argument("--teacher-cfg", type=float, default=2.0)
+    ap.add_argument("--student-cfg", type=float, default=1.0)
+    ap.add_argument("--sigma-shift", type=float, default=3.0,
+                    help="Anima 公式 workflow は 3.0")
+    ap.add_argument("--lora-rank", type=int, default=32)
+    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("--lpips-weight", type=float, default=0.0,
+                    help="LPIPS regularization weight。0 で disable (start 推奨)、後で 0.1-0.5 に")
+    ap.add_argument("--resolution", type=int, default=1024)
+    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)
+    ap.add_argument("--weight-mode", default="uniform", choices=["uniform", "inv_sigma"])
+    ap.add_argument("--neg-prompt", default="",
+                    help="negative prompt for teacher CFG (空文字 = empty conditioning)")
+    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 (DiT + Qwen3 + WanVAE + LLMAdapter)")
+    bundle = build_anima(device=device, dtype=dtype)
+
+    # ----- teacher = base のクローン (frozen, no LoRA) -----
+    print("[setup] cloning DiT for teacher (frozen) and student (wide LoRA)")
+    teacher_transformer = copy.deepcopy(bundle.transformer).to(device=device, dtype=dtype).eval()
+    for p in teacher_transformer.parameters():
+        p.requires_grad = False
+
+    # ----- student = base に wide LoRA を attach -----
+    student_transformer = bundle.transformer  # bundle 側の方を再利用 (メモリ節約)
+    student_transformer = attach_wide_lora(student_transformer, rank=args.lora_rank)
+    student_transformer.to(device=device, dtype=dtype)
+    # base 凍結、LoRA だけ trainable
+    for n, p in student_transformer.named_parameters():
+        p.requires_grad = ("lora_" in n)
+    trainable = sum(p.numel() for p in student_transformer.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in student_transformer.parameters())
+    print(f"[setup] student trainable: {trainable/1e6:.1f}M / {total/1e6:.1f}M")
+    # bundle.transformer は student と同じ object (LoRA wrap 済)。
+    # text_encode/vae_*以外には bundle.transformer 不要なので参照だけ残す。
+    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 for regularization")
+
+    # ----- schedules -----
+    # student schedule: 推論時に使う 8-step grid (t=1 → 0)
+    student_sched = make_schedule(args.student_steps, args.sigma_shift, device=device, dtype=torch.float32)
+    student_ts = student_sched.timesteps  # (N+1,)
+    print(f"[schedule] student t = {student_ts.tolist()}")
+
+    # ----- dataset -----
+    print(f"[data] loading {args.dataset}")
+    dataset = TextOnlyDataset(args.dataset)
+    print(f"   {len(dataset)} captions")
+    loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        collate_fn=text_collate,
+        drop_last=True,
+    )
+
+    # ----- optimizer -----
+    opt = torch.optim.AdamW(
+        [p for p in student_transformer.parameters() if p.requires_grad],
+        lr=args.lr, betas=(0.9, 0.999), weight_decay=args.weight_decay, eps=1e-8,
+    )
+
+    # ----- training loop -----
+    print(f"[train] total_steps={args.total_steps} batch={args.batch_size} "
+          f"teacher={args.teacher_steps} student={args.student_steps} "
+          f"lpips_weight={args.lpips_weight}")
+    log_path = out_dir / "traj_log.jsonl"
+    log_f = open(log_path, "a", buffering=1)
+    t0 = time.time()
+    data_iter = iter(loader)
+
+    # student velocity_fn: バッチ次元 latents + (B,) timesteps + cond -> velocity
+    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)
+
+    # negative cond は 1 度だけエンコード
+    with torch.no_grad():
+        cond_neg = bundle.text_encode([args.neg_prompt or ""])  # (1, 512, 1024)
+    # latent サイズは VAE 8x down 想定 (Anima WanVAE)
+    H_lat = args.resolution // 8
+    W_lat = args.resolution // 8
+
+    for step in range(args.total_steps):
+        try:
+            captions = next(data_iter)
+        except StopIteration:
+            data_iter = iter(loader)
+            captions = next(data_iter)
+
+        B = len(captions)
+
+        # text encode (positive)
+        with torch.no_grad():
+            cond_pos = bundle.text_encode(captions)  # (B, 512, 1024)
+            # neg は B にブロードキャスト
+            cond_neg_b = cond_neg.expand(B, -1, -1).contiguous() if B > 1 else cond_neg
+
+        # 初期 noise (t=1 の状態)
+        init_noise = torch.randn(
+            B, 16, 1, H_lat, W_lat, device=device, dtype=dtype,
+        )
+
+        # ----- teacher rollout (no_grad inside fetch_trajectory) -----
+        teacher_sched = make_schedule(args.teacher_steps, args.sigma_shift,
+                                      device=device, dtype=torch.float32)
+        teacher_sched = snap_to_targets(teacher_sched, student_ts)
+        teacher_transformer.eval()
+        teacher_traj = fetch_trajectory(
+            _teacher_v, init_noise, teacher_sched, student_ts,
+            cond_pos, cond_neg_b, cfg_scale=args.teacher_cfg,
+        )
+
+        # ----- student align loss -----
+        student_transformer.train()
+        l_align = align_trajectory(
+            _student_v, teacher_traj, student_ts,
+            cond_pos, cond_neg_b, cfg_scale_student=args.student_cfg,
+            weight_mode=args.weight_mode,
+        )
+
+        # ----- LPIPS reg (任意) -----
+        if lpips_fn is not None:
+            l_reg = compute_regularization(
+                _student_v, bundle, init_noise,
+                teacher_final_latent=teacher_traj[-1],
+                student_timesteps=student_ts,
+                cond_pos=cond_pos, cond_neg=cond_neg_b,
+                lpips_fn=lpips_fn,
+                cfg_scale=args.student_cfg,
+            )
+        else:
+            l_reg = torch.zeros((), device=device)
+
+        loss = l_align + args.lpips_weight * l_reg
+
+        opt.zero_grad()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(
+            [p for p in student_transformer.parameters() if p.requires_grad],
+            args.grad_clip,
+        )
+        opt.step()
+
+        # ----- log -----
+        if step % args.log_every == 0:
+            metrics = {
+                "step": step,
+                "elapsed": time.time() - t0,
+                "loss": float(loss.detach()),
+                "l_align": float(l_align.detach()),
+                "l_reg": float(l_reg.detach()),
+            }
+            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)
+
+        # ----- checkpoint -----
+        if step > 0 and step % args.sample_every == 0:
+            save_lora_state(student_transformer, out_dir, f"traj_step{step:05d}")
+            print(f"[save] traj_step{step:05d}.safetensors", flush=True)
+            try:
+                import modal
+                modal.Volume.from_name("anima-outputs").commit()
+                print(f"[save] volume committed at step {step}", flush=True)
+            except Exception as e:
+                print(f"[save] volume commit failed: {e}", flush=True)
+
+    # final
+    print("[done] saving final LoRA")
+    save_lora_state(student_transformer, out_dir, "traj_final")
+    log_f.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/distill/traj_loss.py

@@ -0,0 +1,207 @@
+"""
+Trajectory imitation losses for Anima distillation.
+
+Algorithm (DiffSynth-Studio Z-Image trajectory_imitation 由来):
+  1. Teacher (= base Anima, frozen, no LoRA) を CFG=2.0 で 50 step rollout し
+     **student の 8 timesteps と一致する点** で latent を保存
+     → teacher_traj = [x_at_t_1, x_at_t_2, ..., x_at_t_N]  (N=8)
+  2. align loss: 各 student segment [t_i → t_{i+1}] で
+        target_v = (x_{i+1} - x_i) / (t_{i+1} - t_i)
+        student_v = student(x_i, t_i, cond)
+        L_align = MSE(student_v, target_v)
+  3. (optional) regularization: student が自前で 8-step rollout した最終画像と
+     teacher の最終画像を VAE decode して LPIPS で比較
+
+R3GAN (Phase B 失敗の元凶) は**一切使わない**。critic なし、単一ネットワーク。
+"""
+from __future__ import annotations
+from typing import Callable
+
+import torch
+import torch.nn.functional as F
+
+from .anima_loader import AnimaBundle
+from .traj_scheduler import TrajSchedule
+
+
+# ----- helpers -------------------------------------------------------------
+
+def _broadcast_t(t: torch.Tensor, B: int, device, dtype) -> torch.Tensor:
+    """scalar timestep を (B,) に膨らませる。"""
+    if t.dim() == 0:
+        t = t.expand(B)
+    return t.to(device=device, dtype=dtype)
+
+
+def cfg_guided(
+    model_fn: Callable[..., torch.Tensor],
+    x: torch.Tensor,
+    t: torch.Tensor,
+    cond_pos: torch.Tensor,
+    cond_neg: torch.Tensor | None,
+    cfg_scale: float,
+) -> torch.Tensor:
+    """classifier-free guidance を 1 行で。cfg_scale<=1 なら neg を呼ばない。"""
+    v_pos = model_fn(x, t, cond_pos)
+    if cfg_scale <= 1.0 or cond_neg is None:
+        return v_pos
+    v_neg = model_fn(x, t, cond_neg)
+    return v_neg + cfg_scale * (v_pos - v_neg)
+
+
+# ----- teacher rollout (= "fetch trajectory") -------------------------------
+
+@torch.no_grad()
+def fetch_trajectory(
+    teacher_velocity_fn: Callable[..., torch.Tensor],
+    init_noise: torch.Tensor,         # (B, 16, 1, H, W) at t=1
+    teacher_schedule: TrajSchedule,   # snapped to student timesteps
+    student_timesteps: torch.Tensor,  # (N+1,) student grid points
+    cond_pos: torch.Tensor,
+    cond_neg: torch.Tensor | None,
+    cfg_scale: float = 2.0,
+) -> list[torch.Tensor]:
+    """Teacher を CFG'd Euler で rollout し、student_timesteps と一致する点で
+    latent を保存して返す。
+
+    Returns:
+      list of len(student_timesteps) tensors、それぞれ shape = init_noise.shape。
+      i 番目の要素が `latents at t == student_timesteps[i]`。
+    """
+    B = init_noise.size(0)
+    device = init_noise.device
+    dtype = init_noise.dtype
+
+    ts = teacher_schedule.timesteps  # (M+1,) — student 点を含む snap 済
+    student_set = {float(t.item()) for t in student_timesteps}
+
+    latents = init_noise
+    saved: dict[float, torch.Tensor] = {}
+
+    # 初期点 (t=1) が student に含まれていれば保存
+    if float(ts[0].item()) in student_set:
+        saved[float(ts[0].item())] = latents.detach().clone()
+
+    for i in range(len(ts) - 1):
+        t_cur = ts[i]
+        t_next = ts[i + 1]
+        t_in = _broadcast_t(t_cur, B, device, dtype)
+        v = cfg_guided(teacher_velocity_fn, latents, t_in, cond_pos, cond_neg, cfg_scale)
+        # Anima: euler step は forward 方向 (t は減少)
+        dt = (t_next - t_cur).to(device=device, dtype=dtype)
+        latents = latents + dt * v
+        if float(t_next.item()) in student_set:
+            saved[float(t_next.item())] = latents.detach().clone()
+
+    # student_timesteps の順序で並べて返す
+    return [saved[float(t.item())] for t in student_timesteps]
+
+
+# ----- align loss (segment-wise velocity match) -----------------------------
+
+def align_trajectory(
+    student_velocity_fn: Callable[..., torch.Tensor],
+    teacher_traj: list[torch.Tensor],   # length N+1
+    student_timesteps: torch.Tensor,    # (N+1,)
+    cond_pos: torch.Tensor,
+    cond_neg: torch.Tensor | None,
+    cfg_scale_student: float = 1.0,
+    weight_mode: str = "uniform",
+) -> torch.Tensor:
+    """各 segment ごとに student velocity を target velocity と MSE で揃える。
+
+    Anima 蒸留の核。grad は student の LoRA に流す前提なので no_grad で囲まない。
+    teacher_traj は detach 済 (fetch_trajectory が clone してある) なので
+    target に grad は乗らない。
+    """
+    from .traj_scheduler import training_weight as _w
+    losses = []
+    weights = []
+    B = teacher_traj[0].size(0)
+    device = teacher_traj[0].device
+    dtype = teacher_traj[0].dtype
+
+    for i in range(len(student_timesteps) - 1):
+        t_cur = student_timesteps[i]
+        t_next = student_timesteps[i + 1]
+        x_cur = teacher_traj[i]
+        x_next = teacher_traj[i + 1]
+        dt = (t_next - t_cur).to(device=device, dtype=dtype)
+        # target velocity for this segment (teacher's effective velocity averaged over segment)
+        target_v = (x_next - x_cur) / dt
+        # student prediction at t_cur (greedy: assume student sees teacher's noisy state)
+        t_in = _broadcast_t(t_cur, B, device, dtype)
+        pred_v = cfg_guided(
+            student_velocity_fn, x_cur, t_in, cond_pos, cond_neg, cfg_scale_student
+        )
+        loss_i = F.mse_loss(pred_v.float(), target_v.detach().float())
+        losses.append(loss_i)
+        weights.append(_w(t_cur.unsqueeze(0).cpu(), mode=weight_mode).item())
+
+    losses_t = torch.stack(losses)
+    weights_t = torch.tensor(weights, device=losses_t.device, dtype=losses_t.dtype)
+    weights_t = weights_t / weights_t.sum().clamp(min=1e-8)
+    return (losses_t * weights_t).sum()
+
+
+# ----- regularization (student own rollout + LPIPS vs teacher final) -------
+
+def student_rollout(
+    student_velocity_fn: Callable[..., torch.Tensor],
+    init_noise: torch.Tensor,
+    student_timesteps: torch.Tensor,
+    cond_pos: torch.Tensor,
+    cond_neg: torch.Tensor | None,
+    cfg_scale: float = 1.0,
+) -> torch.Tensor:
+    """student が自前で N-step rollout し最終 latent を返す。**grad は通す**。"""
+    B = init_noise.size(0)
+    device = init_noise.device
+    dtype = init_noise.dtype
+    latents = init_noise
+
+    for i in range(len(student_timesteps) - 1):
+        t_cur = student_timesteps[i]
+        t_next = student_timesteps[i + 1]
+        t_in = _broadcast_t(t_cur, B, device, dtype)
+        v = cfg_guided(
+            student_velocity_fn, latents, t_in, cond_pos, cond_neg, cfg_scale
+        )
+        dt = (t_next - t_cur).to(device=device, dtype=dtype)
+        latents = latents + dt * v
+    return latents
+
+
+def compute_regularization(
+    student_velocity_fn: Callable[..., torch.Tensor],
+    bundle: AnimaBundle,
+    init_noise: torch.Tensor,
+    teacher_final_latent: torch.Tensor,
+    student_timesteps: torch.Tensor,
+    cond_pos: torch.Tensor,
+    cond_neg: torch.Tensor | None,
+    lpips_fn,
+    cfg_scale: float = 1.0,
+) -> torch.Tensor:
+    """LPIPS(VAE_decode(student_final), VAE_decode(teacher_final))。
+    student の rollout は grad-through。VAE は no_grad で固定 (decode 時のみ)。
+    """
+    student_final = student_rollout(
+        student_velocity_fn, init_noise, student_timesteps,
+        cond_pos, cond_neg, cfg_scale,
+    )
+    # VAE decode: Anima は (B, 3, T=1, H, W) を返す → T を squeeze
+    # 注: vae_decode は @torch.no_grad で囲まれているため student grad が VAE 入力で切れる
+    # → LPIPS reg を活かすには別 forward が必要。ここでは bundle.vae.model.decode を直接
+    #    呼んで grad を通す (VAE 自体の weight は frozen のままだが入力勾配は流れる)。
+    vae_dtype = next(bundle.vae.model.parameters()).dtype
+    img_student = bundle.vae.model.decode(student_final.to(dtype=vae_dtype), bundle.vae_scale)
+    with torch.no_grad():
+        img_teacher = bundle.vae.model.decode(
+            teacher_final_latent.to(dtype=vae_dtype), bundle.vae_scale
+        )
+    # (B, 3, 1, H, W) -> (B, 3, H, W)、[-1, 1] のまま LPIPS に渡せる
+    img_student = img_student.squeeze(2)
+    img_teacher = img_teacher.squeeze(2)
+    # LPIPS は float32 を要求するので cast
+    return lpips_fn(img_student.float(), img_teacher.float()).mean()