#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""快速测试：3张图×花瓣发型×896分辨率，新提示词"填充遮罩区域的头发"。
预热1次+正式1次。
"""
import base64, json, os, time
from pathlib import Path
import requests

API = "http://127.0.0.1:8187/api/v1/debug/grow-timing"
TOKEN = "dev-shared-secret-2026"
PROMPT = "填充遮罩区域的头发"
OUT = Path("/home/ubuntu/hair/benchmark_out/bench6")
OUT.mkdir(parents=True, exist_ok=True)

IMGS = [
    ("asdf", "/home/ubuntu/hair/image/asdf.jpg"),
    ("qwer", "/home/ubuntu/hair/image/qwer.jpg"),
    ("girl5", "/home/ubuntu/hair/image/girl_img/girl5.jpg"),
]

def call(img_path, save_grown=None, timeout=300):
    data = {"hair_style": "2", "webui_steps": "15", "redraw_max_side": "896", "redraw_prompt": PROMPT}
    t0 = time.perf_counter()
    with open(img_path, "rb") as f:
        r = requests.post(API, headers={"X-Internal-Token": TOKEN},
            files={"image_file": (os.path.basename(img_path), f, "image/jpeg")},
            data=data, timeout=timeout)
    wall = time.perf_counter() - t0
    j = r.json()
    d = j["data"]; hs = d["per_hairstyle"][0]
    if save_grown and hs.get("grown_b64"):
        b = hs["grown_b64"].split(",")[1] if "," in hs["grown_b64"] else hs["grown_b64"]
        open(save_grown, "wb").write(base64.b64decode(b))
    return {"ok": hs.get("ok"), "total_ms": d["total_ms"], "comfy_ms": hs.get("comfyui_redraw_ms"),
            "grown_path": str(save_grown) if save_grown and hs.get("ok") else None}

results = []
for ilabel, ipath in IMGS:
    print(f"预热 {ilabel}...", flush=True)
    call(ipath)
    save = OUT / f"{ilabel}_flower_896.jpg"
    print(f"正式 {ilabel}...", flush=True)
    r = call(ipath, save_grown=save)
    r["img"] = ilabel
    print(f"  -> total={r['total_ms']}ms ok={r['ok']}", flush=True)
    results.append(r)

json.dump({"prompt": PROMPT, "results": results}, open(OUT/"results.json","w"), ensure_ascii=False, indent=2)
print(f"\n✓ 完成 {sum(1 for r in results if r['ok'])}/3", flush=True)