#!/usr/bin/env python3
"""Submit Wan 2.2 Two-Pass I2V workflow to ComfyUI - Fixed v4"""

import json
import urllib.request
import urllib.error
import uuid

COMFY_URL = "http://localhost:8188"
INPUT_IMAGE = "two_men_locker.png"

PROMPT = """Static camera shot. Two shirtless men stand side by side in a locker room with rainbow pride flags and colorful bunting decorating the walls. The man on the left with curly brown hair gently applies lotion to the shoulders and chest of the man on the right with dark curly hair, massaging tenderly. He then presents a beautiful bouquet of rainbow flowers. They touch noses in an eskimo kiss, then the white man places a gentle kiss on the black man's forehead. Pride flags wave softly in the background. The camera remains perfectly still throughout. No zooming, no panning, no camera movement."""

NEGATIVE = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"

workflow = {
    "1": {
        "class_type": "LoadImage",
        "inputs": {
            "image": INPUT_IMAGE,
            "upload": "image"
        }
    },
    "2": {
        "class_type": "ImageResizeKJv2",
        "inputs": {
            "image": ["1", 0],
            "width": 848,
            "height": 480,
            "upscale_method": "lanczos",
            "keep_proportion": "crop",
            "pad_color": "0, 0, 0",
            "crop_position": "center",
            "divisible_by": 32
        }
    },
    "8": {
        "class_type": "WanVideoBlockSwap",
        "inputs": {
            "blocks_to_swap": 20,
            "offload_img_emb": False,
            "offload_txt_emb": False
        }
    },
    "3": {
        "class_type": "WanVideoModelLoader",
        "inputs": {
            "model": "WanVideo/2_2/Wan2_2-I2V-A14B-HIGH_fp8_e4m3fn_scaled_KJ.safetensors",
            "base_precision": "fp16_fast",
            "quantization": "fp8_e4m3fn_scaled",
            "load_device": "offload_device",
            "attention_mode": "sdpa",
            "block_swap_args": ["8", 0]
        }
    },
    "4": {
        "class_type": "WanVideoModelLoader",
        "inputs": {
            "model": "WanVideo/2_2/Wan2_2-I2V-A14B-LOW_fp8_e4m3fn_scaled_KJ.safetensors",
            "base_precision": "fp16_fast",
            "quantization": "fp8_e4m3fn_scaled",
            "load_device": "offload_device",
            "attention_mode": "sdpa",
            "block_swap_args": ["8", 0]
        }
    },
    "5": {
        "class_type": "WanVideoVAELoader",
        "inputs": {
            "model_name": "wanvideo/Wan2_1_VAE_bf16.safetensors",
            "precision": "bf16"
        }
    },
    "6": {
        "class_type": "LoadWanVideoT5TextEncoder",
        "inputs": {
            "model_name": "umt5-xxl-enc-fp8_e4m3fn.safetensors",
            "precision": "bf16",
            "load_device": "offload_device",
            "quantization": "fp8_e4m3fn"
        }
    },
    "7": {
        "class_type": "WanVideoTextEncode",
        "inputs": {
            "positive_prompt": PROMPT,
            "negative_prompt": NEGATIVE,
            "t5": ["6", 0],
            "force_offload": True,
            "use_disk_cache": False
        }
    },
    "9": {
        "class_type": "WanVideoImageToVideoEncode",
        "inputs": {
            "width": ["2", 1],
            "height": ["2", 2],
            "num_frames": 81,
            "noise_aug_strength": 0,
            "start_latent_strength": 1,
            "end_latent_strength": 1,
            "force_offload": True,
            "vae": ["5", 0],
            "start_image": ["2", 0]
        }
    },
    "10": {
        "class_type": "WanVideoSampler",
        "inputs": {
            "model": ["3", 0],
            "image_embeds": ["9", 0],
            "steps": 6,
            "cfg": 1.5,
            "shift": 8,
            "seed": 42,
            "force_offload": True,
            "scheduler": "dpm++_sde",
            "riflex_freq_index": 0,
            "text_embeds": ["7", 0],
            "end_step": 3
        }
    },
    "11": {
        "class_type": "WanVideoSampler",
        "inputs": {
            "model": ["4", 0],
            "image_embeds": ["9", 0],
            "steps": 6,
            "cfg": 1.0,
            "shift": 8,
            "seed": 42,
            "force_offload": True,
            "scheduler": "dpm++_sde",
            "riflex_freq_index": 0,
            "text_embeds": ["7", 0],
            "samples": ["10", 0],
            "start_step": 3
        }
    },
    "12": {
        "class_type": "WanVideoDecode",
        "inputs": {
            "vae": ["5", 0],
            "samples": ["11", 0],
            "enable_vae_tiling": False,
            "enable_tiling": False,
            "tile_x": 272,
            "tile_y": 272,
            "tile_stride_x": 144,
            "tile_stride_y": 128
        }
    },
    "13": {
        "class_type": "VHS_VideoCombine",
        "inputs": {
            "images": ["12", 0],
            "frame_rate": 16,
            "loop_count": 0,
            "filename_prefix": "Wan22_PrideLove",
            "format": "video/h264-mp4",
            "pix_fmt": "yuv420p",
            "crf": 19,
            "save_metadata": True,
            "trim_to_audio": False,
            "pingpong": False,
            "save_output": True
        }
    }
}

def submit_prompt(workflow):
    prompt_id = str(uuid.uuid4())
    data = {"prompt": workflow, "client_id": "selena-agent"}
    req = urllib.request.Request(
        f"{COMFY_URL}/prompt",
        data=json.dumps(data).encode('utf-8'),
        headers={'Content-Type': 'application/json'}
    )
    try:
        with urllib.request.urlopen(req) as response:
            result = json.loads(response.read().decode('utf-8'))
            print(f"✅ Prompt submitted!")
            print(f"   Prompt ID: {result.get('prompt_id', prompt_id)}")
            return result
    except urllib.error.HTTPError as e:
        error_body = e.read().decode('utf-8')
        print(f"❌ Error {e.code}: {error_body[:1000]}")
        return None

if __name__ == "__main__":
    print("🎬 Submitting Wan 2.2 Two-Pass I2V...")
    submit_prompt(workflow)