app.py

import os
import subprocess
import sys

def install_flash_attn():
    """Auto-detect CUDA, PyTorch, Python versions and install the matching pre-built flash-attn wheel."""
    import torch

    # Python version (e.g., "cp310")
    py_major = sys.version_info.major
    py_minor = sys.version_info.minor
    cp_tag = f"cp{py_major}{py_minor}"

    # PyTorch version (e.g., "2.4" from "2.4.0")
    torch_version = torch.__version__.split("+")[0]  # strip +cu121 if present
    torch_major_minor = ".".join(torch_version.split(".")[:2])  # "2.4"

    # CUDA version (e.g., "cu124" from "12.4")
    cuda_version = torch.version.cuda
    if cuda_version is None:
        print("No CUDA detected, skipping flash-attn installation.")
        return
    cuda_major_minor = cuda_version.replace(".", "")  # "124"
    # flash-attn wheels use shortened CUDA tags like "cu12" (just major) or "cu121", "cu124"
    # Check available tags: most wheels use "cu12" for any 12.x
    cuda_tag_short = f"cu{cuda_version.split('.')[0]}"  # "cu12"

    # CXX11 ABI
    cxx11_abi = torch._C._GLIBCXX_USE_CXX11_ABI
    abi_tag = "cxx11abiTRUE" if cxx11_abi else "cxx11abiFALSE"

    # flash-attn version to install
    flash_attn_version = "2.8.3"

    # Construct the wheel filename
    # Example: flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
    wheel_name = (
        f"flash_attn-{flash_attn_version}+"
        f"{cuda_tag_short}torch{torch_major_minor}{abi_tag}-"
        f"{cp_tag}-{cp_tag}-linux_x86_64.whl"
    )

    base_url = f"https://github.com/Dao-AILab/flash-attention/releases/download/v{flash_attn_version}"
    wheel_url = f"{base_url}/{wheel_name}"

    print(f"Detected environment:")
    print(f"  Python:    {py_major}.{py_minor} ({cp_tag})")
    print(f"  PyTorch:   {torch_version} (torch{torch_major_minor})")
    print(f"  CUDA:      {cuda_version} ({cuda_tag_short})")
    print(f"  CXX11 ABI: {cxx11_abi} ({abi_tag})")
    print(f"  Wheel URL: {wheel_url}")

    result = subprocess.run(
        [sys.executable, "-m", "pip", "install", wheel_url],
        capture_output=True,
        text=True,
    )

    if result.returncode != 0:
        print(f"Pre-built wheel failed:\n{result.stderr}")
        print("Falling back to building flash-attn from source (this may take a while)...")
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "flash-attn", "--no-build-isolation"],
            check=True,
        )
    else:
        print("flash-attn installed successfully from pre-built wheel.")
        print(result.stdout)


install_flash_attn()

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import copy
import torch
from torchvision.io import write_video
from torch.utils.data import Dataset
import numpy as np
from pathlib import Path
from hydra import initialize, compose
from hydra.core.global_hydra import GlobalHydra

from b_spline import build_clamped_bspline, equidistant_points_on_spline

torch.set_grad_enabled(False)

from palette import _palette

import gradio as gr
import numpy as np
from scipy import ndimage
from PIL import Image
import os
from pathlib import Path
import cv2

# from sam_segment import predict_masks_with_sam
from segment_anything import SamPredictor, sam_model_registry

from tensor_utils import (
    image_to_pil,
    image_to_np,
    bbox_from_mask,
    draw_bbox_on_image,
    draw_mask_on_image,
    draw_points_on_image,
    draw_lines_on_image,
    trajectory_interpolate,
    dilate_mask,
    dilate_masks,
)

from optimize_utils import (
    MultiTrajectory,
    Trajectory,
)

import sys

from utils.misc import set_seed

from stream_inference_wrapper import StreamInferenceWrapper
from stream_drag_inference_wrapper import StreamDragInferenceWrapper
from utils.dataset import TextDataset

from video_operations import generate_video, optimize_video

# from compute_objmc import visualize_ground_truth_from_trajectory_file


def extract_layer_as_mask(image_editor, layer_index=0):
    if len(image_editor["layers"]) > layer_index:
        layer = image_editor["layers"][layer_index]
        return image_to_np(layer.convert("L")) > 0
    return None


def apply_mask_to_image(
    mask: np.ndarray | None,
    image: np.ndarray | Image.Image,
    mask_color: list[int],
    alpha: float,
) -> None | Image.Image:
    if image is None:
        return None
    if mask is None:
        return image_to_pil(image)
    mask = np.array(mask)
    new_image = draw_mask_on_image(
        image,
        mask,
        mask_color=mask_color,
        alpha=alpha,
    )
    return new_image


def apply_movable_mask_to_image(
    mask: np.ndarray | None,
    image: np.ndarray | Image.Image,
):
    return apply_mask_to_image(
        mask=mask,
        image=image,
        mask_color=(255, 255, 255),
        alpha=0.35,
    )


def apply_target_mask_to_image(
    mask: np.ndarray | None,
    image: np.ndarray | Image.Image,
):
    return apply_mask_to_image(
        mask=mask,
        image=image,
        mask_color=(255, 64, 64),
        alpha=0.5,
    )


