import os
if os.getenv("SPACES_ZERO_GPU"):
    os.system('pip install --upgrade --no-deps spaces')
import spaces
import copy
import gc
import random
import subprocess
import tempfile
import time
import uuid
import warnings

from tqdm import tqdm
import cv2
import numpy as np
import torch
import torch._dynamo
from torch.nn import functional as F
from PIL import Image

import gradio as gr
from diffusers import (
    FlowMatchEulerDiscreteScheduler,
    SASolverScheduler,
    DEISMultistepScheduler,
    UniPCMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
)
from diffusers import HunyuanVideo15ImageToVideoPipeline
from diffusers.utils.export_utils import export_to_video

from torchao.quantization import (
    quantize_,
    Float8DynamicActivationFloat8WeightConfig,
    Int8WeightOnlyConfig,
)
import aoti
import lora_loader

os.environ["TOKENIZERS_PARALLELISM"] = "true"
warnings.filterwarnings("ignore")
IS_ZERO_GPU = bool(os.getenv("SPACES_ZERO_GPU"))

# --- FRAME EXTRACTION JS & LOGIC ---

# JS to grab timestamp from the output video
get_timestamp_js = """
function() {
    // Select the video element specifically inside the component with id 'generated-video'
    const video = document.querySelector('#generated-video video');

    if (video) {
        console.log("Video found! Time: " + video.currentTime);
        return video.currentTime;
    } else {
        console.log("No video element found.");
        return 0;
    }
}
"""


def extract_frame(video_path, timestamp):
    # Safety check: if no video is present
    if not video_path:
        return None

    print(f"Extracting frame at timestamp: {timestamp}")

    cap = cv2.VideoCapture(video_path)

    if not cap.isOpened():
        return None

    # Calculate frame number
    fps = cap.get(cv2.CAP_PROP_FPS)
    target_frame_num = int(float(timestamp) * fps)

    # Cap total frames to prevent errors at the very end of video
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if target_frame_num >= total_frames:
        target_frame_num = total_frames - 1

    # Set position
    cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame_num)
    ret, frame = cap.read()
    cap.release()

    if ret:
        # Convert from BGR (OpenCV) to RGB (Gradio)
        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    return None

# --- END FRAME EXTRACTION LOGIC ---


def clear_vram():
    gc.collect()
    torch.cuda.empty_cache()


# RIFE
if not os.path.exists("RIFEv4.26_0921.zip"):
    print("Downloading RIFE Model...")
    subprocess.run([
        "wget", "-q",
        "https://huggingface.co/thornmaze/RIFE/resolve/main/RIFEv4.26_0921.zip",
        "-O", "RIFEv4.26_0921.zip"
    ], check=True)
    subprocess.run(["unzip", "-o", "RIFEv4.26_0921.zip"], check=True)

from train_log.RIFE_HDv3 import Model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
rife_model = Model()
rife_model.load_model("train_log", -1)
rife_model.eval()


