save code

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:
-    """提交一次生发任务，返回输出 PNG 字节。失败抛异常。"""
+def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT, prompt: str = None) -> bytes:
+    """提交一次生发任务，返回输出 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})