def get_video_last_frame(
    # video: Optional[torch.Tensor],  # None or (t, h, w, c)
    video_path: str,
):
    """
    Loads the last frame from a video.

    Returns:
        Image: The last frame as a PIL Image.
    """
    print(f"Getting last frame from video: {video_path = }")
    if video_path is None:
        return None

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Failed to open video: {video_path}")
        return None
    try:
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if frame_count <= 0:
            print(f"Video has non-positive frame count: {frame_count}")
            cap.release()
            return None

        # Try direct seek to last frame
        target_index = frame_count - 1
        cap.set(cv2.CAP_PROP_POS_FRAMES, target_index)
        ret, frame = cap.read()

        # Fallback: iterate to last frame if random access failed
        if (not ret) or frame is None:
            print("Direct seek failed, iterating through frames...")
            cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
            last_valid = None
            while True:
                ret_i, frame_i = cap.read()
                if not ret_i:
                    break
                last_valid = frame_i
            frame = last_valid

        if frame is None:
            print("Could not retrieve last frame.")
            return None

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        last_frame_image = Image.fromarray(frame)
        return last_frame_image
    except Exception as e:
        print(f"Error extracting last frame: {e}")
        return None
    finally:
        cap.release()


def sam_predict_segmentation(
    sam_predictor: SamPredictor,
    origin_image: Image.Image | np.ndarray,
    restriction_mask: np.ndarray,  # (h, w), bool
    click_points: list[tuple[int, int]],
    previous_sam_logits: np.ndarray | None,  # (3, 256, 256)
):
    # print(f"{restriction_mask.shape = }")

    origin_image_np = image_to_np(origin_image)
    # print(f"{origin_image_np.shape = }")
    sam_predictor.set_image(origin_image_np)

    if previous_sam_logits is not None:
        print(f"{previous_sam_logits.shape = }")
    else:
        print(f"{previous_sam_logits = }")
    masks, scores, logits = sam_predictor.predict(
        point_coords=np.array(click_points),
        point_labels=np.ones((len(click_points),)),
        mask_input=(previous_sam_logits[0:1] if previous_sam_logits is not None else None),
        multimask_output=True,
    )
    # mask: np.ndarray
    # scores: np.ndarray
    # logits: np.ndarray
    # print(f"{masks.shape = }") # (3, 480, 832)
    # print(f"{logits.shape = }") # (3, 256, 256)

    mask = masks[0]
    mask *= restriction_mask

    logits *= cv2.resize(
        restriction_mask.astype(np.uint8),
        dsize=(256, 256),
        interpolation=cv2.INTER_LINEAR,
    )

    return mask, logits


def sam_predict_segmentation_wrapper(
    sam_predictor: SamPredictor,
    original_image: Image.Image | np.ndarray,
    restriction_mask: np.ndarray | None,
    previous_click_points: list[tuple[int, int]],
    previous_sam_logits: np.ndarray | None,
    bypass_sam_model: bool,
    evt: gr.SelectData,
):
    # print(f"{restriction_mask = }")
    original_image = image_to_pil(original_image).convert("RGB")

    if restriction_mask is None:
        labeled_restriction_mask = np.zeros(
            (original_image.height, original_image.width), dtype=np.int32
        )
    else:
        labeled_restriction_mask, _ = ndimage.label(restriction_mask, structure=np.ones((3, 3)))
    # print(f"{labeled_restriction_mask = }")
    current_click_label = labeled_restriction_mask[evt.index[1], evt.index[0]]
    # print(f"{current_click_label = }")

    if current_click_label == 0:
        selected_component_mask = np.zeros_like(labeled_restriction_mask, dtype=bool)
    else:
        selected_component_mask = labeled_restriction_mask == current_click_label
    # print(f"{selected_component_mask = }")

    if bypass_sam_model:
        click_points = [evt.index]
        mask = selected_component_mask
        logits = None
    else:
        click_points = previous_click_points + [evt.index]
        mask, logits = sam_predict_segmentation(
            sam_predictor=sam_predictor,
            origin_image=original_image,
            restriction_mask=selected_component_mask,
            click_points=click_points,
            previous_sam_logits=previous_sam_logits,
        )

    return mask, click_points, logits


def draw_all_sam_masks(image: Image.Image | None, mask_list: list[np.ndarray]):
    if image is None:
        return None
    if len(mask_list) == 0:
        pass
    else:
        for mask_idx, mask in enumerate(mask_list):
            image = draw_mask_on_image(
                image,
                mask,
                mask_color=tuple(_palette[mask_idx + 1]),
                alpha=0.65,
            )
    return image


def draw_sam_mask_wrapper(
    original_image,
    movable_mask,
    current_mask: np.ndarray | None,
    previous_masks: list[np.ndarray],
    click_points: list[tuple[int, int]],
):
    image = apply_movable_mask_to_image(
        image=original_image,
        mask=movable_mask,
    )
    if image is None:
        return None
    image = draw_all_sam_masks(
        image,
        previous_masks + ([current_mask] if current_mask is not None else []),
    )
    image = draw_points_on_image(
        image,
        click_points,
        color=[(0, 255, 0, 255) for l in click_points],
        radius=5,
    )
    return image


def save_sam_masks(
    current_mask: np.ndarray | None,
    previous_masks: list[np.ndarray],
):
    new_masks = previous_masks + ([current_mask] if current_mask is not None else [])
    return None, new_masks, [], None


def select_target_sam_mask(
    masks_list: list[np.ndarray],
    evt: gr.SelectData,
):
    is_match_mask = False
    for mask_index, sam_mask in enumerate(masks_list):
        # check if evt point in sam_mask
        if sam_mask[evt.index[1], evt.index[0]]:
            is_match_mask = True
            break

    if not is_match_mask:
        print(f"Mask not found for {evt.index = }")
        mask_index = -1
    return mask_index


def draw_rotation_trajectory(
    image,
    points,
):
    image = draw_points_on_image(
        image,
        [points[0]],
        color="green",
        radius=15,
    )
    if len(points) > 1:
        image = draw_points_on_image(
            image,
            points[1:],
            color=[
                (
                    255 - int(float(i) / len(points[1:]) * 255.0),
                    64,
                    int(float(i) / len(points[1:]) * 255.0),
                    255,
                )
                for i in range(len(points[1:]))
            ],
            radius=5,
        )
        for point in points[1:]:
            image = draw_lines_on_image(
                image,
                [points[0], point],
                color="green",
                width=3,
            )

    return image


