save code

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)
+        # Poll for completion (5 min timeout, 0.1s interval)
-        for _ in range(150):
+        for _ in range(3000):
-            time.sleep(2)
+            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__":

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")

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)

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)

local_test/index.html

@@ -136,7 +136,13 @@ generateBtn.addEventListener('click', async () => {
 
     generateBtn.disabled = true;
     generateBtn.textContent = '⏳ 生成中...';
-    resultArea.innerHTML = '<div class="loading"><div class="spinner"></div><br>正在调用 ComfyUI 生成，请耐心等待...</div>';
+
+    const startTime = performance.now();
+    resultArea.innerHTML = '<div class="loading"><div class="spinner"></div><br>正在调用 ComfyUI 生成，请耐心等待...<div id="liveTimer" style="margin-top:8px;font-size:15px;color:#aaa">已用时 0.0s</div></div>';
+    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 = `
+            <div style="text-align:center;margin-bottom:12px;color:#4ade80;font-size:16px;font-weight:600">
+                ⏱️ 本次重绘耗时 ${elapsed}s${serverInfo}
+            </div>
             <div class="result-area">
                 <div class="result-item">
                     <h3>原图</h3>
@@ -166,8 +179,10 @@ generateBtn.addEventListener('click', async () => {
             </div>
         `;
     } catch (err) {
-        resultArea.innerHTML = `<div class="error">❌ ${err.message}</div>`;
+        const elapsed = ((performance.now() - startTime) / 1000).toFixed(1);
+        resultArea.innerHTML = `<div class="error">❌ ${err.message}<br><span style="color:#aaa;font-size:13px">（耗时 ${elapsed}s）</span></div>`;
     } finally {
+        clearInterval(timerId);
         generateBtn.disabled = false;
         generateBtn.textContent = '🚀 生成';
     }