#!/usr/bin/env python3
"""风格穿透探针: Ideogram底图作风格供体, Z-Image i2i 换内容
臂A1: S1作init + 唐人街prompt @ d0.75
臂A2: S1作init + 唐人街prompt @ d0.85
臂B : 纯t2i 唐人街prompt (对照, 无风格供体)
判定: 输出是否带S1的摄影质感/色温, 同时内容完全换成唐人街
"""
import json, time, urllib.request

BASE = "http://127.0.0.1:8190"
OUT = "/home/kai/shared/2026-09-09/ideogram_blind4"
SEED = 9909

PROMPT_STREET = ("雨夜中的赛博朋克唐人街夜市，低角度广角镜头，湿漉漉的柏油路面映满霓虹倒影。"
                 "街道两侧密密麻麻的立体霓虹招牌层层叠叠伸向天空：最高处巨大的红色汉字霓虹招牌「龙轩酒家」微微闪烁，"
                 "旁边竖排蓝色霓虹「深淵茶館」与粉色「小籠包」，还有英文灯牌「NOODLE BAR 24H」。"
                 "街上人群撑着透明雨伞穿行，蒸汽从街边小吃摊升腾，电影感冷蓝暖橙对比打光，超精细细节，8K质感")

NEG = "blurry ugly bad, low quality, watermark"

def t2i_nodes():
    return {
      "5": {"class_type": "UNETLoader", "inputs": {"unet_name": "z_image_turbo_bf16.safetensors", "weight_dtype": "default"}},
      "6": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen_3_4b_fp8_mixed.safetensors", "type": "lumina2", "device": "default"}},
      "7": {"class_type": "VAELoader", "inputs": {"vae_name": "z_image_ae.safetensors"}},
      "9": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["6", 0], "text": PROMPT_STREET}},
      "10": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["6", 0], "text": NEG}},
      "13": {"class_type": "EmptySD3LatentImage", "inputs": {"width": 1360, "height": 768, "batch_size": 1}},
      "3": {"class_type": "KSampler", "inputs": {"model": ["5", 0], "seed": SEED, "steps": 9, "cfg": 1.0,
            "sampler_name": "euler", "scheduler": "simple", "positive": ["9", 0], "negative": ["10", 0],
            "latent_image": ["13", 0], "denoise": 1.0}},
      "8": {"class_type": "VAEDecode", "inputs": {"samples": ["3", 0], "vae": ["7", 0]}},
      "17": {"class_type": "SaveImage", "inputs": {"images": ["8", 0], "filename_prefix": "bld4/st_t2i_ctl"}},
    }

def i2i_nodes(denoise):
    return {
      "5": {"class_type": "UNETLoader", "inputs": {"unet_name": "z_image_turbo_bf16.safetensors", "weight_dtype": "default"}},
      "6": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen_3_4b_fp8_mixed.safetensors", "type": "lumina2", "device": "default"}},
      "7": {"class_type": "VAELoader", "inputs": {"vae_name": "z_image_ae.safetensors"}},
      "1": {"class_type": "LoadImage", "inputs": {"image": "s1_ideo.png"}},
      "2": {"class_type": "VAEEncode", "inputs": {"pixels": ["1", 0], "vae": ["7", 0]}},
      "9": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["6", 0], "text": PROMPT_STREET}},
      "10": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["6", 0], "text": NEG}},
      "3": {"class_type": "KSampler", "inputs": {"model": ["5", 0], "seed": SEED, "steps": 9, "cfg": 1.0,
            "sampler_name": "euler", "scheduler": "simple", "positive": ["9", 0], "negative": ["10", 0],
            "latent_image": ["2", 0], "denoise": denoise}},
      "8": {"class_type": "VAEDecode", "inputs": {"samples": ["3", 0], "vae": ["7", 0]}},
      "17": {"class_type": "SaveImage", "inputs": {"images": ["8", 0], "filename_prefix": f"bld4/st_i2i_d{int(denoise*100)}"}},
    }

ARMS = [("st_i2i_d75", i2i_nodes(0.75), 300), ("st_i2i_d85", i2i_nodes(0.85), 300), ("st_t2i_ctl", t2i_nodes(), 300)]

def post(url, payload, timeout=60):
    req = urllib.request.Request(url, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"})
    return json.loads(urllib.request.urlopen(req, timeout=timeout).read())

def get(url, timeout=30):
    return json.loads(urllib.request.urlopen(url, timeout=timeout).read())

timing = {}
for arm, nodes, timeout in ARMS:
    t0 = time.time()
    pid = post(BASE + "/prompt", {"prompt": (nodes() if callable(nodes) else nodes), "client_id": "styleprobe"})["prompt_id"]
    print(f"[{arm}] submitted {pid}", flush=True)
    while time.time() - t0 < timeout:
        time.sleep(5)
        h = get(f"{BASE}/history/{pid}")
        if pid not in h:
            continue
        st = h[pid].get("status", {})
        if st.get("status_str") == "success":
            for node_out in h[pid].get("outputs", {}).values():
                for im in node_out.get("images", []):
                    dur = round(time.time() - t0, 1)
                    vurl = f"{BASE}/view?filename={urllib.request.quote(im['filename'])}&subfolder={urllib.request.quote(im.get('subfolder',''))}&type=output"
                    data = urllib.request.urlopen(vurl, timeout=60).read()
                    open(f"{OUT}/{arm}.png", "wb").write(data)
                    print(f"[{arm}] DONE {dur}s ({len(data)//1024}KB)", flush=True)
                    timing[arm] = dur
            break
        if st.get("status_str") == "error":
            errs = [json.dumps(m[1], ensure_ascii=False)[:250] for m in st.get("messages", []) if m[0] == "execution_error"]
            print(f"[{arm}] ERROR: {errs}", flush=True)
            timing[arm] = None
            break
    else:
        print(f"[{arm}] TIMEOUT", flush=True)
        timing[arm] = None
    json.dump(timing, open(f"{OUT}/styleprobe_timing.json", "w"), ensure_ascii=False, indent=1)

print("ALL:", timing, flush=True)
import sys
sys.exit(0 if all(v for v in timing.values()) else 1)