def draw_translation_trajectory(
    image,
    points,
    control_points: list[tuple[int, int]] = [],
    is_draw_control_points: bool = True,
):
    if len(points) == 1:
        image = draw_points_on_image(
            image,
            points,
            color=[(255, 64, 0, 255)],
            radius=6,
        )
        return image
    if is_draw_control_points and (len(control_points) >= 2):
        image = draw_points_on_image(
            image,
            control_points,
            color=[(0, 255, 0, 255) for _ in control_points],
            radius=3,
        )
        image = draw_lines_on_image(
            image,
            control_points,
            color=[(0, 255, 0, 255) for _ in control_points],
            width=2,
        )
    image = draw_lines_on_image(
        image,
        points,
        color=[
            (
                255 - int(float(i) / len(points[1:]) * 255.0),
                64,
                int(float(i) / len(points[1:]) * 255.0),
                255,
            )
            for i in range(len(points))
        ],
        width=4,
    )
    image = draw_points_on_image(
        image,
        points,
        color=[
            (
                255 - int(float(i) / len(points[1:]) * 255.0),
                64,
                int(float(i) / len(points[1:]) * 255.0),
                255,
            )
            for i in range(len(points))
        ],
        radius=6,
    )

    return image


def draw_all_trajectories(
    image,
    trajectory: MultiTrajectory,
    is_draw_control_points: bool = True,
):
    print(
        f"""
draw_all_trajectories:
"""
    )
    if trajectory.trajectories is None:
        return image
    for traj in trajectory.trajectories:
        if traj.original_trajectory is None:
            continue
        original_traj = traj.original_trajectory
        if original_traj["is_rotation"]:
            image = draw_rotation_trajectory(image, original_traj["points"])
        else:
            image = draw_translation_trajectory(
                image,
                original_traj["points"],
                original_traj.get("control_points", []),
                is_draw_control_points=is_draw_control_points,
            )

    return image


def draw_trajectory_image(
    original_image,
    movable_mask,
    mask_index,
    masks_list: list[np.ndarray],
    trajectory: MultiTrajectory,
    is_draw_bbox: bool = True,
    is_draw_control_points: bool = True,
):
    print(
        f"""
draw_trajectory_image:
    {mask_index = }
"""
    )
    image = apply_movable_mask_to_image(
        mask=movable_mask,
        image=original_image,
    )
    image = draw_all_sam_masks(image, masks_list)
    if (
        (mask_index is not None)
        and (mask_index >= 0)
        and (mask_index < len(masks_list))
        and is_draw_bbox
    ):
        image = draw_bbox_on_image(image, bbox_from_mask(masks_list[mask_index]))
    image = draw_all_trajectories(
        image,
        trajectory,
        is_draw_control_points=is_draw_control_points,
    )
    return image


def update_trajectory(
    trajectory: MultiTrajectory,
    mask_index: int,
    drag_animation_select: str,
    translate_rotate_select: str,
    evt: gr.SelectData,
):
    print(f"update_trajectory")

    # Work on a deep copy so Gradio sees a new object
    trajectory = copy.deepcopy(trajectory)

    if mask_index < 0:
        print(f"Invalid mask_index: {mask_index}")
        return trajectory

    # print(f"{evt.index = }")
    x_center, y_center = evt.index  # evt.value is (x, y)

    clicked_point = (x_center, y_center)
    print(f"{clicked_point = }")

    # Ensure trajectories list is large enough
    while len(trajectory.trajectories) <= mask_index:
        trajectory.trajectories.append(Trajectory())

    existing_traj_obj = trajectory.trajectories[mask_index]
    if existing_traj_obj.original_trajectory is not None:
        current_trajectory = dict(existing_traj_obj.original_trajectory)
    else:
        current_trajectory = {}

    if translate_rotate_select == "Translation":
        current_trajectory["is_rotation"] = False

        # Append clicked control point
        control_points = current_trajectory.get("control_points", [])
        control_points = control_points + [clicked_point]

        # Drag vs Animation behavior
        if drag_animation_select == "Drag":
            # Restrict to last two control points, sample exactly 2 points
            if len(control_points) > 2:
                control_points = [clicked_point]
            num_traj_points = 2
        elif drag_animation_select == "Animation":
            # No restriction on control points, sample N = 1 + 3 * block_number
            num_traj_points = 1 + 3 * int(trajectory.block_number)
        else:
            raise ValueError(f"Invalid drag_animation_select: {drag_animation_select}")

        current_trajectory["control_points"] = control_points

        # Compute trajectory points along BSpline (or pad if not enough controls)
        if len(control_points) < 2:
            sampled_pts = [control_points[0]] * num_traj_points
        else:
            spline = build_clamped_bspline(control_points, degree=3)
            pts = equidistant_points_on_spline(spline, num_points=num_traj_points, grid=8000)
            sampled_pts = [(int(round(px)), int(round(py))) for px, py in pts]

        current_trajectory["points"] = sampled_pts

    elif translate_rotate_select == "Rotation":
        current_trajectory["is_rotation"] = True

        # Initialize if missing, else apply 3-point logic
        if "points" not in current_trajectory or current_trajectory["points"] is None:
            current_trajectory["points"] = [clicked_point]
        else:
            pts = current_trajectory["points"] + [clicked_point]

            # If about to exceed 3, reset to the new point
            if len(pts) > 3:
                current_trajectory["points"] = [clicked_point]
            # If less than 3, just append
            elif len(pts) < 3:
                current_trajectory["points"] = pts
            else:
                # len(pts) == 3: pts[0] is rotation center
                if drag_animation_select == "Animation":
                    first = trajectory_interpolate(pts[1:], scale=int(trajectory.block_number))
                    second = trajectory_interpolate(first, scale=3)
                    current_trajectory["points"] = pts[0:1] + second
                else:
                    # Drag: do not interpolate
                    current_trajectory["points"] = pts
    else:
        raise ValueError("Invalid translation/rotation selection")

    # Update the Trajectory object in-place (recomputes block_trajectories)
    existing_traj_obj.set_original_trajectory(current_trajectory)
    # print(f"{trajectory = }")

    return trajectory


