app.py

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

# 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 cv2
import tempfile

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

import gc

# --- 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 

# Dimension calculation constants
MOD_VALUE = 32
DEFAULT_H_SLIDER_VALUE = 704
DEFAULT_W_SLIDER_VALUE = 1280
NEW_FORMULA_MAX_AREA = 1280.0 * 704.0

SLIDER_MIN_H, SLIDER_MAX_H = 128, 1280
SLIDER_MIN_W, SLIDER_MAX_W = 128, 1280

# 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=False,
    convert_model_dtype=True,
)
print("Pipeline initialized and ready.")

# --- Helper Functions ---

def extract_first_frame_from_video(video_path):
    """
    Extract the first frame from a video file.
    
    Args:
        video_path: Path to the video file
        
    Returns:
        PIL Image of the first frame, or None if extraction fails
    """
    try:
        cap = cv2.VideoCapture(video_path)
        ret, frame = cap.read()
        cap.release()
        
        if ret:
            # Convert BGR to RGB
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            return Image.fromarray(frame_rgb)
        return None
    except Exception as e:
        print(f"Error extracting frame from video: {e}")
        return None

def get_video_dimensions(video_path):
    """
    Get the dimensions of a video file.
    
    Args:
        video_path: Path to the video file
        
    Returns:
        Tuple of (width, height) or None if extraction fails
    """
    try:
        cap = cv2.VideoCapture(video_path)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        cap.release()
        return width, height
    except Exception as e:
        print(f"Error getting video dimensions: {e}")
        return None

