import os
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import subprocess
subprocess.run('pip install flash-attn==2.7.4.post1 --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

# wan2.2-main/gradio_ti2v.py
import gradio as gr
import torch
from huggingface_hub import snapshot_download
from PIL import Image
import random
import numpy as np
import spaces

import wan
from wan.configs import WAN_CONFIGS, SIZE_CONFIGS, MAX_AREA_CONFIGS, SUPPORTED_SIZES
from wan.utils.utils import cache_video

# --- 1. Global Setup and Model Loading ---

print("Starting Gradio App for Wan 2.2 TI2V-5B...")

# Download model snapshots from Hugging Face Hub
repo_id = "Wan-AI/Wan2.2-TI2V-5B"
print(f"Downloading/loading checkpoints for {repo_id}...")
ckpt_dir = snapshot_download(repo_id, local_dir_use_symlinks=False)
print(f"Using checkpoints from {ckpt_dir}")

# Load the model configuration
TASK_NAME = 'ti2v-5B'
cfg = WAN_CONFIGS[TASK_NAME]
FIXED_FPS = 24
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 121 

# Instantiate the pipeline in the global scope
print("Initializing WanTI2V pipeline...")
device = "cuda" if torch.cuda.is_available() else "cpu"
device_id = 0 if torch.cuda.is_available() else -1
pipeline = wan.WanTI2V(
    config=cfg,
    checkpoint_dir=ckpt_dir,
    device_id=device_id,
    rank=0,
    t5_fsdp=False,
    dit_fsdp=False,
    use_sp=False,
    t5_cpu=False,
    init_on_cpu=True,
    convert_model_dtype=True,
)
print("Pipeline initialized and ready.")


# --- 2. Gradio Inference Function ---
@spaces.GPU(duration=80)
def generate_video(
    image,
    prompt,
    size,
    duration_seconds,
    sampling_steps,
    guide_scale,
    shift,
    seed,
    progress=gr.Progress(track_tqdm=True)
):
    """The main function to generate video, called by the Gradio interface."""
    if seed == -1:
        seed = random.randint(0, sys.maxsize)

    input_image = Image.fromarray(image).convert("RGB") if image is not None else None
    
    # Calculate number of frames based on duration
    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)

    video_tensor = pipeline.generate(
        input_prompt=prompt,
        img=input_image,  # Pass None for T2V, Image for I2V
        size=SIZE_CONFIGS[size],
        max_area=MAX_AREA_CONFIGS[size],
        frame_num=num_frames,  # Use calculated frames instead of cfg.frame_num
        shift=shift,
        sample_solver='unipc',
        sampling_steps=int(sampling_steps),
        guide_scale=guide_scale,
        seed=seed,
        offload_model=True
    )

    # Save the video to a temporary file
    video_path = cache_video(
        tensor=video_tensor[None],  # Add a batch dimension
        save_file=None,  # cache_video will create a temp file
        fps=cfg.sample_fps,
        normalize=True,
        value_range=(-1, 1)
    )

    return video_path


# --- 3. Gradio Interface ---
css = ".gradio-container {max-width: 1100px !important} #output_video {height: 500px;} #input_image {height: 500px;}"

with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Wan 2.2 Text/Image-to-Video Demo (ti2v-5B)")
    gr.Markdown("Generate a video from a text prompt. Optionally, provide an initial image to guide the generation (Image-to-Video).")

    with gr.Row():
        with gr.Column(scale=2):
            image_input = gr.Image(type="numpy", label="Input Image (Optional)", elem_id="input_image")
            prompt_input = gr.Textbox(label="Prompt", value="A beautiful waterfall in a lush jungle, cinematic.", lines=3)
            duration_input = gr.Slider(
                minimum=round(MIN_FRAMES_MODEL/FIXED_FPS, 1), 
                maximum=round(MAX_FRAMES_MODEL/FIXED_FPS, 1), 
                step=0.1, 
                value=2.0, 
                label="Duration (seconds)", 
                info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps."
            )
            size_input = gr.Dropdown(label="Output Resolution", choices=list(SUPPORTED_SIZES[TASK_NAME]), value="704*1280")
        with gr.Column(scale=2):
            video_output = gr.Video(label="Generated Video", elem_id="output_video")
            

            with gr.Accordion("Advanced Settings", open=False):
                steps_input = gr.Slider(label="Sampling Steps", minimum=10, maximum=70, value=35, step=1)
                scale_input = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, value=cfg.sample_guide_scale, step=0.1)
                shift_input = gr.Slider(label="Sample Shift", minimum=1.0, maximum=20.0, value=cfg.sample_shift, step=0.1)
                seed_input = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)

            run_button = gr.Button("Generate Video", variant="primary")
            

    example_image_path = os.path.join(os.path.dirname(__file__), "examples/i2v_input.JPG")
    gr.Examples(
        examples=[
            [None, "A cinematic shot of a boat sailing on a calm sea at sunset.", "1280*704", 2.0],
            [example_image_path, "The cat slowly blinks its eyes.", "704*1280", 1.5],
            [None, "Drone footage flying over a futuristic city with flying cars.", "1280*704", 3.0],
        ],
        inputs=[image_input, prompt_input, size_input, duration_input],
        outputs=video_output,
        fn=generate_video,
        cache_examples=False,
    )

    run_button.click(
        fn=generate_video,
        inputs=[image_input, prompt_input, size_input, duration_input, steps_input, scale_input, shift_input, seed_input],
        outputs=video_output
    )

if __name__ == "__main__":
    demo.launch()