def save_trajectory(
    save_dir: Path,
    saved_trajectory: MultiTrajectory,
    original_image: Image.Image,
    current_block_index: int,
    masks: list[np.ndarray],
):
    print(f"save_trajectory")
    print(f"{save_dir = }")
    print(f"{saved_trajectory = }")
    save_dir = Path(save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)

    drag_animation_select = saved_trajectory.drag_or_animation_select or "Drag"
    save_prefix = f"block_{current_block_index}_{drag_animation_select}"

    # Use MultiTrajectory's save method
    saved_trajectory.save(
        save_dir=save_dir,
        prefix=save_prefix,
    )

    # Save the trajectory image
    trajectory_image = draw_trajectory_image(
        original_image=original_image,
        movable_mask=saved_trajectory.movable_mask,
        mask_index=None,
        masks_list=masks,
        trajectory=saved_trajectory,
        is_draw_bbox=False,
        is_draw_control_points=False,
    )
    trajectory_image.save(save_dir / f"{save_prefix}_trajectory.png")


def clear_current_trajectory(
    idx: int,
    trajectory: MultiTrajectory,
):
    trajectory = copy.deepcopy(trajectory)
    """Clear the trajectory at the given mask index."""
    try:
        idx_int = int(idx)
    except Exception:
        return trajectory

    if not trajectory.trajectories:
        return trajectory

    if idx_int < 0 or idx_int >= len(trajectory.trajectories):
        return trajectory

    # Reset this trajectory (keep the mask)
    mask = trajectory.trajectories[idx_int].mask
    trajectory.trajectories[idx_int] = Trajectory(mask=mask)

    return trajectory


def clear_all_trajectories(
    trajectory: MultiTrajectory,
):
    trajectory = copy.deepcopy(trajectory)
    """Clear all trajectories but keep the masks."""
    if trajectory.trajectories is not None:
        for i in range(len(trajectory.trajectories)):
            mask = trajectory.trajectories[i].mask
            trajectory.trajectories[i] = Trajectory(mask=mask)
    return trajectory


def sync_trajectory_masks(saved_trajectory: MultiTrajectory, dilated_masks: list[np.ndarray]):
    """Resize saved_trajectory.trajectories to match the number of dilated masks,
    and update each Trajectory.mask with the corresponding dilated mask."""
    saved_trajectory = copy.deepcopy(saved_trajectory)

    current_len = len(saved_trajectory.trajectories)
    target_len = len(dilated_masks) if dilated_masks else 0

    if target_len > current_len:
        # Expand: append new empty Trajectory objects
        for _ in range(target_len - current_len):
            saved_trajectory.trajectories.append(Trajectory())
    elif target_len < current_len:
        # Shrink: truncate
        saved_trajectory.trajectories = saved_trajectory.trajectories[:target_len]

    # Update each Trajectory.mask
    for i, mask in enumerate(dilated_masks):
        saved_trajectory.trajectories[i].mask = mask

    return saved_trajectory


def add_listeners_to_trajectory(
    saved_trajectory: MultiTrajectory,
    prompt_box: gr.Textbox,
    trajectory_block_number_slider: gr.Slider,
    drag_animation_select: gr.Dropdown,
    movable_area_mask: gr.State,
    dilated_saved_sam_predicted_masks: gr.State,
):
    # Sync prompt into saved_trajectory when prompt_box changes

    def sync_trajectory_prompt(saved_trajectory: MultiTrajectory, prompt: str):
        saved_trajectory.prompt = prompt
        return saved_trajectory

    prompt_box.change(
        fn=sync_trajectory_prompt,
        inputs=[saved_trajectory, prompt_box],
        outputs=saved_trajectory,
        trigger_mode="always_last",
    )

    # Sync block_number into saved_trajectory when trajectory_block_number_slider changes
    def sync_trajectory_block_number(saved_trajectory: MultiTrajectory, block_number: int):
        saved_trajectory.block_number = block_number
        return saved_trajectory

    trajectory_block_number_slider.change(
        fn=sync_trajectory_block_number,
        inputs=[saved_trajectory, trajectory_block_number_slider],
        outputs=saved_trajectory,
        trigger_mode="always_last",
    )

    # Sync drag_or_animation_select into saved_trajectory when drag_animation_select changes
    def sync_trajectory_drag_animation(
        saved_trajectory: MultiTrajectory, drag_animation_select: str
    ):
        saved_trajectory.drag_or_animation_select = drag_animation_select
        return saved_trajectory

    drag_animation_select.change(
        fn=sync_trajectory_drag_animation,
        inputs=[saved_trajectory, drag_animation_select],
        outputs=saved_trajectory,
        trigger_mode="always_last",
    )

    # Sync movable_area_mask into saved_trajectory when it changes
    def sync_trajectory_movable_mask(saved_trajectory: MultiTrajectory, movable_mask):
        saved_trajectory.movable_mask = movable_mask
        return saved_trajectory

    movable_area_mask.change(
        fn=sync_trajectory_movable_mask,
        inputs=[saved_trajectory, movable_area_mask],
        outputs=saved_trajectory,
        trigger_mode="always_last",
    )

    # Sync dilated_saved_sam_predicted_masks into saved_trajectory when it changes
    dilated_saved_sam_predicted_masks.change(
        fn=sync_trajectory_masks,
        inputs=[saved_trajectory, dilated_saved_sam_predicted_masks],
        outputs=saved_trajectory,
        trigger_mode="always_last",
    )


