From be81567261784a08e5fe13d4729b80f7ebe9f410 Mon Sep 17 00:00:00 2001 From: xsl Date: Wed, 17 Jun 2026 23:36:01 +0800 Subject: [PATCH] save code --- app.py | 6 ++++-- hairline/comfyui.py | 10 ++++++++-- hairline/service.py | 11 ++++++----- 3 files changed, 18 insertions(+), 9 deletions(-) diff --git a/app.py b/app.py index 0384154..04eeac8 100644 --- a/app.py +++ b/app.py @@ -504,6 +504,7 @@ async def hair_grow( gender: Optional[str] = Form(default=None, description="性别 male/female（必填）"), beauty_enabled: bool = Form(default=False, description="是否开启美颜（本期不生效）"), use_mask: bool = Form(default=True, description="是否启用 inpaint 遮罩（测试对比用）。false 时用干净原图生成（空遮罩，不烧模板线）"), + prompt: str = Form(default="补充遮罩区域的头发", description="ComfyUI 提示词，会替换工作流节点60的文本"), ): # 1. gender 必填校验（非法/缺失 → 1004） if gender not in ("male", "female"): @@ -530,7 +531,7 @@ async def hair_grow( from hairline.service import generate_grow_results # 预览 + 生发(ComfyUI) 都是阻塞且较慢，放线程池避免卡住事件循环 - items = await run_in_threadpool(generate_grow_results, image, gender, use_mask) + items = await run_in_threadpool(generate_grow_results, image, gender, use_mask, prompt) if items is None: return err(1001, "无法识别人像") @@ -596,6 +597,7 @@ async def hair_grow_b( marked_image_url: Optional[str] = Form(default=None, description="划线图片 URL"), marked_image_base64: Optional[str] = Form(default=None, description="划线图片 base64"), use_mask: bool = Form(default=True, description="是否画发际线（测试对比用）。false 时跳过划线检测、直接送划线图"), + prompt: str = Form(default="补充遮罩区域的头发", description="ComfyUI 提示词，会替换工作流节点60的文本"), ): # 划线图三选一取图（只需这一张） marked_raw, e = await resolve_image_bytes(marked_image_file, marked_image_url, marked_image_base64) @@ -617,7 +619,7 @@ async def hair_grow_b( from fastapi.concurrency import run_in_threadpool from hairline.service import generate_grow_b - res = await run_in_threadpool(generate_grow_b, marked, use_mask) + res = await run_in_threadpool(generate_grow_b, marked, use_mask, prompt) if res["status"] == "no_face": return err(1001, "无法识别人像") if res["status"] == "no_line": diff --git a/hairline/comfyui.py b/hairline/comfyui.py index f8c7422..5c71765 100644 --- a/hairline/comfyui.py +++ b/hairline/comfyui.py @@ -26,6 +26,7 @@ _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 @@ -63,8 +64,11 @@ def _load_workflow() -> dict: return _workflow -def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT) -> bytes: - """提交一次生发任务，返回输出 PNG 字节。失败抛异常。""" +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 目录 @@ -79,6 +83,8 @@ def run(rgba_png_bytes: bytes, timeout: float = COMFY_TIMEOUT) -> bytes: 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}) diff --git a/hairline/service.py b/hairline/service.py index f901792..ab262f0 100644 --- a/hairline/service.py +++ b/hairline/service.py @@ -128,12 +128,13 @@ 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): +def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = True, prompt: str = None): """该性别全部发际线：预览图(白线) + 生发图(ComfyUI)。同步、串行。 use_mask（默认 True）：是否启用 inpaint 遮罩，用于测试对比（同接口3）。 False 时用**干净原图 + 空遮罩**送 ComfyUI（不烧黑色模板线），与模板无关， 故只跑一次 ComfyUI、N 项复用同一张生发图；预览图（白线）仍按各模板生成。 + prompt（默认 None）：ComfyUI 提示词，非 None 时替换工作流节点60文本。 Returns: list[dict] {"hairline_type","order","image_bgr"(预览), "grown_png"(bytes 或 None)}。 无人脸返回 None。某张 ComfyUI 失败时该项 grown_png=None，不影响其余。 """ @@ -151,7 +152,7 @@ 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()) + shared_grown = comfyui.run(buf.getvalue(), prompt=prompt) except Exception as e: # noqa: BLE001 logger.warning("接口2 生发图失败(无遮罩)：%s", e) @@ -170,7 +171,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()) + grown_png = comfyui.run(buf.getvalue(), prompt=prompt) except Exception as e: # noqa: BLE001 单张失败不拖垮整请求 logger.warning("接口2 生发图失败 type=%s：%s", key, e) @@ -211,7 +212,7 @@ def generate_hairline_pngs(image_bgr: np.ndarray, gender: str): return {"images": images, "best_center": best_center} -def generate_grow_b(marked_bgr: np.ndarray, use_mask: bool = True): +def generate_grow_b(marked_bgr: np.ndarray, use_mask: bool = True, prompt: str = None): """接口3：检测医生手绘发际线 → 遮罩 → 送 ComfyUI 生发（仅需划线图一张）。 检测路径只用来**建遮罩**；ComfyUI 输入图用 **marked 原图**（含医生手绘线， @@ -241,7 +242,7 @@ def generate_grow_b(marked_bgr: np.ndarray, use_mask: bool = True): buf = io.BytesIO() compose_comfy_rgba(marked_bgr, mask).save(buf, format="PNG") # marked 原图 + 遮罩 - grown_png = comfyui.run(buf.getvalue()) + grown_png = comfyui.run(buf.getvalue(), prompt=prompt) return {"grown_png": grown_png, "status": "ok"}