diff --git a/local_test/app.py b/local_test/app.py index 81c3b3c..aec326e 100644 --- a/local_test/app.py +++ b/local_test/app.py @@ -4,11 +4,27 @@ 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" @@ -18,6 +34,7 @@ _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", } @@ -43,7 +60,7 @@ def build_workflow(image_filename, prompt_text, seed=None): # Loaders "16": {"class_type": "UNETLoader", "inputs": { "unet_name": "flux2.0/flux-2-klein-9b-fp8.safetensors", - "weight_dtype": "fp8_e4m3fn"}}, + "weight_dtype": "fp8_e4m3fn_fast"}}, "3": {"class_type": "VAELoader", "inputs": { "vae_name": "flux2-vae.safetensors"}}, "61": {"class_type": "CLIPLoader", "inputs": { @@ -136,7 +153,7 @@ def build_workflow(image_filename, prompt_text, seed=None): "1": {"class_type": "BasicScheduler", "inputs": { "model": ["2", 0], "scheduler": "simple", - "steps": 6, + "steps": 4, "denoise": 1}}, "20": {"class_type": "BasicGuider", "inputs": { "model": ["2", 0], @@ -181,20 +198,42 @@ def index(): @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 - image = Image.open(image_file).convert("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 - mask_img = Image.open(mask_file).convert("RGBA") + 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 @@ -229,6 +268,7 @@ def generate(): ) 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"] @@ -241,12 +281,16 @@ def generate(): ) 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) - for _ in range(150): - time.sleep(2) + # 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 ) @@ -254,7 +298,14 @@ def generate(): 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 + 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] @@ -266,16 +317,22 @@ def generate(): params={"filename": filename, "subfolder": subfolder, "type": img_type}, timeout=30, ) - return send_file( + 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: - return jsonify({"error": str(e)}), 500 + log.error("生成失败，未捕获异常:\n%s", traceback.format_exc()) + return jsonify({"error": f"{type(e).__name__}: {e}"}), 500 if __name__ == "__main__": diff --git a/local_test/bench.py b/local_test/bench.py new file mode 100644 index 0000000..0526ec3 --- /dev/null +++ b/local_test/bench.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +"""Benchmark ComfyUI hair-inpaint workflow across model / dtype / steps.""" +import io, time, sys +import requests +import numpy as np +from PIL import Image, ImageFilter +import app as A + +COMFY = "http://127.0.0.1:8188" + + +def prep_and_upload(): + image = Image.open("用来重绘.jpg").convert("RGB") + mask_img = Image.open("用来重绘.png").convert("RGBA") + mask_data = np.max(np.array(mask_img), axis=2) + m = Image.fromarray(mask_data, mode="L") + if m.size != image.size: + m = m.resize(image.size, Image.LANCZOS) + m = m.filter(ImageFilter.GaussianBlur(radius=4)) + alpha = Image.eval(m, lambda x: 255 - x) + r, g, b = image.split() + rgba = Image.merge("RGBA", (r, g, b, alpha)) + buf = io.BytesIO(); rgba.save(buf, format="PNG"); buf.seek(0) + up = requests.post(f"{COMFY}/upload/image", + files={"image": ("hair_input.png", buf, "image/png")}).json() + return up["name"] + + +def run_once(fname, model, dtype, steps): + wf = A.build_workflow(fname, "填充遮罩区域的头发，皮肤加一点磨皮") + wf["16"]["inputs"]["unet_name"] = model + wf["16"]["inputs"]["weight_dtype"] = dtype + wf["1"]["inputs"]["steps"] = steps + r = requests.post(f"{COMFY}/prompt", json={"prompt": wf}).json() + if "prompt_id" not in r: + raise RuntimeError(f"submit failed: {str(r)[:300]}") + pid = r["prompt_id"] + deadline = time.time() + 180 + while time.time() < deadline: + time.sleep(0.1) + h = requests.get(f"{COMFY}/history/{pid}").json() + if pid not in h: + continue + st = h[pid].get("status", {}) + if st.get("status_str") == "error": + for m in st.get("messages", []): + if m[0] == "execution_error": + raise RuntimeError(str(m[1])[:300]) + raise RuntimeError("execution error") + if "17" in h[pid].get("outputs", {}): + ts = {mm[0]: mm[1].get("timestamp") for mm in st["messages"]} + return (ts["execution_success"] - ts["execution_start"]) / 1000.0 + raise TimeoutError("run exceeded 180s") + + +CONFIGS = [ + ("flux2.0/flux-2-klein-9b-fp8.safetensors", "fp8_e4m3fn", 6, "9B fp8 (当前)"), + ("flux2.0/flux-2-klein-9b-fp8.safetensors", "fp8_e4m3fn_fast", 6, "9B fp8-fast"), + ("flux2.0/flux-2-klein-9b-fp8.safetensors", "fp8_e4m3fn_fast", 4, "9B fp8-fast s4"), + ("flux-2-klein-4b-fp8.safetensors", "fp8_e4m3fn", 6, "4B fp8"), + ("flux-2-klein-4b-fp8.safetensors", "fp8_e4m3fn_fast", 6, "4B fp8-fast"), + ("flux-2-klein-4b-fp8.safetensors", "fp8_e4m3fn_fast", 4, "4B fp8-fast s4"), +] + +fname = prep_and_upload() +print("input uploaded:", fname) +print(f"{'配置':<22}{'warmup':>10}{'run1':>10}{'run2':>10}{'best':>10}") +for model, dtype, steps, label in CONFIGS: + times = [] + for i in range(3): # 1 warmup + 2 measured + try: + t = run_once(fname, model, dtype, steps) + except Exception as e: + t = float('nan'); print("ERR", label, e) + times.append(t) + best = min(times[1:]) + print(f"{label:<22}{times[0]:>9.2f}s{times[1]:>9.2f}s{times[2]:>9.2f}s{best:>9.2f}s") diff --git a/local_test/bench2.py b/local_test/bench2.py new file mode 100644 index 0000000..aa72174 --- /dev/null +++ b/local_test/bench2.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""Steps sweep + resolution test on the working 9B fp8-fast config.""" +import io, time +import requests +import numpy as np +from PIL import Image, ImageFilter +import app as A + +COMFY = "http://127.0.0.1:8188" +MODEL = "flux2.0/flux-2-klein-9b-fp8.safetensors" +DTYPE = "fp8_e4m3fn_fast" + + +def upload(scale=1.0): + image = Image.open("用来重绘.jpg").convert("RGB") + mask_img = Image.open("用来重绘.png").convert("RGBA") + if scale != 1.0: + w, h = image.size + w, h = int(w * scale) // 8 * 8, int(h * scale) // 8 * 8 + image = image.resize((w, h), Image.LANCZOS) + mask_data = np.max(np.array(mask_img), axis=2) + m = Image.fromarray(mask_data, mode="L") + if m.size != image.size: + m = m.resize(image.size, Image.LANCZOS) + m = m.filter(ImageFilter.GaussianBlur(radius=4)) + alpha = Image.eval(m, lambda x: 255 - x) + r, g, b = image.split() + rgba = Image.merge("RGBA", (r, g, b, alpha)) + buf = io.BytesIO(); rgba.save(buf, format="PNG"); buf.seek(0) + up = requests.post(f"{COMFY}/upload/image", + files={"image": ("hair_input.png", buf, "image/png")}).json() + return up["name"], image.size + + +def run(fname, steps): + wf = A.build_workflow(fname, "填充遮罩区域的头发，皮肤加一点磨皮") + wf["16"]["inputs"]["unet_name"] = MODEL + wf["16"]["inputs"]["weight_dtype"] = DTYPE + wf["1"]["inputs"]["steps"] = steps + pid = requests.post(f"{COMFY}/prompt", json={"prompt": wf}).json()["prompt_id"] + dl = time.time() + 120 + while time.time() < dl: + time.sleep(0.1) + h = requests.get(f"{COMFY}/history/{pid}").json() + if pid in h and "17" in h[pid].get("outputs", {}): + ts = {m[0]: m[1].get("timestamp") for m in h[pid]["status"]["messages"]} + return (ts["execution_success"] - ts["execution_start"]) / 1000.0 + return float("nan") + + +print("=== 步数扫描 (9B fp8-fast, 原分辨率 1024x775) ===", flush=True) +fname, sz = upload(1.0) +run(fname, 6) # warmup +res = {} +for s in [2, 3, 4, 6, 8]: + t = min(run(fname, s), run(fname, s)) + res[s] = t + print(f" steps={s}: {t:.2f}s", flush=True) +# derive per-step cost & fixed overhead via two points +per = (res[8] - res[2]) / 6 +fixed = res[2] - per * 2 +print(f" -> 每步 ~{per:.3f}s, 固定开销(VAE/编码/colormatch/加载) ~{fixed:.2f}s", flush=True) + +print("\n=== 分辨率影响 (steps=4) ===", flush=True) +for scale in [1.0, 0.75, 0.6]: + fn, s2 = upload(scale) + run(fn, 4) # warmup + t = min(run(fn, 4), run(fn, 4)) + print(f" {s2[0]}x{s2[1]} ({s2[0]*s2[1]/1e6:.2f}MP): {t:.2f}s", flush=True) diff --git a/local_test/bench_quality.py b/local_test/bench_quality.py new file mode 100644 index 0000000..953b0d9 --- /dev/null +++ b/local_test/bench_quality.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""Generate result images at different step counts for quality comparison.""" +import io, time +import requests +import numpy as np +from PIL import Image, ImageFilter +import app as A + +COMFY = "http://127.0.0.1:8188" +MODEL = "flux2.0/flux-2-klein-9b-fp8.safetensors" +DTYPE = "fp8_e4m3fn_fast" +OUT = "output" + +image = Image.open("用来重绘.jpg").convert("RGB") +mask_img = Image.open("用来重绘.png").convert("RGBA") +mask_data = np.max(np.array(mask_img), axis=2) +m = Image.fromarray(mask_data, mode="L") +if m.size != image.size: + m = m.resize(image.size, Image.LANCZOS) +m = m.filter(ImageFilter.GaussianBlur(radius=4)) +alpha = Image.eval(m, lambda x: 255 - x) +r, g, b = image.split() +rgba = Image.merge("RGBA", (r, g, b, alpha)) +buf = io.BytesIO(); rgba.save(buf, format="PNG"); buf.seek(0) +fname = requests.post(f"{COMFY}/upload/image", + files={"image": ("hair_input.png", buf, "image/png")}).json()["name"] + +# fixed seed for fair comparison +SEED = 123456789 +imgs = [] +labels = [] +for steps in [2, 3, 4, 6]: + wf = A.build_workflow(fname, "填充遮罩区域的头发，皮肤加一点磨皮", seed=SEED) + wf["16"]["inputs"]["unet_name"] = MODEL + wf["16"]["inputs"]["weight_dtype"] = DTYPE + wf["1"]["inputs"]["steps"] = steps + pid = requests.post(f"{COMFY}/prompt", json={"prompt": wf}).json()["prompt_id"] + t0 = time.time() + while True: + time.sleep(0.1) + h = requests.get(f"{COMFY}/history/{pid}").json() + if pid in h and "17" in h[pid].get("outputs", {}): + ts = {mm[0]: mm[1].get("timestamp") for mm in h[pid]["status"]["messages"]} + dur = (ts["execution_success"] - ts["execution_start"]) / 1000.0 + info = h[pid]["outputs"]["17"]["images"][0] + data = requests.get(f"{COMFY}/view", params={ + "filename": info["filename"], "subfolder": info.get("subfolder", ""), + "type": info.get("type", "output")}).content + im = Image.open(io.BytesIO(data)).convert("RGB") + imgs.append(im) + labels.append(f"steps={steps} {dur:.2f}s") + print(f"steps={steps}: {dur:.2f}s", flush=True) + break + +# build side-by-side contact sheet +from PIL import ImageDraw +h0 = imgs[0].height +w0 = imgs[0].width +pad = 10 +bar = 28 +sheet = Image.new("RGB", (w0 * len(imgs) + pad * (len(imgs) + 1), + h0 + bar + pad * 2), (26, 26, 46)) +d = ImageDraw.Draw(sheet) +for i, (im, lb) in enumerate(zip(imgs, labels)): + x = pad + i * (w0 + pad) + sheet.paste(im, (x, bar + pad)) + d.text((x + 4, 6), lb, fill=(233, 69, 96)) +sheet.save(f"{OUT}/compare_steps.png") +print("saved:", f"{OUT}/compare_steps.png", flush=True) diff --git a/local_test/index.html b/local_test/index.html index 6c99a37..61c5504 100644 --- a/local_test/index.html +++ b/local_test/index.html @@ -136,7 +136,13 @@ generateBtn.addEventListener('click', async () => { generateBtn.disabled = true; generateBtn.textContent = '⏳ 生成中...'; - resultArea.innerHTML = '

