xslandClaude Opus 4.8
94ad95850e
在发际线预览基础上，每种发际线再出一张「植发3个月」生发图： - hairline/mask.py: headmark 5步法遮罩(额头上部区域∩SegFormer头部=ROI， 取发际线曲线以上闭合区域)；用 hairline_texture_black 渲染黑线替代手绘检测； compose_comfy_rgba 合成 RGBA(alpha=255-mask, 透明=重绘区, 对齐 ComfyUI mask=1-alpha) - hairline/comfyui.py: ComfyUI 客户端(默认8182)，/upload/image+/prompt(改节点26+随机seed) +轮询/history+/view 取回生发图 - hairline/render.py: 抽出 build_overlay_layer 供遮罩取曲线像素 - hairline/service.py: extract_context 一次出 landmarks/parse_map/502点； generate_grow_results 每种=预览+生发图(同步串行N张，单张ComfyUI失败则grown置空不拖垮整请求) - app.py: /hair/grow 返回 results[].grown_image_base64；重活放线程池避免卡事件循环 - add_hair.json 工作流 + hairline_texture_black/ 黑贴图入库 - 测试: test_mask.py(遮罩几何) + test_api mock ComfyUI 验 grown 字段，35 全绿 实测(5090): female 5张生发图同步约18s；预览/生发图人物五官服饰背景保持、黑线已清除。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
111 lines
4.0 KiB
Python
111 lines
4.0 KiB
Python
"""ComfyUI 客户端：用 add_hair.json 工作流跑生发图（Flux-2 inpaint）。
|
||
|
||
worker 不跑 Flux，只把「划线图 + 遮罩」的 RGBA 上传到本机 ComfyUI(默认 8182)，
|
||
替换工作流节点 26 的输入图、随机 seed，提交 /prompt，轮询 /history，取回 /view 输出。
|
||
"""
|
||
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:8182").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")) # 单张出图最长等待(秒)
|
||
|
||
_INPUT_NODE = "26" # LoadImage：外部输入图(含 alpha 遮罩)
|
||
_SEED_NODE = "6" # RandomNoise
|
||
_OUTPUT_NODE = "17" # SaveImage
|
||
|
||
_workflow = 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) -> bytes:
|
||
"""提交一次生发任务，返回输出 PNG 字节。失败抛异常。"""
|
||
client_id = uuid.uuid4().hex
|
||
with httpx.Client(base_url=COMFYUI_URL, timeout=30.0) 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)
|
||
|
||
# 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) 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）")
|