diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..2c7feea4aaed7c8096701c953b9333fa4a9dbc55 100644 --- a/.gitattributes +++ b/.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 diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..134a3bb1731c5c92b5202aa778fd4df7e013295b --- /dev/null +++ b/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. diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b14485f64190d9a28c48311cb152cd1544352246 --- /dev/null +++ b/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 diff --git a/configs/phase1_anima.toml b/configs/phase1_anima.toml new file mode 100644 index 0000000000000000000000000000000000000000..a96a9f72eb2ed029430061aa3a7fbf0434f7e550 --- /dev/null +++ b/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 追加 diff --git a/configs/phase1_dataset.toml b/configs/phase1_dataset.toml new file mode 100644 index 0000000000000000000000000000000000000000..62719b932eec9dbbf6684781c213eeca4465464c --- /dev/null +++ b/configs/phase1_dataset.toml @@ -0,0 +1,34 @@ +# ============================================================================= +# Phase 1 データセット設定 +# 期待構造: /dataset/cleaned/.{png,jpg,webp} と同名の .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 diff --git a/docs/distillation.md b/docs/distillation.md new file mode 100644 index 0000000000000000000000000000000000000000..f3d04bed9c48c1f8be5566cd288f9d2b4cdcc417 --- /dev/null +++ b/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..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 ` | +| `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) を参照。 diff --git a/docs/dmd2.md b/docs/dmd2.md new file mode 100644 index 0000000000000000000000000000000000000000..e386a3dbf1469e81a03554291a67d4cd83207b2e --- /dev/null +++ b/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 アロケート時の混雑度に依存 diff --git a/docs/dmdx.md b/docs/dmdx.md new file mode 100644 index 0000000000000000000000000000000000000000..537e970c74d09a09db25ccc6d40a34cebeb6e1f1 --- /dev/null +++ b/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 した先例あり) diff --git a/docs/migration_log.md b/docs/migration_log.md new file mode 100644 index 0000000000000000000000000000000000000000..b69af1553cdb13d70edb0c0a1de262f9ceaab430 --- /dev/null +++ b/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 --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/月) のみ diff --git a/docs/operations.md b/docs/operations.md new file mode 100644 index 0000000000000000000000000000000000000000..dee3584db2e0a77ad2bb1b1e377677ebe608f6c0 --- /dev/null +++ b/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)。 diff --git a/docs/pcm.md b/docs/pcm.md new file mode 100644 index 0000000000000000000000000000000000000000..960ea0fb0c3946bbb5d3e139ba61c266be9392bb --- /dev/null +++ b/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 構成より軽量 diff --git a/docs/setup.md b/docs/setup.md new file mode 100644 index 0000000000000000000000000000000000000000..590a372a32b2be10d65b0563905ffddd214585be --- /dev/null +++ b/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` デフォルト) +- 多様性 > 美的偏向(モデルの汎用性を保つため) diff --git a/docs/workflow.md b/docs/workflow.md new file mode 100644 index 0000000000000000000000000000000000000000..202e832ae58dcddbfb13de5989747ef56ae47ddc --- /dev/null +++ b/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/" + +# 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/月) も適時間引き。 diff --git a/modal_app.py b/modal_app.py new file mode 100644 index 0000000000000000000000000000000000000000..80f0ddf2c3b315f28ab2d32c5d508f4e04c69e30 --- /dev/null +++ b/modal_app.py @@ -0,0 +1,2850 @@ +""" +rapid_anima — Anima 生成速度向上 (Modal) +======================================== + +CircleStone Labs Anima (2B DiT) の **生成速度を改善** するための一式。 +- B200 sageattention 実 engage (silent fallback 回避 patched wheel) +- batch=8 並列生成 + sage で per-image 高速化 +- 4-8 step 蒸留 LoRA (PCM / Z-Image trajectory / LADD / Reflow / 公式 Turbo merge) +- 旧 fine-tune / quality タグ除去機能も残置 (履歴・派生用) + +主要 entry point: + # 1) base モデル + 互換 path 整備 (1 度だけ) + modal run modal_app.py::download_models + modal run modal_app.py::setup_anima_paths + + # 2) self-distill dataset 生成 (B200 sage + batch=8 並列) + modal run --detach modal_app.py::generate_dataset_parallel # ~$26, 5000 枚, ~28 分 + + # 3) teacher x0 cache (LADD/PCM 用) + modal run --detach modal_app.py::precompute_teacher_x0_cache + + # 4) 蒸留 (例: PCM) + modal run modal_app.py::download_civitai_lora # warm-start 用 Anima Turbo + modal run --detach modal_app.py::train_pcm_distill --warm-lora /models/loras/anima_turbo.safetensors + + # 5) 公式 Turbo との merge (即動く、$0.5) + modal run modal_app.py::merge_turbo_lora +""" + +import modal + +# --------------------------------------------------------------------------- +# Image: CUDA 12.4 base + diffusion-pipe (upstream が torch/deepspeed を pull) +# 方針: torch まわりは pin しない (upstream が頻繁に更新されるため衝突源になる)。 +# attention は torch SDPA (built-in flash-attn 2) を使う。 +# 追加で flash-attn / sageattention を入れたい場合は最後段で +# --no-build-isolation で source build する形に変更可能。 +# +# add_local_dir は Modal の制約で「最後」に置く必要がある (image variant ごとに付加)。 +# --------------------------------------------------------------------------- +_base_image = ( + modal.Image.from_registry( + "nvidia/cuda:12.4.1-devel-ubuntu22.04", add_python="3.11" + ) + .apt_install( + "git", + "wget", + "build-essential", + "libgl1", + "libglib2.0-0", + "ninja-build", + ) + .env( + { + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "NCCL_P2P_DISABLE": "1", + "NCCL_IB_DISABLE": "1", + "HF_HUB_ENABLE_HF_TRANSFER": "1", + "TRANSFORMERS_OFFLINE": "0", + } + ) + .pip_install( + "huggingface_hub[hf_transfer]>=0.27", + "Pillow>=10.0", + "numpy<2", + # trajectory imitation 蒸留 (train_traj.py) の LPIPS regularization 用。 + # --lpips-weight 0 (デフォルト) では import されないが、image build 時に入れておく。 + "lpips>=0.1.4", + # DRaFT+ (train_draftp.py) の HPSv2 reward 用。 + "open_clip_torch>=2.24", + "hpsv2>=1.2", + "ftfy", "regex", + ) + # diffusion-pipe 本体 (torch / deepspeed / transformers などを引っ張る) + # 失敗を黙殺すると runtime で意味不明エラーになるので `|| true` は付けない + # torchvision は diffusion-pipe (models/base.py) が import するため必要 + .run_commands( + "git clone --recurse-submodules " + "https://github.com/tdrussell/diffusion-pipe /workspace/diffusion-pipe", + "cd /workspace/diffusion-pipe && pip install -r requirements.txt", + "pip install --upgrade-strategy only-if-needed torchvision", + # K: sageattention で推論 1.2-1.5x 速 (Anima/Cosmos の torch SDPA を置換) + # --no-build-isolation で torch の binary ABI に合わせて build + # 失敗しても致命的でない (sageattention 使わなければ torch SDPA fallback) ため + # 後ろに `|| echo` を付けて image build が止まらないように + "pip install --no-build-isolation sageattention || echo '[warn] sageattention install failed, falling back to torch SDPA'", + ) +) + +# 訓練・キャプション掃除・モデル DL 用 (ComfyUI 不要) +image = ( + _base_image + .add_local_dir("configs", "/workspace/configs") + .add_local_dir("scripts", "/workspace/scripts") +) + +# 推論・データ生成用 (ComfyUI を上に積む) +comfy_image = ( + _base_image + .run_commands( + "git clone --depth 1 https://github.com/comfyanonymous/ComfyUI /workspace/ComfyUI", + # ComfyUI の依存は diffusion-pipe とほぼ overlap。 + # --upgrade-strategy only-if-needed で torch を下手にいじらせない + "pip install --upgrade-strategy only-if-needed -r /workspace/ComfyUI/requirements.txt", + ) + .add_local_dir("configs", "/workspace/configs") + .add_local_dir("scripts", "/workspace/scripts") +) + +# --------------------------------------------------------------------------- +# B200 (sm_100 / Blackwell) 専用 ComfyUI image +# - CUDA 12.8 + torch 2.7.0+cu128 が必要 (Blackwell Tensor Core を fully 使うため) +# - sageattention 2.2.0 公式 wheel は B200 で silently fallback する +# → darask0/modal_B200_sageattetion_comfyUI の patched wheel を使う +# (sm_100 dispatch を core.py に追加した版) +# - diffusion-pipe は不要 (ComfyUI HTTP API 経由で生成するだけ) +# +# 用途: generate_dataset_chunk など、B200 で大量画像生成する関数。 +# 他の comfy_image 利用関数は CUDA 12.4 base のまま残す (リスク回避)。 +# --------------------------------------------------------------------------- +B200_SAGE_WHEEL_URL = ( + "https://huggingface.co/darask0/modal_B200_sageattetion_comfyUI/" + "resolve/main/wheels/sageattention-2.2.0-cp311-cp311-linux_x86_64.whl" +) +b200_image = ( + modal.Image.from_registry( + "nvidia/cuda:12.8.1-devel-ubuntu22.04", add_python="3.11" + ) + .apt_install( + "git", "wget", "build-essential", "libgl1", "libglib2.0-0", + "ninja-build", "ca-certificates", + ) + .env({ + "HF_HUB_ENABLE_HF_TRANSFER": "1", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + }) + .pip_install( + "torch==2.7.0", "torchvision==0.22.0", "torchaudio==2.7.0", + index_url="https://download.pytorch.org/whl/cu128", + ) + .pip_install( + "huggingface_hub[hf_transfer]>=0.27", + "Pillow>=10.0", "numpy<2", + ) + .run_commands( + # 1) ComfyUI (torch は --upgrade-strategy only-if-needed で 2.7.0 を維持) + "git clone --depth 1 https://github.com/comfyanonymous/ComfyUI /workspace/ComfyUI", + "pip install --upgrade-strategy only-if-needed -r /workspace/ComfyUI/requirements.txt", + # 2) B200 (sm_100) patched sageattention + # --no-deps で torch 2.7.0+cu128 を保持 (上流 wheel が別 torch を pin する場合あり) + f"pip install --no-deps --force-reinstall {B200_SAGE_WHEEL_URL}", + # 3) build-time sanity: sm_100 dispatch が wheel 内 core.py に残ってるか確認 + # GPU なしでも source inspection だけで判定可 + "python -c \"import inspect, sageattention; " + "src = inspect.getsource(sageattention.core.sageattn); " + "assert 'sm100' in src, 'sm_100 dispatch missing in wheel'; " + "print('[sage-check] sm_100 dispatch OK in wheel')\"", + ) + .add_local_dir("configs", "/workspace/configs") + .add_local_dir("scripts", "/workspace/scripts") +) + +# --------------------------------------------------------------------------- +# Volumes: モデル / データセット / 出力を分離 +# --------------------------------------------------------------------------- +models_vol = modal.Volume.from_name("anima-models", create_if_missing=True) +dataset_vol = modal.Volume.from_name("anima-dataset", create_if_missing=True) +output_vol = modal.Volume.from_name("anima-outputs", create_if_missing=True) +# ファインチューニング素材用 (ユーザーがアップロードした +# hakushiMixAnima_v02.safetensors + anima-highres-aesthetic-boost.safetensors を含む) +# self-distill 系では未使用なので空でも mount 成功するよう create_if_missing=True +comfyui_anima_models_vol = modal.Volume.from_name("comfyui-anima-models", create_if_missing=True) + +VOLUMES = { + "/models": models_vol, + "/dataset": dataset_vol, + "/output": output_vol, + "/comfyui_anima_models": comfyui_anima_models_vol, +} + +# HF token (Modal 上の既存 secret 名 `hf_token`、キーは HF_TOKEN) +hf_secret = modal.Secret.from_name("hf_token", required_keys=["HF_TOKEN"]) + +# HF write-permission token (repo create / upload 用、別 secret) +hf_write_secret = modal.Secret.from_name("HF_TOKEN_WRITE") + +# Civitai API key (Anima Turbo LoRA など Civitai 経由のモデル取得用) +civitai_secret = modal.Secret.from_name("civitai_api_key", required_keys=["CIVITAI_API_KEY"]) + +app = modal.App("rapid-anima") + + +# --------------------------------------------------------------------------- +# Step 1: モデルダウンロード (1回だけ実行) +# hf_transfer + 並列で 5-6GB を数分に短縮 +# --------------------------------------------------------------------------- +@app.function( + image=image, + volumes={"/models": models_vol}, + timeout=3600, + secrets=[hf_secret], + cpu=4.0, + memory=8192, +) +def download_models(): + """Anima base + Qwen3 LLM + Qwen-Image VAE を Volume へ取得""" + from concurrent.futures import ThreadPoolExecutor, as_completed + from huggingface_hub import hf_hub_download + import os + import shutil + + repo = "circlestone-labs/Anima" + targets = [ + ("split_files/diffusion_models/anima-preview3-base.safetensors", + "anima-preview3-base.safetensors"), + ("split_files/text_encoders/qwen_3_06b_base.safetensors", + "qwen_3_06b_base.safetensors"), + ("split_files/vae/qwen_image_vae.safetensors", + "qwen_image_vae.safetensors"), + ] + + os.makedirs("/models/checkpoints", exist_ok=True) + token = os.environ.get("HF_TOKEN") or None + + def fetch(repo_path: str, local_name: str) -> str: + out = f"/models/checkpoints/{local_name}" + if os.path.exists(out): + return f"[skip] {local_name} (exists)" + path = hf_hub_download(repo_id=repo, filename=repo_path, token=token) + shutil.copy(path, out) + return f"[done] {out}" + + # 3 並列ダウンロード (hf_transfer が各ファイル内でも multi-connection) + with ThreadPoolExecutor(max_workers=3) as ex: + futures = {ex.submit(fetch, p, n): n for p, n in targets} + for fut in as_completed(futures): + print(fut.result()) + + models_vol.commit() + print("All models ready in /models/checkpoints/") + + +# --------------------------------------------------------------------------- +# anima-models volume の path layout を script 前提に合わせる +# 現状: /models/hf_anima/split_files/{diffusion_models,text_encoders,vae,loras}/* +# 期待: /models/{checkpoints,text_encoders,vae,loras}/* (generate_dataset.py が前提) +# 既存 hf_anima を消さず symlink で互換 path を作る。冪等(再実行で問題なし)。 +# --------------------------------------------------------------------------- +@app.function( + image=image, + volumes={"/models": models_vol}, + timeout=600, + cpu=2.0, + memory=4096, +) +def setup_anima_paths(): + """anima-models 内の hf_anima/split_files/* を /models/{checkpoints,loras}/ に symlink。 + generate_dataset.py::setup_model_symlinks は単一 dir のみ scan するので + diffusion_models / text_encoders / vae を flat に /models/checkpoints/ に集約する + (script 側の classify() がファイル名で適切な ComfyUI subdir に振り分ける)。 + 冪等で何度実行しても安全。""" + import os + from pathlib import Path + + hf_root = Path("/models/hf_anima/split_files") + # (src_subdir, dst_subdir): diffusion_models/text_encoders/vae は全部 checkpoints/ に flat 化 + mappings = [ + ("diffusion_models", "checkpoints"), + ("text_encoders", "checkpoints"), + ("vae", "checkpoints"), + ("loras", "loras"), + ] + created, skipped = 0, 0 + for src_sub, dst_sub in mappings: + src_dir = hf_root / src_sub + dst_dir = Path("/models") / dst_sub + if not src_dir.exists(): + print(f"[skip-src] {src_dir} not found") + continue + dst_dir.mkdir(parents=True, exist_ok=True) + for src in src_dir.glob("*.safetensors"): + dst = dst_dir / src.name + if dst.is_symlink(): + if os.readlink(dst) == str(src): + skipped += 1 + continue + dst.unlink() + elif dst.exists(): + print(f"[real-file] {dst} exists as regular file, skip") + continue + dst.symlink_to(src) + created += 1 + print(f"[symlink] {dst} -> {src}") + + models_vol.commit() + print(f"[done] created={created} skipped={skipped}") + + # 主要 path 確認 (script 前提) + expected = [ + "/models/checkpoints/anima-base-v1.0.safetensors", + "/models/checkpoints/qwen_3_06b_base.safetensors", + "/models/checkpoints/qwen_image_vae.safetensors", + ] + for p in expected: + status = "OK" if os.path.exists(p) else "MISSING" + print(f" [{status}] {p}") + + +# --------------------------------------------------------------------------- +# Step 2: キャプション掃除 (品質/期間/メタタグ除去) +# CPU だけで完結する軽い処理 → 安い CPU 指定でコスト微減 +# --------------------------------------------------------------------------- +@app.function( + image=image, + volumes={"/dataset": dataset_vol}, + timeout=1800, + cpu=2.0, + memory=4096, +) +def clean_captions( + in_subdir: str = "raw", + out_subdir: str = "cleaned", + keep_artist: bool = True, + drop_artist_prob: float = 0.0, +): + """ + /dataset/{in_subdir} の .txt キャプションから品質/期間/メタタグを除去し + /dataset/{out_subdir} に画像と一緒にコピーする。 + + keep_artist=False または drop_artist_prob>0 で artist タグ(@xxx)もコントロール可能。 + """ + import subprocess + cmd = [ + "python", + "/workspace/scripts/clean_captions.py", + "--input", f"/dataset/{in_subdir}", + "--output", f"/dataset/{out_subdir}", + "--drop-artist-prob", str(drop_artist_prob), + ] + if not keep_artist: + cmd.append("--drop-all-artists") + subprocess.run(cmd, check=True) + dataset_vol.commit() + + +# --------------------------------------------------------------------------- +# Step 3: Phase 1 学習 (審美ファインチューン) +# コスト概算 (10k枚 / 3 epoch): +# A100-80GB ($2.50/hr) × 18-22h = $45-55 (推奨/デフォルト) +# H100 80GB ($3.95/hr) × 12-15h = $47-60 (時短) +# L40S 48GB ($1.95/hr) × 28-35h = $55-68 (節約だが結局割高、OOM 注意) +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="A100-80GB", + volumes=VOLUMES, + timeout=24 * 60 * 60, + secrets=[hf_secret], + cpu=8.0, # DataLoader worker 用に余裕 + memory=32768, # 32GB: VAE/text-encoder の load + dataset cache +) +def train_phase1(config_name: str = "phase1_anima.toml"): + """diffusion-pipe で Phase 1 (LoRA aesthetic FT) を実行""" + import subprocess + import os + + os.chdir("/workspace/diffusion-pipe") + + cfg = f"/workspace/configs/{config_name}" + print(f"[train] using config: {cfg}") + + # DeepSpeed 単 GPU 起動 (公式 README 推奨形式) + cmd = [ + "deepspeed", + "--num_gpus=1", + "train.py", + "--deepspeed", + "--config", + cfg, + ] + subprocess.run(cmd, check=True) + output_vol.commit() + print("Phase 1 training complete. Outputs in /output/phase1/") + + +# --------------------------------------------------------------------------- +# Step 4: サンプル生成で Phase 1 を検証 +# 検証用なので一番安い L40S で十分(品質確認だけ、batch 小) +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="L40S", + volumes=VOLUMES, + timeout=1800, + cpu=4.0, + memory=16384, +) +def generate_samples( + lora_path: str = "/output/phase1/latest/adapter_model.safetensors", + prompts_file: str = "/workspace/scripts/eval_prompts.txt", +): + """学習済み LoRA で素の短いプロンプトから画像生成。品質タグなしを意図的に使う。""" + import subprocess + subprocess.run( + [ + "python", + "/workspace/scripts/infer_test.py", + "--lora", lora_path, + "--prompts", prompts_file, + "--out", "/output/phase1_samples", + ], + check=True, + ) + output_vol.commit() + + +# 旧 train_phase2_lcm / train_phase2_dmd2 (スケルトンのみ) は削除。 +# 蒸留の実装は scripts/distill/ + train_sota_distill function を参照。 + + +# --------------------------------------------------------------------------- +# Civitai LoRA ダウンロード(Anima Turbo LoRA など) +# --------------------------------------------------------------------------- +@app.function( + image=image, + volumes={"/models": models_vol}, + timeout=3600, + secrets=[hf_secret], + cpu=2.0, + memory=4096, +) +def download_anima_v1_base( + delete_preview3: bool = False, +): + """Anima base v1.0 を HF Hub から /models/checkpoints/anima-base-v1.0.safetensors に取得。 + delete_preview3=True で旧 preview3-base.safetensors を削除。""" + import os, urllib.request + url = "https://huggingface.co/circlestone-labs/Anima/resolve/main/split_files/diffusion_models/anima-base-v1.0.safetensors" + out_dir = "/models/checkpoints" + os.makedirs(out_dir, exist_ok=True) + out = f"{out_dir}/anima-base-v1.0.safetensors" + if os.path.exists(out): + sz = os.path.getsize(out) / 1e9 + print(f"[skip] {out} exists ({sz:.2f} GB)") + else: + print(f"[download] {url}") + token = os.environ.get("HF_TOKEN", "") + req = urllib.request.Request(url, headers={ + "User-Agent": "modal-anima/1.0", + **({"Authorization": f"Bearer {token}"} if token else {}), + }) + with urllib.request.urlopen(req, timeout=900) as r, open(out, "wb") as f: + size = 0 + while True: + chunk = r.read(8 * 1024 * 1024) + if not chunk: + break + f.write(chunk) + size += len(chunk) + if size % (100 * 1024 * 1024) < (8 * 1024 * 1024): + print(f" ... {size/1e9:.2f} GB") + print(f"[done] {out} ({os.path.getsize(out)/1e9:.2f} GB)") + models_vol.commit() + if delete_preview3: + old = f"{out_dir}/anima-preview3-base.safetensors" + if os.path.exists(old): + os.remove(old) + print(f"[deleted] {old}") + models_vol.commit() + + +@app.function( + image=image, + volumes={"/models": models_vol}, + timeout=1800, + secrets=[civitai_secret], + cpu=2.0, + memory=4096, +) +def download_civitai_lora( + version_id: int = 2877687, # 既定: Anima Turbo LoRA v0.1 + out_name: str = "anima_turbo.safetensors", +): + """ + Civitai の API トークン経由で LoRA をダウンロードして + /models/loras/{out_name} に保存。 + + 例: + # Anima Turbo LoRA を既定設定で + modal run modal_app.py::download_civitai_lora + + # 別の LoRA: version_id は civitai.com/api/download/models/{id} の数字 + modal run modal_app.py::download_civitai_lora \\ + --version-id 1234567 --out-name another.safetensors + """ + import os + import urllib.request + + token = os.environ["CIVITAI_API_KEY"] + url = f"https://civitai.com/api/download/models/{version_id}?token={token}" + out_dir = "/models/loras" + os.makedirs(out_dir, exist_ok=True) + out = f"{out_dir}/{out_name}" + + if os.path.exists(out): + print(f"[skip] {out} (exists, {os.path.getsize(out)/1e6:.1f} MB)") + return + + print(f"[download] version_id={version_id} -> {out}") + req = urllib.request.Request(url, headers={"User-Agent": "modal-anima/1.0"}) + with urllib.request.urlopen(req, timeout=600) as r, open(out, "wb") as f: + size = 0 + while True: + chunk = r.read(1024 * 1024) + if not chunk: + break + f.write(chunk) + size += len(chunk) + if size % (10 * 1024 * 1024) < (1024 * 1024): # ~10MB ごと + print(f" ... {size/1e6:.0f} MB") + print(f"[done] {out} ({size/1e6:.1f} MB)") + models_vol.commit() + + +# --------------------------------------------------------------------------- +# Self-distillation: Anima base 自身で学習用データセットを生成 +# シリアル(A100-80GB): 5,000 枚 ~12.5h, $31 +# パラレル(B200 × 10): 5,000 枚 ~40 分, $44 ← 速度優先時はこっち +# --------------------------------------------------------------------------- +@app.function( + image=comfy_image, + gpu="A100-80GB", + volumes=VOLUMES, + timeout=18 * 60 * 60, # 18h (上限 24h) + secrets=[hf_secret], + cpu=4.0, + memory=32768, +) +def generate_dataset( + prompts_file: str = "/workspace/scripts/gen_prompts.txt", + workflow_file: str = "/workspace/scripts/anima_workflow.json", + out_subdir: str = "raw", + seeds_per_prompt: int = 50, + max_images: int = 0, + start_from: int = 0, + prompt_end: int = 0, + override_steps: int = 0, + override_cfg: float = 0.0, + fixed_aspect: str = "", + no_prefix: bool = False, +): + """ + gen_prompts.txt × N seed で Anima base から画像生成 → /dataset/raw に保存。 + 再開可能 (既存 .png はスキップ)。 + + 例: + # フル実行 (100 prompts × 50 seeds = 5,000 枚, A100, 12.5h, $31) + modal run --detach modal_app.py::generate_dataset + + # 試走 (10 枚だけ) + modal run modal_app.py::generate_dataset --max-images 10 --seeds-per-prompt 5 + + # 範囲指定 (並列実行で chunk を割り当てる用) + modal run modal_app.py::generate_dataset --start-from 0 --prompt-end 10 + """ + import subprocess + out = f"/dataset/{out_subdir}" + cmd = [ + "python", "/workspace/scripts/generate_dataset.py", + "--prompts", prompts_file, + "--workflow", workflow_file, + "--out", out, + "--comfy-dir", "/workspace/ComfyUI", + "--models-dir", "/models/checkpoints", + "--seeds-per-prompt", str(seeds_per_prompt), + "--start-from", str(start_from), + "--prompt-end", str(prompt_end), + ] + if max_images: + cmd += ["--max-images", str(max_images)] + if override_steps > 0: + cmd += ["--override-steps", str(override_steps)] + if override_cfg > 0: + cmd += ["--override-cfg", str(override_cfg)] + if fixed_aspect: + cmd += ["--fixed-aspect", fixed_aspect] + if no_prefix: + cmd += ["--no-prefix"] + subprocess.run(cmd, check=True) + dataset_vol.commit() + + +# --------------------------------------------------------------------------- +# B200 並列 chunk worker (sageattention 実 engage 版) +# image = b200_image (cu128 + torch 2.7.0 + B200 patched sageattention) +# --------------------------------------------------------------------------- +@app.function( + image=b200_image, + gpu="B200", + volumes=VOLUMES, + timeout=4 * 60 * 60, # 1 chunk 4h 上限 + secrets=[hf_secret], + cpu=4.0, + memory=32768, +) +def generate_dataset_chunk( + prompt_start: int, + prompt_end: int, + seeds_per_prompt: int = 50, + out_subdir: str = "raw", + prompts_file: str = "/workspace/scripts/gen_prompts.txt", + workflow_file: str = "/workspace/scripts/anima_workflow.json", + batch_size: int = 8, +): + """並列 worker。prompt_start..prompt_end の範囲のみ処理。 + batch_size=8 は B200 上で 1 submit あたり 8 枚生成 (per-image 大幅高速化)。""" + import subprocess + out = f"/dataset/{out_subdir}" + print(f"[chunk] prompts {prompt_start}..{prompt_end}, " + f"{seeds_per_prompt} seeds each, batch_size={batch_size}") + subprocess.run( + [ + "python", "/workspace/scripts/generate_dataset.py", + "--prompts", prompts_file, + "--workflow", workflow_file, + "--out", out, + "--comfy-dir", "/workspace/ComfyUI", + "--models-dir", "/models/checkpoints", + "--seeds-per-prompt", str(seeds_per_prompt), + "--start-from", str(prompt_start), + "--prompt-end", str(prompt_end), + "--batch-size", str(batch_size), + ], + check=True, + ) + dataset_vol.commit() + print(f"[chunk] done {prompt_start}..{prompt_end}") + + +# --------------------------------------------------------------------------- +# Fine-tune dataset 並列生成 (hakushi base + aesthetic-boost LoRA、Anima 公式設定) +# 60 themes × seeds × random aspect = 大量生成 +# --------------------------------------------------------------------------- +@app.function( + image=comfy_image, + gpu="B200", + volumes=VOLUMES, + timeout=4 * 60 * 60, + secrets=[hf_secret], + cpu=4.0, + memory=32768, +) +def generate_finetune_chunk( + prompt_start: int, + prompt_end: int, + seeds_per_prompt: int = 8, + out_subdir: str = "finetune_raw", + prompts_file: str = "/workspace/scripts/finetune_prompts.txt", + quality_prefix: str = "masterpiece, best quality, score_7, safe", +): + """並列 worker: prompt_start..prompt_end 範囲で hakushi base + aesthetic LoRA から生成。 + comfyui-anima-models volume の hakushiMixAnima_v02 と anima-highres-aesthetic-boost を + /models/{checkpoints,loras} に symlink して既存 generate_dataset.py パイプラインに乗せる。 + quality_prefix を変えれば SFW (safe) / NSFW (nsfw / explicit) を切替可能。""" + import os, subprocess + # comfyui-anima-models volume から /models に symlink + os.makedirs("/models/checkpoints", exist_ok=True) + os.makedirs("/models/loras", exist_ok=True) + src_unet = "/comfyui_anima_models/diffusion_models/hakushiMixAnima_v02.safetensors" + dst_unet = "/models/checkpoints/hakushiMixAnima_v02.safetensors" + if not os.path.lexists(dst_unet): + os.symlink(src_unet, dst_unet) + print(f"[link] {dst_unet} -> {src_unet}") + src_lora = "/comfyui_anima_models/loras/anima-highres-aesthetic-boost.safetensors" + dst_lora = "/models/loras/anima-highres-aesthetic-boost.safetensors" + if not os.path.lexists(dst_lora): + os.symlink(src_lora, dst_lora) + print(f"[link] {dst_lora} -> {src_lora}") + + out = f"/dataset/{out_subdir}" + print(f"[finetune-chunk] prompts {prompt_start}..{prompt_end}, {seeds_per_prompt} seeds, " + f"prefix='{quality_prefix}', file={prompts_file}") + subprocess.run( + [ + "python", "/workspace/scripts/generate_dataset.py", + "--prompts", prompts_file, + "--workflow", "/workspace/scripts/anima_finetune_dataset_workflow.json", + "--out", out, + "--comfy-dir", "/workspace/ComfyUI", + "--models-dir", "/models/checkpoints", + "--loras-dir", "/models/loras", + "--seeds-per-prompt", str(seeds_per_prompt), + "--start-from", str(prompt_start), + "--prompt-end", str(prompt_end), + "--quality-prefix", quality_prefix, + # aspect ratio は random shuffle (ASPECT_RATIOS 7 種から)、--fixed-aspect 渡さない + ], + check=True, + ) + dataset_vol.commit() + print(f"[finetune-chunk] done {prompt_start}..{prompt_end}") + + +@app.local_entrypoint() +def generate_finetune_dataset_parallel( + workers: int = 4, + total_prompts: int = 60, + seeds_per_prompt: int = 8, + out_subdir: str = "finetune_raw", +): + """SFW: 4 B200 並列で fine-tune dataset を生成 (default)。 + + 60 themes × 8 seeds × 7 aspect ratios random = 1920 枚 (ASPECT_RATIOS shuffle 経由)。 + """ + if total_prompts % workers != 0: + print(f"[warn] total_prompts={total_prompts} not divisible by workers={workers}") + chunk_size = (total_prompts + workers - 1) // workers + chunks = [ + (i * chunk_size, min((i + 1) * chunk_size, total_prompts)) + for i in range(workers) + ] + expected_total = total_prompts * seeds_per_prompt + print(f"[parallel-sfw] spawning {workers} B200 workers, ~{expected_total} images total") + print(f"[parallel-sfw] chunks: {chunks}") + handles = [] + for start, end in chunks: + handles.append( + generate_finetune_chunk.spawn( + prompt_start=start, prompt_end=end, + seeds_per_prompt=seeds_per_prompt, out_subdir=out_subdir, + prompts_file="/workspace/scripts/finetune_prompts.txt", + quality_prefix="masterpiece, best quality, score_7, safe", + ) + ) + for h in handles: + h.get() + print(f"[parallel-sfw] all {workers} workers done, dataset at /dataset/{out_subdir}/") + + +@app.local_entrypoint() +def generate_finetune_dataset_sfw_v2_parallel( + workers: int = 4, + total_prompts: int = 60, + seeds_per_prompt: int = 16, + out_subdir: str = "finetune_raw_v2", +): + """SFW v2: 4 B200 並列で fine-tune dataset v2 (animal/mythology/sports/etc 60 themes)。 + + finetune_prompts_v2.txt の 60 新規 themes × 16 seeds × 7 aspect random = ~960 枚。 + 既存 finetune_raw とは別 subdir、訓練時に親 dir 経由で集約取り込み。 + """ + if total_prompts % workers != 0: + print(f"[warn] total_prompts={total_prompts} not divisible by workers={workers}") + chunk_size = (total_prompts + workers - 1) // workers + chunks = [ + (i * chunk_size, min((i + 1) * chunk_size, total_prompts)) + for i in range(workers) + ] + expected_total = total_prompts * seeds_per_prompt + print(f"[parallel-sfw-v2] spawning {workers} B200 workers, ~{expected_total} images total") + handles = [] + for start, end in chunks: + handles.append( + generate_finetune_chunk.spawn( + prompt_start=start, prompt_end=end, + seeds_per_prompt=seeds_per_prompt, out_subdir=out_subdir, + prompts_file="/workspace/scripts/finetune_prompts_v2.txt", + quality_prefix="masterpiece, best quality, score_7, safe", + ) + ) + for h in handles: + h.get() + print(f"[parallel-sfw-v2] all done, /dataset/{out_subdir}/") + + +@app.local_entrypoint() +def generate_finetune_dataset_sfw_v3_parallel( + workers: int = 4, + total_prompts: int = 60, + seeds_per_prompt: int = 16, + out_subdir: str = "finetune_raw_v3", +): + """SFW v3: 4 B200 並列で fine-tune dataset v3 (60 themes #1-60、SFW、職業/世界観多様)。""" + if total_prompts % workers != 0: + print(f"[warn] total_prompts={total_prompts} not divisible by workers={workers}") + chunk_size = (total_prompts + workers - 1) // workers + chunks = [ + (i * chunk_size, min((i + 1) * chunk_size, total_prompts)) + for i in range(workers) + ] + expected_total = total_prompts * seeds_per_prompt + print(f"[parallel-sfw-v3] spawning {workers} B200 workers, ~{expected_total} images total") + handles = [] + for start, end in chunks: + handles.append( + generate_finetune_chunk.spawn( + prompt_start=start, prompt_end=end, + seeds_per_prompt=seeds_per_prompt, out_subdir=out_subdir, + prompts_file="/workspace/scripts/finetune_prompts_v3.txt", + quality_prefix="masterpiece, best quality, score_7, safe", + ) + ) + for h in handles: + h.get() + print(f"[parallel-sfw-v3] all done, /dataset/{out_subdir}/") + + +@app.local_entrypoint() +def generate_finetune_dataset_explicit_v2_parallel( + workers: int = 4, + total_prompts: int = 60, + seeds_per_prompt: int = 16, + out_subdir: str = "finetune_raw_explicit_v2", +): + """EXPLICIT v2: 4 B200 並列で fine-tune dataset の explicit #31-90 (60 themes) を生成。 + + 内容は finetune_prompts_explicit_v2.txt を参照 (consensual framing 調整 8 件あり)。 + 別 subdir 出力、訓練時に親 dir 経由で全 explicit 合算取り込み。 + """ + if total_prompts % workers != 0: + print(f"[warn] total_prompts={total_prompts} not divisible by workers={workers}") + chunk_size = (total_prompts + workers - 1) // workers + chunks = [ + (i * chunk_size, min((i + 1) * chunk_size, total_prompts)) + for i in range(workers) + ] + expected_total = total_prompts * seeds_per_prompt + print(f"[parallel-explicit-v2] spawning {workers} B200 workers, ~{expected_total} images total") + handles = [] + for start, end in chunks: + handles.append( + generate_finetune_chunk.spawn( + prompt_start=start, prompt_end=end, + seeds_per_prompt=seeds_per_prompt, out_subdir=out_subdir, + prompts_file="/workspace/scripts/finetune_prompts_explicit_v2.txt", + quality_prefix="masterpiece, best quality, score_7, explicit", + ) + ) + for h in handles: + h.get() + print(f"[parallel-explicit-v2] all done, /dataset/{out_subdir}/") + + +@app.local_entrypoint() +def generate_finetune_dataset_explicit_parallel( + workers: int = 4, + total_prompts: int = 30, + seeds_per_prompt: int = 8, + out_subdir: str = "finetune_raw_explicit", +): + """EXPLICIT: 4 B200 並列で fine-tune dataset の explicit 性行為含む分を生成。 + + 30 themes × 8 seeds × 7 aspect ratios random = ~960 枚。 + quality_prefix は "...score_7, explicit" に上書き (Anima 公式 safety tag 最上位)。 + 出力は SFW/NSFW と別 subdir に分けて、訓練時に親 dir を指せば自動で全て拾う。 + """ + if total_prompts % workers != 0: + print(f"[warn] total_prompts={total_prompts} not divisible by workers={workers}") + chunk_size = (total_prompts + workers - 1) // workers + chunks = [ + (i * chunk_size, min((i + 1) * chunk_size, total_prompts)) + for i in range(workers) + ] + expected_total = total_prompts * seeds_per_prompt + print(f"[parallel-explicit] spawning {workers} B200 workers, ~{expected_total} images total") + print(f"[parallel-explicit] chunks: {chunks}") + handles = [] + for start, end in chunks: + handles.append( + generate_finetune_chunk.spawn( + prompt_start=start, prompt_end=end, + seeds_per_prompt=seeds_per_prompt, out_subdir=out_subdir, + prompts_file="/workspace/scripts/finetune_prompts_explicit.txt", + quality_prefix="masterpiece, best quality, score_7, explicit", + ) + ) + for h in handles: + h.get() + print(f"[parallel-explicit] all {workers} workers done, dataset at /dataset/{out_subdir}/") + + +@app.local_entrypoint() +def generate_finetune_dataset_nsfw_parallel( + workers: int = 4, + total_prompts: int = 30, + seeds_per_prompt: int = 8, + out_subdir: str = "finetune_raw_nsfw", +): + """NSFW: 4 B200 並列で fine-tune dataset NSFW 分を生成。 + + 30 themes × 8 seeds × 7 aspect ratios random = ~960 枚。 + quality_prefix は "...score_7, nsfw" に上書き、Anima 公式 safety tag 仕様に準拠。 + 出力は SFW と別 subdir に分けて、訓練時に両方読む形 (AnimaImageCaptionDataset は + rglob で recursive なので、親 dir を指せば自動で両 subdir 拾う)。 + """ + if total_prompts % workers != 0: + print(f"[warn] total_prompts={total_prompts} not divisible by workers={workers}") + chunk_size = (total_prompts + workers - 1) // workers + chunks = [ + (i * chunk_size, min((i + 1) * chunk_size, total_prompts)) + for i in range(workers) + ] + expected_total = total_prompts * seeds_per_prompt + print(f"[parallel-nsfw] spawning {workers} B200 workers, ~{expected_total} images total") + print(f"[parallel-nsfw] chunks: {chunks}") + handles = [] + for start, end in chunks: + handles.append( + generate_finetune_chunk.spawn( + prompt_start=start, prompt_end=end, + seeds_per_prompt=seeds_per_prompt, out_subdir=out_subdir, + prompts_file="/workspace/scripts/finetune_prompts_nsfw.txt", + quality_prefix="masterpiece, best quality, score_7, nsfw", + ) + ) + for h in handles: + h.get() + print(f"[parallel-nsfw] all {workers} workers done, dataset at /dataset/{out_subdir}/") + + +# --------------------------------------------------------------------------- +# 4 hyperparam バリアントで LoRA fine-tuning を並列起動 (diffusion-pipe 経由) +# Anima 公式推奨: rank 32 / lr 2e-5 / llm_adapter_lr=0 必須 +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=12 * 60 * 60, + secrets=[hf_secret], + cpu=8.0, + memory=131072, +) +def train_finetune_variant( + variant_name: str, + lora_rank: int = 32, + learning_rate: float = 2e-5, + epochs: int = 3, + drop_artist_prob: float = 0.0, + dataset_dir: str = "/dataset/finetune_raw", +): + """1 つの hyperparam variant で fine-tune。 + phase1_anima.toml をベースに rank / lr / epochs を override。""" + import os, subprocess, tempfile, shutil + # 出力先 + out_dir = f"/output/finetune_{variant_name}" + os.makedirs(out_dir, exist_ok=True) + + # phase1_anima.toml を temp copy + override + base_cfg = "/workspace/configs/phase1_anima.toml" + new_cfg = f"/tmp/finetune_{variant_name}.toml" + txt = open(base_cfg, encoding="utf-8").read() + # 主要 hyperparam を sed-style 置換 (簡易) + import re + txt = re.sub(r"output_dir\s*=.*", f'output_dir = "{out_dir}"', txt) + txt = re.sub(r"epochs\s*=.*", f"epochs = {epochs}", txt) + txt = re.sub(r"^lr\s*=.*", f"lr = {learning_rate}", txt, flags=re.MULTILINE) + txt = re.sub(r"^rank\s*=.*", f"rank = {lora_rank}", txt, flags=re.MULTILINE) + txt = re.sub(r"^alpha\s*=.*", f"alpha = {lora_rank}", txt, flags=re.MULTILINE) + # llm_adapter_lr=0 を保証 (Anima 公式必須) + txt = re.sub(r"llm_adapter_lr\s*=.*", "llm_adapter_lr = 0", txt) + # dataset 配下のパスも置換 + dataset_cfg = "/workspace/configs/phase1_dataset.toml" + new_dataset_cfg = f"/tmp/dataset_{variant_name}.toml" + dtxt = open(dataset_cfg, encoding="utf-8").read() + dtxt = re.sub(r'path\s*=\s*".*"', f'path = "{dataset_dir}"', dtxt) + if drop_artist_prob > 0: + dtxt = re.sub(r"drop_artist_prob\s*=.*", + f"drop_artist_prob = {drop_artist_prob}", dtxt) + open(new_dataset_cfg, "w", encoding="utf-8").write(dtxt) + # dataset config への参照を toml から書き換え + txt = re.sub(r'dataset\s*=\s*".*"', f'dataset = "{new_dataset_cfg}"', txt) + open(new_cfg, "w", encoding="utf-8").write(txt) + print(f"[finetune {variant_name}] config: {new_cfg}") + print(f" rank={lora_rank} lr={learning_rate} epochs={epochs} drop_artist={drop_artist_prob}") + + # diffusion-pipe 起動 + cmd = [ + "deepspeed", "--num_gpus", "1", + "/workspace/diffusion-pipe/train.py", + "--deepspeed", "--config", new_cfg, + ] + print(f"[finetune {variant_name}] {' '.join(cmd)}") + env = {**os.environ, "PYTHONPATH": "/workspace/diffusion-pipe", "PYTHONUNBUFFERED": "1"} + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + output_vol.commit() + print(f"[finetune {variant_name}] done, output at {out_dir}") + + +@app.local_entrypoint() +def train_finetune_4variants(dataset_dir: str = "/dataset/finetune_raw"): + """A/B/C/D の 4 hyperparam variant を並列起動。各 B200 1 台。""" + variants = [ + # (name, rank, lr, epochs, drop_artist_prob) + ("A_baseline", 32, 2e-5, 3, 0.0), # Anima 公式推奨 + ("B_higher_capacity", 64, 2e-5, 3, 0.0), # rank 2x + ("C_safer_lr", 32, 1e-5, 4, 0.0), # lr 半分、epoch 1 増 + ("D_artist_focus", 32, 3e-5, 3, 0.0), # lr 強め (artist drop は 0 固定で stylistic 強化狙い) + ] + print(f"[parallel] spawning {len(variants)} fine-tune variants") + handles = [] + for name, rank, lr, eps, drop in variants: + handles.append( + train_finetune_variant.spawn( + variant_name=name, lora_rank=rank, learning_rate=lr, + epochs=eps, drop_artist_prob=drop, dataset_dir=dataset_dir, + ) + ) + for h in handles: + h.get() + print(f"[parallel] all {len(variants)} variants done") + + +@app.local_entrypoint() +def generate_dataset_parallel( + workers: int = 10, + total_prompts: int = 100, + seeds_per_prompt: int = 50, + out_subdir: str = "raw", + batch_size: int = 8, +): + """ + B200 並列で大量画像生成。batch_size=8 で 1 submission に 8 枚 (per-image 高速化)。 + 旧 batch=1 で 40min/$44 → batch=8 で ~15min/$12 想定。 + 各 worker は total_prompts // workers 件の prompt を担当。 + + 例: + modal run modal_app.py::generate_dataset_parallel + modal run modal_app.py::generate_dataset_parallel --workers 5 --seeds-per-prompt 80 + modal run modal_app.py::generate_dataset_parallel --batch-size 4 # OOM 時 fallback + """ + if total_prompts % workers != 0: + print(f"[warn] total_prompts={total_prompts} not divisible by workers={workers}") + chunk_size = (total_prompts + workers - 1) // workers + chunks = [ + (i * chunk_size, min((i + 1) * chunk_size, total_prompts)) + for i in range(workers) + ] + total_imgs = total_prompts * seeds_per_prompt + print(f"[parallel] spawning {workers} B200 workers, ~{total_imgs} images total, " + f"batch_size={batch_size}") + for s, e in chunks: + print(f" worker: prompts [{s}, {e}) -> ~{(e - s) * seeds_per_prompt} imgs") + + handles = [ + generate_dataset_chunk.spawn( + prompt_start=s, + prompt_end=e, + seeds_per_prompt=seeds_per_prompt, + out_subdir=out_subdir, + batch_size=batch_size, + ) + for s, e in chunks + ] + # 全 worker の完了待ち(順不同で完了するが、全部終わるまでブロック) + for h in handles: + h.get() + print(f"[parallel] all {workers} workers complete") + + +# --------------------------------------------------------------------------- +# Path C: SOTA Distillation (Decoupled DMD2 + R3GAN + TSCD on B200) +# Phase A (8-step) → B (4-step) → C (2-step)、各 3-5h、合計 $80-150 +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=24 * 60 * 60, + secrets=[hf_secret], + cpu=8.0, + memory=131072, # 128GB (3 DiT 複製 + R3GAN D + dataset + buffers) +) +def train_sota_distill( + phase: str = "a", + dataset_path: str = "/dataset/raw", + out_dir: str = "/output/distill", + resume: str = "", + total_steps: int = 3000, + resolution: int = 1024, + gen_lora_only: bool = False, + override_adv_weight: float = -1.0, + override_r3gan_gamma: float = -1.0, +): + """ + Decoupled DMD2 + R3GAN + TSCD で Anima を蒸留。 + Phase A/B/C を順次実行(各 phase ごとに modal run)。 + + 例: + # Phase A (8-step, ~4-5h, ~$25-30) + modal run --detach modal_app.py::train_sota_distill --phase a + + # Phase B (4-step, ~3h, ~$19, Phase A 結果から resume) + modal run --detach modal_app.py::train_sota_distill --phase b \\ + --resume /output/distill/phase_a + + # Phase C (2-step, ~3-4h, ~$19-25) + modal run --detach modal_app.py::train_sota_distill --phase c \\ + --resume /output/distill/phase_b + """ + import os + import subprocess + phase_out = f"{out_dir}/phase_{phase}" + # Phase B/C は LoRA-only がデフォルト(Phase A の蒸留 base 上で LoRA 学習) + if phase in ("b", "c") and not gen_lora_only: + gen_lora_only = True + print(f"[train_sota] phase={phase}: defaulting gen_lora_only=True") + cmd = [ + "python", "/workspace/scripts/distill/train_sota.py", + "--phase", phase, + "--dataset", dataset_path, + "--out", phase_out, + "--total-steps", str(total_steps), + "--resolution", str(resolution), + ] + if resume: + cmd += ["--resume", resume] + if gen_lora_only: + cmd += ["--gen-lora-only"] + if override_adv_weight >= 0: + cmd += ["--override-adv-weight", str(override_adv_weight)] + if override_r3gan_gamma >= 0: + cmd += ["--override-r3gan-gamma", str(override_r3gan_gamma)] + print(f"[train_sota] {' '.join(cmd)}") + # diffusion-pipe の内部 import (`from utils.common import ...` 等) を解決させるため + # PYTHONPATH と cwd を明示。PYTHONUNBUFFERED=1 で print() を行毎 flush し、 + # 進捗ログが modal app logs で即時見えるように。 + env = { + **os.environ, + "PYTHONPATH": "/workspace/diffusion-pipe", + "PYTHONUNBUFFERED": "1", + } + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + output_vol.commit() + + +# --------------------------------------------------------------------------- +# Trajectory imitation distillation (DiffSynth-Studio Z-Image 由来、Anima 移植) +# 単一ネットワーク、critic なし。前回 5 回失敗の R3GAN 不安定性を完全回避する別路線。 +# +# 推奨フロー: +# 1) (一度だけ) Civitai Anima Turbo LoRA を warm-start として download +# modal run modal_app.py::download_civitai_lora +# 2) Smoke test (sign-of-velocity の sanity check, ~$1.5) +# modal run modal_app.py::train_traj_imitation --total-steps 1 \\ +# --teacher-steps 12 --student-steps 8 +# 3) 本番 (warm-start + 2000 step, ~$45) +# modal run --detach modal_app.py::train_traj_imitation \\ +# --warm-lora /models/loras/anima_turbo.safetensors --lpips-weight 0.1 +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=24 * 60 * 60, + secrets=[hf_secret], + cpu=8.0, + memory=131072, # 128GB (teacher + student DiT 2 個 + VAE + LPIPS 余裕) +) +def train_traj_imitation( + dataset_path: str = "/dataset/cleaned", + out_dir: str = "/output/traj", + warm_lora: str = "", # 空なら cold-start + total_steps: int = 2000, + batch_size: int = 1, + teacher_steps: int = 50, + student_steps: int = 8, + teacher_cfg: float = 2.0, + student_cfg: float = 1.0, + sigma_shift: float = 3.0, + lora_rank: int = 32, + lr: float = 1e-4, + lpips_weight: float = 0.0, # smoke では 0、本番で 0.1-0.5 + resolution: int = 1024, + log_every: int = 10, + sample_every: int = 500, + weight_mode: str = "uniform", # "uniform" or "inv_sigma" + neg_prompt: str = "", + seed: int = 42, +): + """ + Anima を trajectory imitation で 8-step 蒸留。warm-start に Civitai Anima Turbo LoRA + を使うのが推奨。詳細は scripts/distill/train_traj.py の docstring 参照。 + """ + import os + import subprocess + cmd = [ + "python", "/workspace/scripts/distill/train_traj.py", + "--dataset", dataset_path, + "--out", out_dir, + "--total-steps", str(total_steps), + "--batch-size", str(batch_size), + "--teacher-steps", str(teacher_steps), + "--student-steps", str(student_steps), + "--teacher-cfg", str(teacher_cfg), + "--student-cfg", str(student_cfg), + "--sigma-shift", str(sigma_shift), + "--lora-rank", str(lora_rank), + "--lr", str(lr), + "--lpips-weight", str(lpips_weight), + "--resolution", str(resolution), + "--log-every", str(log_every), + "--sample-every", str(sample_every), + "--weight-mode", weight_mode, + "--neg-prompt", neg_prompt, + "--seed", str(seed), + ] + if warm_lora: + cmd += ["--warm-lora", warm_lora] + print(f"[train_traj] {' '.join(cmd)}") + # diffusion-pipe の namespace 解決 + 即時 log flush + env = { + **os.environ, + "PYTHONPATH": "/workspace/diffusion-pipe", + "PYTHONUNBUFFERED": "1", + } + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + output_vol.commit() + + +# --------------------------------------------------------------------------- +# AMD Nitro-1 LADD: precompute teacher x0 cache (1 度だけ実行) +# - 全 caption に対し teacher を 20-step CFG=4.5 で回し x0 latent を保存 +# - text embedding も同時に cache → 訓練時の teacher forward 不要 → 高速化 +# B200 で 5000 サンプル precompute は ~30-40 分 ($3-4) +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=6 * 60 * 60, + secrets=[hf_secret], + cpu=4.0, + memory=65536, +) +def precompute_teacher_x0_cache( + dataset_path: str = "/dataset/raw", + out_dir: str = "/dataset/teacher_x0_cache", + num_steps: int = 20, + cfg_scale: float = 4.5, + sigma_shift: float = 3.0, + resolution: int = 768, + max_samples: int = -1, + neg_prompt: str = "", + start_from: int = 0, +): + """LADD 訓練用 teacher x0 + caption embedding cache を作成。""" + import os, subprocess + cmd = [ + "python", "/workspace/scripts/distill/precompute_teacher_x0.py", + "--dataset", dataset_path, "--out", out_dir, + "--num-steps", str(num_steps), "--cfg-scale", str(cfg_scale), + "--sigma-shift", str(sigma_shift), "--resolution", str(resolution), + "--max-samples", str(max_samples), "--start-from", str(start_from), + "--neg-prompt", neg_prompt, + ] + print(f"[precompute] {' '.join(cmd)}") + env = {**os.environ, "PYTHONPATH": "/workspace/diffusion-pipe", "PYTHONUNBUFFERED": "1"} + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + dataset_vol.commit() + + +# --------------------------------------------------------------------------- +# Reflow 用: 同じ precompute だが --save-noise 必須 +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=6 * 60 * 60, + secrets=[hf_secret], + cpu=4.0, + memory=65536, +) +def precompute_for_reflow( + dataset_path: str = "/dataset/raw", + out_dir: str = "/dataset/reflow_cache", + num_steps: int = 20, + cfg_scale: float = 4.5, + sigma_shift: float = 3.0, + resolution: int = 768, + max_samples: int = -1, + neg_prompt: str = "", + start_from: int = 0, +): + """Reflow 用 (noise, x0, emb) triplet を precompute (--save-noise 付き)。""" + import os, subprocess + cmd = [ + "python", "/workspace/scripts/distill/precompute_teacher_x0.py", + "--dataset", dataset_path, "--out", out_dir, + "--num-steps", str(num_steps), "--cfg-scale", str(cfg_scale), + "--sigma-shift", str(sigma_shift), "--resolution", str(resolution), + "--max-samples", str(max_samples), "--start-from", str(start_from), + "--neg-prompt", neg_prompt, + "--save-noise", + ] + print(f"[precompute_reflow] {' '.join(cmd)}") + env = {**os.environ, "PYTHONPATH": "/workspace/diffusion-pipe", "PYTHONUNBUFFERED": "1"} + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + dataset_vol.commit() + + +# --------------------------------------------------------------------------- +# Reflow / InstaFlow 蒸留 (Anima がもともと rectified flow なので最も自然) +# 1 grad-through forward / step、メモリ最安、batch=4-8 で動かせる +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=24 * 60 * 60, + secrets=[hf_secret], + cpu=8.0, + memory=131072, +) +def train_reflow_distill( + cache_dir: str = "/dataset/reflow_cache", + out_dir: str = "/output/reflow", + warm_lora: str = "", + total_steps: int = 8000, + batch_size: int = 4, + grad_accum: int = 2, + resolution: int = 768, + lr: float = 1e-4, + lora_rank: int = 32, + huber_delta: float = 0.03, + lpips_weight: float = 0.1, + lpips_every: int = 4, + t_sampler: str = "u_shape", + log_every: int = 10, + sample_every: int = 500, + seed: int = 42, +): + """Reflow 蒸留。precompute_for_reflow を先に走らせること。""" + import os, subprocess + cmd = [ + "python", "/workspace/scripts/distill/train_reflow.py", + "--cache-dir", cache_dir, "--out", out_dir, + "--total-steps", str(total_steps), + "--batch-size", str(batch_size), "--grad-accum", str(grad_accum), + "--resolution", str(resolution), "--lr", str(lr), + "--lora-rank", str(lora_rank), + "--huber-delta", str(huber_delta), + "--lpips-weight", str(lpips_weight), + "--lpips-every", str(lpips_every), + "--t-sampler", t_sampler, + "--log-every", str(log_every), + "--sample-every", str(sample_every), + "--seed", str(seed), + ] + if warm_lora: + cmd += ["--warm-lora", warm_lora] + print(f"[train_reflow] {' '.join(cmd)}") + env = {**os.environ, "PYTHONPATH": "/workspace/diffusion-pipe", "PYTHONUNBUFFERED": "1"} + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + output_vol.commit() + + +# --------------------------------------------------------------------------- +# PCM (Phased Consistency Model) 蒸留 — SD3-PCM 流派、Anima RF と math 一致 +# 1 grad-through + 3 no_grad forward / step、~60-80 GB on B200 +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=24 * 60 * 60, + secrets=[hf_secret], + cpu=8.0, + memory=131072, +) +def train_pcm_distill( + cache_dir: str = "/dataset/teacher_x0_cache", + out_dir: str = "/output/pcm", + warm_lora: str = "", + total_steps: int = 8000, + batch_size: int = 1, + grad_accum: int = 4, + resolution: int = 768, + num_euler_timesteps: int = 50, + num_phases: int = 4, + sigma_shift: float = 3.0, + w_min: float = 4.0, + w_max: float = 5.0, + w_fixed: float = -1.0, + huber_c: float = 1e-3, + lr: float = 5e-6, + lora_rank: int = 32, + neg_prompt: str = "", + log_every: int = 10, + sample_every: int = 500, + seed: int = 42, +): + """PCM 蒸留 (Anima 移植)。LADD cache の emb/x0 を流用。""" + import os, subprocess + cmd = [ + "python", "/workspace/scripts/distill/train_pcm.py", + "--cache-dir", cache_dir, "--out", out_dir, + "--total-steps", str(total_steps), + "--batch-size", str(batch_size), "--grad-accum", str(grad_accum), + "--resolution", str(resolution), + "--num-euler-timesteps", str(num_euler_timesteps), + "--num-phases", str(num_phases), + "--sigma-shift", str(sigma_shift), + "--w-min", str(w_min), "--w-max", str(w_max), "--w-fixed", str(w_fixed), + "--huber-c", str(huber_c), + "--lr", str(lr), + "--lora-rank", str(lora_rank), + "--neg-prompt", neg_prompt, + "--log-every", str(log_every), + "--sample-every", str(sample_every), + "--seed", str(seed), + ] + if warm_lora: + cmd += ["--warm-lora", warm_lora] + print(f"[train_pcm] {' '.join(cmd)}") + env = {**os.environ, "PYTHONPATH": "/workspace/diffusion-pipe", "PYTHONUNBUFFERED": "1"} + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + output_vol.commit() + + +# --------------------------------------------------------------------------- +# Shortcut Models 蒸留 (Frans et al. 2024) — 単一 LoRA で d 連続値 +# Reflow cache (--save-noise) 必須 +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=24 * 60 * 60, + secrets=[hf_secret], + cpu=8.0, + memory=131072, +) +def train_shortcut_distill( + cache_dir: str = "/dataset/reflow_cache", + out_dir: str = "/output/shortcut", + warm_lora: str = "", + total_steps: int = 2000, + batch_size: int = 4, + grad_accum: int = 2, + denoise_timesteps: int = 128, + bootstrap_every: int = 8, + resolution: int = 768, + lr: float = 2e-5, + lr_d_head: float = 5e-4, + lora_rank: int = 32, + clip_x_bootstrap: float = 4.0, + log_every: int = 10, + sample_every: int = 500, + seed: int = 42, +): + """Shortcut Models 蒸留 (d 連続値、1/2/4/8/128-step 自在)。""" + import os, subprocess + cmd = [ + "python", "/workspace/scripts/distill/train_shortcut.py", + "--cache-dir", cache_dir, "--out", out_dir, + "--total-steps", str(total_steps), + "--batch-size", str(batch_size), "--grad-accum", str(grad_accum), + "--denoise-timesteps", str(denoise_timesteps), + "--bootstrap-every", str(bootstrap_every), + "--resolution", str(resolution), + "--lr", str(lr), "--lr-d-head", str(lr_d_head), + "--lora-rank", str(lora_rank), + "--clip-x-bootstrap", str(clip_x_bootstrap), + "--log-every", str(log_every), "--sample-every", str(sample_every), + "--seed", str(seed), + ] + if warm_lora: + cmd += ["--warm-lora", warm_lora] + print(f"[train_shortcut] {' '.join(cmd)}") + env = {**os.environ, "PYTHONPATH": "/workspace/diffusion-pipe", "PYTHONUNBUFFERED": "1"} + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + output_vol.commit() + + +# --------------------------------------------------------------------------- +# HPSv2 weights download (DRaFT+ 用、1.97 GB、1 度だけ実行) +# --------------------------------------------------------------------------- +@app.function( + image=image, + volumes={"/models": models_vol}, + timeout=1800, + cpu=2.0, + memory=4096, +) +def download_hpsv2_weights(out_name: str = "HPS_v2_compressed.pt"): + """HF Hub から HPSv2 weights を /models/hpsv2/ に取得。""" + import os + from huggingface_hub import hf_hub_download + out_dir = "/models/hpsv2" + os.makedirs(out_dir, exist_ok=True) + dst = f"{out_dir}/{out_name}" + if os.path.exists(dst): + print(f"[skip] {dst} exists ({os.path.getsize(dst)/1e9:.2f} GB)") + return + print(f"[download] xswu/HPSv2/{out_name}") + path = hf_hub_download(repo_id="xswu/HPSv2", filename=out_name, cache_dir="/tmp/hf_cache") + import shutil + shutil.copy(path, dst) + print(f"[done] {dst} ({os.path.getsize(dst)/1e9:.2f} GB)") + models_vol.commit() + + +# --------------------------------------------------------------------------- +# DRaFT+ 品質 fine-tune (Anima 蒸留 LoRA を HPSv2 で reward fit) +# 速度向上 NOT、品質向上のための追加学習。warm-start = 既存 student LoRA 必須。 +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=12 * 60 * 60, + secrets=[hf_secret], + cpu=8.0, + memory=131072, +) +def train_draftp_distill( + dataset_path: str = "/dataset/raw", + out_dir: str = "/output/draftp", + warm_lora: str = "/models/loras/anima_turbo.safetensors", + hps_weights: str = "/models/hpsv2/HPS_v2_compressed.pt", + total_steps: int = 1500, + batch_size: int = 2, + grad_accum: int = 1, + n_student_steps: int = 8, + k_grad: int = 1, + n_lv_samples: int = 2, + resolution: int = 768, + student_cfg: float = 1.0, + sigma_shift: float = 3.0, + lr: float = 1e-4, + kl_coeff: float = 0.2, + lora_rank: int = 32, + log_every: int = 5, + sample_every: int = 200, + seed: int = 42, +): + """DRaFT+ HPSv2 reward fine-tune (student LoRA の品質向上)。""" + import os, subprocess + cmd = [ + "python", "/workspace/scripts/distill/train_draftp.py", + "--dataset", dataset_path, "--out", out_dir, + "--warm-lora", warm_lora, + "--hps-weights", hps_weights, + "--total-steps", str(total_steps), + "--batch-size", str(batch_size), "--grad-accum", str(grad_accum), + "--n-student-steps", str(n_student_steps), + "--K", str(k_grad), "--n-lv-samples", str(n_lv_samples), + "--resolution", str(resolution), "--student-cfg", str(student_cfg), + "--sigma-shift", str(sigma_shift), + "--lr", str(lr), "--kl-coeff", str(kl_coeff), + "--lora-rank", str(lora_rank), + "--log-every", str(log_every), "--sample-every", str(sample_every), + "--seed", str(seed), + ] + print(f"[train_draftp] {' '.join(cmd)}") + env = {**os.environ, "PYTHONPATH": "/workspace/diffusion-pipe", "PYTHONUNBUFFERED": "1"} + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + output_vol.commit() + + +# --------------------------------------------------------------------------- +# SiD2 / SiD-DiT 蒸留 — data-free score identity distillation +# 2 LoRA adapter (student + psi)、D/EMA 不要、ψ warmup 後 θ 活性化 +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=24 * 60 * 60, + secrets=[hf_secret], + cpu=8.0, + memory=131072, +) +def train_sid_distill( + cache_dir: str = "/dataset/teacher_x0_cache", + out_dir: str = "/output/sid", + warm_lora: str = "", + total_outer_steps: int = 8000, + psi_warmup_steps: int = 200, + n_student_steps: int = 4, + batch_size: int = 2, + grad_accum: int = 2, + resolution: int = 768, + teacher_cfg: float = 4.5, + student_cfg: float = 1.0, + alpha: float = 1.2, + mu_t: float = 0.6931, + sigma_t: float = 1.6, + lora_rank: int = 32, + lr_gen: float = 1e-5, + lr_psi: float = 2e-5, + neg_prompt: str = "", + log_every: int = 10, + sample_every: int = 500, + seed: int = 42, +): + """SiD2 / SiD-DiT (data-free)。LADD cache の emb のみ流用、画像不要。""" + import os, subprocess + cmd = [ + "python", "/workspace/scripts/distill/train_sid.py", + "--cache-dir", cache_dir, "--out", out_dir, + "--total-outer-steps", str(total_outer_steps), + "--psi-warmup-steps", str(psi_warmup_steps), + "--n-student-steps", str(n_student_steps), + "--batch-size", str(batch_size), "--grad-accum", str(grad_accum), + "--resolution", str(resolution), + "--teacher-cfg", str(teacher_cfg), "--student-cfg", str(student_cfg), + "--alpha", str(alpha), "--mu-t", str(mu_t), "--sigma-t", str(sigma_t), + "--lora-rank", str(lora_rank), + "--lr-gen", str(lr_gen), "--lr-psi", str(lr_psi), + "--neg-prompt", neg_prompt, + "--log-every", str(log_every), + "--sample-every", str(sample_every), + "--seed", str(seed), + ] + if warm_lora: + cmd += ["--warm-lora", warm_lora] + print(f"[train_sid] {' '.join(cmd)}") + env = {**os.environ, "PYTHONPATH": "/workspace/diffusion-pipe", "PYTHONUNBUFFERED": "1"} + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + output_vol.commit() + + +# --------------------------------------------------------------------------- +# AMD Nitro-1 LADD 蒸留 (PixArt → Anima 移植) +# D backbone = teacher MiniTrainDIT (frozen) + spectral-norm heads (trainable) +# Smooth L1 recon anchor で mean collapse を防ぐ (R3GAN との根本的な違い) +# 推奨: precompute 済 cache を読み、5k step bs=4 grad-accum=4 で ~$36-48 +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=24 * 60 * 60, + secrets=[hf_secret], + cpu=8.0, + memory=131072, +) +def train_ladd_distill( + cache_dir: str = "/dataset/teacher_x0_cache", + out_dir: str = "/output/ladd", + warm_lora: str = "", + total_steps: int = 5000, + batch_size: int = 4, + grad_accum: int = 4, + resolution: int = 768, + recon_lambda: float = 1.0, + lr_g: float = 1e-6, + lr_d: float = 1e-6, + t_d_max: float = 0.75, + lora_rank: int = 32, + block_ids: str = "2,8,14,20,26", + head_hidden: int = 512, + misaligned_pairs_d: bool = True, + log_every: int = 10, + sample_every: int = 500, + seed: int = 42, +): + """LADD 蒸留 (Anima 移植)。precompute_teacher_x0_cache を先に走らせること。""" + import os, subprocess + cmd = [ + "python", "/workspace/scripts/distill/train_ladd.py", + "--cache-dir", cache_dir, "--out", out_dir, + "--total-steps", str(total_steps), + "--batch-size", str(batch_size), "--grad-accum", str(grad_accum), + "--resolution", str(resolution), + "--recon-lambda", str(recon_lambda), + "--lr-g", str(lr_g), "--lr-d", str(lr_d), + "--t-d-max", str(t_d_max), + "--lora-rank", str(lora_rank), + "--block-ids", block_ids, + "--head-hidden", str(head_hidden), + "--log-every", str(log_every), + "--sample-every", str(sample_every), + "--seed", str(seed), + ] + if warm_lora: + cmd += ["--warm-lora", warm_lora] + if misaligned_pairs_d: + cmd += ["--misaligned-pairs-d"] + print(f"[train_ladd] {' '.join(cmd)}") + env = {**os.environ, "PYTHONPATH": "/workspace/diffusion-pipe", "PYTHONUNBUFFERED": "1"} + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + output_vol.commit() + + +# --------------------------------------------------------------------------- +# Cosmos-Predict2.5 公式流 DMD2 蒸留 (Anima の literal upstream 実装由来) +# 2 つの LoRA adapter (student / fake_score) を同じ base に attach、PEFT set_adapter() で切替 +# Alternating: critic × 5 → generator × 1 +# few-step rollout (1-4 step)、grad は最終 step だけ → memory efficient +# 推奨: warm-start に Anima Turbo、5000 outer step、768 解像度 で ~$21 +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=24 * 60 * 60, + secrets=[hf_secret], + cpu=8.0, + memory=131072, +) +def train_dmd2_official_distill( + dataset_path: str = "/dataset/raw", + out_dir: str = "/output/dmd2_official", + warm_lora: str = "", + total_outer_steps: int = 5000, + n_critic_per_gen: int = 5, + n_student_steps: int = 4, + batch_size: int = 1, + resolution: int = 768, + teacher_cfg: float = 3.0, + student_cfg: float = 1.0, + shift: float = 5.0, + lora_rank: int = 32, + lr_gen: float = 5e-6, + lr_critic: float = 1e-5, + log_every: int = 10, + sample_every: int = 500, + neg_prompt: str = "", + seed: int = 42, +): + """ + DMD2 + TrigFlow 蒸留 (cosmos-predict2.5 公式流派の Anima 移植)。 + 詳細は scripts/distill/train_dmd2_official.py の docstring 参照。 + """ + import os + import subprocess + cmd = [ + "python", "/workspace/scripts/distill/train_dmd2_official.py", + "--dataset", dataset_path, + "--out", out_dir, + "--total-outer-steps", str(total_outer_steps), + "--n-critic-per-gen", str(n_critic_per_gen), + "--n-student-steps", str(n_student_steps), + "--batch-size", str(batch_size), + "--resolution", str(resolution), + "--teacher-cfg", str(teacher_cfg), + "--student-cfg", str(student_cfg), + "--shift", str(shift), + "--lora-rank", str(lora_rank), + "--lr-gen", str(lr_gen), + "--lr-critic", str(lr_critic), + "--log-every", str(log_every), + "--sample-every", str(sample_every), + "--neg-prompt", neg_prompt, + "--seed", str(seed), + ] + if warm_lora: + cmd += ["--warm-lora", warm_lora] + print(f"[train_dmd2_official] {' '.join(cmd)}") + env = { + **os.environ, + "PYTHONPATH": "/workspace/diffusion-pipe", + "PYTHONUNBUFFERED": "1", + } + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + output_vol.commit() + + +# --------------------------------------------------------------------------- +# DMDX (ADM = Adversarial Distribution Matching) 蒸留 — arxiv 2507.18569v1 移植 +# DMD2 の reverse-KL grad trick を hinge GAN at t-Δt に置換 (TVD 最小化) +# discriminator は LADD-style (teacher frozen backbone + spectral norm heads) +# precompute_teacher_x0_cache を先に走らせること (real x0 source として使用) +# --------------------------------------------------------------------------- +@app.function( + image=image, + gpu="B200", + volumes=VOLUMES, + timeout=24 * 60 * 60, + secrets=[hf_secret], + cpu=8.0, + memory=131072, +) +def train_dmdx_distill( + cache_dir: str = "/dataset/teacher_x0_cache", + out_dir: str = "/output/dmdx", + warm_lora: str = "", + total_outer_steps: int = 5000, + n_critic_per_gen: int = 2, + n_student_steps: int = 4, + batch_size: int = 1, + resolution: int = 768, + teacher_cfg: float = 4.5, + student_cfg: float = 1.0, + dt_ratio: float = 1.0 / 64, + recon_weight: float = 0.0, + lora_rank: int = 32, + lr_gen: float = 5e-6, + lr_disc: float = 1e-5, + block_ids: str = "2,8,14,20,26", + head_hidden: int = 512, + log_every: int = 10, + sample_every: int = 500, + neg_prompt: str = "", + seed: int = 42, +): + """DMDX ADM-only 蒸留 (Anima 移植)。 + 詳細は scripts/distill/train_dmdx.py / dmdx_loss.py の docstring 参照。""" + import os, subprocess + cmd = [ + "python", "/workspace/scripts/distill/train_dmdx.py", + "--cache-dir", cache_dir, "--out", out_dir, + "--total-outer-steps", str(total_outer_steps), + "--n-critic-per-gen", str(n_critic_per_gen), + "--n-student-steps", str(n_student_steps), + "--batch-size", str(batch_size), + "--resolution", str(resolution), + "--teacher-cfg", str(teacher_cfg), + "--student-cfg", str(student_cfg), + "--dt-ratio", str(dt_ratio), + "--recon-weight", str(recon_weight), + "--lora-rank", str(lora_rank), + "--lr-gen", str(lr_gen), + "--lr-disc", str(lr_disc), + "--block-ids", block_ids, + "--head-hidden", str(head_hidden), + "--log-every", str(log_every), + "--sample-every", str(sample_every), + "--neg-prompt", neg_prompt, + "--seed", str(seed), + ] + if warm_lora: + cmd += ["--warm-lora", warm_lora] + print(f"[train_dmdx] {' '.join(cmd)}") + env = {**os.environ, "PYTHONPATH": "/workspace/diffusion-pipe", "PYTHONUNBUFFERED": "1"} + subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) + output_vol.commit() + + +# --------------------------------------------------------------------------- +# distill LoRA を /models/loras/ に staging (ComfyUI から LoraLoader で読める形に) +# --------------------------------------------------------------------------- +@app.function( + image=image, + volumes=VOLUMES, + timeout=300, + cpu=2.0, + memory=4096, +) +def stage_lora_to_models( + src: str = "/output/distill/phase_b/gen_lora_final.safetensors", + dest_name: str = "phase_b_4step_lora.safetensors", +): + """anima-outputs の LoRA を anima-models/loras/ にコピー。 + LoRA は PEFT 形式 → ComfyUI の LoraLoaderModelOnly は key を `diffusion_model.` + プレフィックスで探すので、必要なら変換する。""" + import os + import shutil + from safetensors.torch import load_file, save_file + + if not os.path.exists(src): + raise SystemExit(f"missing: {src}") + dst = f"/models/loras/{dest_name}" + print(f"[stage] {src} -> {dst}") + sd = load_file(src) + print(f" {len(sd)} keys; sample: {list(sd.keys())[:2]}") + # PEFT の出力 key は 'base_model.model.<...>.lora_A..weight' 形式。 + # ComfyUI の Anima/Cosmos 用 LoraLoader は 'diffusion_model.<...>.lora_down.weight' 等 + # を期待するため、変換する。 + converted = _convert_peft_to_comfy_lora(sd) + print(f" converted to ComfyUI format: {len(converted)} keys; sample: {list(converted.keys())[:2]}") + save_file(converted, dst) + models_vol.commit() + sz = os.path.getsize(dst) / 1e6 + print(f"[done] {dst} ({sz:.1f} MB, {len(converted)} keys)") + + +def _convert_peft_to_comfy_lora(sd: dict) -> dict: + """PEFT 形式 LoRA を ComfyUI(Anima/Cosmos)形式に変換。 + PEFT: 'base_model.model..lora_A..weight' (adapter は default / student 等) + Comfy(Anima): 'diffusion_model..lora_A.weight' ← adapter infix を消すだけ + (Anima Turbo LoRA の実構造を inspect_lora_keys で確認済) + """ + import re + # .lora_A..weight / .lora_B..bias 等を adapter 無しに正規化 + adapter_re = re.compile(r"\.lora_(A|B)\.[^.]+\.(weight|bias)$") + out = {} + for k, v in sd.items(): + nk = k + # PEFT prefix を剥がす + for prefix in ("base_model.model.", "base_model."): + if nk.startswith(prefix): + nk = nk[len(prefix):] + break + # 任意の adapter name (default / student / fake_score 等) を除去 + nk = adapter_re.sub(lambda m: f".lora_{m.group(1)}.{m.group(2)}", nk) + # ComfyUI 用に diffusion_model. を先頭に + if not nk.startswith("diffusion_model."): + nk = "diffusion_model." + nk + out[nk] = v + return out + + +# --------------------------------------------------------------------------- +# Phase A 完了後の比較: gen_final.safetensors を /models/checkpoints に配置 +# --------------------------------------------------------------------------- +@app.function( + image=image, + volumes=VOLUMES, + timeout=600, + cpu=2.0, + memory=8192, +) +def stage_phase_a_to_models( + src: str = "/output/distill/phase_a/gen_final.safetensors", + dest_name: str = "phase_a_distilled.safetensors", + base: str = "/models/checkpoints/anima-preview3-base.safetensors", + add_net_prefix: bool = True, +): + """anima-outputs → anima-models へ copy。 + + 重要: Phase A の save_full_state はバグで integer buffer (RoPE 等) を落としていたため、 + base から欠落 key を merge して完全な ckpt に修復する。 + """ + import os + from safetensors.torch import load_file, save_file + if not os.path.exists(src): + raise SystemExit(f"missing: {src}") + dst = f"/models/checkpoints/{dest_name}" + + print(f"[stage] loading trained: {src}") + trained = load_file(src) + print(f" {len(trained)} keys; sample: {list(trained.keys())[:3]}") + + print(f"[stage] loading base for buffer recovery: {base}") + base_sd = load_file(base) + # base は 'net.' prefix 付き → strip して比較 + base_stripped = {k[4:] if k.startswith("net.") else k: v for k, v in base_sd.items()} + print(f" base {len(base_stripped)} keys") + + # trained に無い key を base から補う + missing = [k for k in base_stripped if k not in trained] + recovered = dict(trained) + for k in missing: + recovered[k] = base_stripped[k] + print(f"[stage] recovered {len(missing)} missing keys from base " + f"(typically integer buffers like RoPE/pos)") + if missing[:5]: + print(f" sample recovered: {missing[:5]}") + + if add_net_prefix: + recovered = {"net." + k: v for k, v in recovered.items()} + print(f"[stage] added 'net.' prefix") + + save_file(recovered, dst) + models_vol.commit() + sz = os.path.getsize(dst) / 1e9 + print(f"[done] {dst} ({sz:.2f} GB, {len(recovered)} keys)") + + +# --------------------------------------------------------------------------- +# 3-way 比較生成 (base / Phase A / Turbo LoRA) × 5 prompt × 2 step counts +# --------------------------------------------------------------------------- +@app.function( + image=comfy_image, + gpu="B200", + volumes=VOLUMES, + timeout=2 * 60 * 60, + secrets=[hf_secret], + cpu=4.0, + memory=32768, +) +def compare_distill_loras( + out_subdir: str = "compare_z_vs_dmd2", + seeds_per_prompt: int = 2, + steps_list: str = "8,4", + base_cfg: float = 4.5, + distill_cfg: float = 1.0, +): + """① Z-Image と ② DMD2 蒸留 LoRA を base と 3-way 比較生成。 + + steps_list で指定した各 step 数で全 3 条件 (base/z_image/dmd2) を生成する。 + base は base_cfg、蒸留 LoRA は distill_cfg、それ以外 (prompt/seed) は全条件同一。 + + 出力 layout (Modal volume anima-dataset): + /dataset/{out_subdir}/ + base_8step_cfg4.5/ base_4step_cfg4.5/ + z_image_8step_cfg1.0/ z_image_4step_cfg1.0/ + dmd2_8step_cfg1.0/ dmd2_4step_cfg1.0/ + + download: + modal volume get anima-dataset {out_subdir}/ ./local_compare/ + """ + import os, json, subprocess, shutil + + out_base = f"/dataset/{out_subdir}" + os.makedirs(out_base, exist_ok=True) + steps_to_run = [int(s) for s in steps_list.split(",")] + + # 1) 蒸留 LoRA を /models/loras/ に staging (なければ) + lora_targets = [ + ("/output/traj_full/traj_final.safetensors", "z_image_traj_final.safetensors"), + ("/output/dmd2_full/dmd2_student_final.safetensors", "dmd2_student_final.safetensors"), + ] + for src, dest in lora_targets: + dst_path = f"/models/loras/{dest}" + if os.path.exists(dst_path): + print(f"[skip stage] {dst_path} exists") + continue + if not os.path.exists(src): + raise SystemExit(f"missing: {src} — training may still be running") + print(f"[stage] {src} -> {dst_path}") + from safetensors.torch import load_file, save_file + sd = load_file(src) + converted = _convert_peft_to_comfy_lora(sd) + save_file(converted, dst_path) + models_vol.commit() + print(f" {len(converted)} keys, {os.path.getsize(dst_path)/1e6:.1f} MB") + + # 2) workflow テンプレ: turbo workflow (LoRA 推論用) の lora_name を書き換える + base_wf = "/workspace/scripts/anima_workflow.json" + turbo_wf = "/workspace/scripts/anima_workflow_turbo.json" + + def _patched_workflow(lora_name: str) -> str: + """turbo workflow を copy して lora_name だけ差し替え、/tmp/wf_.json に保存。""" + with open(turbo_wf) as f: + wf = json.load(f) + for node_id, node in wf.items(): + # _comment などの string value はスキップ + if not isinstance(node, dict): + continue + if node.get("class_type") == "LoraLoaderModelOnly": + node["inputs"]["lora_name"] = lora_name + print(f" patched LoraLoaderModelOnly.lora_name -> {lora_name}") + out_path = f"/tmp/wf_{lora_name.replace('.safetensors','')}.json" + with open(out_path, "w") as f: + json.dump(wf, f) + return out_path + + # (label, workflow, cfg, method_label) + # method_label は README/評価で使う詳細名 (sampler/scheduler は workflow から自動抽出) + conditions = [ + ("base", base_wf, base_cfg, "anima_base_no_lora"), + ("turbo", _patched_workflow("anima_turbo.safetensors"), distill_cfg, "civitai_anima_turbo_v0.1"), + ("z_image", _patched_workflow("z_image_traj_final.safetensors"), distill_cfg, "ours_z_image_trajectory_imitation"), + ("dmd2", _patched_workflow("dmd2_student_final.safetensors"), distill_cfg, "ours_dmd2_trigflow"), + ] + + # 2 軸ループ: steps × conditions + for steps in steps_to_run: + for label, wf, cfg, method in conditions: + out_dir = f"{out_base}/{label}_{steps}step_cfg{cfg}" + cmd = [ + "python", "/workspace/scripts/generate_dataset.py", + "--prompts", "/workspace/scripts/compare_prompts.txt", + "--workflow", wf, + "--out", out_dir, + "--comfy-dir", "/workspace/ComfyUI", + "--models-dir", "/models/checkpoints", + "--seeds-per-prompt", str(seeds_per_prompt), + "--override-steps", str(steps), + "--override-cfg", str(cfg), + "--fixed-aspect", "1024x1024", + "--method-label", method, + ] + print(f"\n=== {label} @ {steps} step CFG {cfg} (method={method}) ===") + subprocess.run(cmd, check=True) + dataset_vol.commit() + print(f"\n[done] images at /dataset/{out_subdir}/") + print(f"download: modal volume get anima-dataset {out_subdir}/ ./local_compare/") + + +# --------------------------------------------------------------------------- +# sageattention の動作検証 + volume への永続化 (image rebuild 短縮用) +# Modal の pip layer は cache されるので普段は不要だが、image 定義改変時の保険。 +# --------------------------------------------------------------------------- +@app.function( + image=image, # ← sageattention 入り base から始まる + gpu="B200", # CUDA kernel 検証のため GPU 必須 + volumes={"/models": models_vol}, + timeout=600, + cpu=4.0, + memory=16384, +) +def cache_sageattention_to_volume(): + """sageattention import + 簡易 CUDA kernel 呼び出し検証 → /models/cache/sageattention/ に .so 保存。""" + import os, shutil, glob, importlib + cache_dir = "/models/cache/sageattention" + os.makedirs(cache_dir, exist_ok=True) + + print("[sage] import test") + try: + import sageattention + print(f"[sage] version: {getattr(sageattention, '__version__', 'unknown')}") + print(f"[sage] path: {sageattention.__file__}") + except ImportError as e: + print(f"[sage] FAILED to import: {e}") + return + + # CUDA kernel 動作確認 (小さい dummy attention) + print("[sage] CUDA kernel smoke test") + try: + import torch + from sageattention import sageattn + B, H, L, D = 1, 8, 64, 64 + q = torch.randn(B, H, L, D, device="cuda", dtype=torch.float16) + k = torch.randn(B, H, L, D, device="cuda", dtype=torch.float16) + v = torch.randn(B, H, L, D, device="cuda", dtype=torch.float16) + out = sageattn(q, k, v) + print(f"[sage] kernel OK, output shape={tuple(out.shape)} dtype={out.dtype}") + except Exception as e: + print(f"[sage] kernel FAILED: {e}") + return + + # site-packages の sageattention ディレクトリを volume に丸ごとコピー + pkg_dir = os.path.dirname(sageattention.__file__) + pkg_name = os.path.basename(pkg_dir) + dst = os.path.join(cache_dir, pkg_name) + if os.path.exists(dst): + shutil.rmtree(dst) + shutil.copytree(pkg_dir, dst) + so_files = glob.glob(f"{dst}/**/*.so", recursive=True) + glob.glob(f"{dst}/**/*.pyd", recursive=True) + print(f"[sage] copied to {dst} ({len(so_files)} compiled binaries)") + for so in so_files[:5]: + print(f" - {so} ({os.path.getsize(so)/1e6:.1f} MB)") + models_vol.commit() + print(f"[sage] saved to volume at {cache_dir}") + + +# --------------------------------------------------------------------------- +# Quick health check: 任意の蒸留 LoRA を 1 枚生成して visual sanity 確認 +# 訓練途中の checkpoint が壊れていないか早期検出する +# --------------------------------------------------------------------------- +@app.function( + image=comfy_image, + gpu="B200", + volumes=VOLUMES, + timeout=30 * 60, + secrets=[hf_secret, civitai_secret], + cpu=4.0, + memory=32768, +) +def quick_check_ckpt( + src_path: str, + label: str, + step_count: int = 4, + cfg: float = 1.0, + sampler_name: str = "", + scheduler: str = "", + out_subdir: str = "ckpt_health_check", +): + """指定 LoRA を stage して 1 枚 verify 生成。出力 path: + /dataset/{out_subdir}/{label}/p0000_s000.png + _summary.json + sampler_name / scheduler 空文字なら workflow JSON のデフォルト (er_sde / simple)。""" + import os, json, subprocess + from safetensors.torch import load_file, save_file + + if not os.path.exists(src_path): + raise SystemExit(f"missing: {src_path}") + + # stage PEFT → ComfyUI 形式 + comfy_name = f"{label}.safetensors" + dst = f"/models/loras/{comfy_name}" + sd = load_file(src_path) + converted = _convert_peft_to_comfy_lora(sd) + save_file(converted, dst) + models_vol.commit() + print(f"[stage] {src_path} -> {dst} ({len(converted)} keys)") + + # workflow patch + base_wf = "/workspace/scripts/anima_verify_workflow.json" + with open(base_wf) as f: + wf = json.load(f) + for nid, node in wf.items(): + if not isinstance(node, dict): + continue + if node.get("class_type") == "LoraLoaderModelOnly": + node["inputs"]["lora_name"] = comfy_name + node["inputs"]["strength_model"] = 1.0 + if node.get("class_type") == "KSampler": + if sampler_name: + node["inputs"]["sampler_name"] = sampler_name + if scheduler: + node["inputs"]["scheduler"] = scheduler + patched = f"/tmp/wf_check_{label}.json" + with open(patched, "w") as f: + json.dump(wf, f) + + sched_tag = f"_{sampler_name}_{scheduler}" if (sampler_name or scheduler) else "" + out_dir = f"/dataset/{out_subdir}/{label}_step{step_count}_cfg{cfg}{sched_tag}" + cmd = [ + "python", "/workspace/scripts/generate_dataset.py", + "--prompts", "/workspace/scripts/verify_prompts.txt", + "--workflow", patched, + "--out", out_dir, + "--comfy-dir", "/workspace/ComfyUI", + "--models-dir", "/models/checkpoints", + "--seeds-per-prompt", "1", + "--base-seed", "42", + "--override-steps", str(step_count), + "--override-cfg", str(cfg), + "--fixed-aspect", "1024x1024", + "--method-label", f"ckpt_{label}", + ] + print(f"[check] {label} @ {step_count} step CFG {cfg}" + f"{' sampler=' + sampler_name if sampler_name else ''}" + f"{' scheduler=' + scheduler if scheduler else ''}") + subprocess.run(cmd, check=True) + dataset_vol.commit() + print(f"[check] done -> {out_dir}/") + + +# --------------------------------------------------------------------------- +# HF Hub upload: 蒸留 LoRA を HF model repo に配布 +# PEFT 原形 + ComfyUI 変換 + README + sample 画像をまとめてアップロード +# --------------------------------------------------------------------------- +@app.function( + image=image, + volumes=VOLUMES, + timeout=30 * 60, + secrets=[hf_write_secret], + cpu=2.0, + memory=8192, +) +def upload_lora_to_hf( + repo_id: str, + peft_path: str = "", + subdir: str = "pcm", + sample_paths: str = "", + readme_path: str = "", + root_readme_path: str = "", + private: bool = False, +): + """指定 LoRA を HF model repo に upload。 + peft_path: /output/pcm_v1/pcm_final.safetensors (空なら LoRA upload skip) + repo_id: darask0/anima-distill-loras + subdir: repo 内サブディレクトリ (pcm / ladd / reflow など複数手法共存用) + sample_paths: comma-separated 'path' または 'path=newname.png' エントリ + readme_path: subdir/README.md の元ファイル (Modal volume path、空なら skip) + root_readme_path: repo 直下 README.md の元ファイル (model card 表示用) + 既存 repo は overwrite、存在しなければ create_repo。""" + import os, json + from pathlib import Path + from safetensors.torch import load_file, save_file + from huggingface_hub import HfApi, create_repo + + token = os.environ.get("HF_TOKEN_WRITE") or os.environ.get("HF_TOKEN") + if not token: + raise SystemExit("no HF token found (HF_TOKEN_WRITE / HF_TOKEN)") + api = HfApi(token=token) + + # 1) repo を確保 (なければ作成) + print(f"[hf] ensuring repo {repo_id} (private={private})") + create_repo(repo_id=repo_id, repo_type="model", token=token, + private=private, exist_ok=True) + + # 1.5) root README (model card) を upload + if root_readme_path: + if not os.path.exists(root_readme_path): + raise SystemExit(f"missing root readme: {root_readme_path}") + print(f"[hf] uploading ROOT README {root_readme_path} -> README.md") + api.upload_file( + path_or_fileobj=root_readme_path, + path_in_repo="README.md", + repo_id=repo_id, repo_type="model", token=token, + ) + + # 2) PEFT 形式そのまま + if peft_path: + peft_dst_name = Path(peft_path).name.replace(".safetensors", "_peft.safetensors") + print(f"[hf] uploading PEFT: {peft_path} -> {subdir}/{peft_dst_name}") + api.upload_file( + path_or_fileobj=peft_path, + path_in_repo=f"{subdir}/{peft_dst_name}", + repo_id=repo_id, repo_type="model", token=token, + ) + + # 3) ComfyUI 形式変換 + upload + sd = load_file(peft_path) + converted = _convert_peft_to_comfy_lora(sd) + comfy_tmp = f"/tmp/{Path(peft_path).stem}_comfy.safetensors" + save_file(converted, comfy_tmp) + comfy_dst_name = Path(peft_path).name.replace(".safetensors", "_comfy.safetensors") + print(f"[hf] uploading ComfyUI: {comfy_tmp} ({len(converted)} keys) -> {subdir}/{comfy_dst_name}") + api.upload_file( + path_or_fileobj=comfy_tmp, + path_in_repo=f"{subdir}/{comfy_dst_name}", + repo_id=repo_id, repo_type="model", token=token, + ) + + # 4) README (subdir/README.md) + if readme_path: + if not os.path.exists(readme_path): + raise SystemExit(f"missing readme: {readme_path}") + print(f"[hf] uploading README {readme_path} -> {subdir}/README.md") + api.upload_file( + path_or_fileobj=readme_path, + path_in_repo=f"{subdir}/README.md", + repo_id=repo_id, repo_type="model", token=token, + ) + + # 5) sample 画像 (entry: "path" or "path=rename.png" 形式) + if sample_paths: + for entry in sample_paths.split(","): + entry = entry.strip() + if not entry: + continue + if "=" in entry: + sp, new_name = entry.split("=", 1) + sp = sp.strip() + new_name = new_name.strip() + else: + sp = entry + new_name = Path(sp).name + if not os.path.exists(sp): + print(f"[hf] sample skip (missing): {sp}") + continue + dst = f"{subdir}/samples/{new_name}" + print(f"[hf] uploading sample: {sp} -> {dst}") + api.upload_file( + path_or_fileobj=sp, + path_in_repo=dst, + repo_id=repo_id, repo_type="model", token=token, + ) + + print(f"[hf] done. https://huggingface.co/{repo_id}/tree/main/{subdir}") + + +# --------------------------------------------------------------------------- +# プロジェクトコード一式を HF model repo に push (rapid-anima リポジトリ用) +# バンドル済 /workspace/{scripts,configs} + volume にアップロード済 README/docs/samples +# を組み合わせて HF にアップロード +# --------------------------------------------------------------------------- +@app.function( + image=image, + volumes=VOLUMES, + timeout=15 * 60, + secrets=[hf_write_secret], + cpu=2.0, + memory=8192, +) +def upload_project_code_to_hf( + repo_id: str = "darask0/rapid-anima", + extras_dir: str = "/dataset/repo_extras", + private: bool = False, + commit_message: str = "Initial commit: rapid-anima distillation codebase", +): + """プロジェクトコード一式を HF model repo にアップロード。 + extras_dir: README.md / docs/ / samples/ / modal_app.py / requirements.txt / LICENSE + を含む Modal volume 上のディレクトリ + バンドル済の /workspace/scripts と /workspace/configs も一緒にアップロード。 + """ + import os, shutil + from pathlib import Path + from huggingface_hub import HfApi, create_repo + + token = os.environ.get("HF_TOKEN_WRITE") or os.environ.get("HF_TOKEN") + if not token: + raise SystemExit("no HF token in env (HF_TOKEN_WRITE / HF_TOKEN)") + api = HfApi(token=token) + + # 1) repo 確保 + print(f"[hf] ensuring repo {repo_id} (private={private})") + create_repo(repo_id=repo_id, repo_type="model", token=token, + private=private, exist_ok=True) + + # 2) アップロード用ツリーを /tmp に組み立てる + upload_root = Path("/tmp/rapid_anima_upload") + if upload_root.exists(): + shutil.rmtree(upload_root) + upload_root.mkdir(parents=True) + + # バンドル済 scripts / configs をコピー (__pycache__ は除外) + def _copy_filtered(src: Path, dst: Path): + def ignore(_dir, names): + return [n for n in names if n == "__pycache__" or n.endswith(".pyc")] + shutil.copytree(src, dst, ignore=ignore) + + _copy_filtered(Path("/workspace/scripts"), upload_root / "scripts") + _copy_filtered(Path("/workspace/configs"), upload_root / "configs") + print(f"[hf] bundled scripts/ ({sum(1 for _ in (upload_root/'scripts').rglob('*'))} entries)") + print(f"[hf] bundled configs/ ({sum(1 for _ in (upload_root/'configs').rglob('*'))} entries)") + + # extras (README / docs / samples / modal_app.py / requirements.txt / LICENSE) + extras = Path(extras_dir) + if not extras.exists(): + raise SystemExit(f"extras_dir not found: {extras_dir}") + for item in extras.iterdir(): + target = upload_root / item.name + if item.is_dir(): + _copy_filtered(item, target) + else: + shutil.copy(item, target) + print(f"[hf] staged extras: {item.name}") + + # 3) upload_folder で一括アップロード + print(f"[hf] uploading folder → {repo_id}") + api.upload_folder( + folder_path=str(upload_root), + repo_id=repo_id, + repo_type="model", + token=token, + commit_message=commit_message, + ignore_patterns=["__pycache__", "*.pyc", ".DS_Store", ".ipynb_checkpoints"], + ) + print(f"[hf] done. https://huggingface.co/{repo_id}") + + +# --------------------------------------------------------------------------- +# 完成済 LoRA をユーザー指定 (Anima_simple ベース) workflow で 2 並列検証 +# 破綻ポイント (2 キャラ / 手指ポーズ / 詳細背景) を含む prompt で品質比較。 +# 全 metadata (sampler, scheduler, step, cfg, gen_time, lora_name) を json sidecar に保存。 +# sageattention は image に同梱済、generate_dataset.py 側で自動有効化。 +# --------------------------------------------------------------------------- +@app.function( + image=comfy_image, + gpu="B200", + volumes=VOLUMES, + timeout=2 * 60 * 60, + secrets=[hf_secret, civitai_secret], + cpu=8.0, + memory=65536, +) +def verify_completed_loras( + out_subdir: str = "verify_completed", + group: str = "A", + seeds_per_prompt: int = 1, + base_seed: int = 42, +): + """完成済 LoRA を Anima_simple workflow で 1 個ずつ生成、metadata 全保存。 + + group="A": base, civitai turbo, ① Z-Image, ② DMD2 (4 conditions、各 8 と 4 step) + group="B": ⑩ DRaFT+ on Z, B DRaFT+ on Turbo, ④ Reflow (3 conditions、各 8 と 4 step) + + base のみ 30 step CFG=4.5 (元の Anima_simple.json の値)、LoRA は 8/4 step CFG=1.0。 + """ + import os, json, subprocess + from safetensors.torch import load_file, save_file + + out_base = f"/dataset/{out_subdir}" + os.makedirs(out_base, exist_ok=True) + + # 1) 完成済 LoRA を staging (PEFT → ComfyUI 形式変換) + stage_targets = [ + ("/output/draftp_full/draftp_final.safetensors", "draftp_on_zimage.safetensors"), + ("/output/draftp_on_turbo/draftp_final.safetensors", "draftp_on_turbo.safetensors"), + ("/output/reflow_full/reflow_final.safetensors", "reflow_final.safetensors"), + ] + for src, dest in stage_targets: + dst = f"/models/loras/{dest}" + if os.path.exists(dst): + continue + if not os.path.exists(src): + print(f"[skip] {src} not found") + continue + sd = load_file(src) + converted = _convert_peft_to_comfy_lora(sd) + save_file(converted, dst) + models_vol.commit() + print(f"[stage] {src} -> {dst} ({len(converted)} keys)") + + # 2) workflow patch helper (lora name 差替え、LoRA なし時は strength=0) + base_wf_path = "/workspace/scripts/anima_verify_workflow.json" + + def _patched(lora_name: str | None, label: str) -> str: + with open(base_wf_path) as f: + wf = json.load(f) + for nid, node in wf.items(): + if not isinstance(node, dict): + continue + if node.get("class_type") == "LoraLoaderModelOnly": + if lora_name is None: + # base 専用: LoRA strength を 0 にして実質 disable + node["inputs"]["strength_model"] = 0.0 + else: + node["inputs"]["lora_name"] = lora_name + node["inputs"]["strength_model"] = 1.0 + path = f"/tmp/verify_wf_{label}.json" + with open(path, "w") as f: + json.dump(wf, f) + return path + + # 3) condition list を group で分割 + # (label, lora_name|None, method_label, steps, cfg) + group_a_conds = [ + ("base_30step", None, "anima_base_no_lora", 30, 4.5), + ("turbo_8step", "anima_turbo.safetensors", "civitai_anima_turbo_v0.1", 8, 1.0), + ("turbo_4step", "anima_turbo.safetensors", "civitai_anima_turbo_v0.1", 4, 1.0), + ("zimage_8step", "z_image_traj_final.safetensors", "ours_z_image_trajectory_imitation", 8, 1.0), + ("zimage_4step", "z_image_traj_final.safetensors", "ours_z_image_trajectory_imitation", 4, 1.0), + ("dmd2_8step", "dmd2_student_final.safetensors", "ours_dmd2_trigflow", 8, 1.0), + ("dmd2_4step", "dmd2_student_final.safetensors", "ours_dmd2_trigflow", 4, 1.0), + ] + group_b_conds = [ + ("draftp_on_z_8step", "draftp_on_zimage.safetensors", "ours_draftp_hpsv2_on_zimage", 8, 1.0), + ("draftp_on_z_4step", "draftp_on_zimage.safetensors", "ours_draftp_hpsv2_on_zimage", 4, 1.0), + ("draftp_on_turbo_8step", "draftp_on_turbo.safetensors", "ours_draftp_hpsv2_on_turbo", 8, 1.0), + ("draftp_on_turbo_4step", "draftp_on_turbo.safetensors", "ours_draftp_hpsv2_on_turbo", 4, 1.0), + ("reflow_8step", "reflow_final.safetensors", "ours_reflow_rfpp", 8, 1.0), + ("reflow_4step", "reflow_final.safetensors", "ours_reflow_rfpp", 4, 1.0), + ] + + if group.upper() == "ALL": + conditions = group_a_conds + group_b_conds + elif group.upper() == "A": + conditions = group_a_conds + else: + conditions = group_b_conds + print(f"[verify] group={group} conditions={len(conditions)}") + + # 4) 各 condition を順次生成 + for label, lora_name, method_label, steps, cfg in conditions: + wf = _patched(lora_name, label) + out_dir = f"{out_base}/{label}_cfg{cfg}" + cmd = [ + "python", "/workspace/scripts/generate_dataset.py", + "--prompts", "/workspace/scripts/verify_prompts.txt", + "--workflow", wf, + "--out", out_dir, + "--comfy-dir", "/workspace/ComfyUI", + "--models-dir", "/models/checkpoints", + "--seeds-per-prompt", str(seeds_per_prompt), + "--base-seed", str(base_seed), + "--override-steps", str(steps), + "--override-cfg", str(cfg), + "--fixed-aspect", "1024x1024", + "--method-label", method_label, + ] + print(f"\n=== {label} | step={steps} cfg={cfg} method={method_label} ===") + subprocess.run(cmd, check=True) + dataset_vol.commit() + print(f"\n[done group {group}] /dataset/{out_subdir}/ に {len(conditions)} 条件生成完了") + + +# --------------------------------------------------------------------------- +# 全 9 条件 (base + civitai turbo + 7 自前蒸留) の比較生成。 +# 各 method の final LoRA が存在する条件のみ自動 staging + 比較対象に追加。 +# --------------------------------------------------------------------------- +@app.function( + image=comfy_image, + gpu="B200", + volumes=VOLUMES, + timeout=4 * 60 * 60, + secrets=[hf_secret, civitai_secret], + cpu=8.0, + memory=65536, +) +def compare_all_methods( + out_subdir: str = "compare_all", + seeds_per_prompt: int = 2, + steps_list: str = "8,4", + base_cfg: float = 4.5, + distill_cfg: float = 1.0, +): + """全手法を 1 つの matrix で比較。final LoRA が存在する method のみ対象に含める。 + + 出力: /dataset/{out_subdir}/