save code

app.py

@@ -504,6 +504,7 @@ async def hair_grow(
     gender: Optional[str] = Form(default=None, description="性别 male/female（必填）"),
     beauty_enabled: bool = Form(default=False, description="是否开启美颜（本期不生效）"),
     use_mask: bool = Form(default=True, description="是否启用 inpaint 遮罩（测试对比用）。false 时用干净原图生成（空遮罩，不烧模板线）"),
+    prompt: str = Form(default="补充遮罩区域的头发", description="ComfyUI 提示词，会替换工作流节点60的文本"),
 ):
     # 1. gender 必填校验（非法/缺失 → 1004）
     if gender not in ("male", "female"):
@@ -530,7 +531,7 @@ async def hair_grow(
         from hairline.service import generate_grow_results
 
         # 预览 + 生发(ComfyUI) 都是阻塞且较慢，放线程池避免卡住事件循环
-        items = await run_in_threadpool(generate_grow_results, image, gender, use_mask)
+        items = await run_in_threadpool(generate_grow_results, image, gender, use_mask, prompt)
         if items is None:
             return err(1001, "无法识别人像")
 
@@ -596,6 +597,7 @@ async def hair_grow_b(
     marked_image_url: Optional[str] = Form(default=None, description="划线图片 URL"),
     marked_image_base64: Optional[str] = Form(default=None, description="划线图片 base64"),
     use_mask: bool = Form(default=True, description="是否画发际线（测试对比用）。false 时跳过划线检测、直接送划线图"),
+    prompt: str = Form(default="补充遮罩区域的头发", description="ComfyUI 提示词，会替换工作流节点60的文本"),
 ):
     # 划线图三选一取图（只需这一张）
     marked_raw, e = await resolve_image_bytes(marked_image_file, marked_image_url, marked_image_base64)
@@ -617,7 +619,7 @@ async def hair_grow_b(
         from fastapi.concurrency import run_in_threadpool
         from hairline.service import generate_grow_b
 
-        res = await run_in_threadpool(generate_grow_b, marked, use_mask)
+        res = await run_in_threadpool(generate_grow_b, marked, use_mask, prompt)
         if res["status"] == "no_face":
             return err(1001, "无法识别人像")
         if res["status"] == "no_line":

hairline/comfyui.py

@@ -26,6 +26,7 @@ _REPO = os.path.dirname(os.path.dirname(__file__))
 _INPUT_NODE = "26"     # LoadImage：外部输入图(含 alpha 遮罩)
 _SEED_NODE = "6"       # RandomNoise
 _OUTPUT_NODE = "17"    # SaveImage
+_PROMPT_NODE = "60"    # JjkText：提示词
 
 _workflow = None
 
@@ -63,8 +64,11 @@ def _load_workflow() -> dict:
     return _workflow
 
 
-def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT) -> bytes:
+def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT, prompt: str = None) -> bytes:
-    """提交一次生发任务，返回输出 PNG 字节。失败抛异常。"""
+    """提交一次生发任务，返回输出 PNG 字节。失败抛异常。
+
+    prompt：非 None 时替换工作流节点60(JjkText)的文本；None 时用工作流内置默认提示词。
+    """
     client_id = uuid.uuid4().hex
     with httpx.Client(base_url=COMFYUI_URL, timeout=30.0, auth=_comfy_auth()) as cli:
         # 1. 上传输入图（含 alpha 遮罩）到 ComfyUI input 目录
@@ -79,6 +83,8 @@ def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT) -> bytes:
         wf = copy.deepcopy(_load_workflow())
         wf[_INPUT_NODE]["inputs"]["image"] = name
         wf[_SEED_NODE]["inputs"]["noise_seed"] = random.randint(0, 2**63 - 1)
+        if prompt is not None:
+            wf[_PROMPT_NODE]["inputs"]["text"] = prompt
 
         # 3. 提交
         r = cli.post("/prompt", json={"prompt": wf, "client_id": client_id})

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"}