xsl/hair
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(接口2): 支持动态切换Flux模型+分辨率 + 模型对比测试脚本

Browse Source 
代码改动:
- comfyui.py: run() 新增 unet_name 参数，提交前自动改写模型节点
  (.gguf→UnetLoaderGGUF, .safetensors→UNETLoader)，并按模型自动同步
  文本编码器(4b→qwen_3_4b, 9b→qwen_3_8b)，避免切换时维度不匹配
- redraw.py: run_redraw() 透传 unet_name
- service.py: generate_grow_results_swap/generate_grow_results 支持
  redraw_max_side(分辨率参数化) 和 unet_name 透传
- app.py: 接口2 新增 flux_model/redraw_max_side 两个 Form 参数(男女路径都加)
- test_interface2.html: 新增 Flux模型/压图长边 下拉选择器
- add_hair.json/0716add-hair-api.json: 工作流默认模型改为 9b

测试脚本:
- benchmark_matrix.py: 4模型×3分辨率×3图×3次 矩阵测试
- benchmark_hairstyle.py: 3图×5发型×10组合 发型对比测试
- benchmark_report.py/benchmark_hairstyle_report.py: HTML报告生成

清理:
- .gitignore: 排除 benchmark_out/、报告HTML、gateway.log、*.bak.*
- 移除 gateway.log 的 git 跟踪
This commit is contained in:
xsl
2026-07-25 16:16:13 +08:00
parent f509fe99b4
commit 94cdfd6de5
13 changed files with 657 additions and 362 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hairline/comfyui.py
+34 -1
View File
@@ -30,6 +30,19 @@ _REPO = os.path.dirname(os.path.dirname(__file__))
_INPUT_NODE = "26" # LoadImage：外部输入图(含 alpha 遮罩)
_SEED_NODE = "6" # RandomNoise
_PROMPT_NODE = "60" # JjkText：提示词
_UNET_NODE = "16" # UNETLoader / UnetLoaderGGUF：Flux 模型加载
_CLIP_NODE = "61" # CLIPLoader：qwen 文本编码器
# Flux 模型 → 配套文本编码器映射。切换 unet 时自动同步编码器，避免维度不匹配。
# 规则：4b 系列配 qwen_3_4b，9b 系列（fp8/GGUF）配 qwen_3_8b_fp8mixed。
def _clip_for_unet(unet_name: str) -> str | None:
"""根据 unet 文件名推断配套的文本编码器文件名；无法推断返回 None。"""
low = unet_name.lower()
if "4b" in low and "9b" not in low:
return "qwen_3_4b.safetensors"
if "9b" in low:
return "qwen_3_8b_fp8mixed.safetensors"
return None
_wf_cache: dict[str, dict] = {} # path → workflow JSON
_wf_output_node: dict[str, str] = {} # path → SaveImage 节点 ID
@@ -86,13 +99,18 @@ def _get_output_node(workflow_path: str | None = None) -> str:
def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT, prompt: str = None,
workflow_path: str | None = None, front: bool = False) -> bytes:
workflow_path: str | None = None, front: bool = False,
unet_name: str | None = None) -> bytes:
"""提交一次生发任务，返回输出 PNG 字节。失败抛异常。
prompt：非 None 时替换工作流节点60(JjkText)的文本；None 时用工作流内置默认提示词。
workflow_path：工作流 JSON 路径，None 则用默认 add_hair.json。
front：True 时任务插到 ComfyUI 队列最前（server 端 "front" 字段，队列号取负）。
接口2 对时延敏感用 True，避免排在接口3/5 的批量任务后面；其余接口保持 False。
unet_name：非 None 时改写工作流里的模型加载节点（节点16），动态切换 Flux 模型。
.safetensors → 保持 UNETLoader 节点类型不变，只替换 unet_name；
.gguf → 自动把节点类型改成 UnetLoaderGGUF（需装 ComfyUI-GGUF 插件）。
None 时用工作流内置默认模型。
"""
path = workflow_path or _WORKFLOW_DEFAULT
output_node = _get_output_node(path)
@@ -120,6 +138,21 @@ def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT, prompt: str = Non
wf[_SEED_NODE]["inputs"]["noise_seed"] = random.randint(0, 2**63 - 1)
if prompt is not None:
wf[_PROMPT_NODE]["inputs"]["text"] = prompt
if unet_name is not None:
node = wf.get(_UNET_NODE)
if node is not None:
# .gguf 需切换到 ComfyUI-GGUF 插件的 UnetLoaderGGUF 节点；
# .safetensors/.ckpt 保持原 UNETLoader 节点类型不变
if unet_name.lower().endswith(".gguf"):
node["class_type"] = "UnetLoaderGGUF"
else:
node["class_type"] = "UNETLoader"
node["inputs"]["unet_name"] = unet_name
# 同步切换配套文本编码器（4b→qwen_3_4b, 9b→qwen_3_8b），避免维度不匹配
clip_node = wf.get(_CLIP_NODE)
clip_name = _clip_for_unet(unet_name)
if clip_node is not None and clip_name is not None:
clip_node["inputs"]["clip_name"] = clip_name
# 诊断：落盘实际提交的工作流 + 输入图，便于和手动 ComfyUI 跑的对比
try:
hairline/redraw.py
+4 -2
View File
@@ -56,7 +56,7 @@ def _process_mask_to_rgba(image_bytes: bytes, mask_bytes: bytes) -> bytes:
def run_redraw(image_bytes: bytes, mask_bytes: bytes,
prompt: str | None = None, timeout: float = 300.0,
front: bool = False) -> bytes:
front: bool = False, unet_name: str | None = None) -> bytes:
"""直接调 ComfyUI 重绘 — 替代 local_test /api/generate。
Args:
@@ -65,6 +65,7 @@ def run_redraw(image_bytes: bytes, mask_bytes: bytes,
prompt: 提示词，None 用默认 "填充遮罩区域的头发，皮肤加一点磨皮，再加一点美颜"
timeout: ComfyUI 超时秒数
front: True 时任务插到 ComfyUI 队列最前（接口2 时延敏感路径用）
unet_name: 非 None 时切换 Flux 模型（如 flux-2-klein-9b-Q5_K_M.gguf），None 用工作流默认
Returns:
重绘后的 PNG 图片字节
@@ -75,4 +76,5 @@ def run_redraw(image_bytes: bytes, mask_bytes: bytes,
"""
rgba_png = _process_mask_to_rgba(image_bytes, mask_bytes)
return comfyui.run(rgba_png, timeout=timeout, prompt=prompt,
workflow_path=_REPAINT_WORKFLOW, front=front)
workflow_path=_REPAINT_WORKFLOW, front=front,
unet_name=unet_name)
hairline/service.py
+26 -13
View File
@@ -54,21 +54,26 @@ _REDRAW_PROMPT = os.getenv("REDRAW_PROMPT", "填充遮罩区域的头发，皮
# 默认压到 896 兜底（ComfyUI ~4s，女性总耗时 9~11s）；追画质可设 REDRAW_MAX_SIDE=1024。
_REDRAW_MAX_SIDE = int(os.getenv("REDRAW_MAX_SIDE", "896"))
def _call_local_redraw(image_png_bytes, mask_png_bytes, timeout=300.0):
def _call_local_redraw(image_png_bytes, mask_png_bytes, timeout=300.0,
max_side=None, unet_name=None):
"""直接调 ComfyUI 重绘（替代原 local_test HTTP 服务）。
传 final 图 + 纯红遮罩 PNG，返回重绘后的 PNG bytes。
失败抛异常（调用方负责 try/except 跳过）。
max_side：送 ComfyUI 前长边压到多少像素，None 用全局默认 _REDRAW_MAX_SIDE。
unet_name：非 None 时切换 Flux 模型，None 用工作流内置默认。
"""
from .redraw import run_redraw
eff_side = _REDRAW_MAX_SIDE if max_side is None else max_side
img = cv2.imdecode(np.frombuffer(image_png_bytes, np.uint8), cv2.IMREAD_UNCHANGED)
scale = 1.0
orig_w = orig_h = 0
if img is not None:
orig_h, orig_w = img.shape[:2]
m = max(orig_h, orig_w)
if _REDRAW_MAX_SIDE > 0 and m > _REDRAW_MAX_SIDE:
scale = _REDRAW_MAX_SIDE / float(m)
if eff_side > 0 and m > eff_side:
scale = eff_side / float(m)
nw, nh = max(1, round(orig_w * scale)), max(1, round(orig_h * scale))
msk = cv2.imdecode(np.frombuffer(mask_png_bytes, np.uint8), cv2.IMREAD_UNCHANGED)
img_s = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_AREA)
@@ -76,10 +81,10 @@ def _call_local_redraw(image_png_bytes, mask_png_bytes, timeout=300.0):
image_png_bytes = cv2.imencode(".png", img_s)[1].tobytes()
mask_png_bytes = cv2.imencode(".png", msk_s)[1].tobytes()
logger.info("接口2女 缩图送 Comfy: %dx%d → %dx%d (max_side=%d)",
orig_w, orig_h, nw, nh, _REDRAW_MAX_SIDE)
orig_w, orig_h, nw, nh, eff_side)
# front=True：接口2 时延敏感，插到 ComfyUI 队列最前，避免排在接口3/5 的批量任务后面
out = run_redraw(image_png_bytes, mask_png_bytes, timeout=timeout,
prompt=_REDRAW_PROMPT, front=True)
prompt=_REDRAW_PROMPT, front=True, unet_name=unet_name)
if scale < 1.0 and out:
out = _upscale_png_to(out, orig_w, orig_h)
return out
@@ -200,7 +205,8 @@ def generate_previews(image_bgr: np.ndarray, gender: str):
def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = True,
prompt: str = None, hair_styles: list[int] | None = None,
workflow_path: str | None = None):
workflow_path: str | None = None,
unet_name: str | None = None):
"""指定发际线类型：发际线透明叠图(白线 RGBA) + 生发图(ComfyUI)。
hair_styles（1-indexed 列表）：指定生成哪几张发际线（按贴图排序）。female: 1..5，male: 1..4。
@@ -236,7 +242,8 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
compose_comfy_rgba(img_s, msk_s).save(buf, format="PNG", compress_level=1)
# front=True：接口2 时延敏感，插到 ComfyUI 队列最前
shared_grown = comfyui.run(buf.getvalue(), prompt=prompt,
workflow_path=workflow_path, front=True)
workflow_path=workflow_path, front=True,
unet_name=unet_name)
if gsc < 1.0 and shared_grown:
shared_grown = _upscale_png_to(shared_grown, w, h)
except Exception as e: # noqa: BLE001
@@ -261,7 +268,8 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
compose_comfy_rgba(m_s, msk_s).save(buf, format="PNG", compress_level=1)
# front=True：接口2 时延敏感，插到 ComfyUI 队列最前
grown_png = comfyui.run(buf.getvalue(), prompt=prompt,
workflow_path=workflow_path, front=True)
workflow_path=workflow_path, front=True,
unet_name=unet_name)
if gsc < 1.0 and grown_png:
grown_png = _upscale_png_to(grown_png, w, h)
except Exception as e: # noqa: BLE001 单张失败不拖垮整请求
@@ -273,7 +281,9 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
def generate_grow_results_swap(image_bgr: np.ndarray, hair_styles: list[int] | None,
redraw_defaults: dict):
redraw_defaults: dict,
redraw_max_side: int | None = None,
unet_name: str | None = None):
"""接口2 女性专用：发际线透明叠图（同 generate_grow_results）+ 换发型重绘图。
grown 图来源（新流程）：对每个选中发型把 female key 映射到 change_hair 的 chang_* hair_id，
@@ -307,16 +317,17 @@ def generate_grow_results_swap(image_bgr: np.ndarray, hair_styles: list[int] | N
# 重绘管线（swapHair + ComfyUI）统一降分辨率：真实照片 swap(SD WebUI)~5s、blend、ComfyUI
# 均随分辨率线性下降。overlay 预览仍用全分辨率；grown_png 最后放大回原尺寸。
eff_side = _REDRAW_MAX_SIDE if redraw_max_side is None else redraw_max_side
redraw_img = image_bgr
hair_mask_redraw = hair_mask_reuse
if _REDRAW_MAX_SIDE > 0 and max(h, w) > _REDRAW_MAX_SIDE:
redraw_img, _rs = _downscale_max_side(image_bgr, _REDRAW_MAX_SIDE)
if eff_side > 0 and max(h, w) > eff_side:
redraw_img, _rs = _downscale_max_side(image_bgr, eff_side)
_nh, _nw = redraw_img.shape[:2]
if hair_mask_redraw is not None:
hair_mask_redraw = cv2.resize(hair_mask_reuse.astype(np.uint8), (_nw, _nh),
interpolation=cv2.INTER_NEAREST).astype(bool)
logger.info("接口2女 管线降分辨率: %dx%d → %dx%d (max_side=%d)",
w, h, _nw, _nh, _REDRAW_MAX_SIDE)
w, h, _nw, _nh, eff_side)
for order, (key, white_path) in items:
white = load_texture_rgba(white_path)
@@ -349,7 +360,9 @@ def generate_grow_results_swap(image_bgr: np.ndarray, hair_styles: list[int] | N
mask_bytes = base64.b64decode(mask_b64)
# 后端直接调 ComfyUI 重绘，返回重绘后的 PNG
_tr0 = _t.perf_counter()
grown_png = _call_local_redraw(final_bytes, mask_bytes)
grown_png = _call_local_redraw(final_bytes, mask_bytes,
max_side=redraw_max_side,
unet_name=unet_name)
_tr1 = _t.perf_counter()
_tm = data.get("timings_ms") or {}
logger.info("接口2女 分段计时 type=%s: swapHair管线=%.2fs (mask=%dms swap=%dms blend=%dms), ComfyUI重绘=%.2fs",