app.py

import os; os.system('pip install --no-deps spaces==0.51.1')
import spaces
import subprocess
import sys
import copy
import random
import tempfile
import warnings
import time
import gc
import uuid
import threading
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 gradio.context import LocalContext
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
import lora_loader
from video_job_api import (
    API_GPU_DURATION_SECONDS,
    DEFAULT_NEGATIVE_PROMPT,
    DEFAULT_PROMPT,
    VideoJobAPI,
    VideoJobRequest,
    VideoJobSettings,
    bind_context_values,
    calculate_dynamic_gpu_duration,
    create_job_api_lifespan,
    resolve_gpu_duration,
)

os.environ["TOKENIZERS_PARALLELISM"] = "true"
warnings.filterwarnings("ignore")
VIDEO_JOB_SETTINGS = VideoJobSettings.from_env()
INFERENCE_SLOT = threading.Lock()
# UI 外层不能再次自动申请 GPU；唯一 ZeroGPU 边界由 run_inference 的动态装饰器负责。
spaces.disable_gradio_auto_wrap()

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

# 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

MODEL_ID = "thornmaze/WAMU_v3_WAN2.2_I2V_LIGHTNING"

LORA_MODELS = []

MAX_DIM = 832
MIN_DIM = 480
SQUARE_DIM = 640
MULTIPLE_OF = 16
MAX_SEED = np.iinfo(np.int32).max

FIXED_FPS = 16
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 321

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)

for i, lora in enumerate(LORA_MODELS):
    name_high_tr = lora["high_tr"].split(".")[0].split("/")[-1] + "Hh"
    name_low_tr = lora["low_tr"].split(".")[0].split("/")[-1] + "Ll"
    
    try: 
        pipe.load_lora_weights(
            lora["repo_id"],
            weight_name=lora["high_tr"],
            adapter_name=name_high_tr
        )
    
        kwargs_lora = {"load_into_transformer_2": True}
        pipe.load_lora_weights(
            lora["repo_id"],
            weight_name=lora["low_tr"],
            adapter_name=name_low_tr,
            **kwargs_lora
        )
    
        pipe.set_adapters([name_high_tr, name_low_tr], adapter_weights=[1.0, 1.0])
    
        pipe.fuse_lora(adapter_names=[name_high_tr], lora_scale=lora["high_scale"], components=["transformer"])
        pipe.fuse_lora(adapter_names=[name_low_tr], lora_scale=lora["low_scale"], components=["transformer_2"])
    
        pipe.unload_lora_weights()

        print(f"Applied: {lora['high_tr']}, hs={lora['high_scale']}/ls={lora['low_scale']}, {i+1}/{len(LORA_MODELS)}") 
    except Exception as e:
        print("Error:", str(e))
        print("Failed LoRA:", name_high_tr)
        pipe.unload_lora_weights()

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='thornmaze/WanTransformer3DModel-sm120-cu130-raa',
)
spaces.aoti_load(
    module=pipe.transformer_2,
    repo_id='thornmaze/WanTransformer3DModel-sm120-cu130-raa',
)

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

default_prompt_i2v = DEFAULT_PROMPT
default_negative_prompt = DEFAULT_NEGATIVE_PROMPT


def model_title():
    return "## Wan 2.2 I2V 14B Lightning — NSFW"


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):
    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,
    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,
    lora_groups,
    progress
):
    """解析本次 ZeroGPU 时长，自定义任务优先，交互调用继续动态估算。

    Args:
        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: 是否为动态估时增加安全余量。
        lora_groups: LoRA 选择，保留回调签名。
        progress: Gradio 进度对象，保留回调签名。

    Returns:
        自定义 API 指定的 GPU 秒数，或原有交互链路的动态估算秒数。
    """
    del (
        processed_last_image,
        prompt,
        negative_prompt,
        guidance_scale_2,
        current_seed,
        scheduler_name,
        flow_shift,
        quality,
        duration_seconds,
        lora_groups,
        progress,
    )

    # 工厂函数保持惰性：API 覆盖存在时不运行任何参数耗时预判。
    return resolve_gpu_duration(
        lambda: calculate_dynamic_gpu_duration(
            image_size=resized_image.size,
            num_frames=num_frames,
            steps=steps,
            guidance_scale=guidance_scale,
            frame_multiplier=frame_multiplier,
            fixed_fps=FIXED_FPS,
            safe_mode=safe_mode,
        )
    )


