#!/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" # 允许浏览器跨域直连本服务（如 hair 项目的测试页）。 # 不引入 flask-cors 依赖，直接在响应头 + OPTIONS 预检里处理。 _CORS_HEADERS = { "Access-Control-Allow-Origin": "*", "Access-Control-Allow-Methods": "POST, OPTIONS", "Access-Control-Allow-Headers": "Content-Type", } @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"}}, "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)