"""ComfyUI 客户端：用 add_hair.json / add_hair2.json 工作流跑生发图（Flux-2 inpaint）。 worker 不跑 Flux，只把「划线图 + 遮罩」的 RGBA 上传到远端 ComfyUI （默认 http://10.60.74.221:8188，可用环境变量 COMFYUI_URL 覆盖）， 替换工作流节点 26 的输入图、随机 seed，提交 /prompt，轮询 /history，取回 /view 输出。 ComfyUI 若开启了 HTTP Basic Auth（user `admin` + 密码），所有请求都带凭据。 支持多工作流：run() 可通过 workflow_path 指定不同工作流 JSON，自动检测 SaveImage 输出节点。 """ from __future__ import annotations import copy import json import os import random import time import uuid import httpx 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"), ) COMFY_TIMEOUT = float(os.getenv("COMFYUI_TIMEOUT", "600")) # 单张出图最长等待(秒) _REPO = os.path.dirname(os.path.dirname(__file__)) _INPUT_NODE = "26" # LoadImage：外部输入图(含 alpha 遮罩) _SEED_NODE = "6" # RandomNoise _PROMPT_NODE = "60" # JjkText：提示词 _wf_cache: dict[str, dict] = {} # path → workflow JSON _wf_output_node: dict[str, str] = {} # path → SaveImage 节点 ID def _comfy_auth(): """ComfyUI Basic Auth 凭据 (user, password)。 user：环境变量 COMFYUI_USER，默认 admin。 password：环境变量 COMFYUI_PASSWORD → worker_config.json.comfyui_password → password.txt。 无密码则返回 None（不带鉴权，兼容未开启 auth 的实例）。 """ user = os.getenv("COMFYUI_USER", "admin") pw = os.getenv("COMFYUI_PASSWORD") if not pw: cfg = os.path.join(_REPO, "worker_config.json") if os.path.isfile(cfg): try: with open(cfg, encoding="utf-8") as f: pw = json.load(f).get("comfyui_password") except Exception: # noqa: BLE001 pw = None if not pw: pwfile = os.path.join(_REPO, "password.txt") if os.path.isfile(pwfile): with open(pwfile, encoding="utf-8") as f: pw = f.read().strip() return (user, pw) if pw else None def _load_workflow(workflow_path: str | None = None) -> dict: """加载工作流 JSON（按路径缓存）。自动检测 SaveImage 节点 ID。""" path = workflow_path or _WORKFLOW_DEFAULT if path not in _wf_cache: with open(path, encoding="utf-8") as f: wf = json.load(f) _wf_cache[path] = wf # 自动检测 SaveImage 输出节点 for node_id, node in wf.items(): if node.get("class_type") == "SaveImage": _wf_output_node[path] = node_id break else: raise ValueError(f"工作流 {path} 中未找到 SaveImage 节点") return _wf_cache[path] def _get_output_node(workflow_path: str | None = None) -> str: """返回指定工作流的 SaveImage 节点 ID。""" path = workflow_path or _WORKFLOW_DEFAULT if path not in _wf_output_node: _load_workflow(path) # 触发检测 return _wf_output_node[path] def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT, prompt: str = None, workflow_path: str | None = None, front: bool = False) -> bytes: """提交一次生发任务，返回输出 PNG 字节。失败抛异常。 prompt：非 None 时替换工作流节点60(JjkText)的文本；None 时用工作流内置默认提示词。 workflow_path：工作流 JSON 路径，None 则用默认 add_hair.json。 front：True 时任务插到 ComfyUI 队列最前（server 端 "front" 字段，队列号取负）。 接口2 对时延敏感用 True，避免排在接口3/5 的批量任务后面；其余接口保持 False。 """ path = workflow_path or _WORKFLOW_DEFAULT output_node = _get_output_node(path) client_id = uuid.uuid4().hex with httpx.Client(base_url=COMFYUI_URL, timeout=30.0, auth=_comfy_auth()) as cli: # 1. 上传输入图（含 alpha 遮罩）到 ComfyUI input 目录 fname = f"hair_{client_id}.png" r = cli.post("/upload/image", files={"image": (fname, rgba_png_bytes, "image/png")}, data={"overwrite": "true", "type": "input"}) r.raise_for_status() up = r.json() name = (up.get("subfolder") + "/" if up.get("subfolder") else "") + up["name"] # 2. 改工作流：节点26 输入图 + 随机 seed wf = copy.deepcopy(_load_workflow(path)) 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 # 诊断：落盘实际提交的工作流 + 输入图，便于和手动 ComfyUI 跑的对比 try: import os as _os _diag = _os.path.join(_os.path.dirname(_os.path.dirname(_os.path.abspath(__file__))), "log", "comfyui_last_submit") _os.makedirs(_diag, exist_ok=True) with open(_os.path.join(_diag, "workflow.json"), "w", encoding="utf-8") as _f: json.dump(wf, _f, ensure_ascii=False, indent=2) with open(_os.path.join(_diag, "input.png"), "wb") as _f: _f.write(rgba_png_bytes) with open(_os.path.join(_diag, "prompt.txt"), "w", encoding="utf-8") as _f: _f.write(prompt if prompt is not None else "(None=用工作流内置默认)") except Exception: # noqa: BLE001 pass # 3. 提交（front=True 时插队到队列最前） payload = {"prompt": wf, "client_id": client_id} if front: payload["front"] = True r = cli.post("/prompt", json=payload) r.raise_for_status() prompt_id = r.json()["prompt_id"] # 4. 轮询 /history deadline = time.time() + timeout outputs = None while time.time() < deadline: hr = cli.get(f"/history/{prompt_id}") hr.raise_for_status() hist = hr.json() if prompt_id in hist: entry = hist[prompt_id] status = entry.get("status", {}) if status.get("status_str") == "error": raise RuntimeError(f"ComfyUI 执行报错: {status}") outputs = entry.get("outputs") if outputs and output_node in outputs: break time.sleep(0.2) if not outputs or output_node not in outputs: raise TimeoutError(f"ComfyUI 出图超时({timeout}s) prompt_id={prompt_id}") # 5. 取回输出图 imgs = outputs[output_node].get("images") or [] if not imgs: raise RuntimeError("ComfyUI 输出无图像") info = imgs[0] vr = cli.get("/view", params={"filename": info["filename"], "subfolder": info.get("subfolder", ""), "type": info.get("type", "output")}) vr.raise_for_status() return vr.content def ping() -> bool: """探测 ComfyUI 是否在线（/system_stats）。""" try: with httpx.Client(base_url=COMFYUI_URL, timeout=3.0, auth=_comfy_auth()) as cli: return cli.get("/system_stats").status_code == 200 except Exception: # noqa: BLE001 return False if __name__ == "__main__": import sys inp = sys.argv[1] if len(sys.argv) > 1 else "tests/output/comfy_input.png" print("ComfyUI:", COMFYUI_URL, "online:", ping()) with open(inp, "rb") as f: png = run(f.read()) out = "tests/output/grown.png" with open(out, "wb") as f: f.write(png) print(f"生发图已存 {out}（{len(png)} bytes）")