def create_generate_video_ui(
    label_root: str | Path,
    text_dataset: Dataset,
    video_path: gr.State,
    stream_drag_inference: StreamDragInferenceWrapper,
    output_dir: str | Path,
    original_image: gr.State,
):
    with gr.Row():
        prompt_index_number = gr.Number(
            label="Step 1: Select Prompt Index Here",
            interactive=True,
            scale=1,
        )
        prompt_box = gr.Textbox(
            label="Prompt",
            interactive=True,
            scale=3,
        )
        save_dir_text_box = gr.Textbox(
            label="Save Directory",
            interactive=False,
            scale=1,
        )
    prompt_index_number.change(
        fn=lambda prompt_index_number: text_dataset[prompt_index_number]["prompts"],
        inputs=prompt_index_number,
        outputs=[
            prompt_box,
        ],
    )
    gr.on(
        triggers=[
            prompt_box.change,
        ],
        fn=lambda prompt_index_number, prompt: str(
            label_root / f"{prompt_index_number:04d}-{prompt[:50].replace(' ', '_')}"
        ),
        inputs=[prompt_index_number, prompt_box],
        outputs=save_dir_text_box,
        trigger_mode="always_last",
    )
    with gr.Row():
        current_block_index_slider = gr.Slider(
            label="Current Start Block Index",
            minimum=0,
            maximum=50,
            value=0,
            step=1,
        )
        generate_block_number_slider = gr.Slider(
            label="Step 2: Select Number of Blocks to Generate",
            minimum=1,
            maximum=50,
            value=2,
            step=1,
        )
    with gr.Row():
        begin_generate_button = gr.Button(
            value="Step 3: Click Here to Begin Generation",
        )
        refresh_video_display_button = gr.Button(value="Refresh Video Display")

    with gr.Row():
        video_display = gr.Video()

    begin_generate_button.click(
        fn=lambda pi, p, sbi, bn: generate_video(
            stream_inference_model=stream_drag_inference,
            prompt_index=pi,
            prompt=p,
            start_block_index=sbi,
            block_number=bn,
            output_dir=output_dir,
        ),
        inputs=[
            prompt_index_number,
            prompt_box,
            current_block_index_slider,
            generate_block_number_slider,
        ],
        outputs=[video_path, current_block_index_slider],
    )
    gr.on(
        triggers=[
            refresh_video_display_button.click,
            video_path.change,
        ],
        fn=lambda video_path: video_path,
        inputs=video_path,
        outputs=video_display,
        trigger_mode="always_last",
    )

    with gr.Row():
        get_last_frame_button = gr.Button(
            value="Get Last Frame (Normally No Need to Click This, In Case the Last Frame Fails to Update due to Gradio Bug)",
        )
    gr.on(
        triggers=[
            video_path.change,
            get_last_frame_button.click,
        ],
        fn=get_video_last_frame,
        inputs=video_path,
        outputs=original_image,
    )

    return (
        prompt_index_number,
        save_dir_text_box,
        prompt_box,
        current_block_index_slider,
        generate_block_number_slider,
    )


def create_movable_area_ui(
    movable_area_mask: gr.State,
    original_image: gr.State,
):

    with gr.Row():
        movable_area_image_editor = gr.ImageEditor(
            label="Step 4: This is Last Frame of Video, Draw Editable Area Here. (Normally This Should Be Large and Cover all Possible Area Where the Object You Want to Move/Animate to)",
            type="pil",
            interactive=True,
            brush=gr.Brush(
                default_size=100,
                colors=[
                    "rgba(0, 0, 255, 0.5)",
                ],
                default_color="auto",
                color_mode="defaults",
            ),
        )
    movable_area_image_editor.change(
        fn=extract_layer_as_mask,
        inputs=movable_area_image_editor,
        outputs=movable_area_mask,
        trigger_mode="always_last",
    )
    original_image.change(
        fn=lambda image: image,
        inputs=original_image,
        outputs=movable_area_image_editor,
        trigger_mode="always_last",
    )
    with gr.Row():
        refresh_movable_area_button = gr.Button(
            value="Refresh Movable Area (Normally No Need to Click This, In Case the Mask Fails to Update due to Gradio Bug)"
        )
    refresh_movable_area_button.click(
        fn=extract_layer_as_mask,
        inputs=movable_area_image_editor,
        outputs=movable_area_mask,
        trigger_mode="always_last",
    )


def create_target_area_ui(
    target_area_mask: gr.State,
    original_image: gr.State,
    movable_area_mask: gr.State,
):
    with gr.Row():
        target_area_image_editor = gr.ImageEditor(
            label="Step 5: Draw Target Area on the Object You Want to Move/Animate (Normally This Should Be a Subset of Editable Area) (Normally This Mask should be Bigger than the Desired Object)",
            type="pil",
            interactive=True,
            brush=gr.Brush(
                default_size=50,
                colors=[
                    "rgba(255, 0, 0, 0.5)",
                ],
                default_color="auto",
                color_mode="defaults",
            ),
        )
    target_area_image_editor.change(
        fn=extract_layer_as_mask,
        inputs=target_area_image_editor,
        outputs=target_area_mask,
        trigger_mode="always_last",
    )
    gr.on(
        triggers=[
            original_image.change,
            movable_area_mask.change,
        ],
        fn=apply_movable_mask_to_image,
        inputs=[
            movable_area_mask,
            original_image,
        ],
        outputs=target_area_image_editor,
        trigger_mode="always_last",
    )

    with gr.Row():
        refresh_target_area_button = gr.Button(
            value="Refresh Target Area (Normally No Need to Click This, In Case the Mask Fails to Update due to Gradio Bug)"
        )
    refresh_target_area_button.click(
        fn=extract_layer_as_mask,
        inputs=target_area_image_editor,
        outputs=target_area_mask,
        trigger_mode="always_last",
    )


