#!/usr/bin/env python3
"""Benchmark ComfyUI hair-inpaint workflow across model / dtype / steps."""
import io, time, sys
import requests
import numpy as np
from PIL import Image, ImageFilter
import app as A

COMFY = "http://127.0.0.1:8188"


def prep_and_upload():
    image = Image.open("用来重绘.jpg").convert("RGB")
    mask_img = Image.open("用来重绘.png").convert("RGBA")
    mask_data = np.max(np.array(mask_img), axis=2)
    m = Image.fromarray(mask_data, mode="L")
    if m.size != image.size:
        m = m.resize(image.size, Image.LANCZOS)
    m = m.filter(ImageFilter.GaussianBlur(radius=4))
    alpha = Image.eval(m, lambda x: 255 - x)
    r, g, b = image.split()
    rgba = Image.merge("RGBA", (r, g, b, alpha))
    buf = io.BytesIO(); rgba.save(buf, format="PNG"); buf.seek(0)
    up = requests.post(f"{COMFY}/upload/image",
                       files={"image": ("hair_input.png", buf, "image/png")}).json()
    return up["name"]


def run_once(fname, model, dtype, steps):
    wf = A.build_workflow(fname, "填充遮罩区域的头发")
    wf["16"]["inputs"]["unet_name"] = model
    wf["16"]["inputs"]["weight_dtype"] = dtype
    wf["1"]["inputs"]["steps"] = steps
    r = requests.post(f"{COMFY}/prompt", json={"prompt": wf}).json()
    if "prompt_id" not in r:
        raise RuntimeError(f"submit failed: {str(r)[:300]}")
    pid = r["prompt_id"]
    deadline = time.time() + 180
    while time.time() < deadline:
        time.sleep(0.1)
        h = requests.get(f"{COMFY}/history/{pid}").json()
        if pid not in h:
            continue
        st = h[pid].get("status", {})
        if st.get("status_str") == "error":
            for m in st.get("messages", []):
                if m[0] == "execution_error":
                    raise RuntimeError(str(m[1])[:300])
            raise RuntimeError("execution error")
        if "17" in h[pid].get("outputs", {}):
            ts = {mm[0]: mm[1].get("timestamp") for mm in st["messages"]}
            return (ts["execution_success"] - ts["execution_start"]) / 1000.0
    raise TimeoutError("run exceeded 180s")


CONFIGS = [
    ("flux2.0/flux-2-klein-9b-fp8.safetensors", "fp8_e4m3fn",      6, "9B fp8 (当前)"),
    ("flux2.0/flux-2-klein-9b-fp8.safetensors", "fp8_e4m3fn_fast", 6, "9B fp8-fast"),
    ("flux2.0/flux-2-klein-9b-fp8.safetensors", "fp8_e4m3fn_fast", 4, "9B fp8-fast s4"),
    ("flux-2-klein-4b-fp8.safetensors",         "fp8_e4m3fn",      6, "4B fp8"),
    ("flux-2-klein-4b-fp8.safetensors",         "fp8_e4m3fn_fast", 6, "4B fp8-fast"),
    ("flux-2-klein-4b-fp8.safetensors",         "fp8_e4m3fn_fast", 4, "4B fp8-fast s4"),
]

fname = prep_and_upload()
print("input uploaded:", fname)
print(f"{'配置':<22}{'warmup':>10}{'run1':>10}{'run2':>10}{'best':>10}")
for model, dtype, steps, label in CONFIGS:
    times = []
    for i in range(3):  # 1 warmup + 2 measured
        try:
            t = run_once(fname, model, dtype, steps)
        except Exception as e:
            t = float('nan'); print("ERR", label, e)
        times.append(t)
    best = min(times[1:])
    print(f"{label:<22}{times[0]:>9.2f}s{times[1]:>9.2f}s{times[2]:>9.2f}s{best:>9.2f}s")