import os; os.system('pip install --upgrade --no-deps spaces')
import spaces
import shutil
import subprocess
import sys
import copy
import random
import tempfile
import warnings
import time
import gc
import uuid
import json
import hashlib
from tqdm import tqdm
import cv2
import numpy as np
import torch
import torch._dynamo
from huggingface_hub import list_models
from torch.nn import functional as F
from PIL import Image

import gradio as gr
from diffusers import (
    FlowMatchEulerDiscreteScheduler,
    SASolverScheduler,
    DEISMultistepScheduler,
    DPMSolverMultistepInverseScheduler,
    UniPCMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
)
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
from diffusers.utils.export_utils import export_to_video

from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig, Int8WeightOnlyConfig
import aoti

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

# if IS_ZERO_GPU:
#     print("Loading...")
#     subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)

# --- 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)
        # Gradio Image component handles Numpy array -> PIL conversion automatically
        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/r3gm/RIFE/resolve/main/RIFEv4.26_0921.zip",
        "-O", "RIFEv4.26_0921.zip"
    ], check=True)
    subprocess.run(["unzip", "-o", "RIFEv4.26_0921.zip"], check=True)

# sys.path.append(os.getcwd())

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, 8)
    Returns:
        List of Numpy Arrays (Height, Width, Channels) - Float32 [0.0, 1.0]
    """
    
    # Handle input shape
    if isinstance(frames_np, list):
        # Convert list of arrays to one big array for easier shape handling if needed, 
        # but here we just grab dims from first frame
        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:
        # Just convert 4D array to list of 3D arrays
        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
    # Load first frame into GPU
    I1 = to_tensor(frames_np[0])

    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


# WAN

ORG_NAME = "TestOrganizationPleaseIgnore"
# MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
MODEL_ID = os.getenv("REPO_ID") or random.choice(
    list(list_models(author=ORG_NAME, filter='diffusers:WanImageToVideoPipeline'))
).modelId
CACHE_DIR = os.path.expanduser("~/.cache/huggingface/")

# Mounted secret JSON with LoRA entries. Each entry may include:
#   repo_id (str): HF repo id, or local directory path starting with /
#   high_tr / low_tr (str, optional): weight path inside repo/dir - independent, not paired
#   high_scale / low_scale (float, optional): fuse scales (default 1.0)
#   trigger / default_prompt / notes (str, optional): shown in Gradio UI
CONFIG_PATH = os.getenv("WAN555_CONFIG_PATH", "/secret/config-wan555/active.json")
# Immutable snapshot of the config that was fused at Space startup.
STARTUP_CONFIG_SNAPSHOT = os.path.join(tempfile.gettempdir(), "wan555_startup_active.json")

_cached_config_hash = None
_cached_loras = None
_startup_config_hash = None
_startup_loras = None


def get_file_hash(path):
    if not os.path.exists(path):
        return None
    with open(path, "rb") as f:
        return hashlib.md5(f.read()).hexdigest()


def load_lora_config(force=False):
    """Load LoRA config from JSON. Reload only when file content hash changes."""
    global _cached_config_hash, _cached_loras

    current_hash = get_file_hash(CONFIG_PATH)
    if current_hash is None:
        print(f"WARNING: LoRA config not found: {CONFIG_PATH}")
        _cached_config_hash = None
        _cached_loras = []
        return _cached_loras

    if not force and current_hash == _cached_config_hash and _cached_loras is not None:
        print("LoRA config unchanged (using cache)")
        return _cached_loras

    print(f"Loading LoRA config from {CONFIG_PATH}")
    try:
        with open(CONFIG_PATH, "r", encoding="utf-8") as f:
            config = json.load(f)
        if not isinstance(config, list):
            raise ValueError("active.json must be a JSON array of LoRA entries")
        _cached_loras = config
        _cached_config_hash = current_hash
        print(f"Loaded {len(config)} LoRA entries")
        return config
    except Exception as e:
        print(f"Failed to load LoRA config: {e}")
        _cached_config_hash = None
        _cached_loras = []
        return _cached_loras


def snapshot_startup_config(config, config_hash):
    """Keep an immutable copy of the LoRA config fused at startup."""
    global _startup_loras, _startup_config_hash
    _startup_loras = copy.deepcopy(config or [])
    _startup_config_hash = config_hash
    try:
        with open(STARTUP_CONFIG_SNAPSHOT, "w", encoding="utf-8") as f:
            json.dump(_startup_loras, f, ensure_ascii=False, indent=2)
        print(f"Saved startup LoRA snapshot: {STARTUP_CONFIG_SNAPSHOT}")
    except Exception as e:
        print(f"WARNING: could not write startup LoRA snapshot: {e}")


def entry_display_name(entry, index):
    for key in ("high_tr", "low_tr"):
        val = entry.get(key)
        if val:
            return str(val).split("/")[-1]
    return entry.get("repo_id", f"LoRA {index + 1}")


def is_local_repo(repo_id):
    return isinstance(repo_id, str) and (
        repo_id.startswith("/")
        or repo_id.startswith("\\")
        or (len(repo_id) > 1 and repo_id[1] == ":")
    )


def apply_loras_to_pipeline(pipe, lora_models):
    """Load and fuse LoRA weights into the pipeline. Modifies pipe in-place."""
    for i, lora in enumerate(lora_models):
        repo_id = lora.get("repo_id")
        if not repo_id:
            print(f"Skipping LoRA entry {i}: missing repo_id")
            continue

        high_tr = lora.get("high_tr") or None
        low_tr = lora.get("low_tr") or None
        if not high_tr and not low_tr:
            print(f"Skipping LoRA entry {i}: neither high_tr nor low_tr set")
            continue

        high_scale = float(lora.get("high_scale", 1.0))
        low_scale = float(lora.get("low_scale", 1.0))
        local = is_local_repo(repo_id)
        name_high_tr = None
        name_low_tr = None

        try:
            adapters_to_set = []

            if high_tr:
                name_high_tr = high_tr.split(".")[0].split("/")[-1] + "Hh"
                load_kwargs = {
                    "weight_name": high_tr,
                    "adapter_name": name_high_tr,
                }
                if local:
                    load_kwargs["local_files_only"] = True
                pipe.load_lora_weights(repo_id, **load_kwargs)
                adapters_to_set.append(name_high_tr)

            if low_tr:
                name_low_tr = low_tr.split(".")[0].split("/")[-1] + "Ll"
                load_kwargs = {
                    "weight_name": low_tr,
                    "adapter_name": name_low_tr,
                    "load_into_transformer_2": True,
                }
                if local:
                    load_kwargs["local_files_only"] = True
                pipe.load_lora_weights(repo_id, **load_kwargs)
                adapters_to_set.append(name_low_tr)

            if adapters_to_set:
                pipe.set_adapters(adapters_to_set, adapter_weights=[1.0] * len(adapters_to_set))

            if name_high_tr:
                pipe.fuse_lora(
                    adapter_names=[name_high_tr],
                    lora_scale=high_scale,
                    components=["transformer"],
                )
            if name_low_tr:
                pipe.fuse_lora(
                    adapter_names=[name_low_tr],
                    lora_scale=low_scale,
                    components=["transformer_2"],
                )

            pipe.unload_lora_weights()

            print(
                f"Applied: high={high_tr}@{high_scale}, low={low_tr}@{low_scale}, "
                f"local={local}, {i + 1}/{len(lora_models)}"
            )
        except Exception as e:
            print("Error:", str(e))
            print("Failed LoRA entry:", entry_display_name(lora, i), "repo:", repo_id)
            try:
                pipe.unload_lora_weights()
            except Exception:
                pass


def refresh_lora_display_config():
    """Re-read active.json for UI display only. Does NOT touch the pipeline."""
    load_lora_config(force=True)
    return format_lora_display()


# Load initial LoRA config and freeze a startup snapshot of what gets fused.
load_lora_config()
snapshot_startup_config(_cached_loras, _cached_config_hash)

MAX_DIM = 832
MIN_DIM = 512
SQUARE_DIM = 768
MULTIPLE_OF = 16
MAX_SEED = np.iinfo(np.int32).max

FIXED_FPS = 16
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 1280

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

SCHEDULER_MAP = {
    "FlowMatchEulerDiscrete": FlowMatchEulerDiscreteScheduler,
    "SASolver": SASolverScheduler,
    "DEISMultistep": DEISMultistepScheduler,
    "DPMSolverMultistepInverse": DPMSolverMultistepInverseScheduler,
    "UniPCMultistep": UniPCMultistepScheduler,
    "DPMSolverMultistep": DPMSolverMultistepScheduler,
    "DPMSolverSinglestep": DPMSolverSinglestepScheduler,
}

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

apply_loras_to_pipeline(pipe, _cached_loras)

# if os.path.exists(CACHE_DIR):
#     shutil.rmtree(CACHE_DIR)
#     print("Deleted Hugging Face cache.")
# else:
#     print("No hub cache found.")

quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
torch._dynamo.reset()
quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
torch._dynamo.reset()
quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
torch._dynamo.reset()

spaces.aoti_load(
    module=pipe.transformer,
    repo_id='cbensimon/WanTransformer3DModel-sm120-cu130-raa',
)
spaces.aoti_load(
    module=pipe.transformer_2,
    repo_id='cbensimon/WanTransformer3DModel-sm120-cu130-raa',
)

# pipe.vae.enable_slicing()
# pipe.vae.enable_tiling()

default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
for _entry in _cached_loras:
    _dp = (_entry.get("default_prompt") or "").strip()
    if _dp:
        default_prompt_i2v = _dp
        break

default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"


def model_title():
    repo_name = MODEL_ID.split('/')[-1].replace("_", " ")
    url = f"https://huggingface.co/{MODEL_ID}"
    return f"## This space is currently running [{repo_name}]({url}) 🐢"


def resize_image(image: Image.Image) -> Image.Image:
    width, height = image.size
    if width == height:
        return image.resize((SQUARE_DIM, SQUARE_DIM), Image.LANCZOS)
    
    aspect_ratio = width / height
    MAX_ASPECT_RATIO = MAX_DIM / MIN_DIM
    MIN_ASPECT_RATIO = MIN_DIM / MAX_DIM

    image_to_resize = image
    if aspect_ratio > MAX_ASPECT_RATIO:
        target_w, target_h = MAX_DIM, MIN_DIM
        crop_width = int(round(height * MAX_ASPECT_RATIO))
        left = (width - crop_width) // 2
        image_to_resize = image.crop((left, 0, left + crop_width, height))
    elif aspect_ratio < MIN_ASPECT_RATIO:
        target_w, target_h = MIN_DIM, MAX_DIM
        crop_height = int(round(width / MIN_ASPECT_RATIO))
        top = (height - crop_height) // 2
        image_to_resize = image.crop((0, top, width, top + crop_height))
    else:
        if width > height:
            target_w = MAX_DIM
            target_h = int(round(target_w / aspect_ratio))
        else:
            target_h = MAX_DIM
            target_w = int(round(target_h * aspect_ratio))

    final_w = round(target_w / MULTIPLE_OF) * MULTIPLE_OF
    final_h = round(target_h / MULTIPLE_OF) * MULTIPLE_OF
    final_w = max(MIN_DIM, min(MAX_DIM, final_w))
    final_h = max(MIN_DIM, min(MAX_DIM, final_h))
    return image_to_resize.resize((final_w, final_h), Image.LANCZOS)


def resize_and_crop_to_match(target_image, reference_image):
    ref_width, ref_height = reference_image.size
    target_width, target_height = target_image.size
    scale = max(ref_width / target_width, ref_height / target_height)
    new_width, new_height = int(target_width * scale), int(target_height * scale)
    resized = target_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    left, top = (new_width - ref_width) // 2, (new_height - ref_height) // 2
    return resized.crop((left, top, left + ref_width, top + ref_height))


def get_num_frames(duration_seconds: float):
    return 1 + int(np.clip(
        int(round(duration_seconds * FIXED_FPS)),
        MIN_FRAMES_MODEL,
        MAX_FRAMES_MODEL,
    ))


def get_inference_duration(
    resized_image,
    processed_last_image,
    prompt,
    steps,
    negative_prompt,
    num_frames,
    guidance_scale,
    guidance_scale_2,
    current_seed,
    scheduler_name,
    flow_shift,
    frame_multiplier,
    quality,
    duration_seconds,
    safe_mode,
    progress
):
    BASE_FRAMES_HEIGHT_WIDTH = 81 * 832 * 624
    BASE_STEP_DURATION = 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

    if guidance_scale > 1:
        gen_time = gen_time * 2.4

    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


@spaces.GPU(duration=get_inference_duration, size='xlarge')
def run_inference(
    resized_image,
    processed_last_image,
    prompt,
    steps,
    negative_prompt,
    num_frames,
    guidance_scale,
    guidance_scale_2,
    current_seed,
    scheduler_name,
    flow_shift,
    frame_multiplier,
    quality,
    duration_seconds,
    safe_mode=False,
    progress=gr.Progress(track_tqdm=True),
):
    scheduler_class = SCHEDULER_MAP.get(scheduler_name)
    if scheduler_class.__name__ != pipe.scheduler.config._class_name or flow_shift != pipe.scheduler.config.get("flow_shift", "shift"):
        config = copy.deepcopy(original_scheduler.config)
        if scheduler_class == FlowMatchEulerDiscreteScheduler:
            config['shift'] = flow_shift
        else:
            config['flow_shift'] = flow_shift
        pipe.scheduler = scheduler_class.from_config(config)

    clear_vram()

    task_name = str(uuid.uuid4())[:8]
    print(f"Generating {num_frames} frames, task: {task_name}, {duration_seconds}, {resized_image.size}")
    start = time.time()
    result = pipe(
        image=resized_image,
        last_image=processed_last_image,
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=resized_image.height,
        width=resized_image.width,
        num_frames=num_frames,
        guidance_scale=float(guidance_scale),
        guidance_scale_2=float(guidance_scale_2),
        num_inference_steps=int(steps),
        generator=torch.Generator(device="cuda").manual_seed(current_seed),
        output_type="np" 
    )
    print("gen time passed:", time.time() - start)
    
    raw_frames_np = result.frames[0]  # Returns (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(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,
    last_image,
    prompt,
    steps=4,
    negative_prompt=default_negative_prompt,
    duration_seconds=MAX_DURATION,
    guidance_scale=1,
    guidance_scale_2=1,
    seed=42,
    randomize_seed=False,
    quality=5,
    scheduler="UniPCMultistep",
    flow_shift=6.0,
    frame_multiplier=16,
    safe_mode=False,
    video_component=True,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generate a video from an input image using the Wan 2.2 14B I2V model with Lightning LoRA.
    This function takes an input image and generates a video animation based on the provided
    prompt and parameters. It uses an FP8 qunatized Wan 2.2 14B Image-to-Video model in with Lightning LoRA
    for fast generation in 4-8 steps.
    Args:
        input_image (PIL.Image): The input image to animate. Will be resized to target dimensions.
        last_image (PIL.Image, optional): The optional last image for the video.
        prompt (str): Text prompt describing the desired animation or motion.
        steps (int, optional): Number of inference steps. More steps = higher quality but slower.
            Defaults to 4. Range: 1-30.
        negative_prompt (str, optional): Negative prompt to avoid unwanted elements.
            Defaults to default_negative_prompt (contains unwanted visual artifacts).
        duration_seconds (float, optional): Duration of the generated video in seconds.
            Defaults to 2. Clamped between MIN_FRAMES_MODEL/FIXED_FPS and MAX_FRAMES_MODEL/FIXED_FPS.
        guidance_scale (float, optional): Controls adherence to the prompt. Higher values = more adherence.
            Defaults to 1.0. Range: 0.0-20.0.
        guidance_scale_2 (float, optional): Controls adherence to the prompt. Higher values = more adherence.
            Defaults to 1.0. Range: 0.0-20.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.
            Defaults to False.
        quality (float, optional): Video output quality. Default is 5. 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 "UniPCMultistep".
        flow_shift (float, optional): The flow shift value for compatible schedulers. Defaults to 6.0.
        frame_multiplier (int, optional): The int value for fps enhancer
        video_component(bool, optional): Show video player in output.
            Defaults to True.
        progress (gr.Progress, optional): Gradio progress tracker. Defaults to gr.Progress(track_tqdm=True).
    Returns:
        tuple: A tuple containing:
            - video_path (str): Path for the video component.
            - video_path (str): Path for the file download component. Attempt to avoid reconversion in video 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)
        - Output dimensions are adjusted to be multiples of MOD_VALUE (32)
        - The function uses GPU acceleration via the @spaces.GPU decorator
        - Generation time varies based on steps and duration (see get_duration function)
    """
    
    if input_image is None:
        raise gr.Error("Please upload an input image.")

    # Refresh display config only. Weight hot-reload is not supported mid-request on ZeroGPU
    # (request-time .to('cuda') is blocked outside startup / @spaces.GPU).
    load_lora_config(force=True)

    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)

    processed_last_image = None
    if last_image:
        processed_last_image = resize_and_crop_to_match(last_image, resized_image)

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

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


