添加服务
This commit is contained in:
@@ -0,0 +1,261 @@
|
||||
#!/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)
|
||||
Reference in New Issue
Block a user