save code

local_test/bench.py

@@ -0,0 +1,77 @@
+#!/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")