@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,
    lora_groups=None,
    progress=gr.Progress(track_tqdm=True),
):
    """在 ZeroGPU 上执行现有 Wan I2V 推理并生成一个临时 MP4。

    Args:
        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: MP4 编码质量。
        duration_seconds: 用于日志和 ZeroGPU 时长估算的视频秒数。
        safe_mode: 是否为 ZeroGPU 估时增加安全余量。
        lora_groups: 要动态加载的 LoRA 精确名称列表。
        progress: Gradio 进度对象。

    Returns:
        生成的单个临时 MP4 路径与截短任务标识。
    """
    task_name = str(uuid.uuid4())[:8]
    video_path = None
    video_ready = False
    lora_attempted = False
    result = None
    raw_frames_np = None
    final_frames = None

    try:
        scheduler_class = SCHEDULER_MAP.get(scheduler_name)
        if scheduler_class is None:
            raise ValueError(f"Unsupported scheduler: {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()
        print(f"Generating {num_frames} frames, task: {task_name}, {duration_seconds}, {resized_image.size}, lora={lora_groups}")
        start = time.time()

        if lora_groups:
            # 从第一项开始加载就视为已污染管线，部分加载失败也必须进入 finally 卸载。
            lora_attempted = True
            try:
                for idx, name in enumerate(lora_groups):
                    if name and name != "(None)":
                        lora_loader.load_lora_to_pipe(pipe, name, adapter_name=f"lora_{idx}")
                print(f"LoRA loaded: {lora_groups}")
            except Exception as exc:
                print(f"LoRA warning: {type(exc).__name__}")
                # 保留原 UI 的降级语义，但不能让部分 LoRA 参与本次或后续推理。
                lora_loader.unload_lora(pipe)
                lora_attempted = False

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

        video_ready = True
        return video_path, task_name
    finally:
        # 无论模型、插帧还是编码在哪一步失败，都恢复可供下一任务复用的全局管线。
        cleanup_error = None
        try:
            if lora_attempted:
                lora_loader.unload_lora(pipe)
        except Exception as exc:
            cleanup_error = exc
        try:
            pipe.scheduler = copy.deepcopy(original_scheduler)
        except Exception as exc:
            cleanup_error = cleanup_error or exc
        result = None
        raw_frames_np = None
        final_frames = None
        try:
            clear_vram()
        except Exception as exc:
            cleanup_error = cleanup_error or exc
        if video_path is not None and (not video_ready or cleanup_error is not None):
            try:
                os.remove(video_path)
            except FileNotFoundError:
                pass
            except OSError as exc:
                print(f"Failed to remove temporary video: {type(exc).__name__}")
        if cleanup_error is not None:
            raise RuntimeError("Inference cleanup failed.") from cleanup_error


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,
    lora_groups=None,
    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.")

    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,
        lora_groups,
        progress,
    )
    print(f"GPU complete: {task_n}")

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


def generate_video_ui(
    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,
    lora_groups=None,
    video_component=True,
    progress=gr.Progress(track_tqdm=True),
):
    """让现有 UI 通过与自定义 API 共用的非阻塞推理槽位生成视频。

    Args:
        input_image: UI 上传的首图。
        last_image: 可选尾图。
        prompt: 正向提示词。
        steps: 推理步数。
        negative_prompt: 负向提示词。
        duration_seconds: 目标视频时长。
        guidance_scale: 高噪声阶段引导强度。
        guidance_scale_2: 低噪声阶段引导强度。
        seed: 固定随机种子。
        randomize_seed: 是否在生成前随机化种子。
        quality: MP4 编码质量。
        scheduler: 调度器名称。
        flow_shift: 调度器流偏移值。
        frame_multiplier: 输出帧率倍数对应的帧率值。
        safe_mode: 是否申请额外 ZeroGPU 运行时间。
        lora_groups: 当前 LoRA 下拉框选择的精确名称列表。
        video_component: 是否把结果同时显示在视频组件中。
        progress: Gradio 进度对象。

    Returns:
        与原 generate_video 一致的视频组件路径、下载路径和实际 seed。
    """
    if not INFERENCE_SLOT.acquire(blocking=False):
        raise gr.Error("The generation service is busy. Please retry shortly.")
    try:
        return generate_video(
            input_image,
            last_image,
            prompt,
            steps,
            negative_prompt,
            duration_seconds,
            guidance_scale,
            guidance_scale_2,
            seed,
            randomize_seed,
            quality,
            scheduler,
            flow_shift,
            frame_multiplier,
            safe_mode,
            lora_groups,
            video_component,
            progress,
        )
    finally:
        INFERENCE_SLOT.release()


