xsl
94cdfd6de5
代码改动: - 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 跟踪
81 lines
3.1 KiB
Python
81 lines
3.1 KiB
Python
"""直接调 ComfyUI 重绘 — 替代 local_test HTTP 服务。
|
||
|
||
将 local_test/app.py 的核心逻辑（遮罩处理 + ComfyUI 调用）提取为 Python 函数，
|
||
不再需要独立 Flask 服务。使用 0716add-hair-api.json 工作流（steps=4）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
import logging
|
||
import os
|
||
|
||
import numpy as np
|
||
from PIL import Image, ImageFilter
|
||
|
||
from . import comfyui
|
||
|
||
logger = logging.getLogger("hair.worker")
|
||
|
||
_DEFAULT_PROMPT = "填充遮罩区域的头发，皮肤加一点磨皮，再加一点美颜"
|
||
_REPO = os.path.dirname(os.path.dirname(__file__))
|
||
_REPAINT_WORKFLOW = os.path.join(_REPO, "0716add-hair-api.json")
|
||
|
||
|
||
def _process_mask_to_rgba(image_bytes: bytes, mask_bytes: bytes) -> bytes:
|
||
"""将分开的 image + mask 处理为 ComfyUI 用的 RGBA PNG bytes。
|
||
|
||
复制 local_test/app.py 的遮罩处理逻辑：
|
||
1. 加载 image 为 RGB
|
||
2. 加载 mask 为 RGBA，取所有通道 max 值（支持红/白/alpha 遮罩）
|
||
3. resize mask 到与 image 一致
|
||
4. 高斯模糊(radius=4) 柔化边缘
|
||
5. alpha = 255 - mask（绘制区=255 → alpha=0 → 重绘区）
|
||
6. 合成 RGBA PNG
|
||
"""
|
||
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
||
mask_img = Image.open(io.BytesIO(mask_bytes)).convert("RGBA")
|
||
mask_arr = np.array(mask_img)
|
||
mask_data = np.max(mask_arr, axis=2) # (H, W) uint8
|
||
|
||
mask_data_img = Image.fromarray(mask_data, mode="L")
|
||
if mask_data_img.size != image.size:
|
||
mask_data_img = mask_data_img.resize(image.size, Image.LANCZOS)
|
||
mask_data_img = mask_data_img.filter(ImageFilter.GaussianBlur(radius=4))
|
||
|
||
# ComfyUI LoadImage: mask = 1.0 - (alpha/255)
|
||
# alpha=0 -> mask=1.0 (inpaint), alpha=255 -> mask=0.0 (keep)
|
||
comfyui_alpha = Image.eval(mask_data_img, lambda x: 255 - x)
|
||
|
||
r, g, b = image.split()
|
||
rgba = Image.merge("RGBA", (r, g, b, comfyui_alpha))
|
||
|
||
buf = io.BytesIO()
|
||
rgba.save(buf, format="PNG")
|
||
return buf.getvalue()
|
||
|
||
|
||
def run_redraw(image_bytes: bytes, mask_bytes: bytes,
|
||
prompt: str | None = None, timeout: float = 300.0,
|
||
front: bool = False, unet_name: str | None = None) -> bytes:
|
||
"""直接调 ComfyUI 重绘 — 替代 local_test /api/generate。
|
||
|
||
Args:
|
||
image_bytes: 人物图片字节（JPG/PNG）
|
||
mask_bytes: 遮罩图片字节（支持红/白/alpha 遮罩格式）
|
||
prompt: 提示词，None 用默认 "填充遮罩区域的头发，皮肤加一点磨皮，再加一点美颜"
|
||
timeout: ComfyUI 超时秒数
|
||
front: True 时任务插到 ComfyUI 队列最前（接口2 时延敏感路径用）
|
||
unet_name: 非 None 时切换 Flux 模型（如 flux-2-klein-9b-Q5_K_M.gguf），None 用工作流默认
|
||
|
||
Returns:
|
||
重绘后的 PNG 图片字节
|
||
|
||
Raises:
|
||
RuntimeError: ComfyUI 执行失败
|
||
TimeoutError: ComfyUI 超时
|
||
"""
|
||
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,
|
||
unet_name=unet_name)
|