#!/usr/bin/env python3
"""LTX Foley 分辨率 A/B 测试: 640×352 vs 1280×704 (视频 latent 冻结, 音频质量应不受影响)"""
import json, time, urllib.request

COMFY = "http://localhost:8188"
CONTAINER = "comfyui-primary"
VIDEO = "h3_design_v1_foley_src.mp4"
NUM_FRAMES = 193

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 young man in a dark charcoal overcoat with thin metal glasses speaks in a dimly lit room, tense suppressed 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):
    start = time.time()
    while time.time() - start < 900:
        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

# 提交 640×352 测试 (低分辨率)
print(f"[{time.strftime('%H:%M:%S')}] 提交 640×352 Foley 测试...", flush=True)
r = submit(build_wf(640, 352, "foley_ab_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 耗时: {t:.0f}s (vs 1280×704 ~510s)", flush=True)
print("TEST_DONE", flush=True)