""" 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}/