"""
WAN 2.2 Multi-Task Video Generation - Bilingual UI
I2V: Lightning 14B (6 steps, FP8+AoT)
T2V: Lightning 14B (4 steps, Lightning LoRA + FP8)
V2V: VACE 14B (mask-based video editing, preserves unmasked regions)
LoRA: from lkzd7/WAN2.2_LoraSet_NSFW (I2V only)
"""
import os

import spaces
import shutil
import subprocess
import copy
import random
import tempfile
import warnings
import time
import gc
import uuid
import traceback
import types
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

def patch_diffusers_torchao_logger_bug():
    """Patch linoytsaban/diffusers@wan22-loras logger init order before importing diffusers."""
    import sys

    rel_path = os.path.join("diffusers", "quantizers", "torchao", "torchao_quantizer.py")
    for root in sys.path:
        if not root:
            continue
        target = os.path.join(root, rel_path)
        if not os.path.exists(target):
            continue
        try:
            with open(target, "r", encoding="utf-8") as file:
                text = file.read()
            logger_line = "logger = logging.get_logger(__name__)"
            call_pos = text.find("_update_torch_safe_globals()")
            logger_pos = text.find(logger_line)
            if call_pos != -1 and logger_pos > call_pos:
                marker = "from ..base import DiffusersQuantizer\n\n"
                text = text.replace(marker, marker + logger_line + "\n\n", 1)
                with open(target, "w", encoding="utf-8") as file:
                    file.write(text)
                print(f"Patched diffusers torchao logger bug: {target}")
            return
        except Exception as exc:
            print(f"Diffusers torchao logger patch failed for {target}: {exc}")


patch_diffusers_torchao_logger_bug()

from diffusers import (
    AutoencoderKLWan,
    FlowMatchEulerDiscreteScheduler,
    WanPipeline,
    SASolverScheduler,
    DEISMultistepScheduler,
    DPMSolverMultistepInverseScheduler,
    UniPCMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
)
from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
from diffusers.pipelines.wan.pipeline_wan_vace import WanVACEPipeline
from diffusers.utils.export_utils import export_to_video
from diffusers.utils import load_video
from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig, Int8WeightOnlyConfig
import aoti
import lora_loader

# SAM2 for video mask generation
from sam2.sam2_video_predictor import SAM2VideoPredictor

os.environ["TOKENIZERS_PARALLELISM"] = "true"
warnings.filterwarnings("ignore")

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

# ============ RIFE ============
get_timestamp_js = """
function() {
    const video = document.querySelector('#generated-video video');
    if (video) { return video.currentTime; }
    return 0;
}
"""

