77 lines
2.7 KiB
Python
77 lines
2.7 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) -> bytes:
|
||
"""直接调 ComfyUI 重绘 — 替代 local_test /api/generate。
|
||
|
||
Args:
|
||
image_bytes: 人物图片字节（JPG/PNG）
|
||
mask_bytes: 遮罩图片字节（支持红/白/alpha 遮罩格式）
|
||
prompt: 提示词，None 用默认 "填充遮罩区域的头发，皮肤加一点磨皮"
|
||
timeout: ComfyUI 超时秒数
|
||
|
||
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)
|