def create_sam_segmentation_ui(
    original_image: gr.State,
    movable_area_mask: gr.State,
    target_area_mask: gr.State,
    sam_predictor: SamPredictor,
    sam_click_points: gr.State,
    sam_saved_logits: gr.State,
    current_sam_predicted_mask: gr.State,
    saved_sam_predicted_masks: gr.State,
    dilated_current_sam_predicted_mask: gr.State,
    dilated_saved_sam_predicted_masks: gr.State,
):
    with gr.Row():
        refresh_sam_segment_click_image_button = gr.Button(
            value="Refresh Target Area Mask Display (Normally No Need to Click This, In Case the Mask Fails to Update due to Gradio Bug)"
        )
    with gr.Row():
        sam_segment_click_image = gr.Image(
            label="Step 6: Click to Perform SAM Segment on Target Area, Segment the Object You Want to Move/Animate. The SAM Mask is Restricted within the Target Area Mask",
            type="pil",
            interactive=True,
        )
    gr.on(
        triggers=[
            original_image.change,
            movable_area_mask.change,
            target_area_mask.change,
            refresh_sam_segment_click_image_button.click,
        ],
        fn=lambda movable_mask, target_mask, image: apply_target_mask_to_image(
            target_mask,
            apply_movable_mask_to_image(
                movable_mask,
                image,
            ),
        ),
        inputs=[
            movable_area_mask,
            target_area_mask,
            original_image,
        ],
        outputs=sam_segment_click_image,
        trigger_mode="always_last",
    )

    with gr.Row():
        dilate_mask_slider = gr.Slider(
            label="Dilate Mask Pixel",
            minimum=0,
            maximum=50,
            value=15,
            step=1,
        )
        bypass_sam_model_check_box = gr.Checkbox(
            label="Bypass SAM Model",
            value=False,
        )

    def sam_predict_segmentation_wrapper_wrapper(
        oi,
        rm,
        pcp,
        psl,
        bs,
        evt: gr.SelectData,
    ):
        return sam_predict_segmentation_wrapper(
            sam_predictor=sam_predictor,
            original_image=oi,
            restriction_mask=rm,
            previous_click_points=pcp,
            previous_sam_logits=psl,
            bypass_sam_model=bs,
            evt=evt,
        )

    sam_segment_click_image.select(
        fn=sam_predict_segmentation_wrapper_wrapper,
        inputs=[
            original_image,
            target_area_mask,
            sam_click_points,
            sam_saved_logits,
            bypass_sam_model_check_box,
        ],
        outputs=[
            current_sam_predicted_mask,
            sam_click_points,
            sam_saved_logits,
        ],
        trigger_mode="always_last",
    )
    gr.on(
        triggers=[
            current_sam_predicted_mask.change,
            dilate_mask_slider.change,
        ],
        fn=dilate_mask,
        inputs=[
            current_sam_predicted_mask,
            dilate_mask_slider,
        ],
        outputs=dilated_current_sam_predicted_mask,
        trigger_mode="always_last",
    )
    gr.on(
        triggers=[
            saved_sam_predicted_masks.change,
            dilate_mask_slider.change,
        ],
        fn=dilate_masks,
        inputs=[
            saved_sam_predicted_masks,
            dilate_mask_slider,
        ],
        outputs=dilated_saved_sam_predicted_masks,
        trigger_mode="always_last",
    )


def create_sam_mask_management_ui(
    original_image: gr.State,
    movable_area_mask: gr.State,
    dilated_current_sam_predicted_mask: gr.State,
    dilated_saved_sam_predicted_masks: gr.State,
    sam_click_points: gr.State,
    current_sam_predicted_mask: gr.State,
    saved_sam_predicted_masks: gr.State,
    sam_saved_logits: gr.State,
):
    with gr.Row():
        save_sam_masks_button = gr.Button(
            value="Step 7: Save the Current SAM Mask",
        )
        cancel_sam_mask_button = gr.Button(value="Cancel Current SAM Mask")
        delete_sam_mask_button = gr.Button(value="Delete All SAM Masks")
    save_sam_masks_button.click(
        fn=save_sam_masks,
        inputs=[
            current_sam_predicted_mask,
            saved_sam_predicted_masks,
        ],
        outputs=[
            current_sam_predicted_mask,
            saved_sam_predicted_masks,
            sam_click_points,
            sam_saved_logits,
        ],
        trigger_mode="always_last",
    )
    with gr.Row():
        sam_segment_display_image = gr.Image(
            label="Step 8: Display the SAM Segmentation, Click to Select Target Object to Create Trajectory",
            type="pil",
            interactive=True,
        )
    gr.on(
        triggers=[
            original_image.change,
            movable_area_mask.change,
            dilated_current_sam_predicted_mask.change,
            dilated_saved_sam_predicted_masks.change,
            sam_click_points.change,
        ],
        fn=draw_sam_mask_wrapper,
        inputs=[
            original_image,
            movable_area_mask,
            dilated_current_sam_predicted_mask,
            dilated_saved_sam_predicted_masks,
            sam_click_points,
        ],
        outputs=sam_segment_display_image,
        trigger_mode="always_last",
    )
    cancel_sam_mask_button.click(
        fn=lambda: (None, [], None),
        outputs=[
            current_sam_predicted_mask,
            sam_click_points,
            sam_saved_logits,
        ],
        trigger_mode="always_last",
    )
    gr.on(
        triggers=[
            # target_area_mask.change,
            delete_sam_mask_button.click,
        ],
        fn=lambda: (None, [], [], None),
        outputs=[
            current_sam_predicted_mask,
            saved_sam_predicted_masks,
            sam_click_points,
            sam_saved_logits,
        ],
        trigger_mode="always_last",
    )
    with gr.Row():
        current_selected_mask_index_number = gr.Number(
            label="Current Selected Mask Index",
            interactive=False,
        )

    sam_segment_display_image.select(
        fn=select_target_sam_mask,
        inputs=[
            saved_sam_predicted_masks,
        ],
        outputs=[
            current_selected_mask_index_number,
        ],
        trigger_mode="always_last",
    )

    return current_selected_mask_index_number


