#!/usr/bin/env python3
"""LTX Foley 分辨率 A/B 测试 case2: dragon_legend 奇幻场景 (环境音丰富)。
640×352 前台快速出, 1280×704 后台对照。同视频同参数同 seed。"""
import json, time, urllib.request, sys, os

COMFY = "http://localhost:8188"
VIDEO = "case2_dragon_src.mp4"
NUM_FRAMES = 193  # 8s @ 24fps + 1 (force_rate=24 重采样)

LTX = {
    "modelName": "ltx-2.3-22b-dev_transformer_only_int8_convrot.safetensors",
    "foleyLoraName": "ltx-2.3-foley-400-steps.safetensors",
    "foleyLoraStrength": 1.0,
    "videoVaeName": "LTX23_video_vae_bf16.safetensors",
    "audioVaeName": "LTX23_audio_vae_bf16.safetensors",
    "clipName1": "gemma_3_12B_it_fp8_scaled.safetensors",
    "clipName2": "ltx-2.3_text_projection_bf16.safetensors",
    "cfg": 2.05, "samplerName": "euler_ancestral_cfg_pp", "schedulerSteps": 30,
    "maxShift": 2.05, "baseShift": 0.95,
    "frameRate": 24,
}
prompt = "A colossal dragon battles in an ancient ruined city, fire and stone debris flying, stormy sky, epic fantasy atmosphere"
negative = "music, melody, song, singing, vocals, score, soundtrack, beat, rhythm bed, instrumental backing, tinny, thin, harsh, clipped, distorted, low bitrate, subtitles, captions, on-screen text"
AUDIO_POS = "strictly diegetic in-world ambient sound, on-location field recording, raw foley, natural room tone, wind, rustle, environmental texture, no scored music, no voices, unscored scene"
AUDIO_NEG = "non-diegetic audio, background music, BGM, soundtrack, musical score, underscore, theme music, cue, instrumentation, instruments, melody, harmony, chord progression, rhythm, beat, pulse, tempo, percussion, drums, bass, orchestra, electronic music, synthesizer, vocal melody, singing, hooks, drops, chorus, verse, bridge, intro, outro, leitmotif, jingle, speech, dialogue, voices, narration, any structured musical composition or arrangement, subtitles, captions, on-screen text, burned-in text, title cards, Chinese characters, hanzi, handwriting, calligraphy, written words"

def build_wf(W, H, prefix):
    latentW, latentH = W // 32, H // 32
    return {
        "3": {"class_type": "UNETLoader", "inputs": {"unet_name": LTX["modelName"], "weight_dtype": "default"}},
        "10": {"class_type": "LTX2LoraLoaderAdvanced", "inputs": {
            "model": ["3", 0], "lora_name": LTX["foleyLoraName"],
            "strength_model": LTX["foleyLoraStrength"], "video": 0.0,
            "video_to_audio": 1.0, "audio": 1.0, "audio_to_video": 0.0, "other": 1.0}},
        "200": {"class_type": "ModelSamplingLTXV", "inputs": {
            "model": ["10", 0], "max_shift": LTX["maxShift"], "base_shift": LTX["baseShift"]}},
        "26": {"class_type": "DualCLIPLoader", "inputs": {
            "clip_name1": LTX["clipName1"], "clip_name2": LTX["clipName2"], "type": "ltxv"}},
        "5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["26", 0], "text": f"{prompt}, {AUDIO_POS}"}},
        "6": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["26", 0], "text": f"{negative}, {AUDIO_NEG}"}},
        "7": {"class_type": "LTXVConditioning", "inputs": {"positive": ["5", 0], "negative": ["6", 0], "frame_rate": LTX["frameRate"]}},
        "31": {"class_type": "VAELoader", "inputs": {"vae_name": LTX["videoVaeName"]}},
        "21": {"class_type": "VAELoader", "inputs": {"vae_name": LTX["audioVaeName"]}},
        "100": {"class_type": "VHS_LoadVideoFFmpeg", "inputs": {
            "video": VIDEO, "force_rate": LTX["frameRate"],
            "custom_width": W, "custom_height": H,
            "frame_load_cap": NUM_FRAMES, "start_time": 0}},
        "101": {"class_type": "VHS_VAEEncodeBatched", "inputs": {"pixels": ["100", 0], "vae": ["31", 0], "per_batch": NUM_FRAMES}},
        "102": {"class_type": "SolidMask", "inputs": {"value": 0, "width": latentW, "height": latentH}},
        "103": {"class_type": "SetLatentNoiseMask", "inputs": {"samples": ["101", 0], "mask": ["102", 0]}},
        "50": {"class_type": "LTXVEmptyLatentAudio", "inputs": {
            "frames_number": NUM_FRAMES, "frame_rate": LTX["frameRate"],
            "batch_size": 1, "audio_vae": ["21", 0]}},
        "23": {"class_type": "LTXVConcatAVLatent", "inputs": {"video_latent": ["103", 0], "audio_latent": ["50", 0]}},
        "15": {"class_type": "RandomNoise", "inputs": {"noise_seed": 42}},
        "27": {"class_type": "LTXVScheduler", "inputs": {
            "steps": LTX["schedulerSteps"], "max_shift": LTX["maxShift"],
            "base_shift": LTX["baseShift"], "stretch": True, "terminal": 0.1}},
        "13": {"class_type": "KSamplerSelect", "inputs": {"sampler_name": LTX["samplerName"]}},
        "37": {"class_type": "CFGGuider", "inputs": {"model": ["200", 0], "positive": ["7", 0], "negative": ["7", 1], "cfg": LTX["cfg"]}},
        "16": {"class_type": "SamplerCustomAdvanced", "inputs": {
            "noise": ["15", 0], "guider": ["37", 0], "sampler": ["13", 0],
            "sigmas": ["27", 0], "latent_image": ["23", 0]}},
        "24": {"class_type": "LTXVSeparateAVLatent", "inputs": {"av_latent": ["16", 0]}},
        "25": {"class_type": "LTXVAudioVAEDecode", "inputs": {"samples": ["24", 1], "audio_vae": ["21", 0]}},
        "110": {"class_type": "SaveAudio", "inputs": {"audio": ["25", 0], "filename_prefix": prefix}},
    }

