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:
xsl
2026-06-22 23:22:25 +08:00
co-authored by Claude
parent 28a6062fea
commit cb1989c042
6 changed files with 333 additions and 39 deletions
+38 -15
View File
@@ -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 Authuser `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]