app.py

import spaces
from huggingface_hub import snapshot_download, hf_hub_download
import os
import subprocess
import importlib, site
from PIL import Image
import uuid
import shutil
import time
import cv2
import json
import gradio as gr
import sys
import gc 
import numpy as np
from gradio_modal import Modal

def draw_points(img, points, labels):
    out = img.copy()
    for p, label in zip(points, labels):
        x, y = int(p[0]), int(p[1])
        color = (0, 255, 0) if label == 1.0 else (255, 0, 0)
        out = cv2.circle(out, (x, y), 10, color, -1)
    return out

def compose_img_mask(img, color_mask, fac: float = 0.5):
    out_f = fac * img / 255 + (1 - fac) * color_mask / 255
    out_u = (255 * out_f).astype("uint8")
    return out_u

BASE = os.path.dirname(os.path.abspath(__file__))
PREPROCESS_DIR = os.path.join(BASE, "wan", "modules", "animate", "preprocess")
sys.path.append(PREPROCESS_DIR)

# Re-discover all .pth/.egg-link files
for sitedir in site.getsitepackages():
    site.addsitedir(sitedir)

# Clear caches so importlib will pick up new modules
importlib.invalidate_caches()

def sh(cmd): subprocess.check_call(cmd, shell=True)



try:

    sh("pip install flash-attn --no-build-isolation")
    # print("Attempting to download and build sam2...")

    # print("download sam")
    # sam_dir = snapshot_download(repo_id="alexnasa/sam2")

    # @spaces.GPU(duration=500)
    # def install_sam():

    #     os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0"
    #     sh(f"cd {sam_dir} && python setup.py build_ext --inplace && pip install -e .")

    # print("install sam")
    # install_sam()
    
    print("Attempting to download")

    print("download sam")
    snapshot_download(repo_id="alexnasa/sam2_C_cpu", local_dir=f"{os.getcwd()}" )

    # tell Python to re-scan site-packages now that the egg-link exists
    import importlib, site; site.addsitedir(site.getsitepackages()[0]); importlib.invalidate_caches()

    print("sam2 installed successfully.")

except Exception as e:
    raise gr.Error("sam2 installation failed")

import torch
from generate import generate, load_model
from preprocess_data import run as run_preprocess
from preprocess_data import load_preprocess_models
from process_pipepline import get_frames
from sam2.build_sam import build_sam2, build_sam2_video_predictor
from sam2.sam2_image_predictor import SAM2ImagePredictor

print(f"Torch version: {torch.__version__}")

os.environ["PROCESSED_RESULTS"] = f"{os.getcwd()}/processed_results"

snapshot_download(repo_id="Wan-AI/Wan2.2-Animate-14B", local_dir="./Wan2.2-Animate-14B")
wan_animate = load_model(True)

rc_mapping = {
    "Video → Ref Image" : False,
    "Video ← Ref Image" : True
}

def prune_prompt_dicts(pts_by_frame: dict, lbs_by_frame: dict):
    """Remove frames with empty/invalid clicks; keep only frames where
    points and labels both exist and lengths match (>0)."""
    clean_pts, clean_lbs = {}, {}
    for k, pts in pts_by_frame.items():
        lbs = lbs_by_frame.get(k, [])
        # normalize weird cases (None, np arrays with size 0, etc.)
        pts_len = len(pts) if pts is not None else 0
        lbs_len = len(lbs) if lbs is not None else 0
        if pts_len > 0 and lbs_len > 0 and pts_len == lbs_len:
            clean_pts[int(k)] = pts
            clean_lbs[int(k)] = lbs
    return clean_pts, clean_lbs


def _ensure_frame_bucket(ps, frame_idx):
    if "selections" not in ps:
        ps["selections"] = {}
    if frame_idx not in ps["selections"]:
        ps["selections"][frame_idx] = {"points": [], "labels": []}

def _get_points_labels_for_frame(ps, frame_idx):
    if "selections" not in ps or frame_idx not in ps["selections"]:
        return [], []
    bucket = ps["selections"][frame_idx]
    return bucket["points"], bucket["labels"]


