save
This commit is contained in:
+63
-11
@@ -29,7 +29,8 @@ import logging
|
||||
logger = logging.getLogger("hair.worker")
|
||||
|
||||
# 接口2 女性发型 key → change_hair hair_id（chang_*）映射：换发型+Flux-2 整帧重绘用。
|
||||
# 与接口12 final 的 5 型一一对应。
|
||||
# 与接口12 final 的 5 型一一对应。female 6/7（bigflower/clasicalflower）无对应 LoRA，
|
||||
# 走与男性一致的原生生发（ComfyUI add_hair）管线，故不在本表。
|
||||
_FEMALE_KEY_TO_CHANG = {
|
||||
"ellipse": "chang_tuoyuan", # 椭圆
|
||||
"flower": "chang_huaban", # 花瓣
|
||||
@@ -38,6 +39,16 @@ _FEMALE_KEY_TO_CHANG = {
|
||||
"wave": "chang_bolang", # 波浪
|
||||
}
|
||||
|
||||
# 发际线贴图显式顺序表：决定 hair_style 序号（1-indexed）。
|
||||
# 不再依赖文件名字母序——字母序会因新增/重命名文件而错位，破坏现有前端/客户端取值。
|
||||
# key 须与 _gender_key 派生结果一致（已去空格）：如 "inverse_arc"（源 man_ inverse_arc.png）、
|
||||
# "Softpetal"（源 man_Soft petal.png，大写 S 保留）。表外未知 key 兜底排到末尾。
|
||||
_HAIRSTYLE_ORDER = {
|
||||
"female": ["ellipse", "flower", "heart", "straight", "wave",
|
||||
"bigflower", "clasicalflower"], # 1..7
|
||||
"male": ["ellipse", "inverse_arc", "m", "straight", "heart", "Softpetal"], # 1..6
|
||||
}
|
||||
|
||||
_REPO = os.path.dirname(os.path.dirname(__file__))
|
||||
_TEXTURE_DIR = os.path.join(_REPO, "hairline_texture")
|
||||
_BLACK_TEXTURE_DIR = os.path.join(_REPO, "hairline_texture_black")
|
||||
@@ -123,10 +134,13 @@ def _gender_key(stem: str):
|
||||
|
||||
|
||||
def get_texture_map(level: str = "middle") -> dict:
|
||||
"""扫描指定档位贴图目录建 {gender: [(key, path)]}，按 key 排序、按档位缓存。
|
||||
"""扫描指定档位贴图目录建 {gender: [(key, path)]}，按显式顺序表排序、按档位缓存。
|
||||
|
||||
level：middle(默认) / high / low，分别对应 hairline_texture[/_high|/_low]。
|
||||
文件名规范化去空格（如 `man_ inverse_arc.png` → key `inverse_arc`）。
|
||||
|
||||
排序依据 _HAIRSTYLE_ORDER：表内 key 按表序、表外未知 key 兜底排到末尾（再按字母序），
|
||||
保证新增/重命名文件不会打乱现有 hair_style 序号。
|
||||
"""
|
||||
if level not in _TEXTURE_DIRS:
|
||||
raise ValueError(f"hairline_level 必须是 middle/high/low，收到 {level!r}")
|
||||
@@ -140,7 +154,9 @@ def get_texture_map(level: str = "middle") -> dict:
|
||||
if gender:
|
||||
mapping[gender].append((key, path))
|
||||
for g in mapping:
|
||||
mapping[g].sort(key=lambda kp: kp[0])
|
||||
order = _HAIRSTYLE_ORDER.get(g, [])
|
||||
idx = {k: i for i, k in enumerate(order)}
|
||||
mapping[g].sort(key=lambda kp: (idx.get(kp[0], len(idx)), kp[0]))
|
||||
_texture_maps[level] = mapping
|
||||
return mapping
|
||||
|
||||
@@ -203,7 +219,7 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
|
||||
workflow_path: str | None = None):
|
||||
"""指定发际线类型：发际线透明叠图(白线 RGBA) + 生发图(ComfyUI)。
|
||||
|
||||
hair_styles（1-indexed 列表）：指定生成哪几张发际线（按贴图排序）。female: 1..5，male: 1..4。
|
||||
hair_styles（1-indexed 列表）：指定生成哪几张发际线（按贴图排序）。female: 1..7，male: 1..6。
|
||||
为 None 时生成全部（兼容旧调用）。
|
||||
use_mask（默认 True）：是否启用 inpaint 遮罩，用于测试对比（同接口3）。
|
||||
False 时用**干净原图 + 空遮罩**送 ComfyUI（不烧黑色模板线）。
|
||||
@@ -273,15 +289,22 @@ def generate_grow_results(image_bgr: np.ndarray, gender: str, use_mask: bool = T
|
||||
|
||||
|
||||
def generate_grow_results_swap(image_bgr: np.ndarray, hair_styles: list[int] | None,
|
||||
redraw_defaults: dict):
|
||||
"""接口2 女性专用：发际线透明叠图（同 generate_grow_results）+ 换发型重绘图。
|
||||
redraw_defaults: dict,
|
||||
redraw_max_side: int | None = None,
|
||||
unet_name: str | None = None,
|
||||
prompt: str | None = None):
|
||||
"""接口2 女性专用：发际线透明叠图（同 generate_grow_results）+ 生发图。
|
||||
|
||||
grown 图来源（新流程）：对每个选中发型把 female key 映射到 change_hair 的 chang_* hair_id，
|
||||
调 face_analysis.hairline_grow.generate_hairline_redraw（= 接口12 final 管线，参数用
|
||||
redraw_defaults）拿到 ④ final（接缝融合基底）+ ⑤-② 纯红遮罩 PNG，再**后端直接调
|
||||
ComfyUI**（0716add-hair-api.json 工作流）完成发际线带重绘，重绘结果作为生发图。
|
||||
生发图来源按发型分两路：
|
||||
- 换发型（chang_* 表内，1..5）：female key → change_hair 的 chang_* hair_id，调
|
||||
face_analysis.hairline_grow.generate_hairline_redraw（= 接口12 final 管线，参数用
|
||||
redraw_defaults）拿到 ④ final（接缝融合基底）+ ⑤-② 纯红遮罩 PNG，再**后端直接调
|
||||
ComfyUI**（0716add-hair-api.json 工作流）完成发际线带重绘。
|
||||
- 原生生发（表外，6/7 bigflower/clasicalflower 等）：无对应 change_hair LoRA，改走与男性
|
||||
一致的原生生发（ComfyUI add_hair），由 _grow_native_one 完成（黑模板 + inpaint 遮罩）。
|
||||
|
||||
overlay 仍是发际线曲线透明层（与 generate_grow_results 完全一致）。
|
||||
prompt 仅用于原生生发分支（换发型分支的提示词由 redraw 流程内部固定）。
|
||||
Returns: list[dict] {"hairline_type","order","overlay","grown_png"(jpg bytes 或 None)}；
|
||||
无人脸返回 None。单个发型换发型/重绘失败时 grown_png=None，不抛异常。
|
||||
"""
|
||||
@@ -325,7 +348,10 @@ def generate_grow_results_swap(image_bgr: np.ndarray, hair_styles: list[int] | N
|
||||
grown_png = None
|
||||
chang_id = _FEMALE_KEY_TO_CHANG.get(key)
|
||||
if chang_id is None:
|
||||
logger.warning("接口2 换发型：female key=%s 无对应 chang_id，跳过生发图", key)
|
||||
# 无对应 change_hair LoRA（如 bigflower/clasicalflower）→ 走与男性一致的原生生发
|
||||
logger.info("接口2 女 key=%s 无 chang_id，走原生生发(ComfyUI add_hair)", key)
|
||||
grown_png = _grow_native_one(image_bgr, ctx, white_path,
|
||||
prompt=prompt, unet_name=unet_name)
|
||||
else:
|
||||
try:
|
||||
import time as _t
|
||||
@@ -394,6 +420,32 @@ def _grow_from_texture(image_bgr: np.ndarray, ctx: dict, white_path: str | None,
|
||||
return None
|
||||
|
||||
|
||||
def _grow_native_one(image_bgr: np.ndarray, ctx: dict, white_path: str,
|
||||
prompt: str | None = None, unet_name: str | None = None):
|
||||
"""对单个发际线做原生生发（ComfyUI add_hair），与男性 generate_grow_results 一致。
|
||||
|
||||
接口2 女性新发型（bigflower/clasicalflower 等）无对应 change_hair LoRA，改走此路径：
|
||||
黑模板 → build_inpaint_mask → 限边（_prep_comfy_input）→ comfyui.run(front=True)。
|
||||
失败返回 None，不抛异常。结果按限边前原图尺寸放大回原尺寸（仅展示对齐）。
|
||||
"""
|
||||
try:
|
||||
h, w = image_bgr.shape[:2]
|
||||
black = load_texture_rgba(_black_texture_path(white_path))
|
||||
marked, mask = build_inpaint_mask(
|
||||
image_bgr, ctx["landmarks"], ctx["parse_map"], ctx["points"], black)
|
||||
m_s, msk_s, gsc = _prep_comfy_input(marked, mask)
|
||||
buf = io.BytesIO()
|
||||
compose_comfy_rgba(m_s, msk_s).save(buf, format="PNG", compress_level=1)
|
||||
# front=True：接口2 时延敏感，插到 ComfyUI 队列最前
|
||||
grown_png = comfyui.run(buf.getvalue(), prompt=prompt, front=True, unet_name=unet_name)
|
||||
if gsc < 1.0 and grown_png:
|
||||
grown_png = _upscale_png_to(grown_png, w, h)
|
||||
return grown_png
|
||||
except Exception as e: # noqa: BLE001 单张失败不拖垮整请求
|
||||
logger.warning("接口2 女原生生发图失败：%s", e)
|
||||
return None
|
||||
|
||||
|
||||
def generate_hairline_pngs(image_bgr: np.ndarray, gender: str,
|
||||
hair_styles: list[int], use_mask: bool = True,
|
||||
prompt: str | None = None,
|
||||
|
||||
Reference in New Issue
Block a user