import os
import subprocess
import sys

# Disable torch.compile / dynamo before any torch import
os.environ["TORCH_COMPILE_DISABLE"] = "1"
os.environ["TORCHDYNAMO_DISABLE"] = "1"

# Install xformers for memory-efficient attention
subprocess.run([sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"], check=False)

# Clone LTX-2 repo and install packages
LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
LTX_REPO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "LTX-2")
LTX_COMMIT_SHA = "ae855f8538843825f9015a419cf4ba5edaf5eec2"

if not os.path.exists(LTX_REPO_DIR):
    print(f"Cloning {LTX_REPO_URL}...")
    os.makedirs(LTX_REPO_DIR)
    subprocess.run(["git", "init", LTX_REPO_DIR], check=True)
    subprocess.run(["git", "remote", "add", "origin", LTX_REPO_URL], cwd=LTX_REPO_DIR, check=True)
    subprocess.run(["git", "fetch", "--depth", "1", "origin", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, check=True)
    subprocess.run(["git", "checkout", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, check=True)

print("Installing ltx-core and ltx-pipelines from cloned repo...")
subprocess.run(
    [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", "-e",
     os.path.join(LTX_REPO_DIR, "packages", "ltx-core"),
     "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines")],
    check=True,
)

# Purge pip cache to save storage space
print("Purging pip cache...")
subprocess.run([sys.executable, "-m", "pip", "cache", "purge"], check=False)

# Delete cloned repository .git folder to save space
git_dir = os.path.join(LTX_REPO_DIR, ".git")
if os.path.exists(git_dir):
    import shutil
    print("Deleting cloned LTX-2 repo .git folder to save space...")
    shutil.rmtree(git_dir, ignore_errors=True)

sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines", "src"))
sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-core", "src"))

import logging
import random
import tempfile
from pathlib import Path

import torch
torch._dynamo.config.suppress_errors = True
torch._dynamo.config.disable = True

import spaces
import gradio as gr
import numpy as np
from huggingface_hub import hf_hub_download, snapshot_download

from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.media_io import encode_video, load_video_conditioning, decode_audio_from_file, get_videostream_metadata
from ltx_pipelines.utils.helpers import (
    encode_prompts,
    cleanup_memory,
    simple_denoising_func,
    denoise_audio_video,
)
from ltx_pipelines.utils import euler_denoising_loop
from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
from ltx_core.components.noisers import GaussianNoiser
from ltx_core.components.diffusion_steps import EulerDiffusionStep
from ltx_core.types import VideoPixelShape, LatentState
from ltx_core.components.protocols import DiffusionStepProtocol
from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.model.upsampler import upsample_video

# Force-patch xformers attention into the LTX attention module.
from ltx_core.model.transformer import attention as _attn_mod
print(f"[ATTN] Before patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}")
try:
    from xformers.ops import memory_efficient_attention as _mea
    _attn_mod.memory_efficient_attention = _mea
    print(f"[ATTN] After patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}")
except Exception as e:
    print(f"[ATTN] xformers patch FAILED: {type(e).__name__}: {e}")

# Disable xformers FA3 dispatch
try:
    from xformers.ops.fmha import _set_use_fa3
    _set_use_fa3(False)
    print("[ATTN] xformers FA3 dispatch disabled (Blackwell-incompatible)")
except Exception as e:
    print(f"[ATTN] FA3 disable FAILED: {type(e).__name__}: {e}")

# FUSE/mmap workaround
import json
import struct
from ltx_core.loader.primitives import StateDict
from ltx_core.loader.sft_loader import SafetensorsStateDictLoader

_SAFETENSORS_DTYPE_MAP = {
    "F64": torch.float64,
    "F32": torch.float32,
    "F16": torch.float16,
    "BF16": torch.bfloat16,
    "F8_E5M2": torch.float8_e5m2,
    "F8_E4M3": torch.float8_e4m3fn,
    "I64": torch.int64,
    "I32": torch.int32,
    "I16": torch.int16,
    "I8": torch.int8,
    "U8": torch.uint8,
    "BOOL": torch.bool,
}

def _patched_load(self, path, sd_ops, device=None):
    sd = {}
    size = 0
    dtype = set()
    device = device or torch.device("cpu")
    model_paths = path if isinstance(path, list) else [path]
    for shard_path in model_paths:
        with open(shard_path, "rb") as f:
            header_len = struct.unpack("<Q", f.read(8))[0]
            header = json.loads(f.read(header_len).decode("utf-8"))
            data_base = 8 + header_len
            for name, meta in header.items():
                if name == "__metadata__":
                    continue
                expected_name = name if sd_ops is None else sd_ops.apply_to_key(name)
                if expected_name is None:
                    continue
                start, end = meta["data_offsets"]
                f.seek(data_base + start)
                buf = f.read(end - start)
                t = torch.frombuffer(
                    bytearray(buf), dtype=_SAFETENSORS_DTYPE_MAP[meta["dtype"]]
                ).reshape(meta["shape"])
                t = t.to(device=device, non_blocking=True, copy=False)
                kvs = (
                    ((expected_name, t),)
                    if sd_ops is None
                    else sd_ops.apply_to_key_value(expected_name, t)
                )
                for key, v in kvs:
                    size += v.nbytes
                    dtype.add(v.dtype)
                    sd[key] = v
    return StateDict(sd=sd, device=device, size=size, dtype=dtype)

SafetensorsStateDictLoader.load = _patched_load
print("[FUSE-PATCH] SafetensorsStateDictLoader.load replaced (chunked-read)")

logging.getLogger().setLevel(logging.INFO)

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

RESOLUTIONS = {
    "high": {"16:9": (1536, 1024), "9:16": (1024, 1536), "1:1": (1024, 1024)},
    "low": {"16:9": (768, 512), "9:16": (512, 768), "1:1": (768, 768)},
}

LTX_MOUNT = "/models/ltx"
GEMMA_MOUNT = "/models/gemma"

import shutil
def print_disk(tag):
    try:
        u = shutil.disk_usage(".")
        print(f"[DISK {tag}] total={u.total/1024**3:.2f}GB, used={u.used/1024**3:.2f}GB, free={u.free/1024**3:.2f}GB")
    except Exception as e:
        print(f"[DISK {tag}] error checking usage: {e}")

# Check if mounts exist
if os.path.exists(LTX_MOUNT) and os.path.exists(GEMMA_MOUNT):
    print("LTX and Gemma mounts detected. Performing fast-path model initialization...")
    mounted_files = os.listdir(LTX_MOUNT)
    distilled_file = next((f for f in mounted_files if "distilled" in f), "ltx-2.3-22b-distilled-1.1.safetensors")
    distilled_checkpoint_path = os.path.join(LTX_MOUNT, distilled_file)
    spatial_upsampler_path = os.path.join(LTX_MOUNT, "ltx-2.3-spatial-upscaler-x2-1.1.safetensors")
    gemma_root = GEMMA_MOUNT

    print("Initializing DistilledPipeline...")
    pipeline = DistilledPipeline(
        distilled_checkpoint_path=distilled_checkpoint_path,
        spatial_upsampler_path=spatial_upsampler_path,
        gemma_root=gemma_root,
        loras=[],
        quantization=QuantizationPolicy.fp8_cast(),
    )
    ledger = pipeline.model_ledger
    print("Preloading models for ZeroGPU...")
    _transformer = ledger.transformer()
    _video_encoder = ledger.video_encoder()
    _video_decoder = ledger.video_decoder()
    _audio_decoder = ledger.audio_decoder()
    _vocoder = ledger.vocoder()
    _spatial_upsampler = ledger.spatial_upsampler()
    _text_encoder = ledger.text_encoder()
    _embeddings_processor = ledger.gemma_embeddings_processor()
else:
    print("Mounts not found. Initiating sequential download and loading to bypass 50GB storage limit...")
    print_disk("startup")
    
    os.makedirs("models", exist_ok=True)
    
    print("1. Downloading Gemma text encoder (24 GB)...")
    gemma_root = snapshot_download(
        repo_id="Lightricks/gemma-3-12b-it-qat-q4_0-unquantized",
        local_dir="models/gemma",
        local_dir_use_symlinks=False
    )
    print_disk("after_gemma_download")

    print("2. Downloading spatial upscaler (1 GB)...")
    spatial_upsampler_path = hf_hub_download(
        repo_id="Lightricks/LTX-2.3",
        filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
        local_dir="models",
        local_dir_use_symlinks=False
    )
    print_disk("after_upscaler_download")

    print("3. Instantiating DistilledPipeline with Gemma and spatial upscaler (using dummy path for base model)...")
    pipeline = DistilledPipeline(
        distilled_checkpoint_path="models/dummy_base.safetensors",
        spatial_upsampler_path=spatial_upsampler_path,
        gemma_root=gemma_root,
        loras=[],
        quantization=QuantizationPolicy.fp8_cast(),
    )
    ledger = pipeline.model_ledger

    print("4. Preloading Gemma and upscaler models...")
    _text_encoder = ledger.text_encoder()
    _embeddings_processor = ledger.gemma_embeddings_processor()
    _spatial_upsampler = ledger.spatial_upsampler()
    print("Gemma and upscaler preloaded in CPU/GPU memory.")

    print("5. Deleting Gemma and upscaler files from disk to free storage space...")
    for f in os.listdir("models/gemma"):
        if f.endswith(".safetensors"):
            os.remove(os.path.join("models/gemma", f))
    if os.path.exists(spatial_upsampler_path):
        os.remove(spatial_upsampler_path)
    print_disk("after_gemma_upscaler_deletion")

    print("6. Downloading base model (29.5 GB)...")
    real_checkpoint_path = hf_hub_download(
        repo_id="Lightricks/LTX-2.3-fp8",
        filename="ltx-2.3-22b-distilled-fp8.safetensors",
        local_dir="models",
        local_dir_use_symlinks=False
    )
    print_disk("after_base_model_download")

    print("7. Rebuilding model builders for base LTX model...")
    ledger.checkpoint_path = real_checkpoint_path
    ledger.gemma_root_path = None  # Prevent searching for deleted Gemma files
    ledger.build_model_builders()

    print("8. Preloading base LTX models...")
    _transformer = ledger.transformer()
    _video_encoder = ledger.video_encoder()
    _video_decoder = ledger.video_decoder()
    _audio_decoder = ledger.audio_decoder()
    _vocoder = ledger.vocoder()
    print("Base LTX models loaded and cached.")

    print("9. Deleting base LTX model weights from disk to free storage...")
    if os.path.exists(real_checkpoint_path):
        os.remove(real_checkpoint_path)
    print_disk("final_cleanup")

# Bind lambda caches to ledger
ledger.transformer = lambda: _transformer
ledger.video_encoder = lambda: _video_encoder
ledger.video_decoder = lambda: _video_decoder
ledger.audio_decoder = lambda: _audio_decoder
ledger.vocoder = lambda: _vocoder
ledger.spatial_upsampler = lambda: _spatial_upsampler
ledger.text_encoder = lambda: _text_encoder
ledger.gemma_embeddings_processor = lambda: _embeddings_processor
print("All models preloaded and mapped successfully!")


def log_memory(tag: str):
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3
        peak = torch.cuda.max_memory_allocated() / 1024**3
        free, total = torch.cuda.mem_get_info()
        print(f"[VRAM {tag}] allocated={allocated:.2f}GB peak={peak:.2f}GB free={free / 1024**3:.2f}GB total={total / 1024**3:.2f}GB")


def detect_aspect_ratio(image) -> str:
    if image is None:
        return "16:9"
    if hasattr(image, "size"):
        w, h = image.size
    elif hasattr(image, "shape"):
        h, w = image.shape[:2]
    else:
        return "16:9"
    ratio = w / h
    candidates = {"16:9": 16 / 9, "9:16": 9 / 16, "1:1": 1.0}
    return min(candidates, key=lambda k: abs(ratio - candidates[k]))


def on_image_upload(image, high_res):
    aspect = detect_aspect_ratio(image)
    tier = "high" if high_res else "low"
    w, h = RESOLUTIONS[tier][aspect]
    return gr.update(value=w), gr.update(value=h)


def on_highres_toggle(image, high_res):
    aspect = detect_aspect_ratio(image)
    tier = "high" if high_res else "low"
    w, h = RESOLUTIONS[tier][aspect]
    return gr.update(value=w), gr.update(value=h)


# VIDEO TO VIDEO INFERENCE
@spaces.GPU(duration=120)
@torch.inference_mode()
def generate_video_to_video(
    input_video: str,
    prompt: str,
    strength: float = 0.6,
    duration: float = 3.0,
    audio_mode: str = "Keep original audio",
    enhance_prompt: bool = False,
    seed: int = 42,
    randomize_seed: bool = True,
    height: int = 512,
    width: int = 768,
    progress=gr.Progress(track_tqdm=True),
):
    try:
        if input_video is None:
            raise ValueError("An input video must be uploaded for Video-to-Video generation.")

        torch.cuda.reset_peak_memory_stats()
        log_memory("V2V start")

        current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
        generator = torch.Generator(device=pipeline.device).manual_seed(current_seed)
        noiser = GaussianNoiser(generator=generator)
        stepper = EulerDiffusionStep()
        dtype = pipeline.dtype

        # Detect original metadata
        try:
            fps, orig_frames, w, h = get_videostream_metadata(input_video)
            print(f"Loaded original video: {orig_frames} frames, {fps} fps, size={w}x{h}")
        except Exception as e:
            print(f"Could not load stream metadata: {e}. Defaulting to 24 FPS.")
            fps = DEFAULT_FRAME_RATE

        frame_rate = float(fps) if fps > 0 else DEFAULT_FRAME_RATE
        num_frames = int(duration * frame_rate) + 1
        num_frames = ((num_frames - 1 + 7) // 8) * 8 + 1

        print(f"Processing V2V: {height}x{width}, target={num_frames} frames ({duration}s), seed={current_seed}")

        # Load video frames for Stage 1 (half resolution)
        video_pixel_stage_1 = load_video_conditioning(
            video_path=input_video,
            height=int(height // 2),
            width=int(width // 2),
            frame_cap=num_frames,
            dtype=dtype,
            device=pipeline.device
        )
        # Handle cases where the video has fewer frames than expected
        F_actual = video_pixel_stage_1.shape[2]
        if F_actual < num_frames:
            num_frames = ((F_actual - 1) // 8) * 8 + 1
            if num_frames < 9:
                num_frames = 9
            video_pixel_stage_1 = video_pixel_stage_1[:, :, :num_frames]
            print(f"Capping frame count to actual video frames: {num_frames}")

        # Load video frames for Stage 2 (full resolution)
        video_pixel_stage_2 = load_video_conditioning(
            video_path=input_video,
            height=int(height),
            width=int(width),
            frame_cap=num_frames,
            dtype=dtype,
            device=pipeline.device
        )
        video_pixel_stage_2 = video_pixel_stage_2[:, :, :num_frames]

        # Encode prompts
        (ctx_p,) = encode_prompts(
            [prompt],
            pipeline.model_ledger,
            enhance_first_prompt=enhance_prompt,
            enhance_prompt_image=None,
        )
        video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding

        # Stage 1: Initial low resolution video denoising
        video_encoder = pipeline.model_ledger.video_encoder()
        transformer = pipeline.model_ledger.transformer()

        # Map strength to starting step in the 8-step distilled schedule
        num_steps = max(1, int(strength * 8))
        start_idx = 8 - num_steps
        stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES[start_idx:]).to(pipeline.device)
        print(f"V2V Stage 1 schedule: {len(stage_1_sigmas)-1} steps, starting at sigma={stage_1_sigmas[0]:.4f}")

        # Encode downscaled video to latents
        stage_1_initial_video_latent = video_encoder(video_pixel_stage_1)

        def denoising_loop(
            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
        ) -> tuple[LatentState, LatentState]:
            return euler_denoising_loop(
                sigmas=sigmas,
                video_state=video_state,
                audio_state=audio_state,
                stepper=stepper,
                denoise_fn=simple_denoising_func(
                    video_context=video_context,
                    audio_context=audio_context,
                    transformer=transformer,
                ),
            )

        stage_1_output_shape = VideoPixelShape(
            batch=1,
            frames=num_frames,
            width=width // 2,
            height=height // 2,
            fps=frame_rate,
        )

        video_state, audio_state = denoise_audio_video(
            output_shape=stage_1_output_shape,
            conditionings=[],
            noiser=noiser,
            sigmas=stage_1_sigmas,
            stepper=stepper,
            denoising_loop_fn=denoising_loop,
            components=pipeline.pipeline_components,
            dtype=dtype,
            device=pipeline.device,
            noise_scale=stage_1_sigmas[0],
            initial_video_latent=stage_1_initial_video_latent,
            initial_audio_latent=None,
        )

        # Stage 2: Upsample and refine
        upscaled_video_latent = upsample_video(
            latent=video_state.latent[:1],
            video_encoder=video_encoder,
            upsampler=pipeline.model_ledger.spatial_upsampler()
        )

        torch.cuda.synchronize()
        cleanup_memory()

        stage_2_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(pipeline.device)
        stage_2_output_shape = VideoPixelShape(
            batch=1,
            frames=num_frames,
            width=width,
            height=height,
            fps=frame_rate
        )

        video_state, audio_state = denoise_audio_video(
            output_shape=stage_2_output_shape,
            conditionings=[],
            noiser=noiser,
            sigmas=stage_2_sigmas,
            stepper=stepper,
            denoising_loop_fn=denoising_loop,
            components=pipeline.pipeline_components,
            dtype=dtype,
            device=pipeline.device,
            noise_scale=stage_2_sigmas[0],
            initial_video_latent=upscaled_video_latent,
            initial_audio_latent=audio_state.latent,
        )

        torch.cuda.synchronize()
        cleanup_memory()

        # VAE decoding
        decoded_video = vae_decode_video(
            video_state.latent,
            pipeline.model_ledger.video_decoder(),
            TilingConfig.default(),
            generator
        )

        # Handle audio mode
        output_audio = None
        if audio_mode == "Keep original audio":
            try:
                original_audio = decode_audio_from_file(
                    path=input_video,
                    device=pipeline.device,
                    start_time=0.0,
                    max_duration=duration,
                )
                output_audio = original_audio
                print("Original audio successfully extracted.")
            except Exception as e:
                print(f"Failed to extract original audio: {e}. Outputting silent or generated audio.")

        if output_audio is None and audio_mode != "No audio":
            decoded_audio = vae_decode_audio(
                audio_state.latent,
                pipeline.model_ledger.audio_decoder(),
                pipeline.model_ledger.vocoder()
            )
            output_audio = decoded_audio
            print("Generated synchronized audio.")

        # Encode and save output video file
        tiling_config = TilingConfig.default()
        video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
        output_path = tempfile.mktemp(suffix=".mp4")

        encode_video(
            video=decoded_video,
            fps=frame_rate,
            audio=output_audio,
            output_path=output_path,
            video_chunks_number=video_chunks_number,
        )

        log_memory("V2V finished")
        return str(output_path), current_seed

    except Exception as e:
        import traceback
        log_memory("V2V error")
        print(f"Error in V2V: {str(e)}\n{traceback.format_exc()}")
        return None, current_seed


# STANDARD GENERATION INFERENCE (Tab 2)
@spaces.GPU(duration=75)
@torch.inference_mode()
def generate_video(
    input_image,
    prompt: str,
    duration: float,
    enhance_prompt: bool = False,
    seed: int = 42,
    randomize_seed: bool = True,
    height: int = 1024,
    width: int = 1536,
    progress=gr.Progress(track_tqdm=True),
):
    try:
        torch.cuda.reset_peak_memory_stats()
        log_memory("T2V start")

        current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
        frame_rate = DEFAULT_FRAME_RATE
        num_frames = int(duration * frame_rate) + 1
        num_frames = ((num_frames - 1 + 7) // 8) * 8 + 1

        print(f"Generating Video: {height}x{width}, {num_frames} frames ({duration}s), seed={current_seed}")

        images = []
        if input_image is not None:
            output_dir = Path("outputs")
            output_dir.mkdir(exist_ok=True)
            temp_image_path = output_dir / f"temp_input_{current_seed}.jpg"
            if hasattr(input_image, "save"):
                input_image.save(temp_image_path)
            else:
                temp_image_path = Path(input_image)
            images = [ImageConditioningInput(path=str(temp_image_path), frame_idx=0, strength=1.0)]

        tiling_config = TilingConfig.default()
        video_chunks_number = get_video_chunks_number(num_frames, tiling_config)

        video, audio = pipeline(
            prompt=prompt,
            seed=current_seed,
            height=int(height),
            width=int(width),
            num_frames=num_frames,
            frame_rate=frame_rate,
            images=images,
            tiling_config=tiling_config,
            enhance_prompt=enhance_prompt,
        )

        output_path = tempfile.mktemp(suffix=".mp4")
        encode_video(
            video=video,
            fps=frame_rate,
            audio=audio,
            output_path=output_path,
            video_chunks_number=video_chunks_number,
        )

        log_memory("T2V finished")
        return str(output_path), current_seed

    except Exception as e:
        import traceback
        log_memory("T2V error")
        print(f"Error in T2V: {str(e)}\n{traceback.format_exc()}")
        return None, current_seed


# GRADIO UI SETUP
with gr.Blocks(title="LTX V2V") as demo:
    gr.Markdown("# LTX V2V: Distilled 22B Video-to-Video & Generation")
    gr.Markdown(
        "Highly efficient video translation (stylization, restyling, editing) and text/image-to-video generation using LTX-2.3. "
        "[[model]](https://huggingface.co/Lightricks/LTX-2.3) "
        "[[code]](https://github.com/Lightricks/LTX-2)"
    )

    with gr.Tabs():
        # TAB 1: Video to Video
        with gr.TabItem("Video-to-Video (V2V)"):
            with gr.Row():
                with gr.Column():
                    v2v_input_video = gr.Video(label="Input Video", sources=["upload"])
                    v2v_prompt = gr.Textbox(
                        label="Prompt",
                        info="Describe the style, aesthetic, actions or changes to apply (e.g. 'Turn the person into a robot', 'Anime style')",
                        value="A cinematic cartoon rendering of the motion, vibrant styling, detailed painting look",
                        lines=3
                    )
                    v2v_strength = gr.Slider(
                        label="Denoising Strength (0.0 = original, 1.0 = completely new)",
                        minimum=0.1,
                        maximum=1.0,
                        value=0.6,
                        step=0.05
                    )
                    with gr.Row():
                        v2v_duration = gr.Slider(label="Duration (seconds)", minimum=1.0, maximum=10.0, value=3.0, step=0.1)
                        v2v_audio_mode = gr.Dropdown(
                            label="Audio Mode",
                            choices=["Keep original audio", "Generate new audio", "No audio"],
                            value="Keep original audio"
                        )
                    
                    v2v_generate_btn = gr.Button("Transform Video", variant="primary", size="lg")

                    with gr.Accordion("Advanced Settings", open=False):
                        v2v_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=42, step=1)
                        v2v_randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                        with gr.Row():
                            v2v_width = gr.Dropdown(label="Width", choices=[512, 768, 1024, 1536], value=768)
                            v2v_height = gr.Dropdown(label="Height", choices=[512, 768, 1024, 1536], value=512)

                with gr.Column():
                    v2v_output_video = gr.Video(label="Transformed Video", autoplay=True)

            v2v_generate_btn.click(
                fn=generate_video_to_video,
                inputs=[
                    v2v_input_video,
                    v2v_prompt,
                    v2v_strength,
                    v2v_duration,
                    v2v_audio_mode,
                    gr.Checkbox(visible=False, value=False),  # enhance_prompt hidden or set False
                    v2v_seed,
                    v2v_randomize_seed,
                    v2v_height,
                    v2v_width
                ],
                outputs=[v2v_output_video, v2v_seed]
            )

        # TAB 2: Text/Image to Video
        with gr.TabItem("Text/Image-to-Video"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(label="Input Image (Optional)", type="pil")
                    t2v_prompt = gr.Textbox(
                        label="Prompt",
                        info="for best results - make it as elaborate as possible",
                        value="Make this image come alive with cinematic motion, smooth animation",
                        lines=3,
                    )
                    with gr.Row():
                        t2v_duration = gr.Slider(label="Duration (seconds)", minimum=1.0, maximum=10.0, value=3.0, step=0.1)
                        with gr.Column():
                            t2v_enhance_prompt = gr.Checkbox(label="Enhance Prompt", value=False)
                            high_res = gr.Checkbox(label="High Resolution", value=True)

                    t2v_generate_btn = gr.Button("Generate Video", variant="primary", size="lg")

                    with gr.Accordion("Advanced Settings", open=False):
                        t2v_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=10, step=1)
                        t2v_randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                        with gr.Row():
                            t2v_width = gr.Number(label="Width", value=1536, precision=0)
                            t2v_height = gr.Number(label="Height", value=1024, precision=0)

                with gr.Column():
                    t2v_output_video = gr.Video(label="Generated Video", autoplay=True)

            # Auto-detect resolution from image
            input_image.change(
                fn=on_image_upload,
                inputs=[input_image, high_res],
                outputs=[t2v_width, t2v_height],
            )
            high_res.change(
                fn=on_highres_toggle,
                inputs=[input_image, high_res],
                outputs=[t2v_width, t2v_height],
            )

            t2v_generate_btn.click(
                fn=generate_video,
                inputs=[
                    input_image, t2v_prompt, t2v_duration, t2v_enhance_prompt,
                    t2v_seed, t2v_randomize_seed, t2v_height, t2v_width,
                ],
                outputs=[t2v_output_video, t2v_seed],
            )

css = """
.fillable{max-width: 1200px !important}
.progress-text {color: white}
"""

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=css)