"""Generate one clip on the RunPod ComfyUI: InfiniteTalk (image+audio) or Wan 2.2 I2V (image, silent).
Usage:
  gen_video.py it  <base_url> <image_name> <audio_name> <prompt> <prefix> <out_dir>
  gen_video.py wan <base_url> <image_name> <frames> <prompt> <prefix> <out_dir>
Inputs must already be uploaded to ComfyUI input folder."""
import json, random, sys, time, urllib.request, urllib.parse, os

UA = {'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) Chrome/128.0', 'Content-Type': 'application/json'}
NEG_EN = ("bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, overall gray, "
          "worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, "
          "poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, messy background, text, watermark")
NEG_ZH = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"

def infinitetalk(image, audio, prompt, prefix, audio_seconds):
    frames = int(audio_seconds * 25) + 30
    return {
     "1": {"class_type":"MultiTalkModelLoader","inputs":{"model":"Wan2_1-InfiniTetalk-Single_fp16.safetensors"}},
     "2": {"class_type":"WanVideoLoraSelect","inputs":{"lora":"lightx2v_I2V_14B_480p_cfg_step_distill_rank64_bf16.safetensors","strength":1.0,"low_mem_load":False,"merge_loras":False}},
     "3": {"class_type":"WanVideoModelLoader","inputs":{"model":"wan2.1_i2v_480p_14B_bf16.safetensors","base_precision":"bf16","quantization":os.environ.get("IT_QUANT","disabled"),"load_device":"main_device","attention_mode":"sageattn","lora":["2",0],"multitalk_model":["1",0]}},
     "4": {"class_type":"WanVideoVAELoader","inputs":{"model_name":"Wan2_1_VAE_bf16.safetensors","precision":"bf16"}},
     "5": {"class_type":"CLIPVisionLoader","inputs":{"clip_name":"clip_vision_h.safetensors"}},
     "6": {"class_type":"LoadImage","inputs":{"image":image}},
     "7": {"class_type":"ImageResizeKJv2","inputs":{"image":["6",0],"width":480,"height":832,"upscale_method":"lanczos","keep_proportion":"crop","pad_color":"0, 0, 0","crop_position":"center","divisible_by":16,"device":"cpu"}},
     "8": {"class_type":"WanVideoClipVisionEncode","inputs":{"clip_vision":["5",0],"image_1":["7",0],"strength_1":1.0,"strength_2":1.0,"crop":"center","combine_embeds":"average","force_offload":True,"tiles":0,"ratio":0.5}},
     "9": {"class_type":"WanVideoImageToVideoMultiTalk","inputs":{"vae":["4",0],"start_image":["7",0],"clip_embeds":["8",0],"width":480,"height":832,"frame_window_size":81,"motion_frame":9,"force_offload":False,"colormatch":"disabled","tiled_vae":False,"mode":"infinitetalk","output_path":""}},
     "10": {"class_type":"LoadAudio","inputs":{"audio":audio}},
     "11": {"class_type":"Wav2VecModelLoader","inputs":{"model":"wav2vec2-chinese-base_fp16.safetensors","base_precision":"fp16","load_device":"main_device"}},
     "12": {"class_type":"MultiTalkWav2VecEmbeds","inputs":{"wav2vec_model":["11",0],"audio_1":["10",0],"normalize_loudness":True,"num_frames":frames,"fps":25.0,"audio_scale":1.0,"audio_cfg_scale":1.0,"multi_audio_type":"para"}},
     "13": {"class_type":"WanVideoTextEncodeCached","inputs":{"model_name":"umt5-xxl-enc-bf16.safetensors","precision":"bf16","positive_prompt":prompt,"negative_prompt":NEG_EN,"quantization":"disabled","use_disk_cache":True,"device":"gpu"}},
     "14": {"class_type":"WanVideoSampler","inputs":{"model":["3",0],"image_embeds":["9",0],"text_embeds":["13",0],"multitalk_embeds":["12",0],"steps":6,"cfg":1.0,"shift":11.0,"seed":random.randint(1, 2**31),"force_offload":True,"scheduler":"dpm++_sde","riflex_freq_index":0,"denoise_strength":1.0,"batched_cfg":False,"rope_function":"comfy","start_step":0,"end_step":-1,"add_noise_to_samples":False}},
     "15": {"class_type":"WanVideoPassImagesFromSamples","inputs":{"samples":["14",0]}},
     "16": {"class_type":"VHS_VideoCombine","inputs":{"images":["15",0],"audio":["10",0],"frame_rate":25,"loop_count":0,"filename_prefix":prefix,"format":"video/h264-mp4","pix_fmt":"yuv420p","crf":17,"save_metadata":True,"trim_to_audio":True,"pingpong":False,"save_output":True}},
    }

