#!/usr/bin/env python3
"""本地开源图像模型横评：同一提示词、同一尺寸，跑 Krea2 / Boogu / Z-Image / Qwen-Image-2.1。

用法: python3 gen_local.py <prompt.txt> <outdir> [--models krea2,boogu-turbo,...] [--width 1344] [--height 768]
"""
import json, os, sys, time, urllib.request, urllib.parse

HOST = "http://192.168.31.31:8189"
OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({}))

TE_QWEN8 = "qwen3vl_8b_fp8_scaled.safetensors"
TE_QWEN4 = "qwen3vl_4b_fp8_scaled.safetensors"

NEG = ("text, watermark, logo, signature, subtitles, letters, lowres, blurry, jpeg artifacts, "
       "extra fingers, fused fingers, bad anatomy, deformed face, plastic skin, modern city, "
       "modern clothes, photo, dslr snapshot, duplicate face, face in the clouds, extra limbs")


def graph(model, prompt, w, h, seed, steps=None):
    """返回 (workflow, 实际步数, 备注)"""
    if model == "krea2":
        st = steps or 9
        wf = {
            "1": {"class_type": "CLIPLoader", "inputs": {"clip_name": TE_QWEN4, "type": "krea2", "device": "default"}},
            "2": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["1", 0]}},
            "5": {"class_type": "UNETLoader", "inputs": {"unet_name": "krea2_turbo_fp8_scaled.safetensors", "weight_dtype": "default"}},
            "7": {"class_type": "VAELoader", "inputs": {"vae_name": "qwen_image_vae.safetensors"}},
            "4": {"class_type": "EmptySD3LatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}},
            "6": {"class_type": "KSampler", "inputs": {"model": ["5", 0], "positive": ["2", 0], "negative": ["3", 0],
                                                       "latent_image": ["4", 0], "seed": seed, "steps": st, "cfg": 1.0,
                                                       "sampler_name": "er_sde", "scheduler": "simple", "denoise": 1.0}},
            "8": {"class_type": "VAEDecode", "inputs": {"samples": ["6", 0], "vae": ["7", 0]}},
            "9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "sweep-krea2", "images": ["8", 0]}},
        }
        wf["3"] = {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["2", 0]}}
        return wf, st, "er_sde/simple cfg1"

    if model.startswith("boogu"):
        turbo = model == "boogu-turbo"
        unet = "boogu_image_turbo_hotfix_nvfp4.safetensors" if turbo else "boogu_image_base_fp8_scaled.safetensors"
        st = steps or (4 if turbo else 30)
        cfg = 1.0 if turbo else 4.0
        sampler = "lcm" if turbo else "dpmpp_2m"
        sched = "sgm_uniform" if turbo else "simple"
        shift = 3.0 if turbo else 3.16
        wf = {
            "1": {"class_type": "UNETLoader", "inputs": {"unet_name": unet, "weight_dtype": "default"}},
            "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": TE_QWEN8, "type": "boogu", "device": "default"}},
            "3": {"class_type": "ModelSamplingAuraFlow", "inputs": {"model": ["1", 0], "shift": shift}},
            "5": {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}},
            "6": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": prompt}},
            "7": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": "" if turbo else NEG}},
            "8": {"class_type": "EmptyLatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}},
            "9": {"class_type": "KSampler", "inputs": {"model": ["3", 0], "seed": seed, "steps": st, "cfg": cfg,
                                                       "sampler_name": sampler, "scheduler": sched, "denoise": 1.0,
                                                       "positive": ["6", 0], "negative": ["7", 0], "latent_image": ["8", 0]}},
            "10": {"class_type": "VAEDecode", "inputs": {"samples": ["9", 0], "vae": ["5", 0]}},
            "11": {"class_type": "SaveImage", "inputs": {"images": ["10", 0], "filename_prefix": "sweep-" + model}},
        }
        if turbo:
            wf["7"] = {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["6", 0]}}
        return wf, st, f"{sampler}/{sched} cfg{cfg} shift{shift}"

    if model == "z-image":
        st = steps or 8
        wf = {
            "1": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen_3_4b.safetensors", "type": "lumina2", "device": "default"}},
            "2": {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}},
            "3": {"class_type": "UNETLoader", "inputs": {"unet_name": "z_image_turbo_bf16.safetensors", "weight_dtype": "default"}},
            "4": {"class_type": "ModelSamplingAuraFlow", "inputs": {"model": ["3", 0], "shift": 3.0}},
            "5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["1", 0], "text": prompt}},
            "6": {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["5", 0]}},
            "7": {"class_type": "EmptySD3LatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}},
            "8": {"class_type": "KSampler", "inputs": {"model": ["4", 0], "positive": ["5", 0], "negative": ["6", 0],
                                                       "latent_image": ["7", 0], "seed": seed, "steps": st, "cfg": 1.0,
                                                       "sampler_name": "res_multistep", "scheduler": "simple", "denoise": 1.0}},
            "9": {"class_type": "VAEDecode", "inputs": {"samples": ["8", 0], "vae": ["2", 0]}},
            "10": {"class_type": "SaveImage", "inputs": {"images": ["9", 0], "filename_prefix": "sweep-z-image"}},
        }
        return wf, st, "res_multistep/simple cfg1 shift3"

    if model == "qwen21":
        st = steps or 25
        wf = {
            "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "qwen_image_2.1_int8_convrot.safetensors", "weight_dtype": "default"}},
            "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen3vl_8b_int8_convrot.safetensors", "type": "qwen_image", "device": "default"}},
            "3": {"class_type": "VAELoader", "inputs": {"vae_name": "qwen_image_2.1_vae_bf16.safetensors"}},
            "4": {"class_type": "TextEncodeQwenImage21", "inputs": {"clip": ["2", 0], "prompt": prompt,
                                                                   "negative_prompt": NEG, "resolution": 1024, "vae": ["3", 0]}},
            "5": {"class_type": "EmptyLatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}},
            "6": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["4", 0], "negative": ["4", 1],
                                                       "latent_image": ["5", 0], "seed": seed, "steps": st, "cfg": 1.0,
                                                       "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0}},
            "7": {"class_type": "VAEDecode", "inputs": {"samples": ["6", 0], "vae": ["3", 0]}},
            "8": {"class_type": "SaveImage", "inputs": {"filename_prefix": "sweep-qwen21", "images": ["7", 0]}},
        }
        return wf, st, "euler/simple cfg1"

    raise SystemExit("unknown model: " + model)


