Files
hair/local_test/bench.py
T
xsl b61ea6f33b feat(接口5): 生发机制对齐接口2 + 统一ComfyUI重绘提示词
接口5改造(生发机制与接口2一致,按性别分流):
- generate_hairline_pngs: 生发图来源从 _grow_from_texture(局部inpaint)
  改为按性别分流——female走 generate_grow_results_swap(swapHair+Flux整帧重绘),
  male走 generate_grow_results(ComfyUI add_hair)
- 新增参数 redraw_max_side/unet_name/v2_defaults(female路径透传)
- 接口5 handler 加 flux_model/redraw_max_side Form参数
- 保留接口5独有输出: 3档叠图(middle/high/low)/中心点/face_measure
- 已验证: female日志出现"接口2女 管线降分辨率max_side=640"+swap+Flux;
  male走add_hair.json; generate_grow_image=false正确跳过生发

统一ComfyUI重绘提示词:
- "填充遮罩区域的头发,皮肤加一点磨皮,再加一点美颜" → "填充遮罩区域的头发"
- 覆盖: _REDRAW_PROMPT/_DEFAULT_PROMPT常量 + app.py各接口Form默认 +
  工作流JSON节点60(add_hair/0716add-hair-api/hair_repaint) + 全部测试页输入框
2026-07-27 23:36:15 +08:00

78 lines
3.1 KiB
Python

#!/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")