def _format_lora_entries(entries, empty_message):
    if not entries:
        return empty_message

    lines = []
    for i, entry in enumerate(entries):
        name = entry_display_name(entry, i)
        trigger = (entry.get("trigger") or "").strip()
        notes = (entry.get("notes") or "").strip()
        dp = (entry.get("default_prompt") or "").strip()
        repo = entry.get("repo_id", "-")
        high = entry.get("high_tr") or "-"
        low = entry.get("low_tr") or "-"
        hs = entry.get("high_scale", 1.0)
        ls = entry.get("low_scale", 1.0)

        lines.append(f"**{name}**")
        lines.append(
            f"weights: high=`{high}` @ {hs}  ·  low=`{low}` @ {ls}"
        )
        lines.append(f"repo: `{repo}`")
        if trigger:
            lines.append(f"trigger: `{trigger}`")
        if dp:
            lines.append(f"default prompt: `{dp}`")
        if notes:
            lines.append(f"notes: {notes}")
        lines.append("")
    return "\n".join(lines).rstrip()


def format_lora_display():
    """Show currently fused LoRAs vs next-restart config from active.json."""
    session_block = _format_lora_entries(
        _startup_loras or [],
        "_No LoRAs fused in this session._",
    )

    lines = [
        "### Currently loaded (this Space session)",
        "These LoRAs were fused at startup and stay active until the Space restarts.",
        "",
        session_block,
        "",
        "### Next restart (`active.json`)",
    ]

    if _startup_config_hash is not None and _cached_config_hash == _startup_config_hash:
        lines.append("Next restart: same LoRA configuration.")
    else:
        lines.append(
            "Next restart, this ZeroGPU Space will load the following LoRA config from "
            f"`{CONFIG_PATH}`:"
        )
        lines.append("")
        lines.append(
            _format_lora_entries(
                _cached_loras or [],
                f"_No LoRAs in `{CONFIG_PATH}`._",
            )
        )

    return "\n".join(lines)


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 Wan 2.2 in just 4-8 steps, fp8 quantization & AoT compilation - 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=3.5, 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, FIXED_FPS*8],
                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):
                last_image_component = gr.Image(type="pil", label="Last Image (Optional)", sources=["upload", "clipboard"])
                negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, info="Used if any 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=30, step=1, value=6, label="Inference Steps")
                guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale - high noise stage", info="Values above 1 increase GPU usage and may take longer to process.")
                guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale 2 - low noise stage")
                scheduler_dropdown = gr.Dropdown(
                    label="Scheduler",
                    choices=list(SCHEDULER_MAP.keys()),
                    value="UniPCMultistep",
                    info="Select a custom scheduler."
                )
                flow_shift_slider = gr.Slider(minimum=0.5, maximum=15.0, step=0.1, value=3.0, label="Flow Shift")
                play_result_video = gr.Checkbox(label="Display result", value=True, interactive=True)
                gr.Markdown(f"[ZeroGPU help, tips and troubleshooting](https://huggingface.co/datasets/{ORG_NAME}/help/blob/main/gpu_help.md)")
                gr.Markdown(  # TestOrganizationPleaseIgnore/wamu-tools
                    "To use a different model, **duplicate this Space** first, then change the `REPO_ID` environment variable. "
                    "[See compatible models here](https://huggingface.co/models?other=diffusers:WanImageToVideoPipeline&sort=trending&search=WAN2.2_I2V_LIGHTNING)."
                )

            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")

    with gr.Accordion("LoRA info", open=False):
        with gr.Row():
            refresh_loras_btn = gr.Button("🔄 Refresh next-restart info", variant="secondary", size="sm")
        lora_info_display = gr.Markdown(value=format_lora_display(), elem_id="lora-info")

        # Display-only: re-read active.json for the "next restart" box.
        # Does not fuse/reload model weights.
        refresh_loras_btn.click(
            fn=refresh_lora_display_config,
            inputs=[],
            outputs=[lora_info_display],
        )

    ui_inputs = [
        input_image_component, last_image_component, prompt_input, steps_slider,
        negative_prompt_input, duration_seconds_input,
        guidance_scale_input, guidance_scale_2_input, seed_input, randomize_seed_checkbox,
        quality_slider, scheduler_dropdown, flow_shift_slider, frame_multi,
        safe_mode_checkbox,
        play_result_video
    ]
    
    generate_button.click(
        fn=generate_video, 
        inputs=ui_inputs, 
        outputs=[video_output, file_output, seed_input, lora_info_display]
    )
    
    # --- 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,
    )