def preprocess_video(input_video_path, duration, session_id=None):
    
    if session_id is None:
        session_id = uuid.uuid4().hex

    output_dir = os.path.join(os.environ["PROCESSED_RESULTS"], session_id)
    os.makedirs(output_dir, exist_ok=True)

    process_video_path = os.path.join(output_dir, 'input_video.mp4')

    clip_and_set_fps(input_video_path, process_video_path, duration_s=duration)

    return process_video_path

def extract_audio_from_video_ffmpeg(video_path, output_wav_path, sample_rate=None):
    """
    Extracts the audio track from a video file and saves it as a WAV file.

    Args:
        video_path (str): Path to the input video file.
        output_wav_path (str): Path to save the extracted WAV file.
        sample_rate (int, optional): Output sample rate (e.g., 16000). 
                                     If None, keep the original.
    """
    cmd = [
        'ffmpeg',
        '-i', video_path,             # Input video
        '-vn',                        # Disable video
        '-acodec', 'pcm_s16le',       # 16-bit PCM (WAV format)
        '-ac', '1',                   # Mono channel (use '2' for stereo)
        '-y',                         # Overwrite output
        '-loglevel', 'error'          # Cleaner output
    ]

    # Only add the sample rate option if explicitly specified
    if sample_rate is not None:
        cmd.extend(['-ar', str(sample_rate)])

    cmd.append(output_wav_path)
    
    try:
        subprocess.run(cmd, check=True, capture_output=True, text=True)
        return True
    except subprocess.CalledProcessError as e:
        return False

    
def combine_video_and_audio_ffmpeg(video_path, audio_path, output_video_path):
    """
    Combines a silent MP4 video with a WAV audio file into a single MP4 with sound.

    Args:
        video_path (str): Path to the silent video file.
        audio_path (str): Path to the WAV audio file.
        output_video_path (str): Path to save the output MP4 with audio.
    """
    cmd = [
        'ffmpeg',
        '-i', video_path,            # Input video
        '-i', audio_path,            # Input audio
        '-c:v', 'copy',              # Copy video without re-encoding
        '-c:a', 'aac',               # Encode audio as AAC (MP4-compatible)
        '-shortest',                 # Stop when the shortest stream ends
        '-y',                        # Overwrite output
        '-loglevel', 'error',
        output_video_path
    ]
    
    try:
        subprocess.run(cmd, check=True, capture_output=True, text=True)
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"ffmpeg failed ({e.returncode}): {e.stderr.strip()}")


def clip_and_set_fps(input_video_path, output_video_path, duration_s=2, target_fps=30):
    """
    Trim to duration_s and (optionally) change FPS, without resizing.
    - If target_fps is None, keeps the original FPS.
    - Re-encodes video when changing FPS for predictable timing.
    """
    vf = []
    if target_fps is not None:
        vf.append(f"fps={target_fps}")
    vf_arg = ",".join(vf) if vf else None

    cmd = [
        "ffmpeg",
        "-nostdin",
        "-hide_banner",
        "-y",
        "-i", input_video_path,
        "-t", str(duration_s),
    ]

    if vf_arg:
        cmd += ["-vf", vf_arg]

    cmd += [
        "-c:v", "libx264",
        "-pix_fmt", "yuv420p",
        "-preset", "veryfast",
        "-crf", "18",
        "-c:a", "aac",          # use aac so MP4 stays compatible
        "-movflags", "+faststart",
        output_video_path,
    ]

    try:
        subprocess.run(cmd, check=True, capture_output=True, text=True)
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"ffmpeg failed ({e.returncode}): {e.stderr.strip()}")


def is_portrait(video_file):

    # Get video information
    cap = cv2.VideoCapture(video_file)
    if not cap.isOpened():
        error_msg = "Cannot open video file"
        gr.Warning(error_msg)
    
    orig_frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    orig_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    orig_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    cap.release()

    return orig_width < orig_height

def calculate_time_required(max_duration_s, rc_bool):

    if max_duration_s == -2:
        return None
    if max_duration_s == -1:
        return 75
    if max_duration_s == 1 or max_duration_s == 2:
        return 120
    elif max_duration_s == 3 or  max_duration_s == 4:
        return 180
    elif max_duration_s == 5 or  max_duration_s == 6:
        return 260
    elif max_duration_s == 7 or  max_duration_s == 8:
        return 330
    elif max_duration_s == 9 or  max_duration_s == 10:
        return 340
    
