# -*- coding: utf-8 -*-
"""预缓存：VAE latent + Text Encoder 输出（冒烟测试数据集）"""
import os, subprocess, sys

os.environ["PYTHONUTF8"] = "1"
os.environ["HF_HUB_OFFLINE"] = "1"

VENV = r"M:\AI\sd\musubi-tuner\.venv\Scripts"
PY = os.path.join(VENV, "python.exe")
MUSUBI = r"M:\AI\sd\musubi-tuner"          # musubi-tuner 根目录（cwd 必须在此，musubi_tuner 包才能导入）
SRC = os.path.join(MUSUBI, "src", "musubi_tuner")
PROJ = r"D:\F\NewI\opencode\daily-workspace\projects\脸部LoRA训练-Qwen-Image"
MODELS = r"M:\AI\sd\models\qwen-edit-2511"

CFG = r"D:\F\NewI\opencode\daily-workspace\projects\脸部LoRA训练-Qwen-Image\config\dataset_active.toml"
VAE = os.path.join(MODELS, "diffusion_pytorch_model.safetensors")
TE = os.path.join(MODELS, "text_encoder", "model-00001-of-00004.safetensors")

steps = [
    ("VAE latent 缓存", [PY, os.path.join(SRC, "qwen_image_cache_latents.py"),
        "--dataset_config", CFG, "--vae", VAE, "--device", "cuda", "--skip_existing",
        "--model_version", "edit-2511"]),
    ("TextEncoder 缓存", [PY, os.path.join(SRC, "qwen_image_cache_text_encoder_outputs.py"),
        "--dataset_config", CFG, "--text_encoder", TE, "--fp8_vl", "--device", "cuda", "--skip_existing",
        "--model_version", "edit-2511"]),
]

for name, cmd in steps:
    print(f"\n===== {name} =====")
    print(" ".join(cmd))
    # cwd=musubi-tuner 根目录：qwen_image_cache_*.py 内部 from musubi_tuner.dataset import config_utils 才能找到包
    r = subprocess.run(cmd, cwd=MUSUBI)
    if r.returncode != 0:
        print(f"[FAIL] {name} exit={r.returncode}")
        sys.exit(r.returncode)
    print(f"[OK] {name}")

print("\n全部缓存完成")