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:
Claude
+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]
|
||||
|
||||
+21
-8
@@ -128,15 +128,19 @@ def generate_previews(image_bgr: np.ndarray, gender: str):
|
||||
return results
|
||||
|
||||
|
||||
def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = True, prompt: str = None):
|
||||
"""该性别全部发际线：预览图(白线) + 生发图(ComfyUI)。同步、串行。
|
||||
def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = True,
|
||||
prompt: str = None, hair_style: int = None,
|
||||
workflow_path: str | None = None):
|
||||
"""指定发际线类型：预览图(白线) + 生发图(ComfyUI)。
|
||||
|
||||
hair_style（1-indexed）：指定生成第几张发际线（按贴图排序）。female: 1..5，male: 1..4。
|
||||
为 None 时生成全部（兼容旧调用）。
|
||||
use_mask（默认 True）：是否启用 inpaint 遮罩，用于测试对比（同接口3）。
|
||||
False 时用**干净原图 + 空遮罩**送 ComfyUI（不烧黑色模板线），与模板无关，
|
||||
故只跑一次 ComfyUI、N 项复用同一张生发图；预览图（白线）仍按各模板生成。
|
||||
False 时用**干净原图 + 空遮罩**送 ComfyUI（不烧黑色模板线）。
|
||||
prompt（默认 None）：ComfyUI 提示词，非 None 时替换工作流节点60文本。
|
||||
workflow_path（默认 None）：ComfyUI 工作流 JSON 路径，None 用默认 add_hair.json。
|
||||
Returns: list[dict] {"hairline_type","order","image_bgr"(预览), "grown_png"(bytes 或 None)}。
|
||||
无人脸返回 None。某张 ComfyUI 失败时该项 grown_png=None，不影响其余。
|
||||
无人脸返回 None。某张 ComfyUI 失败时该项 grown_png=None，不抛异常。
|
||||
"""
|
||||
if gender not in ("male", "female"):
|
||||
raise ValueError(f"gender 必须是 male/female，收到 {gender!r}")
|
||||
@@ -145,6 +149,15 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
|
||||
return None
|
||||
uv, ext_faces = load_ext_mesh()
|
||||
|
||||
textures = get_texture_map()[gender] # [(key, path), ...] 已排序
|
||||
if hair_style is not None:
|
||||
if hair_style < 1 or hair_style > len(textures):
|
||||
raise ValueError(
|
||||
f"hair_style={hair_style} 超出范围，{gender} 可选 1..{len(textures)}")
|
||||
items = [(hair_style, textures[hair_style - 1])]
|
||||
else:
|
||||
items = list(enumerate(textures, start=1))
|
||||
|
||||
# 禁用遮罩：干净原图 + 空遮罩，与模板无关 → 只跑一次 ComfyUI，下面 N 项复用
|
||||
shared_grown = None
|
||||
if not use_mask:
|
||||
@@ -152,12 +165,12 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
|
||||
h, w = image_bgr.shape[:2]
|
||||
buf = io.BytesIO()
|
||||
compose_comfy_rgba(image_bgr, np.zeros((h, w), np.uint8)).save(buf, format="PNG")
|
||||
shared_grown = comfyui.run(buf.getvalue(), prompt=prompt)
|
||||
shared_grown = comfyui.run(buf.getvalue(), prompt=prompt, workflow_path=workflow_path)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("接口2 生发图失败(无遮罩)：%s", e)
|
||||
|
||||
results = []
|
||||
for order, (key, white_path) in enumerate(get_texture_map()[gender], start=1):
|
||||
for order, (key, white_path) in items:
|
||||
white = load_texture_rgba(white_path)
|
||||
preview = render_hairline_overlay(image_bgr, ctx["points"], ext_faces, uv, white)
|
||||
|
||||
@@ -171,7 +184,7 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
|
||||
image_bgr, ctx["landmarks"], ctx["parse_map"], ctx["points"], black)
|
||||
buf = io.BytesIO()
|
||||
compose_comfy_rgba(marked, mask).save(buf, format="PNG")
|
||||
grown_png = comfyui.run(buf.getvalue(), prompt=prompt)
|
||||
grown_png = comfyui.run(buf.getvalue(), prompt=prompt, workflow_path=workflow_path)
|
||||
except Exception as e: # noqa: BLE001 单张失败不拖垮整请求
|
||||
logger.warning("接口2 生发图失败 type=%s：%s", key, e)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user