def post(path, obj, timeout=180):
    req = urllib.request.Request(HOST + path, data=json.dumps(obj).encode(),
                                 headers={"Content-Type": "application/json"})
    return json.load(OPENER.open(req, timeout=timeout))


def get_json(path, timeout=180):
    return json.load(OPENER.open(HOST + path, timeout=timeout))


def main():
    prompt = open(sys.argv[1], encoding="utf-8").read().strip()
    outdir = sys.argv[2]
    models = ["krea2", "boogu-turbo", "boogu-base", "z-image", "qwen21"]
    w, h, seed = 1344, 768, 20260928
    for i, a in enumerate(sys.argv):
        if a == "--models":
            models = sys.argv[i + 1].split(",")
        if a == "--width":
            w = int(sys.argv[i + 1])
        if a == "--height":
            h = int(sys.argv[i + 1])
        if a == "--seed":
            seed = int(sys.argv[i + 1])
    os.makedirs(outdir, exist_ok=True)

    st = get_json("/system_stats", timeout=30)
    print(f"ComfyUI OK vram_free={st['devices'][0]['vram_free']/2**20:.0f}MiB", flush=True)

    results = []
    for m in models:
        wf, steps, note = graph(m, prompt, w, h, seed)
        t0 = time.time()
        try:
            pid = post("/prompt", {"prompt": wf})["prompt_id"]
        except Exception as e:
            print(f"[{m}] 提交失败: {e}", flush=True)
            continue
        print(f"[{m}] {w}x{h} steps={steps} {note} -> {pid}", flush=True)
        hist = None
        while time.time() - t0 < 2400:
            hh = get_json("/history/" + pid)
            if pid in hh:
                hist = hh[pid]
                break
            time.sleep(3)
        el = time.time() - t0
        if hist is None:
            print(f"    TIMEOUT {el:.0f}s", flush=True)
            continue
        if (hist.get("status") or {}).get("status_str") == "error":
            for msg in (hist.get("status") or {}).get("messages", []):
                print("    ERR:", json.dumps(msg, ensure_ascii=False)[:500], flush=True)
            continue
        for _nid, out in (hist.get("outputs") or {}).items():
            for im in out.get("images", []):
                q = urllib.parse.urlencode({"filename": im["filename"], "subfolder": im.get("subfolder", ""), "type": "output"})
                data = OPENER.open(HOST + "/view?" + q, timeout=300).read()
                dst = os.path.join(outdir, f"sweep-{m}.png")
                open(dst, "wb").write(data)
                print(f"    saved {dst} {len(data)} bytes  {el:.1f}s ({el/steps:.2f}s/step)", flush=True)
                results.append({"model": m, "elapsed_s": round(el, 1), "steps": steps, "note": note,
                                "size": f"{w}x{h}", "seed": seed, "file": dst})
    json.dump(results, open(os.path.join(outdir, "_local_results.json"), "w", encoding="utf-8"),
              ensure_ascii=False, indent=1)
    print(f"done {len(results)}/{len(models)}", flush=True)


if __name__ == "__main__":
    main()