import os

# Expandable segments help with transient allocation spikes from video processing.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
# qwen_omni_utils -> librosa -> numba eagerly inits CUDA on import; disable that.
os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")

import spaces  # MUST come before any torch / CUDA-touching import
import torch
import gradio as gr
from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
from qwen_omni_utils import process_mm_info

MODEL_ID = "yaolily/TimeChat-Captioner-GRPO-7B"
MAX_PIXELS = 297920
VIDEO_MAX_PIXELS = 297920

DEFAULT_PROMPT = (
    "Thoroughly describe everything in the video, capturing every detail. "
    "Include as much information from the audio as possible, and ensure that "
    "the descriptions of both audio and video are well-coordinated."
)

print(f"Loading model from {MODEL_ID}...")
processor = Qwen2_5OmniProcessor.from_pretrained(MODEL_ID)
model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
model.disable_talker()
print("Model loaded successfully.")


@spaces.GPU(duration=180)
def caption_video(video_path: str, prompt: str, max_frames: int, fps: float) -> str:
    """Generate a detailed, time-aware audio-visual caption for a multi-scene video.

    Args:
        video_path: Path to the input video file (recommended ~60 s clips).
        prompt: Instruction prompt for the captioner.
        max_frames: Maximum number of frames to sample from the video.
        fps: Frames-per-second sampling rate.
    """
    if video_path is None:
        return "Please upload a video first."

    prompt_text = prompt.strip() if prompt and prompt.strip() else DEFAULT_PROMPT

    conversation = [
        {
            "role": "user",
            "content": [
                {
                    "type": "video",
                    "video": video_path,
                    "max_pixels": MAX_PIXELS,
                    "max_frames": int(max_frames),
                    "fps": float(fps),
                    "video_max_pixels": VIDEO_MAX_PIXELS,
                },
                {
                    "type": "text",
                    "text": prompt_text,
                },
            ],
        },
    ]

    text = processor.apply_chat_template(
        conversation, add_generation_prompt=True, tokenize=False
    )
    audios, images, videos = process_mm_info(
        conversation, use_audio_in_video=True
    )
    inputs = processor(
        text=text,
        audio=audios,
        images=images,
        videos=videos,
        return_tensors="pt",
        padding=True,
        use_audio_in_video=True,
    )
    inputs = inputs.to(model.device).to(model.dtype)

    with torch.inference_mode():
        text_ids = model.generate(
            **inputs,
            use_audio_in_video=True,
            generation_mode="text",
            thinker_max_new_tokens=8192,
            talker_max_new_tokens=8192,
            use_cache=True,
        )

    generated_ids = text_ids[0][inputs.input_ids[0].size(0):]
    response = processor.decode(generated_ids, skip_special_tokens=True)
    return response


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        "# TimeChat-Captioner\n"
        "Generate detailed, time-aware, and structurally coherent audio-visual captions "
        "for multi-scene videos. Upload a short video clip (~60 s recommended) and get a "
        "script-like description with timestamps.\n\n"
        "Model: [yaolily/TimeChat-Captioner-GRPO-7B](https://huggingface.co/yaolily/TimeChat-Captioner-GRPO-7B) | "
        "Paper: [arXiv:2602.08711](https://arxiv.org/abs/2602.08711)"
    )

    with gr.Row():
        video_input = gr.Video(label="Upload Video", sources=["upload"])
        caption_output = gr.Textbox(
            label="Generated Caption",
            lines=20,
            max_lines=50,
        )

    prompt_input = gr.Textbox(
        label="Prompt",
        value=DEFAULT_PROMPT,
        lines=3,
    )

    with gr.Accordion("Advanced Settings", open=False):
        max_frames_slider = gr.Slider(
            label="Max Frames",
            minimum=16,
            maximum=160,
            value=160,
            step=8,
            info="Maximum number of frames sampled from the video. Lower = faster.",
        )
        fps_slider = gr.Slider(
            label="Sampling FPS",
            minimum=0.5,
            maximum=4.0,
            value=2.0,
            step=0.5,
            info="Frames per second to sample. Lower = fewer frames, faster inference.",
        )

    run_btn = gr.Button("Generate Caption", variant="primary")

    run_btn.click(
        fn=caption_video,
        inputs=[video_input, prompt_input, max_frames_slider, fps_slider],
        outputs=caption_output,
        api_name="caption_video",
    )

    gr.Examples(
        examples=[
            ["example_video.mp4", DEFAULT_PROMPT, 160, 2.0],
        ],
        inputs=[video_input, prompt_input, max_frames_slider, fps_slider],
        outputs=caption_output,
        fn=caption_video,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)