save code

hairline/service.py

@@ -128,12 +128,13 @@ def generate_previews(image_bgr: np.ndarray, gender: str):
     return results
 
 
-def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = True):
+def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = True, prompt: str = None):
     """该性别全部发际线：预览图(白线) + 生发图(ComfyUI)。同步、串行。
 
     use_mask（默认 True）：是否启用 inpaint 遮罩，用于测试对比（同接口3）。
       False 时用**干净原图 + 空遮罩**送 ComfyUI（不烧黑色模板线），与模板无关，
       故只跑一次 ComfyUI、N 项复用同一张生发图；预览图（白线）仍按各模板生成。
+    prompt（默认 None）：ComfyUI 提示词，非 None 时替换工作流节点60文本。
     Returns: list[dict] {"hairline_type","order","image_bgr"(预览), "grown_png"(bytes 或 None)}。
     无人脸返回 None。某张 ComfyUI 失败时该项 grown_png=None，不影响其余。
     """
@@ -151,7 +152,7 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
             h, w = image_bgr.shape[:2]
             buf = io.BytesIO()
             compose_comfy_rgba(image_bgr, np.zeros((h, w), np.uint8)).save(buf, format="PNG")
-            shared_grown = comfyui.run(buf.getvalue())
+            shared_grown = comfyui.run(buf.getvalue(), prompt=prompt)
         except Exception as e:  # noqa: BLE001
             logger.warning("接口2 生发图失败(无遮罩)：%s", e)
 
@@ -170,7 +171,7 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
                     image_bgr, ctx["landmarks"], ctx["parse_map"], ctx["points"], black)
                 buf = io.BytesIO()
                 compose_comfy_rgba(marked, mask).save(buf, format="PNG")
-                grown_png = comfyui.run(buf.getvalue())
+                grown_png = comfyui.run(buf.getvalue(), prompt=prompt)
             except Exception as e:  # noqa: BLE001  单张失败不拖垮整请求
                 logger.warning("接口2 生发图失败 type=%s：%s", key, e)
 
@@ -211,7 +212,7 @@ def generate_hairline_pngs(image_bgr: np.ndarray, gender: str):
     return {"images": images, "best_center": best_center}
 
 
-def generate_grow_b(marked_bgr: np.ndarray, use_mask: bool = True):
+def generate_grow_b(marked_bgr: np.ndarray, use_mask: bool = True, prompt: str = None):
     """接口3：检测医生手绘发际线 → 遮罩 → 送 ComfyUI 生发（仅需划线图一张）。
 
     检测路径只用来**建遮罩**；ComfyUI 输入图用 **marked 原图**（含医生手绘线，
@@ -241,7 +242,7 @@ def generate_grow_b(marked_bgr: np.ndarray, use_mask: bool = True):
 
     buf = io.BytesIO()
     compose_comfy_rgba(marked_bgr, mask).save(buf, format="PNG")   # marked 原图 + 遮罩
-    grown_png = comfyui.run(buf.getvalue())
+    grown_png = comfyui.run(buf.getvalue(), prompt=prompt)
     return {"grown_png": grown_png, "status": "ok"}