# Gradio 6 的 MCP 工具名取自回调函数 __name__；保留升级前的 generate_video 工具名。
generate_video_ui.__name__ = "generate_video"


def build_gradio_request(headers: dict[str, str], job_id: str) -> gr.Request:
    """为后台 ZeroGPU 调用重建最小 Gradio 请求对象。

    Args:
        headers: 仅含 ZeroGPU 身份所需字段的筛选后请求头。
        job_id: 用作隔离会话哈希的自定义任务标识。

    Returns:
        可供 spaces.GPU 装饰器读取身份信息的 Gradio 请求。
    """
    return gr.Request(
        username=headers.get("x-gradio-user"),
        session_hash=job_id,
        headers=dict(headers),
        query_params={},
        cookies={},
        path_params={},
        client={"host": "127.0.0.1", "port": 0},
        url="",
    )


def execute_video_job(
    payload: VideoJobRequest,
    input_image: Image.Image,
    last_image: Image.Image | None,
    zero_gpu_headers: dict[str, str],
    job_id: str,
) -> tuple[str, int]:
    """在后台线程恢复 Gradio 上下文并调用现有视频生成链路。

    Args:
        payload: 已通过公开 Schema 校验的命名任务参数。
        input_image: 已安全抓取并解码的首图。
        last_image: 已安全抓取并解码的可选尾图。
        zero_gpu_headers: 仅含短效 ZeroGPU 身份字段的请求头。
        job_id: 用于隔离后台 Gradio 请求上下文的任务标识。

    Returns:
        现有生成函数产生的临时 MP4 明确路径与实际使用的 seed。
    """
    request_context = build_gradio_request(zero_gpu_headers, job_id)

    # 通用 ContextVar 绑定器保证成功或异常时都恢复四个 Gradio 本地上下文。
    with bind_context_values(
        (
            (LocalContext.request, request_context),
            (LocalContext.blocks, demo),
            (LocalContext.in_event_listener, True),
            (LocalContext.event_id, None),
            (API_GPU_DURATION_SECONDS, payload.gpu_duration_seconds),
        )
    ):
        _, video_path, used_seed = generate_video(
            input_image=input_image,
            last_image=last_image,
            prompt=payload.prompt,
            steps=payload.steps,
            negative_prompt=payload.negative_prompt,
            duration_seconds=payload.duration_seconds,
            guidance_scale=payload.guidance_scale,
            guidance_scale_2=payload.guidance_scale_2,
            seed=payload.seed,
            randomize_seed=payload.randomize_seed,
            quality=payload.quality,
            scheduler=payload.scheduler,
            flow_shift=payload.flow_shift,
            frame_multiplier=payload.frame_multiplier,
            safe_mode=payload.safe_mode,
            lora_groups=payload.lora_groups,
            video_component=False,
        )
    return video_path, int(used_seed)


VIDEO_JOB_API = VideoJobAPI(
    settings=VIDEO_JOB_SETTINGS,
    executor=execute_video_job,
    allowed_loras=set(lora_loader.get_lora_choices()),
    inference_slot=INFERENCE_SLOT,
)
JOB_API_LIFESPAN = create_job_api_lifespan(VIDEO_JOB_API)


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")
                lora_dropdown = gr.Dropdown(choices=lora_loader.get_lora_choices(), label="LoRA (NSFW)", multiselect=True, info="Select scenario LoRAs")
                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, 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,
        lora_dropdown,
        play_result_video
    ]
    
    generate_button.click(
        fn=generate_video_ui,
        inputs=ui_inputs,
        outputs=[video_output, file_output, seed_input],
        api_name="generate_video",
        concurrency_limit=1,
    )
    
    # --- 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(default_concurrency_limit=1).launch(
        mcp_server=True,
        css=CSS,
        show_error=True,
        ssr_mode=False,
        app_kwargs={"lifespan": JOB_API_LIFESPAN},
    )