#!/usr/bin/env python3
"""Case08 r2v A/B round 2 at 241f (VDN 362f OOM boundary): VDN -> :8188, LXV -> :8190.
Same seed/refs/res; length 241. Watchdog: SUCCESS only on status_str=success.
"""
import json, subprocess, time, urllib.request

P8188 = "http://localhost:8188"
P8190 = "http://localhost:8190"
PROMPT = open("/home/kai/文档/LTX/validition_v1/08/prompt.txt").read().strip()
SEED, STEPS, W, H, LEN = 20260907, 8, 1216, 672, 241


def graph(prefix, model_chain):
    g = {
        "10": {"class_type": "CLIPLoader", "inputs": {
            "clip_name": "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "type": "minimax"}},
        "11": {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_video_vae_fp16.safetensors"}},
        "12": {"class_type": "UNETLoader", "inputs": {
            "unet_name": "minimax_h3_fl2va_int8_convrot.safetensors", "weight_dtype": "default"}},
        "13": {"class_type": "VAELoader", "inputs": {"vae_name": "minimax_h3_audio_vae_fp32.safetensors"}},
        "20": {"class_type": "MiniMaxH3ReferenceToVideo", "inputs": {
            "clip": ["10", 0], "vae": ["11", 0], "audio_vae": ["13", 0],
            "prompt": PROMPT, "width": W, "height": H, "length": LEN,
            "ref_image_size": "match",
            "ref_images.ref_image_0": ["14", 0],
            "ref_images.ref_image_1": ["141", 0],
            "ref_images.ref_image_2": ["142", 0]}},
        "14": {"class_type": "LoadImage", "inputs": {"image": "ab08_rn.jpg"}},
        "141": {"class_type": "LoadImage", "inputs": {"image": "ab08_wv.jpg"}},
        "142": {"class_type": "LoadImage", "inputs": {"image": "ab08_bg.png"}},
        "30": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "er_sde"}},
        "31": {"class_type": "BasicScheduler", "inputs": {
            "model": ["15", 0], "scheduler": "beta", "steps": STEPS, "denoise": 1.0}},
        "32": {"class_type": "RandomNoise", "inputs": {"noise_seed": SEED, "control_after_generate": "fixed"}},
        "33": {"class_type": "BasicGuider", "inputs": {"model": ["15", 0], "conditioning": ["20", 0]}},
        "34": {"class_type": "SamplerCustomAdvanced", "inputs": {
            "noise": ["32", 0], "guider": ["33", 0], "sampler": ["30", 0],
            "sigmas": ["31", 0], "latent_image": ["20", 1]}},
        "40": {"class_type": "VAEDecode", "inputs": {"samples": ["34", 0], "vae": ["11", 0]}},
        "41": {"class_type": "VAEDecodeAudio", "inputs": {"samples": ["34", 0], "vae": ["13", 0]}},
        "42": {"class_type": "CreateVideo", "inputs": {"images": ["40", 0], "fps": 24, "audio": ["41", 0]}},
        "50": {"class_type": "SaveVideo", "inputs": {
            "video": ["42", 0], "filename_prefix": prefix, "format": "mp4", "codec": "auto"}},
    }
    g["15"] = model_chain
    return {"prompt": g}


VDN_CHAIN = {"class_type": "ApplyVDNH3", "inputs": {
    "model": ["12", 0], "vdn_checkpoint": "vdn-minimax-h3-int8-convrot-comfyui",
    "apply_turbo_adapter": True, "strength": 1.0, "lora_mode": "merge",
    "branch_weights": "stream", "retain_buffers": "off",
    "attention_backend": "flex", "verbose": True}}

LXV_CHAIN = {"class_type": "LoraLoaderModelOnly", "inputs": {
    "model": ["12", 0], "lora_name": "MiniMax-H3-Ref2VA-Acc-8Step.safetensors",
    "strength_model": 1.0}}


def post(port, graph_obj):
    req = urllib.request.Request(port + "/prompt",
        data=json.dumps(graph_obj).encode(), headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.load(r)


def get(port, path):
    try:
        with urllib.request.urlopen(port + path, timeout=10) as r:
            return json.load(r)
    except Exception:
        return None


def state(port, pid):
    h = get(port, f"/history/{pid}")
    if h and pid in h:
        st = h[pid].get("status", {}).get("status_str", "")
        return "SUCCESS" if st == "success" else f"ERROR({st})"
    q = get(port, "/queue")
    if q:
        for items in (q.get("queue_running", []), q.get("queue_pending", [])):
            for it in items:
                if len(it) > 1 and it[1] == pid:
                    return "QUEUED"
    return "GONE"


vid = post(P8188, graph("MiniMaxH3/case08_241_vdn", VDN_CHAIN)).get("prompt_id")
lid = post(P8190, graph("MiniMaxH3/case08_241_lxv", LXV_CHAIN)).get("prompt_id")
print(f"[{time.strftime('%H:%M:%S')}] VDN-241 :8188 -> {vid}", flush=True)
print(f"[{time.strftime('%H:%M:%S')}] LXV-241 :8190 -> {lid}", flush=True)

for i in range(70):
    time.sleep(60)
    sv = state(P8188, vid)
    sl = state(P8190, lid)
    av = subprocess.run(["free", "-g"], capture_output=True, text=True
                        ).stdout.splitlines()[1].split()[6]
    if av.isdigit() and int(av) <= 13:
        print(f"!!! MEMORY CRITICAL avail={av}G", flush=True)
    print(f"[{time.strftime('%H:%M:%S')}] VDN:{sv} LXV:{sl} avail={av}G", flush=True)
    if sv.startswith("ERROR") or sl.startswith("ERROR"):
        print("!!! an arm FAILED", flush=True)
        break
    if sv == "SUCCESS" and sl == "SUCCESS":
        break

print("=== final outputs ===", flush=True)
for c in ("comfyui-primary", "comfyui-secondary"):
    out = subprocess.run(["docker", "exec", c, "sh", "-c",
        "ls -la /root/ComfyUI/output/MiniMaxH3/ 2>/dev/null | grep case08"],
        capture_output=True, text=True).stdout
    print(c, out.strip() or "(none)", flush=True)
print("=== exec times ===", flush=True)
for c in ("comfyui-primary", "comfyui-secondary"):
    out = subprocess.run(["docker", "logs", c, "--since", "90m"], capture_output=True, text=True).stdout
    tails = [l for l in out.splitlines() if "Prompt executed" in l][-2:]
    print(c, tails, flush=True)