xsl/hair
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(接口2/7): 接口2加hair_style参数选单张发型；新增接口7用add_hair2工作流

Browse Source 
接口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
Download Patch File Download Diff File Expand all files Collapse all files
app.py
+131 -5
View File
@@ -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 端生发
# ---------------------------------------------------------------------------
docs/实现说明.md
+13 -4
View File
@@ -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"`。
docs/接口文档.md
+63 -7
View File
@@ -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 工作流 |
---
docs/网关待改动.md
+67
View File
@@ -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`。
hairline/comfyui.py
+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 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]
hairline/service.py
+21 -8
View File
@@ -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)