144 lines
5.4 KiB
Python
144 lines
5.4 KiB
Python
"""ComfyUI 客户端：用 add_hair.json 工作流跑生发图（Flux-2 inpaint）。
|
||
|
||
worker 不跑 Flux，只把「划线图 + 遮罩」的 RGBA 上传到本机 ComfyUI(默认 8188)，
|
||
替换工作流节点 26 的输入图、随机 seed，提交 /prompt，轮询 /history，取回 /view 输出。
|
||
ComfyUI 开启了 HTTP Basic Auth（user `admin` + 密码），所有请求都带凭据。
|
||
"""
|
||
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_PATH = 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
|
||
_OUTPUT_NODE = "17" # SaveImage
|
||
_PROMPT_NODE = "60" # JjkText：提示词
|
||
|
||
_workflow = None
|
||
|
||
|
||
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() -> dict:
|
||
global _workflow
|
||
if _workflow is None:
|
||
with open(WORKFLOW_PATH, encoding="utf-8") as f:
|
||
_workflow = json.load(f)
|
||
return _workflow
|
||
|
||
|
||
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 目录
|
||
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())
|
||
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})
|
||
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(1.0)
|
||
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）")
|