def get_display_time_required(max_duration_s, rc_bool):
    # the 30 seconds extra is just for saftey in case of a unexpected slow down
    return calculate_time_required(max_duration_s, rc_bool)

def update_time_required(max_duration_s, rc_str):

    rc_bool = rc_mapping[rc_str]

    duration_s = get_display_time_required(max_duration_s, rc_bool)
    duration_m = duration_s / 60
    
    return gr.update(value=f"⌚ Zero GPU Required: ~{duration_s}.0s ({duration_m:.1f} mins)")


def get_duration(input_video, max_duration_s, edited_frame, rc_bool, pts_by_frame, lbs_by_frame, session_id, progress):
    return calculate_time_required(max_duration_s, rc_bool)

@spaces.GPU(duration=get_duration)
def _animate(input_video, max_duration_s, edited_frame, rc_bool, pts_by_frame, lbs_by_frame, session_id=None, progress=gr.Progress(track_tqdm=True)):
    if session_id is None:
        session_id = uuid.uuid4().hex

    output_dir = os.path.join(os.environ["PROCESSED_RESULTS"], session_id)
    os.makedirs(output_dir, exist_ok=True)

    preprocess_dir = os.path.join(output_dir, "preprocess_dir")
    os.makedirs(preprocess_dir, exist_ok=True)

    output_video_path = os.path.join(output_dir, 'result.mp4')

    start_preprocess = time.time()

    if is_portrait(input_video):
        w, h = 480, 832
    else:
        w, h = 832, 480

    if max_duration_s == -1:
        if is_portrait(input_video):
            w, h = 360, 640
        else:
            w, h = 640, 360
    elif max_duration_s == -2:
        if is_portrait(input_video):
            w, h = 240, 426
        else:
            w, h = 426, 240

    tag_string = "replace_flag" if rc_bool else "retarget_flag"

    preprocess_model = load_preprocess_models(max_duration_s)

    # NOTE: run_preprocess now receives dicts keyed by frame_idx.
    # Update run_preprocess(...) accordingly in your preprocess_data.py.
    run_preprocess(
        preprocess_model, input_video, edited_frame, preprocess_dir,
        w, h, tag_string, pts_by_frame, lbs_by_frame
    )

    preprocess_time = time.time() - start_preprocess
    print(f"Preprocess took {preprocess_time:.2f} seconds")

    start_generate = time.time()
    generate(wan_animate, preprocess_dir, output_video_path, rc_bool)
    print(f"Generate took {time.time() - start_generate:.2f} seconds")

    gc.collect()
    torch.cuda.empty_cache()
    return output_video_path

def extract_frames(input_video, max_duration_s, session_id=None):
    if not input_video:
        raise gr.Error("Please provide an video")
    
    if session_id is None:
        session_id = uuid.uuid4().hex

    processed_video = preprocess_video(input_video, max_duration_s, session_id)

    if is_portrait(processed_video):
        size = (480, 832)
    else:
        size = (832, 480)
    
    frames = get_frames(processed_video, size)

    output_dir = os.path.join(os.environ["PROCESSED_RESULTS"], session_id)
    os.makedirs(output_dir, exist_ok=True)

    preprocess_dir = os.path.join(output_dir, "preprocess_dir")
    os.makedirs(preprocess_dir, exist_ok=True)

    sam_frame_path = os.path.join(preprocess_dir, "sam_frame.png")
    Image.fromarray(frames[0]).save(sam_frame_path)

    slider = gr.update(minimum=0, maximum=len(frames) - 1, value=0, step=1)

    return sam_frame_path, frames, slider

def set_positive(prompt_state):
    prompt_state["cur_label_val"] = 1.0
    return prompt_state

def set_negative(prompt_state):
    prompt_state["cur_label_val"] = 0.0
    return prompt_state

def set_sam2_mode(prompt_state, mode):
    prompt_state["cur_label_val"] = 1.0 if mode == "✔️Include" else 0.0
    return prompt_state

def set_reset(prompt_state):
    prompt_state["selections"] = {}
    prompt_state["active_frame"] = 0
    prompt_state["cur_label_val"] = 1.0
    return prompt_state, gr.update(value="✔️Include")


