优化为4b 模型

hairline/service.py

@@ -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)