xslandCursor
98b9108837
问题：接口2 与接口3/5 乱序调用时耗时抖动（最差 15~22s）。两个根因： 1. GPU 24G 常驻 21.4G，Flux-2(3.9G) 无法完全驻留显存，每次采样动态换页， 速度随空闲显存波动（2s~8s）； 2. ComfyUI 单队列 FIFO，接口2 排在接口3/5 批量任务后面。 改动： - hairline/comfyui.py: run() 新增 front 参数，/prompt 带 "front": true 插队到队列最前； redraw.py 透传；service.py 接口2 三处调用（女重绘 + 男有/无遮罩）传 front=True， 接口3/5 仍走普通队列。 - add_hair.json / 0716add-hair-api.json: 节点61 CLIPLoader device default→cpu。 qwen CLIP(4G) 不再占显存（文本条件缓存常年命中），ComfyUI 显存 8.8G→4.5G， Flux-2 完全驻留，采样稳定 ~3-5s。代价：换 prompt 后首次请求 CPU 编码 ~11s（一次性）。 - 提示词全局统一为「填充遮罩区域的头发，皮肤加一点磨皮，再加一点美颜」： app.py 4处默认值、service.py _REDRAW_PROMPT、redraw.py _DEFAULT_PROMPT、 4个工作流节点60内置文案、测试页(test_interface2/3/7/12/12_final)、local_test。 任何两个不同 prompt 交替提交都会打爆 CLIP 编码缓存（--cache-classic 只存最近一次）， 之前测试页旧文案与服务端不一致导致交替测试每次 +11s。 - app.py: 接口7 /api/v1/hair/grow-v2 下线（业务弃用；add_hair2.json 的 Klein-9b 会把常驻 Klein-4b 挤出显存）。保留 stub 返回 1007 明确报错，避免裸 404。 实测（1024 档）：接口2女 8.5~10s、接口2男 ~5s、接口3 ~7-10s，交替混跑无尖刺。 Co-authored-by: Cursor <cursoragent@cursor.com> (cherry-picked from ubuntu3090 e7b62f2；已适配 main 分支代码结构：main 无 _REDRAW_PROMPT/_REDRAW_MAX_SIDE 缩图逻辑，front=True 直接加在 _call_local_redraw / generate_grow_results 的调用点；另把 main 独有的 benchmark_*.py 里的 prompt 一并统一) Co-authored-by: Cursor <cursoragent@cursor.com>
79 lines
2.9 KiB
Python
79 lines
2.9 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) -> bytes:
|
||
"""直接调 ComfyUI 重绘 — 替代 local_test /api/generate。
|
||
|
||
Args:
|
||
image_bytes: 人物图片字节（JPG/PNG）
|
||
mask_bytes: 遮罩图片字节（支持红/白/alpha 遮罩格式）
|
||
prompt: 提示词，None 用默认 "填充遮罩区域的头发，皮肤加一点磨皮，再加一点美颜"
|
||
timeout: ComfyUI 超时秒数
|
||
front: True 时任务插到 ComfyUI 队列最前（接口2 时延敏感路径用）
|
||
|
||
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)
|