def get_select_coords(prompt_state, frames, frame_idx, img_path, evt: gr.SelectData):
    img = np.array(frames[frame_idx])

    # y, x from the event
    i = evt.index[1]
    j = evt.index[0]

    # make sure we're editing the correct frame
    prompt_state["active_frame"] = int(frame_idx)
    binary_mask, prompt_state = add_point(prompt_state, frames, frame_idx, i, j)

    colored_mask = np.zeros_like(img)
    colored_mask[binary_mask] = [0, 255, 0]
    out_u = compose_img_mask(img, colored_mask, 0.5)

    pts, lbs = _get_points_labels_for_frame(prompt_state, frame_idx)
    overlaid = draw_points(out_u, pts, lbs)
    return overlaid, prompt_state
 

def add_point(prompt_state, frames, frame_idx, i, j):
    _ensure_frame_bucket(prompt_state, frame_idx)
    prompt_state["selections"][frame_idx]["points"].append([int(j), int(i)])
    prompt_state["selections"][frame_idx]["labels"].append(prompt_state["cur_label_val"])

    pts, lbs = _get_points_labels_for_frame(prompt_state, frame_idx)
    mask, logit, prompt_state = get_sam_mask(
        prompt_state, frames, frame_idx,
        np.array(pts, dtype=np.float32),
        np.array(lbs, dtype=np.int32)
    )
    return mask, prompt_state


def open_modal(input_video, edited_frame):

    if not input_video:
        raise gr.Error("Please provide an video")
    
    if not edited_frame:
        raise gr.Error("Please provide an image")
    
    return Modal(visible=False)

def get_sam_mask(prompt_state, frames, frame_idx, input_points, input_labels):
    """
    :param frame_idx: int
    :param input_points: (N, 2) float32
    :param input_labels: (N,)   int32
    :return: (H, W) boolean mask, (H, W) float32 logits (or None), prompt_state
    """

    model_cfg = "sam2_hiera_l.yaml"
    ckpt_path = "./Wan2.2-Animate-14B/process_checkpoint"
    sam2_checkpoint_path = os.path.join(ckpt_path, 'sam2/sam2_hiera_large.pt')
    
    video_predictor_local = build_sam2_video_predictor(model_cfg, sam2_checkpoint_path, device="cpu")    
    inference_state = video_predictor_local.init_state(images=np.array(frames), device="cpu")

    # Run point prompt
    _, out_obj_ids, out_mask_logits = video_predictor_local.add_new_points(
        inference_state=inference_state,
        frame_idx=int(frame_idx),
        obj_id=0,
        points=input_points,
        labels=input_labels,
    )

    # Get logits for object 0 as numpy
    logits = out_mask_logits[0].detach().cpu().numpy()  # shape could be (H,W) or (1,H,W)
    if logits.ndim == 3:
        logits = logits.squeeze(0)

    # --- Normalize to the frame's size (critical) ---
    H, W = frames[frame_idx].shape[:2]
    if logits.shape != (H, W):
        # Resize logits to (W,H); use nearest for crisp mask boundaries
        # (If you want smoother logits, you can use INTER_LINEAR then threshold)
        resized = cv2.resize(logits, (W, H), interpolation=cv2.INTER_NEAREST)
        logits = resized.astype(np.float32)

    # Threshold -> boolean mask
    mask = (logits > 0.0)

    return mask, logits, prompt_state


