From cb1989c042d801f82f15f23d66a3dabbb941bb8e Mon Sep 17 00:00:00 2001 From: xsl Date: Mon, 22 Jun 2026 23:22:25 +0800 Subject: [PATCH] =?UTF-8?q?feat(=E6=8E=A5=E5=8F=A32/7):=20=E6=8E=A5?= =?UTF-8?q?=E5=8F=A32=E5=8A=A0hair=5Fstyle=E5=8F=82=E6=95=B0=E9=80=89?= =?UTF-8?q?=E5=8D=95=E5=BC=A0=E5=8F=91=E5=9E=8B=EF=BC=9B=E6=96=B0=E5=A2=9E?= =?UTF-8?q?=E6=8E=A5=E5=8F=A37=E7=94=A8add=5Fhair2=E5=B7=A5=E4=BD=9C?= =?UTF-8?q?=E6=B5=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 接口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 --- app.py | 136 ++++++++++++++++++++++++++++++++++++++++++-- docs/实现说明.md | 17 ++++-- docs/接口文档.md | 70 ++++++++++++++++++++--- docs/网关待改动.md | 67 ++++++++++++++++++++++ hairline/comfyui.py | 53 ++++++++++++----- hairline/service.py | 29 +++++++--- 6 files changed, 333 insertions(+), 39 deletions(-) diff --git a/app.py b/app.py index 744da69..ec42902 100644 --- a/app.py +++ b/app.py @@ -446,8 +446,7 @@ async def face_measure( summary="接口2 C端生发", tags=["生发"], description=f""" -输入用户正面照 + **性别**，返回该性别对应的多张「建议发际线预览图」（本期为 -**发际线曲线叠加在原照片上的预览图**，非最终文生图生发图）。每个方案包含： +输入用户正面照 + **性别** + **发型序号**，返回指定发际线类型的预览图与生发图。每个方案包含： - 预览图（worker 返回 `image_base64`，网关落盘后改写为 `image_url`） - 发际线类型 `hairline_type`（英文 key） - 顺序 `order`（本期固定 `1..N`，不排序） @@ -460,6 +459,8 @@ async def face_measure( - **gender**（必填）：`male` / `female`。决定返回的贴图集合（female 5 张 / male 4 张）。 非法或缺失返回 `1004`。 +- **hair_style**（必填）：`int`，发型序号。`female`：1=ellipse, 2=flower, 3=heart, 4=straight, 5=wave； + `male`：1=ellipse, 2=inverse_arc, 3=m, 4=straight。越界返回 `1007`。 - **beauty_enabled**：本期保留但不生效。 `hairline_type` 取值：`ellipse` / `flower` / `heart` / `straight` / `wave`（female）， @@ -477,7 +478,6 @@ async def face_measure( "data": { "results": [ {"image_base64": "iVBORw0KGgo...", "hairline_type": "ellipse", "order": 1}, - {"image_base64": "iVBORw0KGgo...", "hairline_type": "flower", "order": 2}, ] }, } @@ -502,6 +502,7 @@ async def hair_grow( image_url: Optional[str] = Form(default=None, description="图片 URL"), image_base64: Optional[str] = Form(default=None, description="图片 base64（需带 data:image/...;base64, 前缀）"), gender: Optional[str] = Form(default=None, description="性别 male/female（必填）"), + hair_style: Optional[int] = Form(default=None, description="发型序号（必填）。female:1-5 male:1-4"), beauty_enabled: bool = Form(default=False, description="是否开启美颜（本期不生效）"), use_mask: bool = Form(default=True, description="是否启用 inpaint 遮罩（测试对比用）。false 时用干净原图生成（空遮罩，不烧模板线）"), prompt: str = Form(default="补充遮罩区域的头发", description="ComfyUI 提示词，会替换工作流节点60的文本"), @@ -510,7 +511,12 @@ async def hair_grow( if gender not in ("male", "female"): return err(1004, "gender 必填且只能为 male / female") - # 2. 三选一取图 + # 2. hair_style 必填校验（越界 → 1007） + max_styles = {"female": 5, "male": 4}[gender] + if hair_style is None or not isinstance(hair_style, int) or hair_style < 1 or hair_style > max_styles: + return err(1007, f"hair_style 必填且为 1..{max_styles} 的整数，收到 {hair_style!r}") + + # 3. 三选一取图 raw, e = await resolve_image_bytes(image_file, image_url, image_base64) if e is not None: return e @@ -531,7 +537,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, prompt) + items = await run_in_threadpool(generate_grow_results, image, gender, use_mask, prompt, hair_style) if items is None: return err(1001, "无法识别人像") @@ -550,6 +556,126 @@ async def hair_grow( return err(1007, f"处理失败：{ex}") +# --------------------------------------------------------------------------- +# 接口 7：C 端生发 v2（add_hair2.json 工作流） +# --------------------------------------------------------------------------- + +_WORKFLOW2_PATH = os.path.join(os.path.dirname(__file__), "add_hair2.json") + + +@app.post( + "/api/v1/hair/grow-v2", + summary="接口7 C端生发 v2（add_hair2 工作流）", + tags=["生发"], + description=f""" +输入用户正面照 + **性别** + **发型序号**，使用 add_hair2.json 工作流生成指定发际线类型的预览图与生发图。 +功能与接口2 完全一致，仅 ComfyUI 工作流不同。 + +{_image_fields_desc} + +图片同时支持 `multipart/form-data` 文件上传（字段名 `image_file`）。 + +--- + +- **gender**（必填）：`male` / `female`。决定返回的贴图集合（female 5 张 / male 4 张）。 + 非法或缺失返回 `1004`。 +- **hair_style**（必填）：`int`，发型序号。`female`：1=ellipse, 2=flower, 3=heart, 4=straight, 5=wave； + `male`：1=ellipse, 2=inverse_arc, 3=m, 4=straight。越界返回 `1007`。 +- **beauty_enabled**：本期保留但不生效。 + +`hairline_type` 取值：`ellipse` / `flower` / `heart` / `straight` / `wave`（female）， +`ellipse` / `m` / `straight` / `inverse_arc`（male）。 +""", + responses={ + 200: { + "description": "成功", + "content": { + "application/json": { + "example": { + "code": 0, + "message": "success", + "request_id": "mock-request-id", + "data": { + "results": [ + {"image_base64": "iVBORw0KGgo...", "hairline_type": "ellipse", "order": 1}, + ] + }, + } + } + }, + }, + 400: { + "description": "参数错误 / 图片识别失败", + "content": { + "application/json": { + "examples": { + "图片参数错误": {"value": {"code": 1007, "message": "图片参数错误：必须且只能传 image_file / image_url / image_base64 其中一个", "request_id": "x", "data": None}}, + "非正面照": {"value": {"code": 1003, "message": "角度问题，请上传正面照", "request_id": "x", "data": None}}, + } + } + }, + }, + }, +) +async def hair_grow_v2( + image_file: Optional[UploadFile] = File(default=None, description="上传图片文件（JPG/PNG，≤ 1 MB）"), + image_url: Optional[str] = Form(default=None, description="图片 URL"), + image_base64: Optional[str] = Form(default=None, description="图片 base64（需带 data:image/...;base64, 前缀）"), + gender: Optional[str] = Form(default=None, description="性别 male/female（必填）"), + hair_style: Optional[int] = Form(default=None, description="发型序号（必填）。female:1-5 male:1-4"), + 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"): + return err(1004, "gender 必填且只能为 male / female") + + # 2. hair_style 必填校验（越界 → 1007） + max_styles = {"female": 5, "male": 4}[gender] + if hair_style is None or not isinstance(hair_style, int) or hair_style < 1 or hair_style > max_styles: + return err(1007, f"hair_style 必填且为 1..{max_styles} 的整数，收到 {hair_style!r}") + + # 3. 三选一取图 + raw, e = await resolve_image_bytes(image_file, image_url, image_base64) + if e is not None: + return e + if len(raw) > MAX_FILE_BYTES: + return err(1006, "文件超出 1 MB 限制") + + image = cv2.imdecode(np.frombuffer(raw, np.uint8), cv2.IMREAD_COLOR) + if image is None: + return err(1008, "图片格式不支持（仅 JPG / PNG）") + + h, w = image.shape[:2] + short_side, long_side = min(w, h), max(w, h) + if short_side < MIN_SHORT_SIDE or long_side < MIN_LONG_SIDE: + return err(1002, "人像分辨率过低") + + try: + from fastapi.concurrency import run_in_threadpool + from hairline.service import generate_grow_results + + # 预览 + 生发(ComfyUI) 都是阻塞且较慢，放线程池避免卡住事件循环 + items = await run_in_threadpool(generate_grow_results, image, gender, use_mask, prompt, hair_style, _WORKFLOW2_PATH) + if items is None: + return err(1001, "无法识别人像") + + results = [] + for p in items: + results.append({ + "image_base64": _jpg_b64(p["image_bgr"]), # 预览图 JPG + "grown_image_base64": (_png_to_jpg_b64(p["grown_png"]) # 生发图 JPG + if p["grown_png"] else None), + "hairline_type": p["hairline_type"], + "order": p["order"], + }) + return ok({"results": results}) + except Exception as ex: # noqa: BLE001 + logger.exception("接口7 处理异常") + return err(1007, f"处理失败：{ex}") + + # --------------------------------------------------------------------------- # 接口 3：B 端生发 # --------------------------------------------------------------------------- diff --git a/docs/实现说明.md b/docs/实现说明.md index 9af7ccd..79fdb32 100644 --- a/docs/实现说明.md +++ b/docs/实现说明.md @@ -13,7 +13,7 @@ ``` 客户端 ──HTTPS──> 外网网关(gateway/) ──HTTP(X-Internal-Token)──> worker(GPU 机, app.py) │ 薄代理 + 落盘改URL │ 跑算法(本地模型/ComfyUI) - └ 接口4 本机直接调豆包(不转发) └ 接口1/2/3/5 + └ 接口4 本机直接调豆包(不转发) └ 接口1/2/3/5/7 ``` - **worker**（`app.py` + `face_analysis/` + `hairline/`）：跑真正的算法，**纯本地、无外网依赖**。 @@ -32,6 +32,7 @@ | 2 | `results[].image_base64` / `results[].grown_image_base64`(可空) | `results[].image_url` / `results[].grown_image_url` | | 3 | `hair_growth_image_base64`(可空) | `hair_growth_image_url` | | 5 | `hairline_images[].image_base64` | `hairline_images[].image_url` | +| 7 | `results[].image_base64` / `results[].grown_image_base64`(可空) | `results[].image_url` / `results[].grown_image_url` | | 4 | （网关本机产出，无图片字段，`features` 为 JSON 字符串） | — | > 实现建议：递归遍历 data，凡 key 以 `_base64` 结尾就落盘改 `_url`，自动覆盖嵌套/新增字段。 @@ -40,7 +41,7 @@ ### 错误码 -`1001` 无法识别人像 ｜ `1002` 分辨率过低 ｜ `1003` 非正面 ｜ `1004` gender 必填/非法（接口2/5）｜ +`1001` 无法识别人像 ｜ `1002` 分辨率过低 ｜ `1003` 非正面 ｜ `1004` gender 必填/非法（接口2/5/7）｜ `1006` >1MB ｜ `1007` 图片参数错误（0或多个）/未预期异常 ｜ `1008` 格式不支持。 --- @@ -55,11 +56,19 @@ - 门槛可配：`MIN_SHORT_SIDE`/`MIN_LONG_SIDE`(默认600/800)、姿态阈值 `FRONTAL_*_THR`(默认30°)。 ### 接口2 C端生发 `/api/v1/hair/grow`（worker）—— 预览 + 生发图 -- **做什么**：正面照 + `gender`(必填) → 该性别全部发际线(female5/male4) 各一组：**预览图**(发际线叠在照片上) + **生发后图**(植发3个月效果)。 +- **做什么**：正面照 + `gender`(必填) + `hair_style`(必填，int) → 指定发际线类型 **1 组**：**预览图**(发际线叠在照片上) + **生发后图**(植发3个月效果)。 - **怎么实现**（`hairline/`）：移植 head3d——MediaPipe(Tasks) + SegFormer 分割 + 17 锚点射线检测 → 502 点 mesh， 按 `face_ext.obj` 的 UV 把发际线贴图渲染到额头(预览)。生发：黑贴图渲染遮罩 → 调本机 **ComfyUI 8182** 的 `add_hair.json`(Flux-2) 出图。**关键坑**：obj 是重排序，需 `INDEX_MAP_468` 把 MP 序→OBJ 序。 - 同步、一次 N 张(~18s)。返回 `results[].image_base64` + `grown_image_base64`。 + 单张请求(~5s)。返回 `results[].image_base64` + `grown_image_base64`。 +- `hair_style` 映射：female 1=ellipse 2=flower 3=heart 4=straight 5=wave；male 1=ellipse 2=inverse_arc 3=m 4=straight。 + +### 接口7 C端生发 v2 `/api/v1/hair/grow-v2`（worker）—— 接口2同款，add_hair2 工作流 +- **做什么**：与接口 2 完全一致（正面照 + `gender` + `hair_style` → 1 组预览+生发图）。 +- **与接口 2 的唯一区别**：ComfyUI 工作流使用 `add_hair2.json`（Flux-2 Klein 9b），输入/遮罩节点同为 26， + SaveImage 输出节点为 75（自动检测）。其他参数、响应结构、错误码**完全相同**。 +- **网关改动**：在 `gateway/app.py` 新增路由 `POST /api/v1/hair/grow-v2`，转发到 worker 同路径即可（盲转发， + base64→URL 改写逻辑无需改动，数组内图片字段已覆盖）。详见 [`网关待改动.md`](网关待改动.md)。 ### 接口3 B端生发 `/api/v1/hair/grow-b`（worker）—— 马克笔发际线 - **做什么**：医生用马克笔在额头画好发际线，**只传这一张划线图** → 检测线 → 生发图。输出 `hair_growth_image_url` + `hairline_type="custom"`。 diff --git a/docs/接口文档.md b/docs/接口文档.md index 3549e2a..718537c 100644 --- a/docs/接口文档.md +++ b/docs/接口文档.md @@ -19,6 +19,7 @@ | 3 B 端生发 | POST | `/api/v1/hair/grow-b` | | 4 用户特征 | POST | `/api/v1/face/features` | | 5 发际线 PNG 生成 | POST | `/api/v1/hairline/generate` | +| 7 C 端生发 v2 | POST | `/api/v1/hair/grow-v2` | --- @@ -177,9 +178,9 @@ ## 接口 2：C 端生发接口 -**说明**：输入用户正面照 + 性别，按性别对应的发际线类型，逐张把建议发际线渲染到照片上，输出多张方案。 +**说明**：输入用户正面照 + 性别 + 发型序号，按指定发际线类型渲染一张预览图 + 生发图。 -> **每个方案返回两张图**：`image_url`=「原照片 + 发际线曲线叠加的**预览图**」；`grown_image_url`= +> **每个请求返回两张图**：`image_url`=「原照片 + 发际线曲线叠加的**预览图**」；`grown_image_url`= > 经 ComfyUI/Flux 的「植发 3 个月**生发后图片**」。两者均已实现，实现简述见 [`实现说明.md`](实现说明.md)。 **请求**：`POST /api/v1/hair/grow` @@ -191,12 +192,13 @@ | 参数 | 类型 | 必填 | 说明 | |------|------|------|------| | gender | string | **是** | 性别：`male` / `female`。决定使用的发际线贴图集合 | +| hair_style | int | **是** | 发型序号，1-indexed。female：1=ellipse, 2=flower, 3=heart, 4=straight, 5=wave；male：1=ellipse, 2=inverse_arc, 3=m, 4=straight。越界返回 `1007` | | beauty_enabled | bool | 否 | 生发图是否带美颜效果，默认 false（当前阶段不生效） | -| use_mask | bool | 否 | 是否启用 inpaint 遮罩，默认 `true`。`false` 时用干净原图生成（空遮罩、不烧模板黑线），供测试对比；此时 N 张结果共用同一张生发图 | +| use_mask | bool | 否 | 是否启用 inpaint 遮罩，默认 `true`。`false` 时用干净原图生成（空遮罩、不烧模板黑线），供测试对比 | ### 输出（data） -`results`：发际线方案数组，**数量 = 该性别的发际线类型数**（`female` 5 个 / `male` 4 个）。每个元素： +`results`：发际线方案数组，**含 1 个元素**（指定发型）。每个元素： | 字段 | 类型 | 说明 | |------|------|------| @@ -205,8 +207,8 @@ | hairline_type | string | 发际线类型 key：`ellipse`/`flower`/`heart`/`straight`/`wave`（female），`ellipse`/`m`/`straight`/`inverse_arc`（male） | | order | int | 排序序号（当前阶段固定 `1..N`，按贴图顺序，暂不计算合适度） | -> ⚠️ 生发图由本机 ComfyUI（Flux-2，端口 8182）生成，**一次请求生成全部 N 张、同步返回**， -> 单请求耗时可达数分钟，调用方超时需放大。worker 侧返回 `image_base64` / `grown_image_base64`， +> ⚠️ 生发图由本机 ComfyUI（Flux-2，端口 8182）生成，**一次请求生成指定发型的 1 张、同步返回**。 +> worker 侧返回 `image_base64` / `grown_image_base64`， > 网关落盘后改写为上表的 `image_url` / `grown_image_url`。 ### 响应示例 @@ -361,15 +363,69 @@ --- +## 接口 7：C 端生发 v2 接口 + +**说明**：功能与[接口 2](#接口-2c-端生发接口)完全一致，仅 ComfyUI 工作流不同——使用 `add_hair2.json` 替代 `add_hair.json`。 + +**请求**：`POST /api/v1/hair/grow-v2` + +### 输入 + +与接口 2 完全相同。图片参数见「通用约定 → 图片传参字段」。专属参数： + +| 参数 | 类型 | 必填 | 说明 | +|------|------|------|------| +| gender | string | **是** | 性别：`male` / `female`。决定使用的发际线贴图集合 | +| hair_style | int | **是** | 发型序号，1-indexed。female：1=ellipse, 2=flower, 3=heart, 4=straight, 5=wave；male：1=ellipse, 2=inverse_arc, 3=m, 4=straight。越界返回 `1007` | +| beauty_enabled | bool | 否 | 生发图是否带美颜效果，默认 false（当前阶段不生效） | +| use_mask | bool | 否 | 是否启用 inpaint 遮罩，默认 `true`。`false` 时用干净原图生成（空遮罩、不烧模板黑线） | + +### 输出（data） + +与接口 2 完全相同。`results`：发际线方案数组，**含 1 个元素**。每个元素： + +| 字段 | 类型 | 说明 | +|------|------|------| +| image_url | string | 方案**预览图** URL（发际线曲线叠加图） | +| grown_image_url | string | **生发后图片** URL（ComfyUI/Flux「植发 3 个月」效果图） | +| hairline_type | string | 发际线类型 key | +| order | int | 排序序号 | + +> ⚠️ 与接口 2 的区别：本接口使用 `add_hair2.json` 工作流（Flux-2 Klein 9b），输入/遮罩节点同为 26， +> SaveImage 输出节点为 75。 + +### 响应示例 + +```json +{ + "code": 0, + "message": "success", + "request_id": "mock-request-id", + "data": { + "results": [ + { + "image_url": "https://hair.xiangsilian.com/static/sample.jpg", + "grown_image_url": "https://hair.xiangsilian.com/static/sample.jpg", + "hairline_type": "ellipse", + "order": 1 + } + ] + } +} +``` + +--- + ## 汇总：输入输出一览 | 接口 | 输入 | 主要输出 | |------|------|----------| | 1 四庭七眼测量 | 用户照片 | 标注 PNG（无人物）+ 四庭/七眼厘米数值与坐标 | -| 2 C 端生发 | 用户照片 | 生发后图片 + 多张发际线（带类型与排序） | +| 2 C 端生发 | 用户照片 | 生发后图片 + 指定发际线预览（单张） | | 3 B 端生发 | 划线图片 | 最合适发际线图片 + 生发后图片 | | 4 用户特征 | 用户照片 | 6 个用户特征字段（脸形/眉形/年龄/动静/性别/基因风格） | | 5 发际线 PNG | 用户照片 | N 张发际线 PNG + 最合适发际线面部中间点坐标 | +| 7 C 端生发 v2 | 用户照片 + gender + hair_style | 同接口2，使用 add_hair2.json 工作流 | --- diff --git a/docs/网关待改动.md b/docs/网关待改动.md index d72ee0b..9038c43 100644 --- a/docs/网关待改动.md +++ b/docs/网关待改动.md @@ -44,6 +44,73 @@ filename = f"{uuid.uuid4().hex}.{ext}" --- +## 5. 🔲【新增】接口7 C端生发 v2（`/api/v1/hair/grow-v2`） + +**背景**：worker 侧已新增接口 7，功能与接口 2 完全一致，区别仅在于 ComfyUI 工作流使用 `add_hair2.json`（而非 `add_hair.json`）。 + +**网关需新增一个路由**，代理转发到 worker 的同路径： + +```python +# gateway/app.py + +@app.post("/api/v1/hair/grow-v2", tags=["生发"]) +async def hair_grow_v2(request: Request): + """接口7：C端生发 v2（add_hair2 工作流）""" + return await _proxy(request, "/api/v1/hair/grow-v2") +``` + +**无需额外改动**： +- 请求：multipart/form-data，参数与接口 2 完全相同（`image_file/url/base64` 三选一 + `gender` + `hair_style` + `beauty_enabled` + `use_mask` + `prompt`），网关盲转发即可 +- 响应：结构与接口 2 完全一致，`results[].image_base64` / `results[].grown_image_base64` 经现有 `rewrite_base64_to_url` 自动改写为 URL +- base64→URL：数组内图片字段递归改写已覆盖，无需修改 + +### 入参（与接口 2 一致） + +| 参数 | 类型 | 必填 | 说明 | +|------|------|------|------| +| image_file / image_url / image_base64 | — | **三选一** | 用户正面照 | +| gender | string | **是** | `male` / `female` | +| hair_style | int | **是** | 发型序号。female: 1–5，male: 1–4 | +| beauty_enabled | bool | 否 | 美颜开关（本期不生效） | +| use_mask | bool | 否 | 默认 `true`，`false` 跳过遮罩 | +| prompt | string | 否 | ComfyUI 提示词 | + +### 出参（与接口 2 一致） + +```json +{ + "code": 0, + "message": "success", + "request_id": "gw-xxxxxxxx", + "data": { + "results": [ + { + "image_url": "https://hair.xiangsilian.com/static/annotations/xxx.jpg", + "grown_image_url": "https://hair.xiangsilian.com/static/annotations/xxx.jpg", + "hairline_type": "ellipse", + "order": 1 + } + ] + } +} +``` + +### worker 侧信息 + +| 项目 | 值 | +|------|-----| +| worker 路径 | `/api/v1/hair/grow-v2` | +| 工作流文件 | `add_hair2.json` | +| 输入/遮罩节点 | 26（LoadImage） | +| 输出节点 | 75（SaveImage，自动检测） | +| 提示词节点 | 60（JjkText） | + +### 🔲【可选·仅影响 /docs】OpenAPI 表单声明 + +若想让网关 `/docs` 展示准确，在 `gateway/app.py` 新增 `_GROW_V2_FORMS`（或复用 `_GROW_FORMS` 并补充 `gender`/`hair_style` 字段），然后将路由函数签名改为显式声明 Form 参数（参考接口 4 的写法）。不改也不影响实际转发。 + +--- + ## 已经做好、无需再动的 - **接口4 在网关本机实现**（调豆包，不转发 worker）——已完成；config 里配 `ark`。 diff --git a/hairline/comfyui.py b/hairline/comfyui.py index 5c71765..507cf34 100644 --- a/hairline/comfyui.py +++ b/hairline/comfyui.py @@ -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] diff --git a/hairline/service.py b/hairline/service.py index ab262f0..3e8fdc6 100644 --- a/hairline/service.py +++ b/hairline/service.py @@ -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)