xsl
b61ea6f33b
接口5改造（生发机制与接口2一致，按性别分流）： - generate_hairline_pngs: 生发图来源从 _grow_from_texture(局部inpaint) 改为按性别分流——female走 generate_grow_results_swap(swapHair+Flux整帧重绘)， male走 generate_grow_results(ComfyUI add_hair) - 新增参数 redraw_max_side/unet_name/v2_defaults(female路径透传) - 接口5 handler 加 flux_model/redraw_max_side Form参数 - 保留接口5独有输出: 3档叠图(middle/high/low)/中心点/face_measure - 已验证: female日志出现"接口2女 管线降分辨率max_side=640"+swap+Flux; male走add_hair.json; generate_grow_image=false正确跳过生发 统一ComfyUI重绘提示词: - "填充遮罩区域的头发，皮肤加一点磨皮，再加一点美颜" → "填充遮罩区域的头发" - 覆盖: _REDRAW_PROMPT/_DEFAULT_PROMPT常量 + app.py各接口Form默认 + 工作流JSON节点60(add_hair/0716add-hair-api/hair_repaint) + 全部测试页输入框
81 lines
3.0 KiB
Python
81 lines
3.0 KiB
Python
"""直接调 ComfyUI 重绘 — 替代 local_test HTTP 服务。
|
||
|
||
将 local_test/app.py 的核心逻辑（遮罩处理 + ComfyUI 调用）提取为 Python 函数，
|
||
不再需要独立 Flask 服务。使用 0716add-hair-api.json 工作流（steps=4）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
import logging
|
||
import os
|
||
|
||
import numpy as np
|
||
from PIL import Image, ImageFilter
|
||
|
||
from . import comfyui
|
||
|
||
logger = logging.getLogger("hair.worker")
|
||
|
||
_DEFAULT_PROMPT = "填充遮罩区域的头发"
|
||
_REPO = os.path.dirname(os.path.dirname(__file__))
|
||
_REPAINT_WORKFLOW = os.path.join(_REPO, "0716add-hair-api.json")
|
||
|
||
|
||
def _process_mask_to_rgba(image_bytes: bytes, mask_bytes: bytes) -> bytes:
|
||
"""将分开的 image + mask 处理为 ComfyUI 用的 RGBA PNG bytes。
|
||
|
||
复制 local_test/app.py 的遮罩处理逻辑：
|
||
1. 加载 image 为 RGB
|
||
2. 加载 mask 为 RGBA，取所有通道 max 值（支持红/白/alpha 遮罩）
|
||
3. resize mask 到与 image 一致
|
||
4. 高斯模糊(radius=4) 柔化边缘
|
||
5. alpha = 255 - mask（绘制区=255 → alpha=0 → 重绘区）
|
||
6. 合成 RGBA PNG
|
||
"""
|
||
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
||
mask_img = Image.open(io.BytesIO(mask_bytes)).convert("RGBA")
|
||
mask_arr = np.array(mask_img)
|
||
mask_data = np.max(mask_arr, axis=2) # (H, W) uint8
|
||
|
||
mask_data_img = Image.fromarray(mask_data, mode="L")
|
||
if mask_data_img.size != image.size:
|
||
mask_data_img = mask_data_img.resize(image.size, Image.LANCZOS)
|
||
mask_data_img = mask_data_img.filter(ImageFilter.GaussianBlur(radius=4))
|
||
|
||
# ComfyUI LoadImage: mask = 1.0 - (alpha/255)
|
||
# alpha=0 -> mask=1.0 (inpaint), alpha=255 -> mask=0.0 (keep)
|
||
comfyui_alpha = Image.eval(mask_data_img, lambda x: 255 - x)
|
||
|
||
r, g, b = image.split()
|
||
rgba = Image.merge("RGBA", (r, g, b, comfyui_alpha))
|
||
|
||
buf = io.BytesIO()
|
||
rgba.save(buf, format="PNG")
|
||
return buf.getvalue()
|
||
|
||
|
||
def run_redraw(image_bytes: bytes, mask_bytes: bytes,
|
||
prompt: str | None = None, timeout: float = 300.0,
|
||
front: bool = False, unet_name: str | None = None) -> bytes:
|
||
"""直接调 ComfyUI 重绘 — 替代 local_test /api/generate。
|
||
|
||
Args:
|
||
image_bytes: 人物图片字节（JPG/PNG）
|
||
mask_bytes: 遮罩图片字节（支持红/白/alpha 遮罩格式）
|
||
prompt: 提示词，None 用默认 "填充遮罩区域的头发"
|
||
timeout: ComfyUI 超时秒数
|
||
front: True 时任务插到 ComfyUI 队列最前（接口2 时延敏感路径用）
|
||
unet_name: 非 None 时切换 Flux 模型（如 flux-2-klein-9b-Q5_K_M.gguf），None 用工作流默认
|
||
|
||
Returns:
|
||
重绘后的 PNG 图片字节
|
||
|
||
Raises:
|
||
RuntimeError: ComfyUI 执行失败
|
||
TimeoutError: ComfyUI 超时
|
||
"""
|
||
rgba_png = _process_mask_to_rgba(image_bytes, mask_bytes)
|
||
return comfyui.run(rgba_png, timeout=timeout, prompt=prompt,
|
||
workflow_path=_REPAINT_WORKFLOW, front=front,
|
||
unet_name=unet_name)
|