def animate_scene(input_video, max_duration_s, edited_frame, rc_str,
                  prompt_state=None, session_id=None,
                  progress=gr.Progress(track_tqdm=True)):

    # Build dicts keyed by frame index, or empty if none
    if not prompt_state or "selections" not in prompt_state:
        pts_by_frame, lbs_by_frame = {}, {}
    else:
        pts_by_frame = {int(k): v["points"] for k, v in prompt_state["selections"].items()}
        lbs_by_frame = {int(k): v["labels"] for k, v in prompt_state["selections"].items()}

    # 🧹 prune empties/invalids
    pts_by_frame, lbs_by_frame = prune_prompt_dicts(pts_by_frame, lbs_by_frame)

    # If nothing left, treat as vanilla run (no prompts)
    if not pts_by_frame:
        pts_by_frame, lbs_by_frame = {}, {}


    if not input_video:
        raise gr.Error("Please provide an video")
    if not edited_frame:
        raise gr.Error("Please provide an image")

    if session_id is None:
        session_id = uuid.uuid4().hex

    input_video = preprocess_video(input_video, max_duration_s, session_id)
    rc_bool = rc_mapping[rc_str]

    output_dir = os.path.join(os.environ["PROCESSED_RESULTS"], session_id)
    os.makedirs(output_dir, exist_ok=True)

    input_audio_path = os.path.join(output_dir, 'input_audio.wav')
    audio_extracted = extract_audio_from_video_ffmpeg(input_video, input_audio_path)

    edited_frame_png = os.path.join(output_dir, 'edited_frame.png')
    Image.open(edited_frame).save(edited_frame_png)

    print(f'{session_id} inference started with {rc_str} and duration of {max_duration_s}')

    try:
        
        output_video_path = _animate(
            input_video, max_duration_s, edited_frame_png, rc_bool,
            pts_by_frame, lbs_by_frame, session_id, progress
        )
    except Exception as e:
        err = str(e).lower()
        print(f"{session_id} failed due to {err}")

        if "zerogpu quotas" not in err and "pro gpu quota" not in err:
            raise

        if max_duration_s > 2:
            raise

        max_duration_s = -1

        try:
            pts_by_frame, lbs_by_frame = {}, {}
            
            output_video_path = _animate(
                input_video, max_duration_s, edited_frame_png, rc_bool,
                pts_by_frame, lbs_by_frame, session_id, progress
            )
        except Exception as e:
            err = str(e).lower()
            print(f"{session_id} failed due to {err}")

            if "unlogged" in err:
                max_duration_s = -2

                gr.Info("Sign up for free or login to get at least 2 generations for free")

                output_video_path = _animate(
                    input_video, max_duration_s, edited_frame_png, rc_bool,
                    pts_by_frame, lbs_by_frame, session_id, progress
                )
            else:
                raise   

    final_video_path = os.path.join(output_dir, 'final_result.mp4')

    preprocess_dir = os.path.join(output_dir, "preprocess_dir")
    pose_video = os.path.join(preprocess_dir, 'src_pose.mp4')

    if rc_bool:
        mask_video = os.path.join(preprocess_dir, 'src_mask.mp4')
        bg_video   = os.path.join(preprocess_dir, 'src_bg.mp4')
        face_video = os.path.join(preprocess_dir, 'src_face.mp4')
    else:
        mask_video = pose_video
        bg_video   = pose_video
        face_video = pose_video

    if audio_extracted:
        combine_video_and_audio_ffmpeg(output_video_path, input_audio_path, final_video_path)
    else:
        final_video_path = output_video_path

    print(f"task for {session_id} finalised")
    return final_video_path, pose_video, bg_video, mask_video, face_video

css = """
    #col-container {
        margin: 0 auto;
        max-width: 1600px;
    }

    #modal-container {
    width: 100vw;            /* Take full viewport width */
    height: 100vh;           /* Take full viewport height (optional) */
    display: flex;           
    justify-content: center; /* Center content horizontally */
    align-items: center;     /* Center content vertically if desired */
    }

    #modal-content {
    width: 100%;
    max-width: 700px;         /* Limit content width */
    margin: 0 auto;
    border-radius: 8px;
    padding: 1.5rem;
    }

    #step-column {
        padding: 10px;
        border-radius: 8px;
        box-shadow: var(--card-shadow);
        margin: 10px;
    }

    #col-showcase {
        margin: 0 auto;
        max-width: 1100px;
    }

    .button-gradient {
        background: linear-gradient(45deg, rgb(255, 65, 108), rgb(255, 75, 43), rgb(255, 155, 0), rgb(255, 65, 108)) 0% 0% / 400% 400%;
        border: none;
        padding: 14px 28px;
        font-size: 16px;
        font-weight: bold;
        color: white;
        border-radius: 10px;
        cursor: pointer;
        transition: 0.3s ease-in-out;
        animation: 2s linear 0s infinite normal none running gradientAnimation;
        box-shadow: rgba(255, 65, 108, 0.6) 0px 4px 10px;
    }

    .toggle-container {
    display: inline-flex;
    background-color: #ffd6ff;  /* light pink background */
    border-radius: 9999px;
    padding: 4px;
    position: relative;
    width: fit-content;
    font-family: sans-serif;
    }

    .toggle-container input[type="radio"] {
    display: none;
    }

    .toggle-container label {
    position: relative;
    z-index: 2;
    flex: 1;
    text-align: center;
    font-weight: 700;
    color: #4b2ab5; /* dark purple text for unselected */
    padding: 6px 22px;
    border-radius: 9999px;
    cursor: pointer;
    transition: color 0.25s ease;
    }

    /* Moving highlight */
    .toggle-highlight {
    position: absolute;
    top: 4px;
    left: 4px;
    width: calc(50% - 4px);
    height: calc(100% - 8px);
    background-color: #4b2ab5; /* dark purple background */
    border-radius: 9999px;
    transition: transform 0.25s ease;
    z-index: 1;
    }

    /* When "True" is checked */
    #true:checked ~ label[for="true"] {
    color: #ffd6ff; /* light pink text */
    }

    /* When "False" is checked */
    #false:checked ~ label[for="false"] {
    color: #ffd6ff; /* light pink text */
    }

    /* Move highlight to right side when False is checked */
    #false:checked ~ .toggle-highlight {
    transform: translateX(100%);
    }
    """