def submit(wf):
    req = urllib.request.Request(
        f"{COMFY}/prompt",
        data=json.dumps({"prompt": wf}).encode(),
        headers={'Content-Type': 'application/json'})
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.load(r)

def wait(pid, label, timeout=900):
    start = time.time()
    while time.time() - start < timeout:
        try:
            with urllib.request.urlopen(f"{COMFY}/history/{pid}", timeout=15) as r:
                d = json.load(r)
            if pid in d:
                st = d[pid].get('status', {})
                if st.get('completed'):
                    files = []
                    for nid, out in d[pid].get('outputs', {}).items():
                        for k in ('gifs','videos','audio','images'):
                            if k in out:
                                for item in out[k]:
                                    files.append(item.get('filename'))
                    elapsed = time.time() - start
                    print(f"✅ {label} DONE in {elapsed:.0f}s ({elapsed/60:.1f}min): {files}", flush=True)
                    return elapsed, files
                for m in st.get('messages', []):
                    if m[0] == 'execution_error':
                        print(f"❌ {label} ERROR: {m[1].get('exception_message','')[:300]}", flush=True)
                        return None, None
        except Exception as e:
            print(f"  poll err: {e}", flush=True)
        time.sleep(10)
    print(f"❌ {label} TIMEOUT", flush=True)
    return None, None

which = sys.argv[1] if len(sys.argv) > 1 else "640"
if which == "640":
    print(f"[{time.strftime('%H:%M:%S')}] 提交 640×352 case2 Foley...", flush=True)
    r = submit(build_wf(640, 352, "case2_foley_640x352"))
    if 'error' in r and r['error']:
        print(f"❌ 提交失败: {json.dumps(r['error'])[:400]}", flush=True)
    else:
        pid = r['prompt_id']
        print(f"  prompt_id: {pid}", flush=True)
        t, files = wait(pid, "640×352")
        if files:
            print(f"  ⏱ 640×352 case2 耗时: {t:.0f}s", flush=True)
else:
    print(f"[{time.strftime('%H:%M:%S')}] 提交 1280×704 case2 Foley (对照)...", flush=True)
    r = submit(build_wf(1280, 704, "case2_foley_1280x704"))
    if 'error' in r and r['error']:
        print(f"❌ 提交失败: {json.dumps(r['error'])[:400]}", flush=True)
    else:
        pid = r['prompt_id']
        print(f"  prompt_id: {pid}", flush=True)
        t, files = wait(pid, "1280×704", timeout=1500)
        if files:
            print(f"  ⏱ 1280×704 case2 耗时: {t:.0f}s", flush=True)
print("CASE2_" + which + "_DONE", flush=True)