def extract_frame(video_path, timestamp):
    if not video_path:
        return None
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None
    fps = cap.get(cv2.CAP_PROP_FPS)
    target_frame_num = int(float(timestamp) * fps)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if target_frame_num >= total_frames:
        target_frame_num = total_frames - 1
    cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame_num)
    ret, frame = cap.read()
    cap.release()
    if ret:
        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    return None

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)

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):
    if isinstance(frames_np, list):
        T = len(frames_np)
        H, W, C = frames_np[0].shape
    else:
        T, H, W, C = frames_np.shape
    if multiplier < 2:
        return list(frames_np) if isinstance(frames_np, np.ndarray) else frames_np
    n_interp = multiplier - 1
    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)
    def to_tensor(frame_np):
        t = torch.from_numpy(frame_np).to(device)
        t = t.permute(2, 0, 1).unsqueeze(0)
        return F.pad(t, padding).half()
    def from_tensor(tensor):
        t = tensor[0, :, :H, :W]
        return t.permute(1, 2, 0).float().cpu().numpy()
    def make_inference(I0, I1, n):
        if rife_model.version >= 3.9:
            return [rife_model.inference(I0, I1, (i+1) * 1. / (n+1), scale) for i in range(n)]
        else:
            middle = rife_model.inference(I0, I1, scale)
            if n == 1: return [middle]
            first_half = make_inference(I0, middle, n//2)
            second_half = make_inference(middle, I1, n//2)
            return [*first_half, middle, *second_half] if n % 2 else [*first_half, *second_half]
    output_frames = []
    I1 = to_tensor(frames_np[0])
    with tqdm(total=T-1, desc="Interpolating", unit="frame") as pbar:
        for i in range(T - 1):
            I0 = I1
            output_frames.append(from_tensor(I0))
            I1 = to_tensor(frames_np[i+1])
            for mid in make_inference(I0, I1, n_interp):
                output_frames.append(from_tensor(mid))
            if (i + 1) % 50 == 0:
                pbar.update(50)
        pbar.update((T-1) % 50)
        output_frames.append(from_tensor(I1))
    del I0, I1
    torch.cuda.empty_cache()
    return output_frames

# ============ Config ============
FIXED_FPS = 16
MAX_FRAMES_MODEL = 241  # ~15s@16fps, requires more VRAM/time
MAX_SEED = np.iinfo(np.int32).max

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

default_negative_prompt = (
    "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, "
    "still image, overall gray, worst quality, low quality, JPEG artifacts, ugly, incomplete, "
    "extra fingers, poorly drawn hands, poorly drawn face, deformed, disfigured, "
    "malformed limbs, fused fingers, still frame, messy background, three legs, "
    "many people in background, walking backwards, watermark, text, signature"
)

# ============ Load I2V Pipeline (Lightning, AoT compiled) ============
print("Loading I2V Pipeline (Lightning 14B)...")
i2v_pipe = WanImageToVideoPipeline.from_pretrained(
    "TestOrganizationPleaseIgnore/WAMU_v2_WAN2.2_I2V_LIGHTNING",
    torch_dtype=torch.bfloat16,
)
i2v_original_scheduler = copy.deepcopy(i2v_pipe.scheduler)

def attach_manual_torchao_quantizer(model):
    """Expose manual torchao quantization config so PEFT can inject LoRA into quantized Wan layers."""
    model.hf_quantizer = types.SimpleNamespace(
        quantization_config=types.SimpleNamespace(
            get_apply_tensor_subclass=lambda: Float8DynamicActivationFloat8WeightConfig()
        )
    )

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

i2v_pipe.enable_model_cpu_offload()
print("I2V pipeline ready (FP8 quantized, CPU offload)")

for vae_memory_method in ("enable_slicing", "enable_tiling"):
    if hasattr(i2v_pipe.vae, vae_memory_method):
        getattr(i2v_pipe.vae, vae_memory_method)()

# ============ T2V Pipeline (on-demand, 14B + Wan22 Lightning LoRA) ============
# Use T2V-A14B + Wan22 Lightning LoRA (separate HIGH/LOW for dual transformer)
# Load on-demand with CPU offload to avoid OOM alongside I2V
T2V_MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
T2V_LORA_REPO = "Kijai/WanVideo_comfy"
T2V_LORA_HIGH = "LoRAs/Wan22-Lightning/Wan22_A14B_T2V_HIGH_Lightning_4steps_lora_250928_rank128_fp16.safetensors"
T2V_LORA_LOW = "LoRAs/Wan22-Lightning/Wan22_A14B_T2V_LOW_Lightning_4steps_lora_250928_rank64_fp16.safetensors"
t2v_pipe = None
t2v_ready = False

def load_t2v_pipeline():
    """Load T2V 14B + Lightning LoRA on-demand with CPU offload."""
    global t2v_pipe, t2v_ready

    if t2v_pipe is not None and t2v_ready:
        print("T2V pipeline reused from memory")
        return t2v_pipe

    print("Loading T2V Pipeline (14B + Lightning LoRA) first time...")

    # Move I2V components to CPU to make room
    i2v_pipe.to('cpu')
    clear_vram()

    t2v_vae = AutoencoderKLWan.from_pretrained(T2V_MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
    t2v_pipe = WanPipeline.from_pretrained(
        T2V_MODEL_ID,
        transformer=WanTransformer3DModel.from_pretrained(
            'linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
            subfolder='transformer',
            torch_dtype=torch.bfloat16,
        ),
        transformer_2=WanTransformer3DModel.from_pretrained(
            'linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
            subfolder='transformer_2',
            torch_dtype=torch.bfloat16,
        ),
        vae=t2v_vae,
        torch_dtype=torch.bfloat16,
    )

    # Load and fuse Lightning LoRAs (HIGH for transformer, LOW for transformer_2)
    print("Fusing Lightning LoRA HIGH (transformer)...")
    from safetensors.torch import load_file
    from huggingface_hub import hf_hub_download

    # Download LoRA files
    high_path = hf_hub_download(T2V_LORA_REPO, T2V_LORA_HIGH)
    low_path = hf_hub_download(T2V_LORA_REPO, T2V_LORA_LOW)

    # Load HIGH LoRA into transformer
    t2v_pipe.load_lora_weights(high_path, adapter_name="lightning_high")
    t2v_pipe.set_adapters(["lightning_high"], adapter_weights=[1.0])
    t2v_pipe.fuse_lora(adapter_names=["lightning_high"], lora_scale=1.0, components=["transformer"])
    t2v_pipe.unload_lora_weights()

    # Load LOW LoRA into transformer_2
    print("Fusing Lightning LoRA LOW (transformer_2)...")
    t2v_pipe.load_lora_weights(low_path, adapter_name="lightning_low", load_into_transformer_2=True)
    t2v_pipe.set_adapters(["lightning_low"], adapter_weights=[1.0])
    t2v_pipe.fuse_lora(adapter_names=["lightning_low"], lora_scale=1.0, components=["transformer_2"])
    t2v_pipe.unload_lora_weights()

    # Use model CPU offload — only one component on GPU at a time
    t2v_pipe.enable_model_cpu_offload()

    t2v_ready = True
    print("T2V pipeline ready (14B + Lightning + CPU offload)")
    return t2v_pipe

def unload_t2v_pipeline():
    """Clear VRAM after T2V; I2V uses CPU offload so no manual .to('cuda') needed."""
    clear_vram()
    print("T2V done, VRAM cleared (I2V uses CPU offload)")

# Keep cache for on-demand T2V loading

# ============ SAM2 Video Segmentation ============
sam2_predictor = None

def get_sam2_predictor():
    global sam2_predictor
    if sam2_predictor is None:
        print("Loading SAM2.1 hiera-large...")
        sam2_predictor = SAM2VideoPredictor.from_pretrained("facebook/sam2.1-hiera-large")
        print("SAM2 loaded")
    return sam2_predictor

def extract_first_frame_from_video(video_path):
    """Extract first frame from video as PIL Image."""
    cap = cv2.VideoCapture(video_path)
    ret, frame = cap.read()
    cap.release()
    if ret:
        return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    return None

def video_to_frames_dir(video_path, max_frames=None):
    """Extract video frames to a temp directory for SAM2."""
    tmp_dir = tempfile.mkdtemp(prefix="sam2_frames_")
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS) or 16
    idx = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if max_frames and idx >= max_frames:
            break
        cv2.imwrite(os.path.join(tmp_dir, f"{idx:05d}.jpg"), frame)
        idx += 1
    cap.release()
    print(f"Extracted {idx} frames to {tmp_dir} (fps={fps:.1f})")
    return tmp_dir, idx, fps

@spaces.GPU(duration=300)
def generate_mask_video(video_path, points_json, num_frames_limit=None):
    """Generate mask video using SAM2 from user-clicked points."""
    import json

    if not video_path:
        raise gr.Error("请先上传视频 / Upload a video first")
    if not points_json or points_json.strip() == "[]":
        raise gr.Error("请在视频第一帧上点击要编辑的区域 / Click on the area to edit")

    points_data = json.loads(points_json)
    if not points_data:
        raise gr.Error("没有标记点 / No points marked")

    # Extract frames
    frames_dir, total_frames, fps = video_to_frames_dir(video_path, max_frames=num_frames_limit)

    predictor = get_sam2_predictor()

    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
        state = predictor.init_state(video_path=frames_dir)

        # Add points (all on frame 0)
        pos_points = []
        neg_points = []
        for p in points_data:
            if p.get("label", 1) == 1:
                pos_points.append([p["x"], p["y"]])
            else:
                neg_points.append([p["x"], p["y"]])

        all_points = pos_points + neg_points
        all_labels = [1] * len(pos_points) + [0] * len(neg_points)

        points_np = np.array(all_points, dtype=np.float32)
        labels_np = np.array(all_labels, dtype=np.int32)

        _, _, _ = predictor.add_new_points_or_box(
            state,
            frame_idx=0,
            obj_id=1,
            points=points_np,
            labels=labels_np,
        )

        # Propagate through video
        all_masks = {}
        for frame_idx, obj_ids, masks in predictor.propagate_in_video(state):
            # masks shape: (num_objects, 1, H, W)
            mask = (masks[0, 0] > 0.0).cpu().numpy().astype(np.uint8) * 255
            all_masks[frame_idx] = mask

    # Build mask video
    out_path = os.path.join(tempfile.mkdtemp(), "mask_video.mp4")
    # Get frame size from first mask
    first_mask = all_masks[0]
    h, w = first_mask.shape
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(out_path, fourcc, fps, (w, h), isColor=False)
    for i in range(total_frames):
        if i in all_masks:
            writer.write(all_masks[i])
        elif all_masks:
            # Use nearest available mask
            nearest = min(all_masks.keys(), key=lambda k: abs(k - i))
            writer.write(all_masks[nearest])
    writer.release()

    # Cleanup frames dir
    shutil.rmtree(frames_dir, ignore_errors=True)

    print(f"Mask video generated: {out_path} ({total_frames} frames, {w}x{h})")
    return out_path

# ============ V2V Pipeline (VACE 14B, on-demand) ============
VACE_MODEL_ID = "Wan-AI/Wan2.1-VACE-14B-diffusers"
v2v_pipe = None
v2v_ready = False

def load_v2v_pipeline():
    """Load VACE 14B pipeline on-demand for mask-based video editing."""
    global v2v_pipe, v2v_ready

    # Move I2V to CPU to free GPU
    i2v_pipe.to('cpu')
    clear_vram()

    if v2v_pipe is not None and v2v_ready:
        v2v_pipe.to('cuda')
        print("VACE pipeline restored to GPU")
        return v2v_pipe

    print("Loading VACE 14B Pipeline first time (this downloads ~75GB)...")

    v2v_vae = AutoencoderKLWan.from_pretrained(VACE_MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
    v2v_pipe = WanVACEPipeline.from_pretrained(
        VACE_MODEL_ID,
        vae=v2v_vae,
        torch_dtype=torch.bfloat16,
    )
    v2v_pipe.scheduler = UniPCMultistepScheduler.from_config(v2v_pipe.scheduler.config, flow_shift=5.0)

    # Quantize to fit in A100 80GB
    quantize_(v2v_pipe.text_encoder, Int8WeightOnlyConfig())
    major, minor = torch.cuda.get_device_capability()
    if (major > 8) or (major == 8 and minor >= 9):
        quantize_(v2v_pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
    else:
        quantize_(v2v_pipe.transformer, Int8WeightOnlyConfig())

    v2v_pipe.to('cuda')

    v2v_ready = True
    print("VACE 14B pipeline ready (quantized, on GPU)")
    return v2v_pipe

def unload_v2v_pipeline():
    """Move V2V to CPU and clear VRAM; I2V uses CPU offload so no manual .to('cuda') needed."""
    global v2v_pipe
    if v2v_pipe is not None:
        v2v_pipe.to('cpu')
    clear_vram()
    print("VACE → CPU, VRAM cleared (I2V uses CPU offload)")

def load_video_frames_and_masks(video_path, mask_path, num_frames, target_h, target_w):
    """Load source video frames and mask video frames for VACE."""
    # Load source video frames as PIL Images
    src_frames = load_video(video_path)[:num_frames]
    print(f"Loaded {len(src_frames)} source frames (original size: {src_frames[0].size if src_frames else 'N/A'})")

    # Load mask video frames
    mask_frames_raw = load_video(mask_path)[:num_frames]

    # Convert mask to L mode (white=edit, black=keep) — don't resize, let pipeline handle it
    masks = []
    for mf in mask_frames_raw:
        gray = mf.convert("L")
        masks.append(gray)
    print(f"Loaded {len(masks)} mask frames")

    # Pad or trim to match
    while len(masks) < len(src_frames):
        masks.append(masks[-1] if masks else Image.new("L", src_frames[0].size, 0))
    while len(src_frames) < len(masks):
        src_frames.append(src_frames[-1] if src_frames else Image.new("RGB", (target_w, target_h), (128, 128, 128)))

    frame_count = min(len(src_frames), len(masks))
    src_frames = src_frames[:frame_count]
    masks = masks[:frame_count]

    return src_frames, masks

# ============ Utils ============
def resize_image(image, max_dim=832, min_dim=480, square_dim=640, multiple_of=16):
    width, height = image.size
    if width == height:
        return image.resize((square_dim, square_dim), Image.LANCZOS)
    aspect_ratio = width / height
    max_ar = max_dim / min_dim
    min_ar = min_dim / max_dim
    if aspect_ratio > max_ar:
        crop_width = int(round(height * max_ar))
        left = (width - crop_width) // 2
        image = image.crop((left, 0, left + crop_width, height))
        target_w, target_h = max_dim, min_dim
    elif aspect_ratio < min_ar:
        crop_height = int(round(width / min_ar))
        top = (height - crop_height) // 2
        image = image.crop((0, top, width, top + crop_height))
        target_w, target_h = min_dim, max_dim
    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 = max(min_dim, min(max_dim, round(target_w / multiple_of) * multiple_of))
    final_h = max(min_dim, min(max_dim, round(target_h / multiple_of) * multiple_of))
    return image.resize((final_w, final_h), Image.LANCZOS)

def resize_and_crop_to_match(target_image, reference_image):
    ref_w, ref_h = reference_image.size
    tgt_w, tgt_h = target_image.size
    scale = max(ref_w / tgt_w, ref_h / tgt_h)
    new_w, new_h = int(tgt_w * scale), int(tgt_h * scale)
    resized = target_image.resize((new_w, new_h), Image.Resampling.LANCZOS)
    left, top = (new_w - ref_w) // 2, (new_h - ref_h) // 2
    return resized.crop((left, top, left + ref_w, top + ref_h))

def get_num_frames(duration_seconds):
    raw = int(round(duration_seconds * FIXED_FPS))
    raw = ((raw - 1) // 4) * 4 + 1
    return int(np.clip(raw, 9, MAX_FRAMES_MODEL))

def extract_video_path(input_video):
    if input_video is None:
        return None
    if isinstance(input_video, str):
        return input_video
    if isinstance(input_video, dict):
        # Gradio 5.x format: {'video': filepath, ...} or {'name': filepath, ...} or {'path': filepath}
        return input_video.get("video", input_video.get("path", input_video.get("name", None)))
    # Could be a Gradio VideoData object
    if hasattr(input_video, 'video'):
        return input_video.video
    if hasattr(input_video, 'path'):
        return input_video.path
    if hasattr(input_video, 'name'):
        return input_video.name
    return str(input_video)

def extract_first_frame(video_input):
    path = extract_video_path(video_input)
    if not path or not os.path.exists(path):
        return None
    cap = cv2.VideoCapture(path)
    ret, frame = cap.read()
    cap.release()
    if ret:
        return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    return None

# ============ Inference ============
def estimate_gpu_duration(
    task_type, input_image, input_video, mask_video, prompt, negative_prompt,
    duration_seconds, steps, guidance_scale, guidance_scale_2,
    current_seed, scheduler_name, flow_shift, frame_multiplier,
    quality, last_image_input, lora_groups, *args, **kwargs,
):
    task = str(task_type or "")
    duration = float(duration_seconds or 3)
    step_count = int(steps or 6)
    num_frames = get_num_frames(duration)

    if "V2V" in task:
        return 600
    if "T2V" in task:
        return 600

    frame_factor = max(1, int(frame_multiplier or FIXED_FPS) // FIXED_FPS)
    estimate = 40 + step_count * 8 + max(0, num_frames - 9) * 2.5
    if float(guidance_scale or 1) > 1 or float(guidance_scale_2 or 1) > 1:
        estimate *= 1.2
    if frame_factor > 1:
        estimate += num_frames * (frame_factor - 1) * 0.3
    if lora_groups:
        lora_count = len(lora_groups) if isinstance(lora_groups, (list, tuple)) else 1
        estimate += 10 + max(0, lora_count - 1) * 5
    return int(max(60, min(600, estimate)))


@spaces.GPU(duration=estimate_gpu_duration)
def run_inference(
    task_type, input_image, input_video, mask_video, prompt, negative_prompt,
    duration_seconds, steps, guidance_scale, guidance_scale_2,
    current_seed, scheduler_name, flow_shift, frame_multiplier,
    quality, last_image_input, lora_groups,
    progress=gr.Progress(track_tqdm=True),
):
    clear_vram()
    num_frames = get_num_frames(duration_seconds)
    task_id = str(uuid.uuid4())[:8]
    print(f"Task: {task_id}, type={task_type}, duration={duration_seconds}s, frames={num_frames}")
    start = time.time()

    if "T2V" in task_type:
        # ====== T2V: 14B + Lightning LoRA (4 steps, dual guidance) ======
        t2v_steps = max(int(steps), 4)
        print(f"T2V: steps={t2v_steps}, guidance={guidance_scale}/{guidance_scale_2}, frames={num_frames}")

        pipe = load_t2v_pipeline()
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=480,
            width=832,
            num_frames=num_frames,
            guidance_scale=float(guidance_scale),
            guidance_scale_2=float(guidance_scale_2),
            num_inference_steps=t2v_steps,
            generator=torch.Generator(device="cpu").manual_seed(int(current_seed)),
            output_type="np",
        )
        unload_t2v_pipeline()

    else:
        # ====== I2V / V2V ======
        if "V2V" in task_type:
            # V2V uses VACE pipeline — mask-based video editing
            print(f"V2V VACE: input_video type={type(input_video)}, value={input_video}")
            video_path = extract_video_path(input_video)
            if not video_path or not os.path.exists(video_path):
                raise gr.Error("Upload a source video for V2V / V2V请上传原视频")

            # Get mask video path
            mask_path = extract_video_path(mask_video)
            if not mask_path or not os.path.exists(mask_path):
                raise gr.Error("Upload a mask video for V2V / V2V请上传遮罩视频（黑白视频，白色=编辑区域）")

            target_h, target_w = 480, 832
            src_frames, masks = load_video_frames_and_masks(video_path, mask_path, num_frames, target_h, target_w)

            # Ensure num_frames satisfies (n-1) % 4 == 0 for VACE
            n = len(src_frames)
            n = (n - 1) // 4 * 4 + 1
            n = max(n, 5)
            src_frames = src_frames[:n]
            masks = masks[:n]

            # Load VACE pipeline
            pipe = load_v2v_pipeline()
            v2v_steps = max(int(steps), 20)
            print(f"V2V VACE: steps={v2v_steps}, guidance={guidance_scale}, frames={len(src_frames)}")

            result = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                video=src_frames,
                mask=masks,
                height=target_h,
                width=target_w,
                num_frames=len(src_frames),
                guidance_scale=max(float(guidance_scale), 5.0),
                num_inference_steps=v2v_steps,
                generator=torch.Generator(device="cuda").manual_seed(int(current_seed)),
                output_type="np",
            )
            unload_v2v_pipeline()

        else:
            # ====== I2V ======
            if input_image is None:
                raise gr.Error("Upload an image / 请上传图片")

            scheduler_class = SCHEDULER_MAP.get(scheduler_name)
            if scheduler_class and scheduler_class.__name__ != i2v_pipe.scheduler.config._class_name:
                config = copy.deepcopy(i2v_original_scheduler.config)
                if scheduler_class == FlowMatchEulerDiscreteScheduler:
                    config['shift'] = flow_shift
                else:
                    config['flow_shift'] = flow_shift
                i2v_pipe.scheduler = scheduler_class.from_config(config)

            loaded_lora_adapters = []
            lora_plan = lora_loader.get_lora_blend_plan(lora_groups)
            if lora_plan:
                try:
                    print(f"LoRA blend plan: {lora_plan}")
                    for idx, entry in enumerate(lora_plan):
                        loaded_lora_adapters.extend(
                            lora_loader.load_lora_to_pipe(
                                i2v_pipe,
                                entry["name"],
                                adapter_name=f"lora_{idx}",
                                weight=entry["weight"],
                            )
                        )
                    if loaded_lora_adapters:
                        lora_loader.activate_loras(i2v_pipe, loaded_lora_adapters)
                except Exception as e:
                    details = traceback.format_exc()
                    print(details)
                    lora_loader.unload_lora(i2v_pipe)
                    raise gr.Error(f"LoRA failed: {type(e).__name__}: {str(e)[:700]}")

            lora_loaded = bool(loaded_lora_adapters)

            resized_image = resize_image(input_image)
            processed_last = None
            if last_image_input:
                processed_last = resize_and_crop_to_match(last_image_input, resized_image)

            print(f"I2V: size={resized_image.size}, steps={int(steps)}, guidance={guidance_scale}/{guidance_scale_2}")

            try:
                result = i2v_pipe(
                    image=resized_image,
                    last_image=processed_last,
                    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(int(current_seed)),
                    attention_kwargs={"scale": 1.0} if lora_loaded else None,
                    output_type="np",
                )
            except RuntimeError as e:
                if "NVML" in str(e) or "out of memory" in str(e).lower() or "CUDACachingAllocator" in str(e):
                    print(f"OOM with {num_frames} frames, retrying with reduced frames...")
                    clear_vram()
                    reduced_frames = max(9, ((num_frames - 1) // 2 // 4) * 4 + 1)
                    print(f"Retrying with {reduced_frames} frames")
                    result = i2v_pipe(
                        image=resized_image,
                        last_image=processed_last,
                        prompt=prompt,
                        negative_prompt=negative_prompt,
                        height=resized_image.height,
                        width=resized_image.width,
                        num_frames=reduced_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(int(current_seed)),
                        attention_kwargs={"scale": 1.0} if lora_loaded else None,
                        output_type="np",
                    )
                else:
                    raise
            except Exception as e:
                details = traceback.format_exc()
                print(details)
                raise gr.Error(f"I2V failed: {type(e).__name__}: {str(e)[:700]}\n{details[-1200:]}")
            finally:
                if lora_loaded:
                    lora_loader.unload_lora(i2v_pipe)
                    clear_vram()

    raw_frames = result.frames[0]
    elapsed = time.time() - start
    print(f"Generation took {elapsed:.1f}s ({len(raw_frames)} frames)")

    frame_factor = frame_multiplier // FIXED_FPS
    if frame_factor > 1:
        rife_model.device()
        rife_model.flownet = rife_model.flownet.half()
        final_frames = interpolate_bits(raw_frames, multiplier=int(frame_factor))
    else:
        final_frames = list(raw_frames)
    final_fps = FIXED_FPS * max(1, frame_factor)

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name
    export_to_video(final_frames, video_path, fps=final_fps, quality=quality)
    return video_path, task_id

# ============ Generate ============
def generate_video(
    task_type, input_image, input_video, mask_video, prompt,
    lora_groups, duration_seconds, frame_multiplier,
    steps, guidance_scale, guidance_scale_2,
    negative_prompt, quality, seed, randomize_seed,
    scheduler, flow_shift, last_image, display_result,
    progress=gr.Progress(track_tqdm=True),
):
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt / 请输入提示词")
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    try:
        video_path, task_id = run_inference(
            task_type, input_image, input_video, mask_video, prompt, negative_prompt,
            duration_seconds, steps, guidance_scale, guidance_scale_2,
            current_seed, scheduler, flow_shift, frame_multiplier,
            quality, last_image, lora_groups,
        )
    except Exception as e:
        err_msg = str(e)
        if "GPU task aborted" in err_msg or "aborted" in err_msg.lower():
            raise gr.Error("GPU task timed out. Try reducing duration, steps, or output FPS. / GPU超时，请缩短时长、减少步数或降低输出帧率。")
        raise
    print(f"Done: {task_id}")
    return (video_path if display_result else None), video_path, current_seed

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

with gr.Blocks(theme=gr.themes.Soft(), css=CSS, delete_cache=(3600, 10800)) as demo:
    gr.Markdown("## WAN 2.2 Multi-Task Video Generation / 多任务视频生成")
    gr.Markdown("#### I2V (Lightning 8-step) · T2V (Lightning 14B 4-step) · V2V (VACE 14B) · Multi-LoRA")
    gr.Markdown("---")

    task_type = gr.Radio(
        choices=[
            "I2V (图生视频 / Image-to-Video)",
            "T2V (文生视频 / Text-to-Video)",
            "V2V (视频生视频 / Video-to-Video)",
        ],
        value="I2V (图生视频 / Image-to-Video)",
        label="Task Type / 任务类型",
    )

    with gr.Row():
        with gr.Column():
            with gr.Group():
                input_image = gr.Image(type="pil", label="Input Image / 输入图片 (I2V)", sources=["upload", "clipboard"])
            with gr.Group():
                input_video = gr.Video(label="Source Video / 原视频 (V2V)", sources=["upload"], visible=False, interactive=True)
            with gr.Group():
                mask_video = gr.Video(label="Mask Video / 遮罩视频 (V2V, 白色=编辑区域)", sources=["upload"], visible=False, interactive=True)
            v2v_guide = gr.Markdown(
                value="""### 📖 V2V 使用方法 / How to use V2V

**步骤1**: 上传原视频到 "Source Video"
**步骤2**: 点击下方"提取第一帧"按钮
**步骤3**: 在第一帧上**点击**要编辑的区域（绿色=编辑，红色=排除）
**步骤4**: 点击"生成遮罩"按钮，SAM2 会自动追踪整个视频
**步骤5**: 填写 Prompt 描述目标效果，点击生成

💡 也可以直接上传自己的遮罩视频（白色=编辑区域）
""",
                visible=False,
            )
            with gr.Group(visible=False) as v2v_mask_tools:
                first_frame_display = gr.Image(label="第一帧预览 / First Frame (点击标记区域)", type="pil", interactive=False)
                points_store = gr.State(value=[])
                points_display = gr.Textbox(label="标记点 / Points", value="无标记 / No points", interactive=False)
                with gr.Row():
                    point_mode = gr.Radio(choices=["include (编辑)", "exclude (排除)"], value="include (编辑)", label="点击模式")
                with gr.Row():
                    extract_frame_btn = gr.Button("📷 提取第一帧 / Extract First Frame", variant="secondary")
                    gen_mask_btn = gr.Button("🎭 生成遮罩 / Generate Mask (SAM2)", variant="primary")
                    clear_points_btn = gr.Button("🗑️ 清除标记 / Clear Points")

            prompt_input = gr.Textbox(
                label="Prompt / 提示词", value="",
                placeholder="Describe the video... / 描述你想生成的视频...", lines=3,
            )
            duration_slider = gr.Slider(
                minimum=0.5, maximum=15, step=0.5, value=3,
                label="Duration / 时长 (seconds/秒)",
                info="Max ~15s (241 frames @16fps) / 最大约15秒",
            )
            frame_multi = gr.Dropdown(choices=[16, 32, 64], value=16, label="Output FPS / 输出帧率", info="RIFE interpolation / RIFE插帧")

            with gr.Accordion("⚙️ Advanced Settings / 高级设置", open=False):
                last_image = gr.Image(type="pil", label="Last Frame / 末帧 (Optional)", sources=["upload", "clipboard"])
                negative_prompt_input = gr.Textbox(label="Negative Prompt / 负面提示词", value=default_negative_prompt, lines=3)
                with gr.Row():
                    steps_slider = gr.Slider(minimum=1, maximum=50, step=1, value=8, label="Steps / 步数", info="I2V LoRA: 8-12 | T2V: 4-8 | V2V: 25-50")
                    quality_sl = gr.Slider(minimum=1, maximum=10, step=1, value=8, label="Quality / 质量")
                with gr.Row():
                    guidance_h = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance High / 引导(高噪声)")
                    guidance_l = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1.0, label="Guidance Low / 引导(低噪声)")
                with gr.Row():
                    scheduler_dd = gr.Dropdown(choices=list(SCHEDULER_MAP.keys()), value="UniPCMultistep", label="Scheduler / 调度器")
                    flow_shift_sl = gr.Slider(minimum=0.5, maximum=15.0, step=0.1, value=3.0, label="Flow Shift / 流偏移")
                with gr.Row():
                    seed_sl = gr.Slider(minimum=0, maximum=MAX_SEED, step=1, value=42, label="Seed / 种子")
                    random_seed_cb = gr.Checkbox(label="Random / 随机", value=True)
                lora_dd = gr.Dropdown(
                    choices=lora_loader.get_lora_choices(),
                    label="LoRA blend (I2V only / 仅I2V)",
                    multiselect=True,
                    max_choices=lora_loader.MAX_LORAS,
                    info="Up to 3. First action LoRA is primary; helper/style LoRAs are auto-weighted lower.",
                )
                display_cb = gr.Checkbox(label="Display / 显示", value=True)

            generate_btn = gr.Button("🎬 Generate / 生成视频", variant="primary", size="lg")

        with gr.Column():
            video_output = gr.Video(label="Generated Video / 生成的视频", autoplay=True, sources=["upload"], show_download_button=True, show_share_button=True, interactive=False, elem_id="generated-video")
            with gr.Row():
                grab_frame_btn = gr.Button("📸 Use Frame / 使用帧", variant="secondary")
                timestamp_box = gr.Number(value=0, label="Timestamp", visible=False, elem_id="hidden-timestamp")
            file_output = gr.File(label="Download / 下载")

    def update_task_ui(task):
        is_v2v = "V2V" in task
        is_t2v = "T2V" in task
        if is_t2v:
            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(value=4), gr.update(value=1.0), gr.update(value=1.0)
        elif is_v2v:
            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(value=30), gr.update(value=5.0), gr.update(value=1.0)
        else:
            return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(value=8), gr.update(value=1.0), gr.update(value=1.0)

    task_type.change(update_task_ui, inputs=[task_type], outputs=[input_image, input_video, mask_video, v2v_guide, v2v_mask_tools, steps_slider, guidance_h, guidance_l])

    # V2V mask generation callbacks
    def on_extract_first_frame(video):
        vpath = extract_video_path(video)
        if not vpath or not os.path.exists(vpath):
            raise gr.Error("请先上传视频 / Upload video first")
        frame = extract_first_frame_from_video(vpath)
        if frame is None:
            raise gr.Error("无法提取第一帧 / Failed to extract first frame")
        return frame, [], "无标记 / No points"

    def on_click_frame(img, points, mode, evt: gr.SelectData):
        if img is None:
            return img, points, "请先提取第一帧 / Extract first frame first"
        x, y = evt.index
        label = 1 if "include" in mode else 0
        points.append({"x": x, "y": y, "label": label})
        # Draw points on image
        display_img = img.copy()
        draw = __import__('PIL').ImageDraw.Draw(display_img)
        for p in points:
            color = (0, 255, 0) if p["label"] == 1 else (255, 0, 0)
            r = 8
            draw.ellipse([p["x"]-r, p["y"]-r, p["x"]+r, p["y"]+r], fill=color, outline="white", width=2)
        info = f"{len([p for p in points if p['label']==1])} include, {len([p for p in points if p['label']==0])} exclude"
        return display_img, points, info

    def on_clear_points(original_video):
        vpath = extract_video_path(original_video)
        if vpath and os.path.exists(vpath):
            frame = extract_first_frame_from_video(vpath)
            return frame, [], "无标记 / No points"
        return None, [], "无标记 / No points"

    def on_generate_mask(video, points):
        import json
        vpath = extract_video_path(video)
        if not vpath:
            raise gr.Error("请先上传视频 / Upload video first")
        if not points:
            raise gr.Error("请先在第一帧上点击标记 / Click on first frame to mark areas")
        mask_path = generate_mask_video(vpath, json.dumps(points))
        return mask_path

    extract_frame_btn.click(fn=on_extract_first_frame, inputs=[input_video], outputs=[first_frame_display, points_store, points_display])
    first_frame_display.select(fn=on_click_frame, inputs=[first_frame_display, points_store, point_mode], outputs=[first_frame_display, points_store, points_display])
    clear_points_btn.click(fn=on_clear_points, inputs=[input_video], outputs=[first_frame_display, points_store, points_display])
    gen_mask_btn.click(fn=on_generate_mask, inputs=[input_video, points_store], outputs=[mask_video])
    generate_btn.click(
        fn=generate_video,
        inputs=[task_type, input_image, input_video, mask_video, prompt_input, lora_dd, duration_slider, frame_multi,
                steps_slider, guidance_h, guidance_l, negative_prompt_input, quality_sl, seed_sl, random_seed_cb,
                scheduler_dd, flow_shift_sl, last_image, display_cb],
        outputs=[video_output, file_output, seed_sl],
    )
    grab_frame_btn.click(fn=None, inputs=None, outputs=[timestamp_box], js=get_timestamp_js)
    timestamp_box.change(fn=extract_frame, inputs=[video_output, timestamp_box], outputs=[input_image])

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