save code

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)