# -*- coding: utf-8 -*-
"""
LoRA checkpoint 验证工具
========================
加载 edit-2511 底模 + 训练好的 LoRA，用多组提示词生成对比图，
验证 LoRA 是否生效、效果如何。

用法:
  python validate_lora.py                    # 验证最新 checkpoint
  python validate_lora.py --ckpt <路径>      # 验证指定 checkpoint
  python validate_lora.py --all              # 验证 output/checkpoints 下所有
"""
import glob
import os
import subprocess
import sys

sys.stdout.reconfigure(encoding="utf-8", errors="replace")

VENV = r"D:\AI\sd\musubi-tuner\.venv\Scripts"
PY = os.path.join(VENV, "python.exe")
SRC = r"D:\AI\sd\musubi-tuner\src\musubi_tuner"
MODELS = r"D:\AI\sd\models\qwen-edit-2511"
PROJ = r"D:\F\NewI\opencode\daily-workspace\projects\脸部LoRA训练-Qwen-Image"

TRIGGER = "lm_face_v1"
CKPT_DIR = os.path.join(PROJ, "output", "checkpoints")
OUT_DIR = os.path.join(PROJ, "output", "验证")

# 验证提示词：同一触发词 + 不同场景，看脸是否稳定一致
PROMPTS = [
    ("证件照", f"{TRIGGER}, professional headshot, business attire, neutral expression, studio lighting, plain white background"),
    ("日常", f"{TRIGGER}, natural lifestyle portrait, by a window with soft diffused sunlight, warm tones"),
    ("户外全身", f"{TRIGGER}, full body shot, standing in a park, natural daylight, casual clothing"),
    ("半侧脸", f"{TRIGGER}, three-quarter view portrait, soft golden hour light, shallow depth of field"),
]

BASE_CMD = [
    PY, os.path.join(SRC, "qwen_image_generate_image.py"),
    "--dit", os.path.join(MODELS, "transformer", "diffusion_pytorch_model-00001-of-00005.safetensors"),
    "--vae", os.path.join(MODELS, "diffusion_pytorch_model.safetensors"),
    "--text_encoder", os.path.join(MODELS, "text_encoder", "model-00001-of-00004.safetensors"),
    "--model_version", "edit-2511",
    "--fp8", "--fp8_scaled", "--blocks_to_swap", "24",
    "--infer_steps", "25",
    "--image_size", "1024", "1024",
    "--seed", "42",
]


def pick_ckpts(args):
    if args.ckpt:
        return [args.ckpt]
    if args.all and os.path.isdir(CKPT_DIR):
        return sorted(glob.glob(os.path.join(CKPT_DIR, "*.safetensors")))
    ckpts = sorted(glob.glob(os.path.join(CKPT_DIR, "*.safetensors")))
    return [ckpts[-1]] if ckpts else []


def main():
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default=None, help="指定 checkpoint 路径")
    ap.add_argument("--all", action="store_true", help="验证所有 checkpoint")
    args = ap.parse_args()

    ckpts = pick_ckpts(args)
    if not ckpts:
        print(f"[WARN] 没有找到 checkpoint: {CKPT_DIR}")
        sys.exit(1)

    os.makedirs(OUT_DIR, exist_ok=True)
    print(f"找到 {len(ckpts)} 个 checkpoint，开始验证（每个 4 张图）...")

    for ckpt in ckpts:
        name = os.path.splitext(os.path.basename(ckpt))[0]
        print(f"\n===== {name} =====")
        for tag, prompt in PROMPTS:
            out = os.path.join(OUT_DIR, f"{name}_{tag}.png")
            cmd = BASE_CMD + ["--lora_weight", ckpt, "--lora_multiplier", "0.8",
                              "--prompt", prompt, "--save_path", out]
            print(f"  生成 [{tag}] ...")
            r = subprocess.run(cmd)
            if r.returncode == 0:
                print(f"    OK -> {out}")
            else:
                print(f"    FAIL exit={r.returncode}")

    print(f"\n验证完成，图片在 {OUT_DIR}")


if __name__ == "__main__":
    main()