#!/usr/bin/env python3
"""Qwen-Image-2.1 API smoke test for ComfyUI (GPU machine localhost:8189).

usage:
  qwen21_run.py t2i  <promptfile> <prefix> [steps] [w] [h] [seed]
  qwen21_run.py edit <promptfile> <prefix> <input_image_name> [steps] [w] [h] [seed]
"""
import json, sys, time, threading, subprocess, urllib.request, os

HOST = "http://127.0.0.1:8189"
UNET = "qwen_image_2.1_int8_convrot.safetensors"
CLIP = "qwen3vl_8b_int8_convrot.safetensors"
VAE = "qwen_image_2.1_vae_bf16.safetensors"


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


def get(path):
    return json.load(urllib.request.urlopen(HOST + path, timeout=120))


def vram_samples(stop, out):
    while not stop.is_set():
        try:
            s = subprocess.run(["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
                               capture_output=True, text=True, timeout=5).stdout.strip()
            out.append(int(s))
        except Exception:
            pass
        time.sleep(1)


def main():
    mode = sys.argv[1]
    prompt = open(sys.argv[2]).read().strip()
    prefix = sys.argv[3]
    steps = int(sys.argv[4]) if len(sys.argv) > 4 else 25
    w = int(sys.argv[5]) if len(sys.argv) > 5 else 1024
    h = int(sys.argv[6]) if len(sys.argv) > 6 else 1024
    seed = int(sys.argv[7]) if len(sys.argv) > 7 else 12345

    wf = {
        "1": {"class_type": "UNETLoader", "inputs": {"unet_name": UNET, "weight_dtype": "default"}},
        "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": CLIP, "type": "qwen_image", "device": "default"}},
        "3": {"class_type": "VAELoader", "inputs": {"vae_name": VAE}},
        "4": {"class_type": "TextEncodeQwenImage21",
              "inputs": {"clip": ["2", 0], "prompt": prompt, "negative_prompt": "", "resolution": 1024,
                         "vae": ["3", 0]}},
        "6": {"class_type": "KSampler",
              "inputs": {"model": ["1", 0], "positive": ["4", 0], "negative": ["4", 1], "seed": seed,
                         "steps": steps, "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": prefix, "images": ["7", 0]}},
    }

    if mode == "t2i":
        wf["5"] = {"class_type": "EmptyLatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}}
        wf["6"]["inputs"]["latent_image"] = ["5", 0]
    elif mode == "edit":
        imgs = sys.argv[8:]
        for i, im in enumerate(imgs, start=1):
            nid = str(100 + i)
            wf[nid] = {"class_type": "LoadImage", "inputs": {"image": im, "upload": "image"}}
            wf["4"]["inputs"]["images.image_%d" % i] = [nid, 0]   # autogrow reference slot
        wf["6"]["inputs"]["latent_image"] = ["4", 2]   # canvas from encode latent (image_1 size)
    else:
        raise SystemExit("mode must be t2i|edit")

    stop = threading.Event()
    vram = []
    t = threading.Thread(target=vram_samples, args=(stop, vram), daemon=True)
    t.start()

    t0 = time.time()
    r = post("/prompt", {"prompt": wf})
    pid = r["prompt_id"]
    print("prompt_id:", pid, "queued_at", time.strftime("%H:%M:%S"), flush=True)

    hist = None
    while time.time() - t0 < 1800:
        h_ = get("/history/" + pid)
        if pid in h_:
            hist = h_[pid]
            break
        time.sleep(3)
    elapsed = time.time() - t0
    stop.set()

    if hist is None:
        print("TIMEOUT after %.1fs" % elapsed)
        sys.exit(1)

    status = hist.get("status", {})
    print("elapsed: %.1fs (%.1fs/step @%d steps)  status=%s" % (
        elapsed, elapsed / steps, steps, status.get("status_str")))
    print("vram_used MiB: start=%d peak=%d end=%d" % (vram[0] if vram else -1, max(vram) if vram else -1,
                                                      vram[-1] if vram else -1))
    if status.get("status_str") == "error":
        for m in status.get("messages", []):
            print("ERR:", json.dumps(m)[:800])
        sys.exit(2)

    outdir = "/home/zyw/ComfyUI/output"
    files = []
    for node_id, out in hist.get("outputs", {}).items():
        for im in out.get("images", []):
            p = os.path.join(outdir, im.get("subfolder", ""), im["filename"])
            files.append(p)
            print("OUTPUT:", p, os.path.getsize(p), "bytes", flush=True)
    print("files:", len(files))


if __name__ == "__main__":
    main()