xsl
5bcaeb2594
测试结论: 纯生发提示词(不做皮肤处理)效果最佳。 - 替换所有位置的提示词默认值(原"填充遮罩区域的头发,皮肤加一点磨皮,再加一点美颜"等) - 涉及25个文件: app.py/hairline/service.py/hairline/redraw.py/工作流json/测试页/benchmark脚本/local_test - _REDRAW_PROMPT / _DEFAULT_PROMPT / 各接口Form默认值 / 测试页输入框默认值 全部统一
340 lines
12 KiB
Python
340 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
"""Hair inpainting service - calls ComfyUI workflow with image + mask."""
|
||
import io
|
||
import json
|
||
import time
|
||
import random
|
||
import logging
|
||
import traceback
|
||
import requests
|
||
from flask import Flask, request, jsonify, send_file
|
||
import numpy as np
|
||
from PIL import Image, ImageFilter
|
||
|
||
# 让 PIL 支持 iPhone 的 HEIC/HEIF 照片（浏览器 accept="image/*" 会允许选中它们）。
|
||
try:
|
||
from pillow_heif import register_heif_opener
|
||
register_heif_opener()
|
||
_HEIF_OK = True
|
||
except Exception:
|
||
_HEIF_OK = False
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||
)
|
||
log = logging.getLogger("hair")
|
||
|
||
app = Flask(__name__)
|
||
COMFYUI_URL = "http://127.0.0.1:8188"
|
||
|
||
# 允许浏览器跨域直连本服务（如 hair 项目的测试页）。
|
||
# 不引入 flask-cors 依赖，直接在响应头 + OPTIONS 预检里处理。
|
||
_CORS_HEADERS = {
|
||
"Access-Control-Allow-Origin": "*",
|
||
"Access-Control-Allow-Methods": "POST, OPTIONS",
|
||
"Access-Control-Allow-Headers": "Content-Type",
|
||
"Access-Control-Expose-Headers": "X-Generate-Time",
|
||
}
|
||
|
||
|
||
@app.after_request
|
||
def _add_cors(resp):
|
||
for k, v in _CORS_HEADERS.items():
|
||
resp.headers[k] = v
|
||
return resp
|
||
|
||
|
||
@app.route("/api/generate", methods=["OPTIONS"])
|
||
def _generate_preflight():
|
||
"""CORS 预检：浏览器 POST 前会先发 OPTIONS。"""
|
||
return ("", 204)
|
||
|
||
|
||
def build_workflow(image_filename, prompt_text, seed=None):
|
||
"""Build ComfyUI API workflow from the 0716add-hair.json structure."""
|
||
if seed is None:
|
||
seed = random.randint(0, 2**53)
|
||
|
||
return {
|
||
# Loaders
|
||
"16": {"class_type": "UNETLoader", "inputs": {
|
||
"unet_name": "flux2.0/flux-2-klein-9b-fp8.safetensors",
|
||
"weight_dtype": "fp8_e4m3fn_fast"}},
|
||
"3": {"class_type": "VAELoader", "inputs": {
|
||
"vae_name": "flux2-vae.safetensors"}},
|
||
"61": {"class_type": "CLIPLoader", "inputs": {
|
||
"clip_name": "qwen_3_8b_fp8mixed.safetensors",
|
||
"type": "flux2",
|
||
"device": "default"}},
|
||
|
||
# Input image (with mask in alpha channel)
|
||
"26": {"class_type": "LoadImage", "inputs": {
|
||
"image": image_filename}},
|
||
|
||
# Prompt
|
||
"60": {"class_type": "JjkText", "inputs": {
|
||
"text": prompt_text}},
|
||
"22": {"class_type": "CLIPTextEncode", "inputs": {
|
||
"clip": ["61", 0],
|
||
"text": ["60", 0]}},
|
||
|
||
# Image size
|
||
"31": {"class_type": "easy imageSize", "inputs": {
|
||
"image": ["26", 0]}},
|
||
|
||
# Mask processing: fill holes -> convert to image -> scale -> back to mask
|
||
"33": {"class_type": "Mask Fill Holes", "inputs": {
|
||
"masks": ["26", 1]}},
|
||
"36": {"class_type": "Convert Masks to Images", "inputs": {
|
||
"masks": ["33", 0]}},
|
||
"39": {"class_type": "ImageScale", "inputs": {
|
||
"image": ["36", 0],
|
||
"upscale_method": "nearest-exact",
|
||
"width": ["31", 0],
|
||
"height": ["31", 1],
|
||
"crop": "disabled"}},
|
||
"37": {"class_type": "Image To Mask", "inputs": {
|
||
"image": ["39", 0],
|
||
"method": "intensity"}},
|
||
|
||
# Scale image+mask by aspect ratio
|
||
"32": {"class_type": "LayerUtility: ImageScaleByAspectRatio V2", "inputs": {
|
||
"image": ["26", 0],
|
||
"mask": ["37", 0],
|
||
"aspect_ratio": "custom",
|
||
"proportional_width": ["31", 0],
|
||
"proportional_height": ["31", 1],
|
||
"fit": "letterbox",
|
||
"method": "lanczos",
|
||
"round_to_multiple": "8",
|
||
"scale_to_side": "None",
|
||
"scale_to_length": 1024,
|
||
"background_color": "#000000"}},
|
||
|
||
# Preview (pass_through=true, just passes the image through)
|
||
"44": {"class_type": "ImageAndMaskPreview", "inputs": {
|
||
"image": ["32", 0],
|
||
"mask": ["32", 1],
|
||
"mask_opacity": 1,
|
||
"mask_color": "FFFF00",
|
||
"pass_through": True}},
|
||
|
||
# Get size of scaled image
|
||
"14": {"class_type": "GetImageSize+", "inputs": {
|
||
"image": ["44", 0]}},
|
||
|
||
# VAE encode the image
|
||
"13": {"class_type": "VAEEncode", "inputs": {
|
||
"pixels": ["44", 0],
|
||
"vae": ["3", 0]}},
|
||
|
||
# Flux model setup
|
||
"2": {"class_type": "ModelSamplingFlux", "inputs": {
|
||
"model": ["16", 0],
|
||
"max_shift": 1.15,
|
||
"base_shift": 0.5,
|
||
"width": ["14", 0],
|
||
"height": ["14", 1]}},
|
||
"19": {"class_type": "FluxGuidance", "inputs": {
|
||
"conditioning": ["22", 0],
|
||
"guidance": 1}},
|
||
"5": {"class_type": "ReferenceLatent", "inputs": {
|
||
"conditioning": ["19", 0],
|
||
"latent": ["13", 0]}},
|
||
|
||
# Empty latent for sampling
|
||
"7": {"class_type": "EmptySD3LatentImage", "inputs": {
|
||
"width": ["14", 0],
|
||
"height": ["14", 1],
|
||
"batch_size": 1}},
|
||
|
||
# Scheduler & guider
|
||
"1": {"class_type": "BasicScheduler", "inputs": {
|
||
"model": ["2", 0],
|
||
"scheduler": "simple",
|
||
"steps": 4,
|
||
"denoise": 1}},
|
||
"20": {"class_type": "BasicGuider", "inputs": {
|
||
"model": ["2", 0],
|
||
"conditioning": ["5", 0]}},
|
||
|
||
# Noise & sampler
|
||
"6": {"class_type": "RandomNoise", "inputs": {
|
||
"noise_seed": seed}},
|
||
"8": {"class_type": "KSamplerSelect", "inputs": {
|
||
"sampler_name": "euler"}},
|
||
"9": {"class_type": "SamplerCustomAdvanced", "inputs": {
|
||
"noise": ["6", 0],
|
||
"guider": ["20", 0],
|
||
"sampler": ["8", 0],
|
||
"sigmas": ["1", 0],
|
||
"latent_image": ["7", 0]}},
|
||
|
||
# VAE decode
|
||
"10": {"class_type": "VAEDecode", "inputs": {
|
||
"samples": ["9", 0],
|
||
"vae": ["3", 0]}},
|
||
|
||
# Color match with original image
|
||
"62": {"class_type": "ColorMatch", "inputs": {
|
||
"image_ref": ["26", 0],
|
||
"image_target": ["10", 0],
|
||
"method": "mkl",
|
||
"strength": 1,
|
||
"multithread": True}},
|
||
|
||
# Save result
|
||
"17": {"class_type": "SaveImage", "inputs": {
|
||
"images": ["62", 0],
|
||
"filename_prefix": "hair_inpaint"}},
|
||
}
|
||
|
||
|
||
@app.route("/")
|
||
def index():
|
||
return send_file("index.html")
|
||
|
||
|
||
@app.route("/api/generate", methods=["POST"])
|
||
def generate():
|
||
t_start = time.time()
|
||
try:
|
||
if "image" not in request.files or "mask" not in request.files:
|
||
msg = f"缺少上传文件: files={list(request.files.keys())}"
|
||
log.warning(msg)
|
||
return jsonify({"error": msg}), 400
|
||
image_file = request.files["image"]
|
||
mask_file = request.files["mask"]
|
||
prompt_text = request.form.get("prompt", "填充遮罩区域的头发")
|
||
log.info(
|
||
"收到请求: image=%s mask=%s prompt=%r",
|
||
image_file.filename, mask_file.filename, prompt_text,
|
||
)
|
||
|
||
# Load original image as RGB
|
||
try:
|
||
image = Image.open(image_file).convert("RGB")
|
||
except Exception as e:
|
||
log.error("无法解码人物图片 %s: %s", image_file.filename, e)
|
||
hint = "" if _HEIF_OK else "（当前不支持 HEIC）"
|
||
return jsonify({
|
||
"error": f"无法识别人物图片格式{hint}，请改用 JPG/PNG: {e}"
|
||
}), 400
|
||
|
||
# Load mask and extract mask data from ALL channels (R, G, B, A)
|
||
# This handles different mask formats:
|
||
# - Red mask (R=255 where drawn): user-provided PNG
|
||
# - White mask (R=G=B=255 where drawn): frontend canvas
|
||
# - Alpha mask (A=255 where drawn): transparent brush
|
||
try:
|
||
mask_img = Image.open(mask_file).convert("RGBA")
|
||
except Exception as e:
|
||
log.error("无法解码遮罩图片 %s: %s", mask_file.filename, e)
|
||
return jsonify({
|
||
"error": f"无法识别遮罩图片格式，请改用 JPG/PNG: {e}"
|
||
}), 400
|
||
mask_arr = np.array(mask_img)
|
||
# Use max of all channels: 255 where any color/alpha is drawn, 0 where empty
|
||
mask_data = np.max(mask_arr, axis=2) # (H, W) uint8
|
||
|
||
# Ensure mask matches image size
|
||
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)
|
||
|
||
# Apply slight blur for soft edges (similar to ComfyUI's brush)
|
||
mask_data_img = mask_data_img.filter(ImageFilter.GaussianBlur(radius=4))
|
||
|
||
# ComfyUI LoadImage: mask = 1.0 - (alpha/255)
|
||
# So alpha=0 -> mask=1.0 (inpaint), alpha=255 -> mask=0.0 (keep)
|
||
# We want: drawn area (mask_data=255) -> inpaint -> alpha=0
|
||
# undrawn area (mask_data=0) -> keep -> alpha=255
|
||
comfyui_alpha = Image.eval(mask_data_img, lambda x: 255 - x)
|
||
|
||
# Combine into RGBA (split RGB into separate channels first)
|
||
r, g, b = image.split()
|
||
rgba = Image.merge("RGBA", (r, g, b, comfyui_alpha))
|
||
|
||
# Upload to ComfyUI
|
||
img_bytes = io.BytesIO()
|
||
rgba.save(img_bytes, format="PNG")
|
||
img_bytes.seek(0)
|
||
|
||
upload_resp = requests.post(
|
||
f"{COMFYUI_URL}/upload/image",
|
||
files={"image": ("hair_input.png", img_bytes, "image/png")},
|
||
timeout=30,
|
||
)
|
||
upload_data = upload_resp.json()
|
||
if "name" not in upload_data:
|
||
log.error("ComfyUI 上传图片失败: %s", upload_data)
|
||
return jsonify({"error": f"Upload failed: {upload_data}"}), 500
|
||
image_filename = upload_data["name"]
|
||
|
||
# Build and queue workflow
|
||
workflow = build_workflow(image_filename, prompt_text)
|
||
prompt_resp = requests.post(
|
||
f"{COMFYUI_URL}/prompt",
|
||
json={"prompt": workflow},
|
||
timeout=30,
|
||
)
|
||
prompt_data = prompt_resp.json()
|
||
if "error" in prompt_data:
|
||
log.error(
|
||
"ComfyUI /prompt 校验失败: error=%s node_errors=%s",
|
||
prompt_data.get("error"), prompt_data.get("node_errors"),
|
||
)
|
||
return jsonify({"error": json.dumps(prompt_data["error"], ensure_ascii=False)}), 500
|
||
prompt_id = prompt_data["prompt_id"]
|
||
|
||
# Poll for completion (5 min timeout, 0.1s interval)
|
||
for _ in range(3000):
|
||
time.sleep(0.1)
|
||
history_resp = requests.get(
|
||
f"{COMFYUI_URL}/history/{prompt_id}", timeout=10
|
||
)
|
||
history_data = history_resp.json()
|
||
if prompt_id in history_data:
|
||
status = history_data[prompt_id].get("status", {})
|
||
if status.get("status_str") == "error":
|
||
log.error(
|
||
"ComfyUI 工作流执行失败: %s",
|
||
json.dumps(status, ensure_ascii=False),
|
||
)
|
||
return jsonify({
|
||
"error": "Workflow execution failed",
|
||
"detail": status.get("messages", status),
|
||
}), 500
|
||
outputs = history_data[prompt_id].get("outputs", {})
|
||
if "17" in outputs: # SaveImage node
|
||
image_info = outputs["17"]["images"][0]
|
||
filename = image_info["filename"]
|
||
subfolder = image_info.get("subfolder", "")
|
||
img_type = image_info.get("type", "output")
|
||
view_resp = requests.get(
|
||
f"{COMFYUI_URL}/view",
|
||
params={"filename": filename, "subfolder": subfolder, "type": img_type},
|
||
timeout=30,
|
||
)
|
||
elapsed = time.time() - t_start
|
||
log.info("重绘完成，服务端耗时 %.2fs", elapsed)
|
||
resp = send_file(
|
||
io.BytesIO(view_resp.content), mimetype="image/png"
|
||
)
|
||
resp.headers["X-Generate-Time"] = f"{elapsed:.2f}"
|
||
return resp
|
||
|
||
return jsonify({"error": "Timeout: workflow did not complete in 5 minutes"}), 500
|
||
|
||
except requests.ConnectionError:
|
||
log.error("无法连接 ComfyUI @ %s", COMFYUI_URL)
|
||
return jsonify({"error": "Cannot connect to ComfyUI at " + COMFYUI_URL + ". Is it running?"}), 503
|
||
except Exception as e:
|
||
log.error("生成失败，未捕获异常:\n%s", traceback.format_exc())
|
||
return jsonify({"error": f"{type(e).__name__}: {e}"}), 500
|
||
|
||
|
||
if __name__ == "__main__":
|
||
app.run(host="0.0.0.0", port=8899, debug=False)
|