def wan22(image, frames, prompt, prefix):
    seed = random.randint(1, 2**48)
    W, H = 720, 1280
    return {
     "1": {"class_type":"UNETLoader","inputs":{"unet_name":"wan2.2_i2v_high_noise_14B_fp16.safetensors","weight_dtype":"default"}},
     "2": {"class_type":"UNETLoader","inputs":{"unet_name":"wan2.2_i2v_low_noise_14B_fp16.safetensors","weight_dtype":"default"}},
     "3": {"class_type":"LoraLoaderModelOnly","inputs":{"model":["1",0],"lora_name":"i2v_lightx2v_high_noise_model.safetensors","strength_model":1.0}},
     "4": {"class_type":"LoraLoaderModelOnly","inputs":{"model":["2",0],"lora_name":"i2v_lightx2v_low_noise_model.safetensors","strength_model":1.0}},
     "5": {"class_type":"ModelSamplingSD3","inputs":{"model":["3",0],"shift":5.0}},
     "6": {"class_type":"ModelSamplingSD3","inputs":{"model":["4",0],"shift":5.0}},
     "7": {"class_type":"CLIPLoader","inputs":{"clip_name":"umt5_xxl_fp8_e4m3fn_scaled.safetensors","type":"wan","device":"default"}},
     "8": {"class_type":"CLIPTextEncode","inputs":{"clip":["7",0],"text":prompt}},
     "9": {"class_type":"CLIPTextEncode","inputs":{"clip":["7",0],"text":NEG_ZH}},
     "10": {"class_type":"VAELoader","inputs":{"vae_name":"wan_2.1_vae.safetensors"}},
     "11": {"class_type":"LoadImage","inputs":{"image":image}},
     "12": {"class_type":"ImageScale","inputs":{"image":["11",0],"upscale_method":"lanczos","width":W,"height":H,"crop":"center"}},
     "13": {"class_type":"WanImageToVideo","inputs":{"positive":["8",0],"negative":["9",0],"vae":["10",0],"start_image":["12",0],"width":W,"height":H,"length":frames,"batch_size":1}},
     "14": {"class_type":"KSamplerAdvanced","inputs":{"model":["5",0],"add_noise":"enable","noise_seed":seed,"steps":8,"cfg":1.0,"sampler_name":"euler","scheduler":"simple","start_at_step":0,"end_at_step":4,"return_with_leftover_noise":"enable","positive":["13",0],"negative":["13",1],"latent_image":["13",2]}},
     "15": {"class_type":"KSamplerAdvanced","inputs":{"model":["6",0],"add_noise":"disable","noise_seed":seed,"steps":8,"cfg":1.0,"sampler_name":"euler","scheduler":"simple","start_at_step":4,"end_at_step":10000,"return_with_leftover_noise":"disable","positive":["13",0],"negative":["13",1],"latent_image":["14",0]}},
     "16": {"class_type":"VAEDecode","inputs":{"samples":["15",0],"vae":["10",0]}},
     "17": {"class_type":"VHS_VideoCombine","inputs":{"images":["16",0],"frame_rate":16,"loop_count":0,"filename_prefix":prefix,"format":"video/h264-mp4","pix_fmt":"yuv420p","crf":17,"save_metadata":True,"trim_to_audio":False,"pingpong":False,"save_output":True}},
    }

def call(base, path, data=None):
    r = urllib.request.Request(base + path, data=json.dumps(data).encode() if data is not None else None, headers=UA)
    return json.load(urllib.request.urlopen(r, timeout=60))

def main():
    kind, base = sys.argv[1], sys.argv[2]
    if kind == 'it':
        image, audio, prompt, prefix, out_dir, secs = sys.argv[3], sys.argv[4], sys.argv[5], sys.argv[6], sys.argv[7], float(sys.argv[8])
        graph = infinitetalk(image, audio, prompt, prefix, secs)
    else:
        image, frames, prompt, prefix, out_dir = sys.argv[3], int(sys.argv[4]), sys.argv[5], sys.argv[6], sys.argv[7]
        graph = wan22(image, frames, prompt, prefix)
    try:
        res = call(base, '/prompt', {"prompt": graph, "client_id": "claude"})
    except urllib.error.HTTPError as e:
        print('SUBMIT_HTTP_ERROR', e.code, e.read().decode()[:2000]); sys.exit(1)
    if res.get('node_errors'):
        print('NODE_ERRORS', json.dumps(res['node_errors'])[:2000]); sys.exit(1)
    pid = res['prompt_id']; t0 = time.time(); fails = 0
    print(f'{prefix} SUBMITTED {pid}', flush=True)
    while True:
        time.sleep(10)
        try:
            h = call(base, f'/history/{pid}'); fails = 0
        except Exception as e:
            fails += 1
            if fails >= 18: print(f'{prefix} FAILED_UNREACHABLE {e}'); sys.exit(1)
            continue
        st = h.get(pid, {}).get('status', {})
        if st.get('status_str') in ('success', 'error'):
            break
        if time.time() - t0 > 3600: print(f'{prefix} FAILED_TIMEOUT'); sys.exit(1)
    secs_taken = int(time.time() - t0)
    if st['status_str'] == 'error':
        for m in st.get('messages', []):
            if m[0] == 'execution_error':
                print(f"{prefix} ERROR {m[1].get('node_type')}: {m[1].get('exception_message','')[:1500]}")
        sys.exit(1)
    files = [f for o in h[pid].get('outputs', {}).values() for arr in o.values() if isinstance(arr, list)
             for f in arr if isinstance(f, dict) and str(f.get('filename', '')).endswith('.mp4')]
    if not files: print(f'{prefix} FAILED_NO_MP4'); sys.exit(1)
    f = sorted(files, key=lambda x: '-audio' not in x['filename'])[0]
    q = urllib.parse.urlencode({'filename': f['filename'], 'subfolder': f.get('subfolder', ''), 'type': 'output'})
    os.makedirs(out_dir, exist_ok=True)
    dest = os.path.join(out_dir, prefix + '.mp4')
    r = urllib.request.Request(base + '/view?' + q, headers={'User-Agent': UA['User-Agent']})
    open(dest, 'wb').write(urllib.request.urlopen(r, timeout=300).read())
    print(f'{prefix} DONE in {secs_taken}s -> {dest} ({os.path.getsize(dest)} bytes)')

main()