def _calculate_new_dimensions_wan(pil_image, mod_val, calculation_max_area,
                                 min_slider_h, max_slider_h,
                                 min_slider_w, max_slider_w,
                                 default_h, default_w):
    orig_w, orig_h = pil_image.size
    if orig_w <= 0 or orig_h <= 0:
        return default_h, default_w

    aspect_ratio = orig_h / orig_w
    
    calc_h = round(np.sqrt(calculation_max_area * aspect_ratio))
    calc_w = round(np.sqrt(calculation_max_area / aspect_ratio))

    calc_h = max(mod_val, (calc_h // mod_val) * mod_val)
    calc_w = max(mod_val, (calc_w // mod_val) * mod_val)
    
    new_h = int(np.clip(calc_h, min_slider_h, (max_slider_h // mod_val) * mod_val))
    new_w = int(np.clip(calc_w, min_slider_w, (max_slider_w // mod_val) * mod_val))
    
    return new_h, new_w

def handle_media_upload_for_dims_wan(uploaded_media, current_h_val, current_w_val):
    """
    Handle image or video upload and calculate appropriate dimensions.
    
    Args:
        uploaded_media: The uploaded file (can be image or video path)
        current_h_val: Current height slider value
        current_w_val: Current width slider value
        
    Returns:
        Tuple of (gr.update for height, gr.update for width, first frame as numpy array or None)
    """
    if uploaded_media is None:
        return (gr.update(value=DEFAULT_H_SLIDER_VALUE), 
                gr.update(value=DEFAULT_W_SLIDER_VALUE),
                None)
    
    try:
        pil_image = None
        
        # Check if it's a video file
        if isinstance(uploaded_media, str) and uploaded_media.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.webm')):
            # Extract first frame from video
            pil_image = extract_first_frame_from_video(uploaded_media)
            if pil_image is None:
                gr.Warning("Could not extract frame from video")
                return (gr.update(value=DEFAULT_H_SLIDER_VALUE), 
                        gr.update(value=DEFAULT_W_SLIDER_VALUE),
                        None)
        else:
            # Handle as image
            if hasattr(uploaded_media, 'shape'):  # numpy array
                pil_image = Image.fromarray(uploaded_media).convert("RGB")
            elif isinstance(uploaded_media, str):  # file path
                pil_image = Image.open(uploaded_media).convert("RGB")
            else:  # PIL Image
                pil_image = uploaded_media
        
        # Calculate dimensions
        new_h, new_w = _calculate_new_dimensions_wan(
            pil_image, MOD_VALUE, NEW_FORMULA_MAX_AREA,
            SLIDER_MIN_H, SLIDER_MAX_H, SLIDER_MIN_W, SLIDER_MAX_W,
            DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE
        )
        
        # Convert PIL image to numpy array for display
        display_image = np.array(pil_image)
        
        return gr.update(value=new_h), gr.update(value=new_w), display_image
        
    except Exception as e:
        print(f"Error in handle_media_upload_for_dims_wan: {e}")
        gr.Warning("Error processing uploaded file")
        return (gr.update(value=DEFAULT_H_SLIDER_VALUE), 
                gr.update(value=DEFAULT_W_SLIDER_VALUE),
                None)

def get_duration(video_input,
                 image_preview,
                 prompt, 
                 height,
                 width,
                 duration_seconds, 
                 sampling_steps, 
                 guide_scale, 
                 shift, 
                 seed,
                 progress):
    """Calculate dynamic GPU duration based on parameters."""
    return sampling_steps * 15

# --- 2. Gradio Inference Function ---
@spaces.GPU(duration=get_duration)
def generate_video(
    video_input,
    image_preview,
    prompt,
    height,
    width,
    duration_seconds,
    sampling_steps=38,
    guide_scale=cfg.sample_guide_scale,
    shift=cfg.sample_shift,
    seed=42,
    progress=gr.Progress(track_tqdm=True)
):
    """
    Generate a video from text prompt and optional image/video using the Wan 2.2 TI2V model.
    
    Args:
        video_input: Optional input video file path
        image_preview: Preview image (numpy array) extracted from video or uploaded image
        prompt: Text prompt describing the desired video
        height: Target video height in pixels
        width: Target video width in pixels
        duration_seconds: Desired video duration in seconds
        sampling_steps: Number of denoising steps for video generation
        guide_scale: Guidance scale for classifier-free guidance
        shift: Sample shift parameter for the model
        seed: Random seed for reproducibility (-1 for random)
        progress: Gradio progress tracker
        
    Returns:
        Path to the generated video file
    """
    if seed == -1:
        seed = random.randint(0, sys.maxsize)

    # Ensure dimensions are multiples of MOD_VALUE
    target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
    target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)

    input_image = None
    
    # Process video input if provided
    if video_input is not None:
        if isinstance(video_input, str) and video_input.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.webm')):
            input_image = extract_first_frame_from_video(video_input)
        else:
            # Fallback to image preview
            if image_preview is not None:
                input_image = Image.fromarray(image_preview).convert("RGB")
    elif image_preview is not None:
        # Use image preview if no video input
        input_image = Image.fromarray(image_preview).convert("RGB")
    
    # Resize image to match target dimensions if we have an input image
    if input_image is not None:
        input_image = input_image.resize((target_w, target_h))
    
    # Calculate number of frames based on duration
    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)

    # Create size string for the pipeline
    size_str = f"{target_h}*{target_w}"

    video_tensor = pipeline.generate(
        input_prompt=prompt,
        img=input_image,  # Pass None for T2V, Image for I2V
        size=SIZE_CONFIGS.get(size_str, (target_h, target_w)),
        max_area=MAX_AREA_CONFIGS.get(size_str, target_h * target_w),
        frame_num=num_frames,
        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)
    )
    del video_tensor
    gc.collect()
    return video_path


# --- 3. Gradio Interface ---
css = ".gradio-container {max-width: 1200px !important; margin: 0 auto} #output_video {height: 500px;} #image_preview {height: 400px;}"

with gr.Blocks(css=css, theme=gr.themes.Soft(), delete_cache=(60, 900)) as demo:
    gr.Markdown("# Wan 2.2 TI2V 5B - Video/Image to Video")
    gr.Markdown("Generate high quality videos using **Wan 2.2 5B Text-Image-to-Video model** with support for video input. [[model]](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B), [[paper]](https://arxiv.org/abs/2503.20314)")

    with gr.Row():
        with gr.Column(scale=2):
            video_input = gr.Video(
                label="Upload Video or Image (optional - blank for text-to-video)",
                sources=["upload"],
            )
            image_preview = gr.Image(
                type="numpy", 
                label="Preview (first frame will be extracted from video)",
                elem_id="image_preview",
                interactive=False
            )
            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."
            )
            
            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    height_input = gr.Slider(
                        minimum=SLIDER_MIN_H, 
                        maximum=SLIDER_MAX_H, 
                        step=MOD_VALUE, 
                        value=DEFAULT_H_SLIDER_VALUE, 
                        label=f"Output Height (multiple of {MOD_VALUE})"
                    )
                    width_input = gr.Slider(
                        minimum=SLIDER_MIN_W, 
                        maximum=SLIDER_MAX_W, 
                        step=MOD_VALUE, 
                        value=DEFAULT_W_SLIDER_VALUE, 
                        label=f"Output Width (multiple of {MOD_VALUE})"
                    )
                steps_input = gr.Slider(label="Sampling Steps", minimum=10, maximum=50, value=38, 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)

        with gr.Column(scale=2):
            video_output = gr.Video(label="Generated Video", elem_id="output_video")
            run_button = gr.Button("Generate Video", variant="primary")
            
    # Add video/image upload handler
    video_input.upload(
        fn=handle_media_upload_for_dims_wan,
        inputs=[video_input, height_input, width_input],
        outputs=[height_input, width_input, image_preview]
    )
    
    video_input.clear(
        fn=lambda: (gr.update(value=DEFAULT_H_SLIDER_VALUE), 
                   gr.update(value=DEFAULT_W_SLIDER_VALUE),
                   None),
        inputs=[],
        outputs=[height_input, width_input, image_preview]
    )

    example_image_path = os.path.join(os.path.dirname(__file__), "examples/i2v_input.JPG")
    gr.Examples(
        examples=[
            [example_image_path, "The cat removes the glasses from its eyes.", 1088, 800, 1.5],
            [None, "A cinematic shot of a boat sailing on a calm sea at sunset.", 704, 1280, 2.0],
            [None, "Drone footage flying over a futuristic city with flying cars.", 704, 1280, 2.0],
        ],
        inputs=[video_input, prompt_input, height_input, width_input, duration_input],
        outputs=video_output,
        fn=generate_video,
        cache_examples="lazy",
    )

    run_button.click(
        fn=generate_video,
        inputs=[
            video_input, 
            image_preview, 
            prompt_input, 
            height_input, 
            width_input, 
            duration_input, 
            steps_input, 
            scale_input, 
            shift_input, 
            seed_input
        ],
        outputs=video_output
    )

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