262 lines
9.6 KiB
Python
262 lines
9.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Hair inpainting service - calls ComfyUI workflow with image + mask."""
|
|
import io
|
|
import json
|
|
import time
|
|
import random
|
|
import requests
|
|
from flask import Flask, request, jsonify, send_file
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter
|
|
|
|
app = Flask(__name__)
|
|
COMFYUI_URL = "http://127.0.0.1:8188"
|
|
|
|
|
|
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"}},
|
|
"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": 6,
|
|
"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():
|
|
try:
|
|
image_file = request.files["image"]
|
|
mask_file = request.files["mask"]
|
|
prompt_text = request.form.get("prompt", "填充遮罩区域的头发，皮肤加一点磨皮")
|
|
|
|
# Load original image as RGB
|
|
image = Image.open(image_file).convert("RGB")
|
|
|
|
# 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
|
|
mask_img = Image.open(mask_file).convert("RGBA")
|
|
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:
|
|
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:
|
|
return jsonify({"error": json.dumps(prompt_data["error"], ensure_ascii=False)}), 500
|
|
prompt_id = prompt_data["prompt_id"]
|
|
|
|
# Poll for completion (5 min timeout)
|
|
for _ in range(150):
|
|
time.sleep(2)
|
|
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":
|
|
return jsonify({"error": "Workflow execution failed"}), 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,
|
|
)
|
|
return send_file(
|
|
io.BytesIO(view_resp.content), mimetype="image/png"
|
|
)
|
|
|
|
return jsonify({"error": "Timeout: workflow did not complete in 5 minutes"}), 500
|
|
|
|
except requests.ConnectionError:
|
|
return jsonify({"error": "Cannot connect to ComfyUI at " + COMFYUI_URL + ". Is it running?"}), 503
|
|
except Exception as e:
|
|
return jsonify({"error": str(e)}), 500
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run(host="0.0.0.0", port=8899, debug=False)
|