@torch.no_grad()
def interpolate_bits(frames_np, multiplier=2, scale=1.0):
    """
    Interpolation maintaining Numpy Float 0-1 format.
    Args:
        frames_np: Numpy Array (Time, Height, Width, Channels) - Float32 [0.0, 1.0]
        multiplier: int (2, 4)
    Returns:
        List of Numpy Arrays (Height, Width, Channels) - Float32 [0.0, 1.0]
    """

    # Handle input shape
    if isinstance(frames_np, list):
        T = len(frames_np)
        H, W, C = frames_np[0].shape
    else:
        T, H, W, C = frames_np.shape

    # 1. No Interpolation Case
    if multiplier < 2:
        if isinstance(frames_np, np.ndarray):
            return list(frames_np)
        return frames_np

    n_interp = multiplier - 1

    # Pre-calc padding for RIFE (requires dimensions divisible by 32/scale)
    tmp = max(128, int(128 / scale))
    ph = ((H - 1) // tmp + 1) * tmp
    pw = ((W - 1) // tmp + 1) * tmp
    padding = (0, pw - W, 0, ph - H)

    # Helper: Numpy (H, W, C) Float -> Tensor (1, C, H, W) Half
    def to_tensor(frame_np):
        # frame_np is float32 0-1
        t = torch.from_numpy(frame_np).to(device)
        # HWC -> CHW
        t = t.permute(2, 0, 1).unsqueeze(0)
        return F.pad(t, padding).half()

    # Helper: Tensor (1, C, H, W) Half -> Numpy (H, W, C) Float
    def from_tensor(tensor):
        # Crop padding
        t = tensor[0, :, :H, :W]
        # CHW -> HWC
        t = t.permute(1, 2, 0)
        # Keep as float32, range 0-1
        return t.float().cpu().numpy()

    def make_inference(I0, I1, n):
        if rife_model.version >= 3.9:
            res = []
            for i in range(n):
                res.append(rife_model.inference(I0, I1, (i + 1) * 1. / (n + 1), scale))
            return res
        else:
            middle = rife_model.inference(I0, I1, scale)
            if n == 1:
                return [middle]
            first_half = make_inference(I0, middle, n=n // 2)
            second_half = make_inference(middle, I1, n=n // 2)
            if n % 2:
                return [*first_half, middle, *second_half]
            else:
                return [*first_half, *second_half]

    output_frames = []

    # Process Frames
    I1 = to_tensor(frames_np[0])
    mid_tensors = []

    total_steps = T - 1

    with tqdm(total=total_steps, desc="Interpolating", unit="frame") as pbar:

        for i in range(total_steps):
            I0 = I1
            # Add original frame to output
            output_frames.append(from_tensor(I0))

            # Load next frame
            I1 = to_tensor(frames_np[i + 1])

            # Generate intermediate frames
            mid_tensors = make_inference(I0, I1, n_interp)

            # Append intermediate frames
            for mid in mid_tensors:
                output_frames.append(from_tensor(mid))

            if (i + 1) % 50 == 0:
                pbar.update(50)
        pbar.update(total_steps % 50)

        # Add the very last frame
        output_frames.append(from_tensor(I1))

    # Cleanup
    del I0, I1, mid_tensors
    torch.cuda.empty_cache()

    return output_frames


# HUNYUANVIDEO 1.5

# Step-distilled I2V: 8-12 steps, CFG 1.0, flow shift 7.0, 480p bucket base (target_size=640).
# Alternatives (set MODEL_ID):
#   hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_distilled       (50 steps, CFG 1.0, shift 5.0)
#   hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v                 (50 steps, CFG 6.0, shift 5.0)
#   hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v_distilled       (50 steps, CFG 1.0, shift 7.0, target_size=960)
MODEL_ID = os.getenv(
    "MODEL_ID",
    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled",
)

# Attention backend. H100/H200 -> _flash_3_hub, A100/4090 -> flash_hub, other -> sage_hub.
ATTN_BACKEND = os.getenv("ATTN_BACKEND", "_flash_3_hub")

# Optional prebuilt AoT-Inductor packages for HunyuanVideo15TransformerBlock etc.
# Unset by default: no public package repo exists for this architecture yet.
AOTI_REPO = os.getenv("AOTI_REPO")
AOTI_VARIANT = os.getenv("AOTI_VARIANT")

QUANTIZE = os.getenv("QUANTIZE", "1" if IS_ZERO_GPU else "0") == "1"

MAX_SEED = np.iinfo(np.int32).max

# HunyuanVideo 1.5 is a 24 fps model; the VAE compresses time 4x, so num_frames must be 4k+1.
FIXED_FPS = 24
MIN_FRAMES_MODEL = 25
MAX_FRAMES_MODEL = int(os.getenv("MAX_FRAMES", "121"))  # 121 = the trained 5s length

MIN_DURATION = round(MIN_FRAMES_MODEL / FIXED_FPS, 1)
MAX_DURATION = round(MAX_FRAMES_MODEL / FIXED_FPS, 1)

DEFAULT_STEPS = int(os.getenv("DEFAULT_STEPS", "8"))
DEFAULT_SHIFT = float(os.getenv("DEFAULT_SHIFT", "7.0"))
DEFAULT_GUIDANCE = float(os.getenv("DEFAULT_GUIDANCE", "1.0"))

# Only flow-matching schedulers make sense for this model. FlowMatchEulerDiscrete is what the
# checkpoint ships with; the multistep solvers are run in their flow-prediction mode and are
# experimental here, especially on the meanflow step-distilled weights.
SCHEDULER_MAP = {
    "FlowMatchEulerDiscrete": FlowMatchEulerDiscreteScheduler,
    "UniPCMultistep": UniPCMultistepScheduler,
    "DPMSolverMultistep": DPMSolverMultistepScheduler,
    "DPMSolverSinglestep": DPMSolverSinglestepScheduler,
    "DEISMultistep": DEISMultistepScheduler,
    "SASolver": SASolverScheduler,
}

pipe = HunyuanVideo15ImageToVideoPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
).to('cuda')
original_scheduler = copy.deepcopy(pipe.scheduler)

try:
    pipe.transformer.set_attention_backend(ATTN_BACKEND)
    print(f"Attention backend: {ATTN_BACKEND}")
except Exception as e:
    print(f"Attention backend '{ATTN_BACKEND}' unavailable, using default: {e}")

# Fuse any `fuse_at_startup` LoRAs from the catalog before quantization: fused weights survive
# the fp8 conversion, runtime adapters may not.
try:
    lora_loader.fuse_startup_loras(pipe)
except Exception as e:
    print("Startup LoRA fusion skipped:", e)

if QUANTIZE:
    # Qwen2.5-VL text encoder -> int8 weight only, DiT -> fp8 dynamic activations.
    # The SigLIP image encoder, ByT5 and the VAE are small enough to leave alone.
    try:
        quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
        torch._dynamo.reset()
    except Exception as e:
        print("text_encoder quantization skipped:", e)
    try:
        quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
        torch._dynamo.reset()
    except Exception as e:
        print("transformer quantization skipped:", e)

if AOTI_REPO:
    try:
        aoti.aoti_blocks_load(pipe.transformer, AOTI_REPO, variant=AOTI_VARIANT)
        print(f"AoTI blocks loaded from {AOTI_REPO}")
    except Exception as e:
        print("AoTI load skipped:", e)

pipe.vae.enable_tiling()

default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
default_negative_prompt = "overexposed, low quality, blurry details, subtitles, watermark, static, still frame, jpeg artifacts, deformed, disfigured, extra fingers, malformed hands, malformed face, cluttered background"


def model_title():
    return "## HunyuanVideo 1.5 I2V 8.3B — Fast Preview"


def bucket_size(image: Image.Image):
    """Resolution the pipeline will pick for this image (aspect-ratio bucket around target_size)."""
    height, width = pipe.video_processor.calculate_default_height_width(
        height=image.size[1], width=image.size[0], target_size=pipe.target_size
    )
    return width, height


def resize_image(image: Image.Image) -> Image.Image:
    """Center-crop/resize to the exact bucket the pipeline would choose, so the cost estimate
    below and the actual generation agree. Passing the result back is idempotent."""
    width, height = bucket_size(image)
    return pipe.video_processor.resize(image, height=height, width=width, resize_mode="crop")


def get_num_frames(duration_seconds: float):
    raw = int(round(duration_seconds * FIXED_FPS))
    raw = max(MIN_FRAMES_MODEL, min(MAX_FRAMES_MODEL, raw))
    return ((raw - 1) // 4) * 4 + 1


def get_inference_duration(
    resized_image,
    prompt,
    steps,
    negative_prompt,
    num_frames,
    guidance_scale,
    current_seed,
    scheduler_name,
    flow_shift,
    frame_multiplier,
    quality,
    duration_seconds,
    safe_mode,
    lora_groups,
    lora_scale,
    custom_lora,
    progress
):
    # Calibrated on the 480p step-distilled checkpoint at its base config (121 frames, 704x480).
    # Re-tune BASE_STEP_DURATION from the "gen time passed" logs on your own hardware.
    BASE_FRAMES_HEIGHT_WIDTH = 121 * 704 * 480
    BASE_STEP_DURATION = float(os.getenv("BASE_STEP_DURATION", "3.5"))
    width, height = resized_image.size
    factor = num_frames * width * height / BASE_FRAMES_HEIGHT_WIDTH
    step_duration = BASE_STEP_DURATION * factor ** 1.5
    gen_time = int(steps) * step_duration

    # guidance > 1 turns CFG back on -> two transformer passes per step
    if guidance_scale > 1:
        gen_time = gen_time * 2.0

    frame_factor = frame_multiplier // FIXED_FPS
    if frame_factor > 1:
        total_out_frames = (num_frames * frame_factor) - num_frames
        inter_time = (total_out_frames * 0.02)
        gen_time += inter_time

    total_time = 15 + gen_time
    if safe_mode:
        total_time = total_time * 1.30

    return total_time


def _apply_scheduler(scheduler_name, flow_shift):
    scheduler_class = SCHEDULER_MAP.get(scheduler_name, FlowMatchEulerDiscreteScheduler)

    if scheduler_class is FlowMatchEulerDiscreteScheduler:
        current = pipe.scheduler
        if current.config._class_name == "FlowMatchEulerDiscreteScheduler" and \
                float(current.config.get("shift", -1)) == float(flow_shift):
            return
        config = copy.deepcopy(original_scheduler.config)
        config["shift"] = float(flow_shift)
        pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(config)
        return

    try:
        pipe.scheduler = scheduler_class.from_config({
            "num_train_timesteps": 1000,
            "prediction_type": "flow_prediction",
            "use_flow_sigmas": True,
            "flow_shift": float(flow_shift),
        })
    except Exception as e:
        print(f"Scheduler '{scheduler_name}' failed ({e}); falling back to FlowMatchEulerDiscrete.")
        config = copy.deepcopy(original_scheduler.config)
        config["shift"] = float(flow_shift)
        pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(config)


def _apply_guidance(guidance_scale):
    """HunyuanVideo 1.5 takes CFG through a guider object, not a __call__ argument."""
    try:
        if float(pipe.guider.config.guidance_scale) != float(guidance_scale):
            pipe.guider = pipe.guider.new(guidance_scale=float(guidance_scale))
    except Exception as e:
        print("Could not update guider:", e)


@spaces.GPU(duration=get_inference_duration, size='xlarge')
def run_inference(
    resized_image,
    prompt,
    steps,
    negative_prompt,
    num_frames,
    guidance_scale,
    current_seed,
    scheduler_name,
    flow_shift,
    frame_multiplier,
    quality,
    duration_seconds,
    safe_mode=False,
    lora_groups=None,
    lora_scale=1.0,
    custom_lora="",
    progress=gr.Progress(track_tqdm=True),
):
    _apply_scheduler(scheduler_name, flow_shift)
    _apply_guidance(guidance_scale)

    clear_vram()

    task_name = str(uuid.uuid4())[:8]
    print(f"Generating {num_frames} frames, task: {task_name}, {duration_seconds}, {resized_image.size}, lora={lora_groups}")
    start = time.time()

    lora_loaded = False
    try:
        lora_loaded = lora_loader.load_loras_to_pipe(
            pipe, lora_groups, custom_lora, scale=float(lora_scale)
        )
    except Exception as e:
        print(f"LoRA warning: {e}")
        lora_loader.unload_lora(pipe)

    result = pipe(
        image=resized_image,
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_frames=num_frames,
        num_inference_steps=int(steps),
        generator=torch.Generator(device="cuda").manual_seed(current_seed),
        output_type="np",
    )

    if lora_loaded:
        lora_loader.unload_lora(pipe)

    print("gen time passed:", time.time() - start)

    raw_frames_np = result.frames[0]  # (T, H, W, C) float32
    pipe.scheduler = original_scheduler

    frame_factor = frame_multiplier // FIXED_FPS
    if frame_factor > 1:
        start = time.time()
        print(f"Processing frames (RIFE Multiplier: {frame_factor}x)...")
        rife_model.device()
        rife_model.flownet = rife_model.flownet.half()
        final_frames = interpolate_bits(raw_frames_np, multiplier=int(frame_factor))
        print("Interpolation time passed:", time.time() - start)
    else:
        final_frames = list(raw_frames_np)

    final_fps = FIXED_FPS * int(max(1, frame_factor))

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name

    start = time.time()
    with tqdm(total=3, desc="Rendering Media", unit="clip") as pbar:
        pbar.update(2)
        export_to_video(final_frames, video_path, fps=final_fps, quality=quality)
        pbar.update(1)
    print(f"Export time passed, {final_fps} FPS:", time.time() - start)

    return video_path, task_name


def generate_video(
    input_image,
    prompt,
    steps=DEFAULT_STEPS,
    negative_prompt=default_negative_prompt,
    duration_seconds=MAX_DURATION,
    guidance_scale=DEFAULT_GUIDANCE,
    seed=42,
    randomize_seed=False,
    quality=6,
    scheduler="FlowMatchEulerDiscrete",
    flow_shift=DEFAULT_SHIFT,
    frame_multiplier=FIXED_FPS,
    safe_mode=False,
    lora_groups=None,
    lora_scale=1.0,
    custom_lora="",
    video_component=True,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generate a video from an input image using HunyuanVideo 1.5 I2V (8.3B).

    This function takes an input image and generates a video animation based on the provided
    prompt and parameters. It uses an fp8-quantized HunyuanVideo 1.5 image-to-video model; with
    the step-distilled checkpoint 8-12 steps are enough.

    Args:
        input_image (PIL.Image): The input image to animate. Cropped to the closest aspect-ratio
            bucket around the model's target size (640px for 480p, 960px for 720p).
        prompt (str): Text prompt describing the desired animation or motion.
        steps (int, optional): Number of inference steps. Defaults to 8. Range: 1-50.
            The step-distilled checkpoint is tuned for 8 or 12; the plain checkpoints want 50.
        negative_prompt (str, optional): Negative prompt to avoid unwanted elements.
            Only used when guidance_scale > 1 (CFG is off on the distilled checkpoints).
        duration_seconds (float, optional): Duration of the generated video in seconds.
            Clamped between MIN_FRAMES_MODEL/FIXED_FPS and MAX_FRAMES_MODEL/FIXED_FPS.
        guidance_scale (float, optional): Classifier-free guidance scale, applied through the
            pipeline's guider. Defaults to 1.0 (disabled, one transformer pass per step).
            Values above 1 double the generation time. Range: 0.0-10.0.
        seed (int, optional): Random seed for reproducible results. Defaults to 42.
            Range: 0 to MAX_SEED (2147483647).
        randomize_seed (bool, optional): Whether to use a random seed instead of the provided seed.
        quality (float, optional): Video output quality. Uses variable bit rate.
            Highest quality is 10, lowest is 1.
        scheduler (str, optional): The name of the scheduler to use for inference.
            Defaults to "FlowMatchEulerDiscrete", which is what the checkpoint ships with.
        flow_shift (float, optional): The flow shift value. Defaults to 7.0 for the 480p
            step-distilled checkpoint (5.0 for the plain 480p ones).
        frame_multiplier (int, optional): Target fps; extra frames are produced by RIFE.
        lora_groups (list, optional): LoRA entries from the catalog to apply.
        lora_scale (float, optional): Weight applied to the selected LoRAs.
        custom_lora (str, optional): Extra LoRA as "repo_id" or "repo_id:filename".
        video_component (bool, optional): Show video player in output. Defaults to True.
        progress (gr.Progress, optional): Gradio progress tracker.

    Returns:
        tuple: A tuple containing:
            - video_path (str): Path for the video component.
            - video_path (str): Path for the file download component.
            - current_seed (int): The seed used for generation.

    Raises:
        gr.Error: If input_image is None (no image uploaded).

    Note:
        - Frame count is calculated as duration_seconds * FIXED_FPS (24) rounded to 4k+1
        - Output dimensions come from the model's aspect-ratio buckets, not from sliders
        - The function uses GPU acceleration via the @spaces.GPU decorator
        - Generation time varies based on steps and duration (see get_inference_duration)
    """

    if input_image is None:
        raise gr.Error("Please upload an input image.")

    num_frames = get_num_frames(duration_seconds)
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    resized_image = resize_image(input_image)

    video_path, task_n = run_inference(
        resized_image,
        prompt,
        steps,
        negative_prompt,
        num_frames,
        guidance_scale,
        current_seed,
        scheduler,
        flow_shift,
        frame_multiplier,
        quality,
        duration_seconds,
        safe_mode,
        lora_groups,
        lora_scale,
        custom_lora,
        progress,
    )
    print(f"GPU complete: {task_n}")

    return (video_path if video_component else None), video_path, current_seed


CSS = """
#hidden-timestamp {
    opacity: 0;
    height: 0px;
    width: 0px;
    margin: 0px;
    padding: 0px;
    overflow: hidden;
    position: absolute;
    pointer-events: none;
}
"""


with gr.Blocks(delete_cache=(3600, 10800)) as demo:
    gr.Markdown(model_title())
    gr.Markdown(
        "Run HunyuanVideo 1.5 image-to-video in 8-12 steps, fp8 quantization - "
        "compatible with 🧨 diffusers and ZeroGPU"
    )

    with gr.Row():
        with gr.Column():
            input_image_component = gr.Image(type="pil", label="Input Image", sources=["upload", "clipboard"])
            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v)
            duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=MAX_DURATION, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
            frame_multi = gr.Dropdown(
                choices=[FIXED_FPS, FIXED_FPS * 2, FIXED_FPS * 4],
                value=FIXED_FPS,
                label="Video Fluidity (Frames per Second)",
                info="Extra frames will be generated using flow estimation, which estimates motion between frames to make the video smoother."
            )
            safe_mode_checkbox = gr.Checkbox(
                label="🛠️ Safe Mode",
                value=True,
                info="Requests 30% extra processing time to try to prevent unfinished tasks when the server is busy."
            )
            with gr.Accordion("Advanced Settings", open=False):
                negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, info="Used only if Guidance Scale > 1.", lines=3)
                quality_slider = gr.Slider(minimum=1, maximum=10, step=1, value=6, label="Video Quality", info="If set to 10, the generated video may be too large and won't play in the Gradio preview.")
                seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
                randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
                steps_slider = gr.Slider(minimum=1, maximum=50, step=1, value=DEFAULT_STEPS, label="Inference Steps", info="8 or 12 for the step-distilled checkpoint, 50 for the plain ones.")
                guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=DEFAULT_GUIDANCE, label="Guidance Scale (CFG)", info="1.0 = off. Values above 1 double GPU time and enable the negative prompt.")
                scheduler_dropdown = gr.Dropdown(
                    label="Scheduler",
                    choices=list(SCHEDULER_MAP.keys()),
                    value="FlowMatchEulerDiscrete",
                    info="FlowMatchEulerDiscrete is what the checkpoint ships with; the rest run in flow-prediction mode and are experimental."
                )
                flow_shift_slider = gr.Slider(minimum=0.5, maximum=15.0, step=0.1, value=DEFAULT_SHIFT, label="Flow Shift", info="7.0 for the 480p step-distilled / 720p checkpoints, 5.0 for plain 480p.")
                lora_dropdown = gr.Dropdown(choices=lora_loader.get_lora_choices(), label="LoRA", multiselect=True, info="Entries from loras.json / LORA_CATALOG. HunyuanVideo 1.5 LoRAs only.")
                lora_scale_slider = gr.Slider(minimum=0.0, maximum=2.0, step=0.05, value=1.0, label="LoRA Scale")
                custom_lora_input = gr.Textbox(label="Custom LoRA", value="", placeholder="repo_id or repo_id:file.safetensors", info="Any Hub repo holding a HunyuanVideo 1.5 LoRA.")
                play_result_video = gr.Checkbox(label="Display result", value=True, interactive=True)

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

        with gr.Column():
            # ASSIGNED elem_id="generated-video" so JS can find it
            video_output = gr.Video(label="Generated Video", autoplay=True, sources=["upload"], buttons=["download", "share"], interactive=True, elem_id="generated-video")

            # --- Frame Grabbing UI ---
            with gr.Row():
                grab_frame_btn = gr.Button("📸 Use Current Frame as Input", variant="secondary")
                timestamp_box = gr.Number(value=0, label="Timestamp", visible=True, elem_id="hidden-timestamp")
            # -------------------------

            file_output = gr.File(label="Download Video")

    ui_inputs = [
        input_image_component, prompt_input, steps_slider,
        negative_prompt_input, duration_seconds_input,
        guidance_scale_input, seed_input, randomize_seed_checkbox,
        quality_slider, scheduler_dropdown, flow_shift_slider, frame_multi,
        safe_mode_checkbox,
        lora_dropdown, lora_scale_slider, custom_lora_input,
        play_result_video
    ]

    generate_button.click(
        fn=generate_video,
        inputs=ui_inputs,
        outputs=[video_output, file_output, seed_input]
    )

    # --- Frame Grabbing Events ---
    # 1. Click button -> JS runs -> puts time in hidden number box
    grab_frame_btn.click(
        fn=None,
        inputs=None,
        outputs=[timestamp_box],
        js=get_timestamp_js
    )

    # 2. Hidden number box changes -> Python runs -> puts frame in Input Image
    timestamp_box.change(
        fn=extract_frame,
        inputs=[video_output, timestamp_box],
        outputs=[input_image_component]
    )

if __name__ == "__main__":
    demo.queue().launch(
        mcp_server=True,
        css=CSS,
        show_error=True,
    )