| """ |
| 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 |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| _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", |
| |
| |
| "lpips>=0.1.4", |
| |
| "open_clip_torch>=2.24", |
| "hpsv2>=1.2", |
| "ftfy", "regex", |
| ) |
| |
| |
| |
| .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", |
| |
| |
| |
| |
| "pip install --no-build-isolation sageattention || echo '[warn] sageattention install failed, falling back to torch SDPA'", |
| ) |
| ) |
|
|
| |
| image = ( |
| _base_image |
| .add_local_dir("configs", "/workspace/configs") |
| .add_local_dir("scripts", "/workspace/scripts") |
| ) |
|
|
| |
| comfy_image = ( |
| _base_image |
| .run_commands( |
| "git clone --depth 1 https://github.com/comfyanonymous/ComfyUI /workspace/ComfyUI", |
| |
| |
| "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_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( |
| |
| "git clone --depth 1 https://github.com/comfyanonymous/ComfyUI /workspace/ComfyUI", |
| "pip install --upgrade-strategy only-if-needed -r /workspace/ComfyUI/requirements.txt", |
| |
| |
| f"pip install --no-deps --force-reinstall {B200_SAGE_WHEEL_URL}", |
| |
| |
| "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") |
| ) |
|
|
| |
| |
| |
| 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) |
| |
| |
| |
| 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_secret = modal.Secret.from_name("hf_token", required_keys=["HF_TOKEN"]) |
|
|
| |
| hf_write_secret = modal.Secret.from_name("HF_TOKEN_WRITE") |
|
|
| |
| civitai_secret = modal.Secret.from_name("civitai_api_key", required_keys=["CIVITAI_API_KEY"]) |
|
|
| app = modal.App("rapid-anima") |
|
|
|
|
| |
| |
| |
| |
| @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}" |
|
|
| |
| 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/") |
|
|
|
|
| |
| |
| |
| |
| |
| |
| @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") |
| |
| 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}") |
|
|
| |
| 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}") |
|
|
|
|
| |
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| @app.function( |
| image=image, |
| gpu="A100-80GB", |
| volumes=VOLUMES, |
| timeout=24 * 60 * 60, |
| secrets=[hf_secret], |
| cpu=8.0, |
| memory=32768, |
| ) |
| 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}") |
|
|
| |
| 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/") |
|
|
|
|
| |
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
|
|
|
|
| |
| |
| |
| @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, |
| 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): |
| print(f" ... {size/1e6:.0f} MB") |
| print(f"[done] {out} ({size/1e6:.1f} MB)") |
| models_vol.commit() |
|
|
|
|
| |
| |
| |
| |
| |
| @app.function( |
| image=comfy_image, |
| gpu="A100-80GB", |
| volumes=VOLUMES, |
| timeout=18 * 60 * 60, |
| 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() |
|
|
|
|
| |
| |
| |
| |
| @app.function( |
| image=b200_image, |
| gpu="B200", |
| volumes=VOLUMES, |
| timeout=4 * 60 * 60, |
| 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}") |
|
|
|
|
| |
| |
| |
| |
| @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 |
| |
| 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, |
| |
| ], |
| 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}/") |
|
|
|
|
| |
| |
| |
| |
| @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) |
|
|
| |
| base_cfg = "/workspace/configs/phase1_anima.toml" |
| new_cfg = f"/tmp/finetune_{variant_name}.toml" |
| txt = open(base_cfg, encoding="utf-8").read() |
| |
| 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) |
| |
| txt = re.sub(r"llm_adapter_lr\s*=.*", "llm_adapter_lr = 0", txt) |
| |
| 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) |
| |
| 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}") |
|
|
| |
| 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 = [ |
| |
| ("A_baseline", 32, 2e-5, 3, 0.0), |
| ("B_higher_capacity", 64, 2e-5, 3, 0.0), |
| ("C_safer_lr", 32, 1e-5, 4, 0.0), |
| ("D_artist_focus", 32, 3e-5, 3, 0.0), |
| ] |
| 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 |
| ] |
| |
| for h in handles: |
| h.get() |
| print(f"[parallel] all {workers} workers complete") |
|
|
|
|
| |
| |
| |
| |
| @app.function( |
| image=image, |
| gpu="B200", |
| volumes=VOLUMES, |
| timeout=24 * 60 * 60, |
| secrets=[hf_secret], |
| cpu=8.0, |
| memory=131072, |
| ) |
| 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}" |
| |
| 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)}") |
| |
| |
| |
| env = { |
| **os.environ, |
| "PYTHONPATH": "/workspace/diffusion-pipe", |
| "PYTHONUNBUFFERED": "1", |
| } |
| subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) |
| output_vol.commit() |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| @app.function( |
| image=image, |
| gpu="B200", |
| volumes=VOLUMES, |
| timeout=24 * 60 * 60, |
| secrets=[hf_secret], |
| cpu=8.0, |
| memory=131072, |
| ) |
| def train_traj_imitation( |
| dataset_path: str = "/dataset/cleaned", |
| out_dir: str = "/output/traj", |
| warm_lora: str = "", |
| 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, |
| resolution: int = 1024, |
| log_every: int = 10, |
| sample_every: int = 500, |
| weight_mode: str = "uniform", |
| 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)}") |
| |
| env = { |
| **os.environ, |
| "PYTHONPATH": "/workspace/diffusion-pipe", |
| "PYTHONUNBUFFERED": "1", |
| } |
| subprocess.run(cmd, check=True, cwd="/workspace/diffusion-pipe", env=env) |
| output_vol.commit() |
|
|
|
|
| |
| |
| |
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
| |
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
| |
| |
| |
| |
| @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() |
|
|
|
|
| |
| |
| |
| @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]}") |
| |
| |
| |
| 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.<module>.lora_A.<adapter>.weight' (adapter は default / student 等) |
| Comfy(Anima): 'diffusion_model.<module>.lora_A.weight' ← adapter infix を消すだけ |
| (Anima Turbo LoRA の実構造を inspect_lora_keys で確認済) |
| """ |
| import re |
| |
| adapter_re = re.compile(r"\.lora_(A|B)\.[^.]+\.(weight|bias)$") |
| out = {} |
| for k, v in sd.items(): |
| nk = k |
| |
| for prefix in ("base_model.model.", "base_model."): |
| if nk.startswith(prefix): |
| nk = nk[len(prefix):] |
| break |
| |
| nk = adapter_re.sub(lambda m: f".lora_{m.group(1)}.{m.group(2)}", nk) |
| |
| if not nk.startswith("diffusion_model."): |
| nk = "diffusion_model." + nk |
| out[nk] = v |
| return out |
|
|
|
|
| |
| |
| |
| @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_stripped = {k[4:] if k.startswith("net.") else k: v for k, v in base_sd.items()} |
| print(f" base {len(base_stripped)} keys") |
|
|
| |
| 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)") |
|
|
|
|
| |
| |
| |
| @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(",")] |
|
|
| |
| 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") |
|
|
| |
| 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_<name>.json に保存。""" |
| with open(turbo_wf) as f: |
| wf = json.load(f) |
| for node_id, node in wf.items(): |
| |
| 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 |
|
|
| |
| |
| 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"), |
| ] |
|
|
| |
| 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/") |
|
|
|
|
| |
| |
| |
| |
| @app.function( |
| image=image, |
| gpu="B200", |
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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}") |
|
|
|
|
| |
| |
| |
| |
| @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}") |
|
|
| |
| 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)") |
|
|
| |
| 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}/") |
|
|
|
|
| |
| |
| |
| |
| @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) |
|
|
| |
| 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) |
|
|
| |
| 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, |
| ) |
|
|
| |
| 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, |
| ) |
|
|
| |
| 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, |
| ) |
|
|
| |
| 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, |
| ) |
|
|
| |
| 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}") |
|
|
|
|
| |
| |
| |
| |
| |
| @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) |
|
|
| |
| 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) |
|
|
| |
| upload_root = Path("/tmp/rapid_anima_upload") |
| if upload_root.exists(): |
| shutil.rmtree(upload_root) |
| upload_root.mkdir(parents=True) |
|
|
| |
| 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 = 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}") |
|
|
| |
| 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}") |
|
|
|
|
| |
| |
| |
| |
| |
| |
| @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) |
|
|
| |
| 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)") |
|
|
| |
| 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: |
| |
| 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 |
|
|
| |
| |
| 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)}") |
|
|
| |
| 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)} 条件生成完了") |
|
|
|
|
| |
| |
| |
| |
| @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}/<label>_<steps>step_cfg<cfg>/ |
| 画像 + _times.jsonl + _summary.json (method_label, sampler, scheduler, lora 等) |
| """ |
| 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) |
| steps_to_run = [int(s) for s in steps_list.split(",")] |
|
|
| |
| candidates = [ |
| ("z_image", "/output/traj_full/traj_final.safetensors", "z_image_traj_final.safetensors", "ours_z_image_trajectory_imitation"), |
| ("dmd2", "/output/dmd2_full/dmd2_student_final.safetensors", "dmd2_student_final.safetensors", "ours_dmd2_trigflow"), |
| ("ladd", "/output/ladd_full/ladd_student_final.safetensors", "ladd_student_final.safetensors", "ours_ladd"), |
| ("reflow", "/output/reflow_full/reflow_final.safetensors", "reflow_final.safetensors", "ours_reflow_rfpp"), |
| ("pcm", "/output/pcm_full/pcm_final.safetensors", "pcm_final.safetensors", "ours_pcm_phased_consistency"), |
| ("sid", "/output/sid_full/sid_student_final.safetensors", "sid_student_final.safetensors", "ours_sid2"), |
| ("shortcut", "/output/shortcut_full/shortcut_final.safetensors", "shortcut_final.safetensors", "ours_shortcut_models"), |
| ("draftp", "/output/draftp_full/draftp_final.safetensors", "draftp_final.safetensors", "ours_draftp_hpsv2_reward"), |
| ] |
|
|
| |
| available = [] |
| for label, src, comfy_name, method in candidates: |
| dst = f"/models/loras/{comfy_name}" |
| if not os.path.exists(src): |
| print(f"[skip] {label}: {src} not found (training not complete?)") |
| continue |
| if not os.path.exists(dst): |
| print(f"[stage] {src} -> {dst}") |
| sd = load_file(src) |
| converted = _convert_peft_to_comfy_lora(sd) |
| save_file(converted, dst) |
| models_vol.commit() |
| available.append((label, comfy_name, method)) |
| print(f"[ready] {len(available)} distill methods + base + turbo") |
|
|
| |
| base_wf = "/workspace/scripts/anima_workflow.json" |
| turbo_wf = "/workspace/scripts/anima_workflow_turbo.json" |
|
|
| def _patched_workflow(lora_name: str) -> str: |
| with open(turbo_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"] = 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 |
|
|
| |
| conditions = [ |
| ("base", base_wf, base_cfg, "anima_base_no_lora"), |
| ("turbo", _patched_workflow("anima_turbo.safetensors"), distill_cfg, "civitai_anima_turbo_v0.1"), |
| ] |
| for label, comfy_name, method in available: |
| conditions.append((label, _patched_workflow(comfy_name), distill_cfg, method)) |
| print(f"[conditions] total {len(conditions)} × {len(steps_to_run)} step = {len(conditions)*len(steps_to_run)} runs") |
|
|
| |
| 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_all/") |
|
|
|
|
| @app.function( |
| image=comfy_image, |
| gpu="B200", |
| volumes=VOLUMES, |
| timeout=2 * 60 * 60, |
| secrets=[hf_secret, civitai_secret], |
| cpu=8.0, |
| memory=65536, |
| ) |
| def compare_phase_a( |
| steps_list: str = "8,12", |
| out_subdir: str = "compare", |
| seeds_per_prompt: int = 1, |
| ): |
| """5 prompt × 3 model × N steps の比較画像を /dataset/{out_subdir}/ に出力。 |
| |
| 使い方: |
| modal run modal_app.py::stage_phase_a_to_models # まず gen を models に staging |
| modal run modal_app.py::compare_phase_a # 30 枚生成 |
| """ |
| import os |
| import subprocess |
|
|
| steps_to_run = [int(s) for s in steps_list.split(",")] |
| out_base = f"/dataset/{out_subdir}" |
| os.makedirs(out_base, exist_ok=True) |
|
|
| conditions = [ |
| |
| ("base", "/workspace/scripts/anima_workflow.json", 4.5, False), |
| ("phase_a", "/workspace/scripts/anima_workflow_phase_a.json", 1.0, False), |
| ("turbo", "/workspace/scripts/anima_workflow_turbo.json", 1.0, False), |
| ] |
|
|
| for steps in steps_to_run: |
| for label, wf, cfg, no_prefix in conditions: |
| out_dir = f"{out_base}/{label}_{steps}step" |
| 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", |
| ] |
| if no_prefix: |
| cmd.append("--no-prefix") |
| print(f"\n=== {label} @ {steps} step CFG {cfg} ===") |
| subprocess.run(cmd, check=True) |
| dataset_vol.commit() |
|
|
|
|
| |
| |
| |
| |
| @app.function( |
| image=image, |
| volumes=VOLUMES, |
| timeout=1800, |
| cpu=2.0, |
| memory=16384, |
| ) |
| def merge_turbo_lora( |
| turbo_url: str = "", |
| phase1_lora: str = "/output/phase1/latest/adapter_model.safetensors", |
| out_name: str = "anima_phase1_plus_turbo.safetensors", |
| alpha_phase1: float = 1.0, |
| alpha_turbo: float = 1.0, |
| ): |
| """ |
| Phase 1 LoRA + 既存 Turbo LoRA を 1 ファイルに合成して /output/merged/ に保存。 |
| |
| 使い方: |
| # 1) Civitai から Turbo LoRA を一度だけ取得 (URL は Civitai のダウンロード直リン) |
| modal run modal_app.py::merge_turbo_lora \\ |
| --turbo-url "https://civitai.com/api/download/models/<id>" |
| |
| weight 演算: |
| out_lora = alpha_phase1 * phase1_lora + alpha_turbo * turbo_lora |
| shape/key が完全一致しないキーはスキップしてログ出力。 |
| """ |
| import os |
| import urllib.request |
| from safetensors.torch import load_file, save_file |
|
|
| os.makedirs("/output/merged", exist_ok=True) |
| turbo_path = "/models/loras/anima_turbo.safetensors" |
| os.makedirs("/models/loras", exist_ok=True) |
|
|
| if turbo_url: |
| print(f"[download] {turbo_url} -> {turbo_path}") |
| urllib.request.urlretrieve(turbo_url, turbo_path) |
| elif not os.path.exists(turbo_path): |
| raise SystemExit( |
| f"Turbo LoRA not found. Pass --turbo-url to download, " |
| f"or upload to {turbo_path} via `modal volume put anima-models ...`" |
| ) |
|
|
| print(f"[load] phase1={phase1_lora}") |
| print(f"[load] turbo ={turbo_path}") |
| sd_p1 = load_file(phase1_lora) |
| sd_tb = load_file(turbo_path) |
|
|
| merged, skipped, mismatched = {}, [], [] |
| for k, v in sd_p1.items(): |
| if k in sd_tb and sd_tb[k].shape == v.shape: |
| merged[k] = alpha_phase1 * v + alpha_turbo * sd_tb[k] |
| elif k in sd_tb: |
| mismatched.append((k, tuple(v.shape), tuple(sd_tb[k].shape))) |
| merged[k] = alpha_phase1 * v |
| else: |
| skipped.append(k) |
| merged[k] = alpha_phase1 * v |
| |
| for k, v in sd_tb.items(): |
| if k not in sd_p1: |
| merged[k] = alpha_turbo * v |
|
|
| out = f"/output/merged/{out_name}" |
| save_file(merged, out) |
| output_vol.commit() |
|
|
| print(f"[merge] phase1 keys={len(sd_p1)} turbo keys={len(sd_tb)} -> merged={len(merged)}") |
| if mismatched: |
| print(f"[merge] shape mismatch (kept phase1 only): {len(mismatched)} keys") |
| for k, s1, s2 in mismatched[:5]: |
| print(f" - {k}: phase1{s1} vs turbo{s2}") |
| if skipped: |
| print(f"[merge] phase1-only keys (no turbo): {len(skipped)}") |
| print(f"[done] {out}") |
|
|
|
|
| |
| |
| |
| @app.function(image=image, volumes=VOLUMES, timeout=300, cpu=2.0, memory=4096) |
| def inspect_lora_keys( |
| path1: str = "/models/loras/anima_turbo.safetensors", |
| path2: str = "/models/loras/phase_b_4step_lora.safetensors", |
| ): |
| """2 つの LoRA の key 命名を比較。Phase B output が壊れている件の debug 用。""" |
| from safetensors.torch import load_file |
| for label, path in (("turbo", path1), ("phase_b", path2)): |
| sd = load_file(path) |
| print(f"\n=== {label} ({path}) ===") |
| print(f" total keys: {len(sd)}") |
| for k in list(sd.keys())[:6]: |
| print(f" {k} shape={tuple(sd[k].shape)}") |
| |
| for pat in ("lora_down", "lora_up", "lora_A", "lora_B", "alpha"): |
| n = sum(1 for k in sd if pat in k) |
| print(f" contains '{pat}': {n} keys") |
|
|
|
|
| @app.function(image=image, volumes=VOLUMES, timeout=120, cpu=1.0, memory=2048) |
| def status(): |
| """Volume の中身とサイズを表示""" |
| import subprocess |
| for label, path in (("models", "/models"), ("dataset", "/dataset"), ("output", "/output")): |
| print(f"\n=== /{label} ===") |
| subprocess.run(["du", "-sh", path], check=False) |
| subprocess.run(["find", path, "-maxdepth", "3", "-type", "f", |
| "-printf", "%s\t%p\n"], check=False) |
|
|
|
|
| @app.function(image=image, volumes=VOLUMES, timeout=600, cpu=1.0, memory=2048) |
| def cleanup_checkpoints(keep_latest_n: int = 2, dry_run: bool = True): |
| """ |
| /output/phase1/ 配下の古いチェックポイントを削除して Volume 課金を抑える。 |
| diffusion-pipe は epoch ごとに保存するので、本番は最新 2 個残せば十分。 |
| |
| 例: modal run modal_app.py::cleanup_checkpoints --no-dry-run |
| """ |
| import os |
| import shutil |
|
|
| for root in ("/output/phase1", "/output/phase2_lcm", "/output/phase2_dmd2"): |
| if not os.path.isdir(root): |
| continue |
| |
| subs = sorted( |
| [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))], |
| key=lambda d: os.path.getmtime(os.path.join(root, d)), |
| reverse=True, |
| ) |
| to_keep = set(subs[:keep_latest_n] + ["latest"]) |
| for d in subs: |
| if d in to_keep: |
| continue |
| full = os.path.join(root, d) |
| size_mb = sum( |
| os.path.getsize(os.path.join(dp, f)) |
| for dp, _, fns in os.walk(full) for f in fns |
| ) / 1e6 |
| if dry_run: |
| print(f"[dry-run] would delete {full} ({size_mb:.0f} MB)") |
| else: |
| shutil.rmtree(full) |
| print(f"[deleted] {full} ({size_mb:.0f} MB)") |
| if not dry_run: |
| output_vol.commit() |
|
|
|
|
| @app.local_entrypoint() |
| def main(): |
| """`modal run modal_app.py` で叩いた時の既定動作: ヘルプ表示""" |
| print(__doc__) |
|
|
