#!/usr/bin/env python3
"""Misst Krea-2-Generierungszeiten ueber die ComfyUI-API.

Vergleicht NVFP4 gegen FP8, jeweils mit und ohne das ilore-Style-LoRA.
Erster Lauf je Variante ist Aufwaermen (Modell laden) und zaehlt nicht.
"""
import json, time, urllib.request, urllib.error, statistics, argparse, uuid

HOST = "127.0.0.1:8188"

def build(model, lora, strength, w, h, steps, prompt, seed):
    g = {
        "1":  {"class_type": "UNETLoader",  "inputs": {"unet_name": model, "weight_dtype": "default"}},
        "13": {"class_type": "CLIPLoader",  "inputs": {"clip_name": "qwen3vl_4b_fp8_scaled.safetensors",
                                                       "type": "krea2", "device": "default"}},
        "4":  {"class_type": "VAELoader",   "inputs": {"vae_name": "qwen_image_vae.safetensors"}},
        "6":  {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["13", 0]}},
        "8":  {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["6", 0]}},
        "10": {"class_type": "EmptyLatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}},
        "3":  {"class_type": "VAEDecode",   "inputs": {"samples": ["2", 0], "vae": ["4", 0]}},
        "9":  {"class_type": "SaveImage",   "inputs": {"images": ["3", 0], "filename_prefix": "bench"}},
    }
    model_src = ["1", 0]
    if lora:
        g["20"] = {"class_type": "LoraLoaderModelOnly",
                   "inputs": {"model": ["1", 0], "lora_name": lora, "strength_model": strength}}
        model_src = ["20", 0]
    g["2"] = {"class_type": "KSampler", "inputs": {
        "model": model_src, "positive": ["6", 0], "negative": ["8", 0], "latent_image": ["10", 0],
        "seed": seed, "steps": steps, "cfg": 1.0,
        "sampler_name": "er_sde", "scheduler": "simple", "denoise": 1.0}}
    return g

def run(graph, cid):
    body = json.dumps({"prompt": graph, "client_id": cid}).encode()
    req = urllib.request.Request(f"http://{HOST}/prompt", body, {"Content-Type": "application/json"})
    pid = json.load(urllib.request.urlopen(req))["prompt_id"]
    t0 = time.time()
    while True:
        with urllib.request.urlopen(f"http://{HOST}/history/{pid}") as r:
            hist = json.load(r)
        if pid in hist:
            st = hist[pid].get("status", {})
            if st.get("status_str") == "error":
                raise RuntimeError(json.dumps(st)[:400])
            return time.time() - t0
        time.sleep(0.05)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--runs", type=int, default=5)
    ap.add_argument("--steps", type=int, default=8)
    ap.add_argument("--size", default="1024x1024")
    ap.add_argument("--lora", default="ilore_style_krea2_1000.safetensors")
    ap.add_argument("--strength", type=float, default=0.8)
    ap.add_argument("--prompt", default="a weathered iron battleaxe lying on plain white background")
    a = ap.parse_args()
    w, h = (int(x) for x in a.size.split("x"))
    cid = str(uuid.uuid4())

    variants = [
        ("NVFP4  ohne LoRA", "krea2_turbo_nvfp4.safetensors",      None),
        ("NVFP4  mit LoRA",  "krea2_turbo_nvfp4.safetensors",      a.lora),
        ("FP8    ohne LoRA", "krea2_turbo_fp8_scaled.safetensors", None),
        ("FP8    mit LoRA",  "krea2_turbo_fp8_scaled.safetensors", a.lora),
    ]
    print(f"{w}x{h}, {a.steps} steps, {a.runs} Laeufe je Variante (+1 Warmup)\n")
    results = {}
    for label, model, lora in variants:
        try:
            g = build(model, lora, a.strength, w, h, a.steps, a.prompt, 0)
            run(g, cid)                                   # Warmup
            ts = []
            for i in range(a.runs):
                g = build(model, lora, a.strength, w, h, a.steps, a.prompt, i + 1)
                ts.append(run(g, cid))
            med = statistics.median(ts)
            results[label] = med
            print(f"  {label:20s}  Median {med:6.2f}s   (min {min(ts):.2f} / max {max(ts):.2f})")
        except Exception as e:
            print(f"  {label:20s}  FEHLER: {str(e)[:200]}")
    if "NVFP4  mit LoRA" in results and "FP8    mit LoRA" in results:
        sp = results["FP8    mit LoRA"] / results["NVFP4  mit LoRA"]
        print(f"\n  NVFP4 ist {sp:.2f}x schneller als FP8 (mit LoRA)")

if __name__ == "__main__":
    main()