feat(接口2/7): 接口2加hair_style参数选单张发型;新增接口7用add_hair2工作流
接口2 变更: - 新增必填 hair_style(int) 参数,按序号只生成一张(不再全量) - female:1-5 male:1-4,越界返回1007 接口7 新增: - POST /api/v1/hair/grow-v2,功能与接口2一致 - 使用 add_hair2.json 工作流(Flux-2 Klein 9b) - SaveImage输出节点自动检测(75) comfyui.py 重构: - run() 支持 workflow_path 参数,多工作流按路径缓存 - SaveImage 输出节点自动检测,不再硬编码 - 输入/种子/提示词节点ID两个工作流相同(26/6/60) 文档: - 接口文档、实现说明、网关待改动 三份同步更新 - 网关只需加一行路由,base64→URL改写无需改动 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
+38
-15
@@ -1,8 +1,10 @@
|
||||
"""ComfyUI 客户端:用 add_hair.json 工作流跑生发图(Flux-2 inpaint)。
|
||||
"""ComfyUI 客户端:用 add_hair.json / add_hair2.json 工作流跑生发图(Flux-2 inpaint)。
|
||||
|
||||
worker 不跑 Flux,只把「划线图 + 遮罩」的 RGBA 上传到本机 ComfyUI(默认 8188),
|
||||
替换工作流节点 26 的输入图、随机 seed,提交 /prompt,轮询 /history,取回 /view 输出。
|
||||
ComfyUI 开启了 HTTP Basic Auth(user `admin` + 密码),所有请求都带凭据。
|
||||
|
||||
支持多工作流:run() 可通过 workflow_path 指定不同工作流 JSON,自动检测 SaveImage 输出节点。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -16,7 +18,7 @@ import uuid
|
||||
import httpx
|
||||
|
||||
COMFYUI_URL = os.getenv("COMFYUI_URL", "http://127.0.0.1:8188").rstrip("/")
|
||||
WORKFLOW_PATH = os.getenv(
|
||||
_WORKFLOW_DEFAULT = os.getenv(
|
||||
"ADD_HAIR_WORKFLOW",
|
||||
os.path.join(os.path.dirname(os.path.dirname(__file__)), "add_hair.json"),
|
||||
)
|
||||
@@ -25,10 +27,10 @@ _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
|
||||
_wf_cache: dict[str, dict] = {} # path → workflow JSON
|
||||
_wf_output_node: dict[str, str] = {} # path → SaveImage 节点 ID
|
||||
|
||||
|
||||
def _comfy_auth():
|
||||
@@ -56,19 +58,40 @@ def _comfy_auth():
|
||||
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 _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 run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT, prompt: str = None) -> bytes:
|
||||
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) -> bytes:
|
||||
"""提交一次生发任务,返回输出 PNG 字节。失败抛异常。
|
||||
|
||||
prompt:非 None 时替换工作流节点60(JjkText)的文本;None 时用工作流内置默认提示词。
|
||||
workflow_path:工作流 JSON 路径,None 则用默认 add_hair.json。
|
||||
"""
|
||||
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 目录
|
||||
@@ -80,7 +103,7 @@ def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT, prompt: str = Non
|
||||
name = (up.get("subfolder") + "/" if up.get("subfolder") else "") + up["name"]
|
||||
|
||||
# 2. 改工作流:节点26 输入图 + 随机 seed
|
||||
wf = copy.deepcopy(_load_workflow())
|
||||
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:
|
||||
@@ -104,14 +127,14 @@ def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT, prompt: str = Non
|
||||
if status.get("status_str") == "error":
|
||||
raise RuntimeError(f"ComfyUI 执行报错: {status}")
|
||||
outputs = entry.get("outputs")
|
||||
if outputs and _OUTPUT_NODE in outputs:
|
||||
if outputs and output_node in outputs:
|
||||
break
|
||||
time.sleep(1.0)
|
||||
if not outputs or _OUTPUT_NODE not in outputs:
|
||||
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 []
|
||||
imgs = outputs[output_node].get("images") or []
|
||||
if not imgs:
|
||||
raise RuntimeError("ComfyUI 输出无图像")
|
||||
info = imgs[0]
|
||||
|
||||
Reference in New Issue
Block a user