def set_input_image(prompt_state, frames, frame_index, session_id):
    output_dir = os.path.join(os.environ["PROCESSED_RESULTS"], session_id)
    os.makedirs(output_dir, exist_ok=True)

    preprocess_dir = os.path.join(output_dir, "preprocess_dir")
    os.makedirs(preprocess_dir, exist_ok=True)

    sam_frame_path = os.path.join(preprocess_dir, "sam_frame.png")

    # clamp/check
    if frame_index is None or frame_index < 0 or (frames and frame_index >= len(frames)):
        frame_index = 0
    frame_index = int(frame_index)

    # mark active frame and ensure bucket exists
    prompt_state["active_frame"] = frame_index
    _ensure_frame_bucket(prompt_state, frame_index)

    # base image
    img = np.array(frames[frame_index])

    # fetch any saved selections for this frame
    pts, lbs = _get_points_labels_for_frame(prompt_state, frame_index)

    if len(pts) > 0:
        # recompute mask from saved clicks on this frame
        mask, _logits, _ = get_sam_mask(
            prompt_state,
            frames,
            frame_index,
            np.array(pts, dtype=np.float32),
            np.array(lbs, dtype=np.int32),
        )

        # overlay mask + points
        colored_mask = np.zeros_like(img)
        colored_mask[mask.astype(bool)] = [0, 255, 0]
        out_u = compose_img_mask(img, colored_mask, 0.5)
        overlaid = draw_points(out_u, pts, lbs)
        Image.fromarray(overlaid).save(sam_frame_path)
    else:
        # no selections yet on this frame → show plain frame
        Image.fromarray(img).save(sam_frame_path)

    return sam_frame_path, prompt_state


def open_modal(input_video, edited_frame):

    if not input_video:
        raise gr.Error("Please provide an video")
    
    if not edited_frame:
        raise gr.Error("Please provide an image")
    
    return Modal(visible=True)

def log_change(log_source, session_id, meta_data = None):

    if not meta_data:
        print(f'{session_id} changed {log_source}')
    else:
        print(f'{session_id} changed {log_source} with {meta_data}')

def start_session(request: gr.Request):

    return request.session_hash

def cleanup(request: gr.Request):

    sid = request.session_hash
    
    if sid:
        print(f"{sid} left")
        d1 = os.path.join(os.environ["PROCESSED_RESULTS"], sid)
        shutil.rmtree(d1, ignore_errors=True)

