优化为4b 模型

hairline/redraw.py

@@ -0,0 +1,76 @@
+"""直接调 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) -> bytes:
+    """直接调 ComfyUI 重绘 — 替代 local_test /api/generate。
+
+    Args:
+        image_bytes: 人物图片字节（JPG/PNG）
+        mask_bytes: 遮罩图片字节（支持红/白/alpha 遮罩格式）
+        prompt: 提示词，None 用默认 "填充遮罩区域的头发，皮肤加一点磨皮"
+        timeout: ComfyUI 超时秒数
+
+    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)