正在调用 ComfyUI 生成，请耐心等待...
'; + + const startTime = performance.now(); + resultArea.innerHTML = '

正在调用 ComfyUI 生成，请耐心等待...
已用时 0.0s
'; + const liveTimer = document.getElementById('liveTimer'); + const timerId = setInterval(() => { + if (liveTimer) liveTimer.textContent = '已用时 ' + ((performance.now() - startTime) / 1000).toFixed(1) + 's'; + }, 100); try { const formData = new FormData(); @@ -152,8 +158,15 @@ generateBtn.addEventListener('click', async () => { const resultBlob = await resp.blob(); const resultUrl = URL.createObjectURL(resultBlob); + const elapsed = ((performance.now() - startTime) / 1000).toFixed(1); + // 服务端纯推理耗时（若返回该响应头） + const serverTime = resp.headers.get('X-Generate-Time'); + const serverInfo = serverTime ? `，服务端推理 ${parseFloat(serverTime).toFixed(1)}s` : ''; resultArea.innerHTML = ` +
+ ⏱️ 本次重绘耗时 ${elapsed}s${serverInfo} +

原图

@@ -166,8 +179,10 @@ generateBtn.addEventListener('click', async () => {
`; } catch (err) { - resultArea.innerHTML = `
❌ ${err.message}
`; + const elapsed = ((performance.now() - startTime) / 1000).toFixed(1); + resultArea.innerHTML = `
❌ ${err.message}
（耗时 ${elapsed}s）
`; } finally { + clearInterval(timerId); generateBtn.disabled = false; generateBtn.textContent = '🚀 生成'; } diff --git a/local_test/output/compare_hairline.png b/local_test/output/compare_hairline.png new file mode 100644 index 0000000..7c0fbb1 Binary files /dev/null and b/local_test/output/compare_hairline.png differ diff --git a/local_test/output/compare_steps.png b/local_test/output/compare_steps.png new file mode 100644 index 0000000..d83711d Binary files /dev/null and b/local_test/output/compare_steps.png differ diff --git a/local_test/output/result.png b/local_test/output/result.png index e2a8c8d..4c0d5f9 100644 Binary files a/local_test/output/result.png and b/local_test/output/result.png differ