with gr.Blocks(css=css, title="Wan 2.2 Animate --replace", theme=gr.themes.Ocean()) as demo:
        
    session_state = gr.State()

    prompt_state = gr.State({
        "selections": {  # frame_idx -> {"points": [[x,y],...], "labels": [0/1,...]}
            # e.g. 0: {"points": [], "labels": []}
        },
        "active_frame": 0,  # which frame the user is editing/adding points to
        "cur_label_val": 1.0,
    })

    frames = gr.State()

    demo.load(start_session, outputs=[session_state])

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            gr.HTML(
                """
                <div style="text-align: center;">
                    <p style="font-size:16px; display: inline; margin: 0;">
                        <strong>Wan2.2-Animate-14B </strong>
                    </p>
                    <a href="https://huggingface.co/Wan-AI/Wan2.2-Animate-14B" style="display: inline-block; vertical-align: middle; margin-left: 0.5em;">
                        [Model]
                    </a>
                    <p style="font-size:16px; display: inline; margin: 0;">
                        -- HF Space By:
                    </p>
                    <a href="https://huggingface.co/alexnasa" style="display: inline-block; vertical-align: middle; margin-left: 0.5em;">
                        <img src="https://img.shields.io/badge/🤗-Follow Me-yellow.svg">
                    </a>
                </div>
                """
            )
        with gr.Row():
            with gr.Column(elem_id="step-column"):
                gr.HTML("""
                <div>
                    <span style="font-size: 24px;">1. Upload a Video</span><br>
                </div>
                """)
                input_video = gr.Video(label="Input Video", height=512)

                max_duration_slider = gr.Slider(2, 10, 2, step=1, label="Max Duration")

                gr.Examples(
                    examples=[
                        
                        [
                            "./examples/martialart.mp4", 
                        ],
                        
                        [
                            "./examples/test_example.mp4", 
                        ],
                 
                        [
                            "./examples/dream.mp4", 
                        ],

                    ],
                    inputs=[input_video],
                    cache_examples=False,
                )


            with gr.Column(elem_id="step-column"):
                gr.HTML("""
                <div>
                    <span style="font-size: 24px;">2. Upload a Ref Image</span><br>
                </div>
                """)
                edited_frame = gr.Image(label="Ref Image", type="filepath", height=512)

                default_replace_string = "Video ← Ref Image"
                replace_character_string = gr.Radio(
                    ["Video → Ref Image", "Video ← Ref Image"], value=default_replace_string, show_label=False
                )


                gr.Examples(
                    examples=[

                        [
                            "./examples/ali.png", 
                        ],

                        [
                            "./examples/james.png", 
                        ],


                        [
                            "./examples/amber.png", 
                        ],

                        [
                            "./examples/ella.png", 
                        ],

                    ],
                    inputs=[edited_frame],
                    cache_examples=False,
                )

            with gr.Column(elem_id="step-column"):
                gr.HTML("""
                <div>
                    <span style="font-size: 24px;">3. Wan Animate it!</span><br>
                </div>
                """)
                output_video = gr.Video(label="Edited Video", height=512)
                duration_s = get_display_time_required(2, default_replace_string)
                duration_m = duration_s / 60
                
                time_required = f"⌚ Zero GPU Required: ~{duration_s}.0s ({duration_m:.1f} mins)"

                time_required = gr.Text(value=time_required, show_label=False, visible=True)
                
                action_button = gr.Button("Wan Animate 🦆", variant='primary', elem_classes="button-gradient")

                with gr.Accordion("Advanced 🎭", open=False, visible=True):
                    with gr.Row():
                        adv_masking = gr.Button("Wan Animate PRO ⚙️")
                    with gr.Row():
                        pose_video = gr.Video(label="Pose Video")
                        bg_video = gr.Video(label="Background Video")
                    with gr.Row():
                        face_video = gr.Video(label="Face Video")
                        mask_video = gr.Video(label="Mask Video")
                        

        with gr.Row():
            with gr.Column(elem_id="col-showcase"):

                gr.Examples(
                    examples=[
                       
                        [
                            "./examples/okay.mp4", 
                            2,
                            "./examples/amber.png",     
                            "Video ← Ref Image"
                        ],

                        [
                            "./examples/kat.mp4", 
                            2,
                            "./examples/cat.png",      
                            "Video → Ref Image"
                        ],

                        [
                            "./examples/kat.mp4", 
                            2,
                            "./examples/anime.png",      
                            "Video → Ref Image"
                        ],

                        [
                            "./examples/rachel.mp4", 
                            2,
                            "./examples/james.png",     
                            "Video ← Ref Image"
                        ],

                        [
                            "./examples/test_example.mp4", 
                            2,
                            "./examples/ella.png",     
                            "Video ← Ref Image"
                        ],

                        [
                            "./examples/paul.mp4", 
                            2,
                            "./examples/man.png",      
                            "Video → Ref Image"
                        ],

                        [
                            "./examples/desi.mp4", 
                            2,
                            "./examples/desi.png",     
                            "Video ← Ref Image"
                        ],

                    ],
                    inputs=[input_video, max_duration_slider, edited_frame, replace_character_string],
                    outputs=[output_video, pose_video, bg_video, mask_video, face_video],
                    fn=animate_scene,
                    cache_examples=True,
                )

    with Modal(visible=False, elem_id="modal-container") as modal:
        with gr.Column(elem_id="modal-content"):
            with gr.Row():
                gr.Markdown("## WAN Animate PRO")
                gr.Markdown("Choose a frame and adjust the relevant mask, by including and excluding points, to propagate across frames. If no mask is selected, one will be chosen automatically.")
                gr.Markdown("")
            with gr.Row():
                max_duration_slider_2 = gr.Slider(2, 10, 2, step=1, label="Max Duration", visible=True)
                frame_index = gr.Slider(label="Frame Index", minimum=0, maximum=0, value=0, step=1,)
            with gr.Row():
                first_frame = gr.Image(type="filepath", height=384, interactive=False)
            with gr.Row():

                with gr.Column():
                    sam2_mode_select = gr.Radio(
                        ["✔️Include", "❌Exclude"], value="✔️Include", show_label=False
                    )

                with gr.Column():

                    confirm_sam2_points = gr.Button("Wan Animate 🦆", variant='primary', elem_classes="button-gradient")
                    cancel_button = gr.Button("Cancel")

    action_button.click(fn=animate_scene, inputs=[input_video, max_duration_slider, edited_frame, replace_character_string, prompt_state, session_state], outputs=[output_video, pose_video, bg_video, mask_video, face_video])
    replace_character_string.change(update_time_required, inputs=[max_duration_slider, replace_character_string], outputs=[time_required]) \
        .then(log_change, inputs=[gr.State("mode"), session_state, replace_character_string])

    def sync_slider(val):
        return val
    
    max_duration_slider_2.change(sync_slider, max_duration_slider_2, max_duration_slider)
    max_duration_slider.change(log_change, inputs=[gr.State("slider"), session_state, max_duration_slider]) \
        .then(update_time_required, inputs=[max_duration_slider, replace_character_string], outputs=[time_required]) \
        .then(extract_frames, inputs=[input_video, max_duration_slider, session_state], outputs=[first_frame, frames, frame_index])
    
    input_video.change(log_change, inputs=[gr.State("video"), session_state]) \
        .then(set_reset, inputs=[prompt_state], outputs=[prompt_state, sam2_mode_select])
    edited_frame.change(log_change, inputs=[gr.State("ref image"), session_state])        

    adv_masking.click(log_change, inputs=[gr.State("advanced masking"), session_state]) \
        .then(set_reset, inputs=[prompt_state], outputs=[prompt_state, sam2_mode_select]) \
        .then(open_modal, [input_video, edited_frame], modal) \
        .then(extract_frames, inputs=[input_video, max_duration_slider, session_state], outputs=[first_frame, frames, frame_index])

    confirm_sam2_points.click(lambda: Modal(visible=False), None, modal) \
                        .then(fn=animate_scene, inputs=[input_video, max_duration_slider, edited_frame, replace_character_string, prompt_state, session_state], outputs=[output_video, pose_video, bg_video, mask_video, face_video]) \
                        .then(set_reset, inputs=[prompt_state], outputs=[prompt_state, sam2_mode_select])
    
    sam2_mode_select.change(set_sam2_mode, [prompt_state, sam2_mode_select], [prompt_state])

    cancel_button.click(set_reset, inputs=[prompt_state], outputs=[prompt_state, sam2_mode_select]) \
        .then(lambda: Modal(visible=False), None, modal)

    frame_index.release(set_input_image, [prompt_state, frames, frame_index, session_state], [first_frame, prompt_state])
    first_frame.select(get_select_coords, [prompt_state, frames, frame_index, first_frame], [first_frame, prompt_state])



if __name__ == "__main__":
    demo.queue()
    demo.unload(cleanup)
    demo.launch(ssr_mode=False, share=True)