def create_trajectory_display_ui(
    original_image: gr.State,
    movable_area_mask: gr.State,
    dilated_saved_sam_predicted_masks: gr.State,
    saved_trajectory: gr.State,
    current_selected_mask_index_number: gr.State,
):
    with gr.Row():
        trajectory_block_number_slider = gr.Slider(
            label="Step 9: Select Number of Trajectory Blocks (For Animation Only, More Blocks Means Longer Animation, For Drag, This Should be 1)",
            minimum=1,
            maximum=10,
            value=1,
            step=1,
        )
    with gr.Row():
        drag_animation_select = gr.Dropdown(
            choices=["Drag", "Animation"],
            label="Step 10: Select Drag or Animation",
        )
        translate_rotate_select = gr.Dropdown(
            choices=["Translation", "Rotation"],
            label="Step 11: Select Translation or Rotation",
        )

    with gr.Row():
        trajectory_display_image = gr.Image(
            label="Step 12: Click on the Object in the Image to Create Trajectory. The Translation Trajectory is Controlled by Bspline Interpolation. The Rotation Trajectory is Controlled by 3 Points",
            type="pil",
            interactive=False,
        )
    gr.on(
        triggers=[
            original_image.change,
            movable_area_mask.change,
            current_selected_mask_index_number.change,
            dilated_saved_sam_predicted_masks.change,
            saved_trajectory.change,
        ],
        fn=draw_trajectory_image,
        inputs=[
            original_image,
            movable_area_mask,
            current_selected_mask_index_number,
            dilated_saved_sam_predicted_masks,
            saved_trajectory,
        ],
        outputs=trajectory_display_image,
        trigger_mode="always_last",
    )

    trajectory_display_image.select(
        fn=update_trajectory,
        inputs=[
            saved_trajectory,
            current_selected_mask_index_number,
            drag_animation_select,
            translate_rotate_select,
        ],
        outputs=saved_trajectory,
    )

    return drag_animation_select, trajectory_block_number_slider


def create_trajectory_management_ui(
    save_dir_text_box: gr.Textbox,
    original_image: gr.State,
    current_block_index_slider: gr.Slider,
    saved_trajectory: gr.State,
    dilated_saved_sam_predicted_masks: gr.State,
    current_selected_mask_index_number: gr.Number,
):
    with gr.Row():
        save_trajectory_button = gr.Button(
            value="Step 13: Save Trajectory",
        )
        delete_current_trajectory_button = gr.Button(value="Delete Current Trajectory")
        delete_all_trajectory_button = gr.Button(value="Delete All Trajectories")
    save_trajectory_button.click(
        fn=save_trajectory,
        inputs=[
            save_dir_text_box,
            saved_trajectory,
            original_image,
            current_block_index_slider,
            dilated_saved_sam_predicted_masks,
        ],
    )
    delete_current_trajectory_button.click(
        fn=clear_current_trajectory,
        inputs=[current_selected_mask_index_number, saved_trajectory],
        outputs=[saved_trajectory],
    )
    delete_all_trajectory_button.click(
        fn=clear_all_trajectories,
        inputs=[saved_trajectory],
        outputs=[saved_trajectory],
    )


