优化为4b 模型

2026-07-18 15:30:31 +08:00
parent 659c037270
commit b58cd4c441
289 changed files with 9367 additions and 558 deletions
@@ -18,7 +18,7 @@ import uuid

 import httpx

-COMFYUI_URL = os.getenv("COMFYUI_URL", "http://10.60.74.221:8188").rstrip("/")
+COMFYUI_URL = os.getenv("COMFYUI_URL", "http://127.0.0.1:8188").rstrip("/")
 _WORKFLOW_DEFAULT = os.getenv(
    "ADD_HAIR_WORKFLOW",
    os.path.join(os.path.dirname(os.path.dirname(__file__)), "add_hair.json"),
@@ -145,7 +145,7 @@ def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT, prompt: str = Non
                outputs = entry.get("outputs")
                if outputs and output_node in outputs:
                    break
-            time.sleep(1.0)
+            time.sleep(0.2)
        if not outputs or output_node not in outputs:
            raise TimeoutError(f"ComfyUI 出图超时({timeout}s) prompt_id={prompt_id}")

@@ -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)
@@ -40,27 +40,14 @@ _REPO = os.path.dirname(os.path.dirname(__file__))
 _TEXTURE_DIR = os.path.join(_REPO, "hairline_texture")
 _BLACK_TEXTURE_DIR = os.path.join(_REPO, "hairline_texture_black")

-# 外部重绘服务（local_test，0716add-hair.json 工作流）。接口2 female 用它替代原 Flux-2 重绘。
-# 重绘服务已独立到远程机器；可用 HAIR_LOCAL_REDRAW_URL 覆盖。
-# gpu_worker 同内网走 10.60.74.221，外网（本地开发机）需走公网 117.50.183.232。
-_LOCAL_REDRAW_URL = os.getenv("HAIR_LOCAL_REDRAW_URL", "http://10.60.74.221:8899").rstrip("/")
-
-
 def _call_local_redraw(image_png_bytes, mask_png_bytes, timeout=300.0):
-    """调外部重绘服务（local_test /api/generate）：传 final 图 + 纯红遮罩 PNG，
-    返回重绘后的 PNG bytes。失败抛异常（调用方负责 try/except 跳过）。
+    """直接调 ComfyUI 重绘（替代原 local_test HTTP 服务）。
+
+    传 final 图 + 纯红遮罩 PNG，返回重绘后的 PNG bytes。
+    失败抛异常（调用方负责 try/except 跳过）。
    """
-    import requests
-    files = {
-        "image": ("final.jpg", image_png_bytes, "image/jpeg"),
-        "mask": ("mask.png", mask_png_bytes, "image/png"),
-    }
-    resp = requests.post(_LOCAL_REDRAW_URL + "/api/generate", files=files, timeout=timeout)
-    resp.raise_for_status()
-    if "image/" not in resp.headers.get("Content-Type", ""):
-        # 服务返回了 JSON 错误
-        raise RuntimeError(f"local_test 返回非图片: {resp.text[:200]}")
-    return resp.content
+    from .redraw import run_redraw
+    return run_redraw(image_png_bytes, mask_png_bytes, timeout=timeout)

 # 发际线贴图档位：middle=默认(hairline_texture/)，high/low 各自独立文件夹。
 _TEXTURE_DIRS = {
@@ -69,9 +56,8 @@ _TEXTURE_DIRS = {
    "low": os.path.join(_REPO, "hairline_texture_low"),
 }

-# ⚠️ 本 worker 是 RTX 5090(sm_120)，torch 2.2.2(cu121) 只编到 sm_90，CUDA 跑算子会报
-#    "no kernel image"。SegFormer 默认走 CPU（~2.5s/张）。换 torch cu128 后可设 SEG_DEVICE=cuda。
-_SEG_DEVICE = os.getenv("SEG_DEVICE", "cpu")
+# torch 2.7.1+cu128 已支持 RTX 5090 (sm_120)，SegFormer 走 GPU（~0.05s/张）
+_SEG_DEVICE = os.getenv("SEG_DEVICE", "cuda")

 _landmarker = None
 _parser = None
@@ -241,20 +227,24 @@ def generate_grow_results_swap(image_bgr: np.ndarray, hair_styles: list[int] | N

    grown 图来源（新流程）：对每个选中发型把 female key 映射到 change_hair 的 chang_* hair_id，
    调 face_analysis.hairline_grow.generate_hairline_redraw（= 接口12 final 管线，参数用
-    redraw_defaults）拿到 ④ final（接缝融合基底）+ ⑤-② 纯红遮罩 PNG，再**后端调外部重绘服务
-    local_test**（0716add-hair.json 工作流）完成发际线带重绘，重绘结果作为生发图。
+    redraw_defaults）拿到 ④ final（接缝融合基底）+ ⑤-② 纯红遮罩 PNG，再**后端直接调
+    ComfyUI**（0716add-hair-api.json 工作流）完成发际线带重绘，重绘结果作为生发图。

    overlay 仍是发际线曲线透明层（与 generate_grow_results 完全一致）。
    Returns: list[dict] {"hairline_type","order","overlay","grown_png"(jpg bytes 或 None)}；
             无人脸返回 None。单个发型换发型/重绘失败时 grown_png=None，不抛异常。
    """
    from face_analysis.hairline_grow import generate_hairline_redraw, NoFaceError
+    from face_analysis.head_mask import SEGFORMER_HAIR

    ctx = extract_context(image_bgr)
    if ctx is None:
        return None
    uv, ext_faces = load_ext_mesh()

+    # 复用 extract_context 已算好的 SegFormer parse_map，避免 generate_hairline_redraw 内部重复分割
+    hair_mask_reuse = (ctx["parse_map"] == SEGFORMER_HAIR)
+
    textures = get_texture_map()["female"]  # [(key, path), ...] 已排序
    if hair_styles is not None:
        items = [(s, textures[s - 1]) for s in hair_styles]
@@ -273,7 +263,7 @@ def generate_grow_results_swap(image_bgr: np.ndarray, hair_styles: list[int] | N
            logger.warning("接口2 换发型：female key=%s 无对应 chang_id，跳过生发图", key)
        else:
            try:
-                data = generate_hairline_redraw(image_bgr, chang_id, **redraw_defaults)
+                data = generate_hairline_redraw(image_bgr, chang_id, hair_mask=hair_mask_reuse, **redraw_defaults)
                steps = data.get("steps") or {}
                # ④ final（接缝融合基底）+ ⑤-② 纯红遮罩 PNG
                final_b64 = steps.get("final_base64") or ""
@@ -289,7 +279,7 @@ def generate_grow_results_swap(image_bgr: np.ndarray, hair_styles: list[int] | N
                        mask_b64 = mask_b64.split(",", 1)[1]
                    final_bytes = base64.b64decode(final_b64)
                    mask_bytes = base64.b64decode(mask_b64)
-                    # 后端调外部重绘服务（local_test），返回重绘后的 PNG
+                    # 后端直接调 ComfyUI 重绘，返回重绘后的 PNG
                    grown_png = _call_local_redraw(final_bytes, mask_bytes)
                    if grown_png is None:
                        logger.warning("接口2 换发型：type=%s 重绘结果为空", key)