def create_ui(
    text_dataset: Dataset,
    label_root: str | Path,
    output_dir: str | Path,
    sam_predictor: SamPredictor,
    stream_drag_inference: StreamDragInferenceWrapper,
):
    with gr.Blocks() as demo:
        video_path = gr.State(value=None)
        original_image = gr.State(value=None)
        movable_area_mask = gr.State(value=None)
        target_area_mask = gr.State(value=None)

        sam_click_points = gr.State(value=[])
        sam_saved_logits = gr.State(value=None)
        saved_sam_predicted_masks = gr.State(value=[])
        current_sam_predicted_mask = gr.State(value=None)

        dilated_current_sam_predicted_mask = gr.State(value=None)
        dilated_saved_sam_predicted_masks = gr.State(value=[])

        saved_trajectory = gr.State(value=MultiTrajectory())

        (
            prompt_index_number,
            save_dir_text_box,
            prompt_box,
            current_block_index_slider,
            generate_block_number_slider,
        ) = create_generate_video_ui(
            label_root=label_root,
            text_dataset=text_dataset,
            video_path=video_path,
            stream_drag_inference=stream_drag_inference,
            output_dir=output_dir,
            original_image=original_image,
        )

        create_movable_area_ui(movable_area_mask, original_image)
        create_target_area_ui(target_area_mask, original_image, movable_area_mask)
        create_sam_segmentation_ui(
            original_image=original_image,
            movable_area_mask=movable_area_mask,
            target_area_mask=target_area_mask,
            sam_predictor=sam_predictor,
            sam_click_points=sam_click_points,
            sam_saved_logits=sam_saved_logits,
            current_sam_predicted_mask=current_sam_predicted_mask,
            saved_sam_predicted_masks=saved_sam_predicted_masks,
            dilated_current_sam_predicted_mask=dilated_current_sam_predicted_mask,
            dilated_saved_sam_predicted_masks=dilated_saved_sam_predicted_masks,
        )

        current_selected_mask_index_number = create_sam_mask_management_ui(
            original_image=original_image,
            movable_area_mask=movable_area_mask,
            dilated_current_sam_predicted_mask=dilated_current_sam_predicted_mask,
            dilated_saved_sam_predicted_masks=dilated_saved_sam_predicted_masks,
            sam_click_points=sam_click_points,
            current_sam_predicted_mask=current_sam_predicted_mask,
            saved_sam_predicted_masks=saved_sam_predicted_masks,
            sam_saved_logits=sam_saved_logits,
        )

        drag_animation_select, trajectory_block_number_slider = create_trajectory_display_ui(
            original_image=original_image,
            movable_area_mask=movable_area_mask,
            dilated_saved_sam_predicted_masks=dilated_saved_sam_predicted_masks,
            saved_trajectory=saved_trajectory,
            current_selected_mask_index_number=current_selected_mask_index_number,
        )
        create_trajectory_management_ui(
            save_dir_text_box=save_dir_text_box,
            original_image=original_image,
            current_block_index_slider=current_block_index_slider,
            saved_trajectory=saved_trajectory,
            dilated_saved_sam_predicted_masks=dilated_saved_sam_predicted_masks,
            current_selected_mask_index_number=current_selected_mask_index_number,
        )

        add_listeners_to_trajectory(
            saved_trajectory=saved_trajectory,
            prompt_box=prompt_box,
            trajectory_block_number_slider=trajectory_block_number_slider,
            drag_animation_select=drag_animation_select,
            movable_area_mask=movable_area_mask,
            dilated_saved_sam_predicted_masks=dilated_saved_sam_predicted_masks,
        )

        with gr.Row():
            begin_optimize_button = gr.Button(
                value="Step 14: Click Here to Begin Optimize, Wait for a Moment and the Dragged/Animated Video will be Displayed Above",
            )
        begin_optimize_button.click(
            fn=lambda pi, sbi, st: optimize_video(
                stream_drag_inference_model=stream_drag_inference,
                output_dir=output_dir,
                prompt_index=pi,
                start_block_index=sbi,
                multi_trajectory=st,
            ),
            inputs=[
                prompt_index_number,
                current_block_index_slider,
                saved_trajectory,
            ],
            outputs=[
                video_path,
                current_block_index_slider,
            ],
        )
        with gr.Row():
            clear_all_button = gr.Button(
                value="Step 15: Remember to Click Here to Clear All Before Generation/Editing on Next Video, Otherwise the Previous KV Cache will Affect the Generation/Editing of Next Video",
            )

        def clear_all():
            stream_drag_inference.reset()

            return (
                0,
                None,
                None,
                None,
                None,
                [],
                None,
                [],
                None,
                MultiTrajectory(),
            )

        clear_all_button.click(
            fn=clear_all,
            outputs=[
                current_block_index_slider,
                video_path,
                original_image,
                movable_area_mask,
                target_area_mask,
                sam_click_points,
                sam_saved_logits,
                saved_sam_predicted_masks,
                current_sam_predicted_mask,
                saved_trajectory,
            ],
        )

    return demo


def download_required_files():
    from huggingface_hub import snapshot_download
    import urllib.request

    # 1. Download "checkpoints" directory from gdhe17/Self-Forcing
    if not os.path.exists("checkpoints"):
        print("Downloading checkpoints from gdhe17/Self-Forcing...")
        snapshot_download(
            repo_id="gdhe17/Self-Forcing",
            allow_patterns=["checkpoints/*"],
            local_dir=".",
        )

    # 2. Download Wan-AI/Wan2.1-T2V-1.3B and place in wan_models/Wan2.1-T2V-1.3B
    wan_model_dir = os.path.join("wan_models", "Wan2.1-T2V-1.3B")
    if not os.path.exists(wan_model_dir):
        print("Downloading Wan-AI/Wan2.1-T2V-1.3B...")
        os.makedirs("wan_models", exist_ok=True)
        snapshot_download(
            repo_id="Wan-AI/Wan2.1-T2V-1.3B",
            local_dir=wan_model_dir,
        )

    # 3. Download SAM ViT-H checkpoint
    sam_checkpoint_path = "sam_vit_h_4b8939.pth"
    if not os.path.exists(sam_checkpoint_path):
        print("Downloading SAM ViT-H checkpoint...")
        urllib.request.urlretrieve(
            "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
            sam_checkpoint_path,
        )

def main():
    download_required_files()
    sam_model = sam_model_registry["vit_h"](checkpoint="./sam_vit_h_4b8939.pth")
    sam_model.to(device="cuda")
    sam_predictor = SamPredictor(sam_model)

    SEED = 42

    text_dataset = TextDataset(prompt_path="prompts/MovieGenVideoBench_extended.txt")

    if GlobalHydra.instance().is_initialized():
        GlobalHydra.instance().clear()

    config_dir = "configs"
    stream_config_name = "self_forcing_dmd_vsink_stream_drag"
    with initialize(version_base=None, config_path=config_dir):
        stream_config = compose(config_name=stream_config_name)
    print(f"{stream_config = }")

    stream_drag_inference = StreamDragInferenceWrapper(
        stream_model_config=stream_config,
        checkpoint_path="./checkpoints/self_forcing_dmd.pt",
        total_generate_block_number=36,
        use_ema=True,
        seed=SEED,
    )
    label_save_dir = Path("./saved_labels")
    label_save_dir = label_save_dir / f"{stream_config_name}-seed{SEED}"
    label_save_dir.mkdir(parents=True, exist_ok=True)

    output_save_dir = Path("outputs-editing")
    output_save_dir = output_save_dir / f"{stream_config_name}-seed{SEED}"
    output_save_dir.mkdir(parents=True, exist_ok=True)

    demo = create_ui(
        text_dataset=text_dataset,
        label_root=label_save_dir,
        output_dir=output_save_dir,
        sam_predictor=sam_predictor,
        stream_drag_inference=stream_drag_inference,
    )
    demo.launch(server_name="0.0.0.0")


if __name__ == "__main__":
    main()