import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / any CUDA-touching import
import torch
import gradio as gr
from pathlib import Path

# Monkey-patch safetensors to load on CPU, bypassing ZeroGPU's CUDA interception
import safetensors.torch as _st
_orig_safe_load_file = _st.load_file
def _patched_safe_load_file(filename, device="cpu", **kwargs):
    return _orig_safe_load_file(filename, device="cpu", **kwargs)
_st.load_file = _patched_safe_load_file
# Also patch in peft's save_and_load module so it uses our patched version
import peft.utils.save_and_load as _peft_sl
_peft_sl.safe_load_file = _patched_safe_load_file

from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
from peft import PeftModel
from qwen_vl_utils import process_vision_info

BASE_MODEL_ID = "Qwen/Qwen3-VL-8B-Instruct"
LORA_ADAPTER_ID = "RainChow/PinpointQA-Qwen3-VL-8B-Instruct-LoRA"

print("Loading base model on CPU...")
model = Qwen3VLForConditionalGeneration.from_pretrained(
    BASE_MODEL_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
)

print("Loading and merging LoRA adapter on CPU...")
model = PeftModel.from_pretrained(model, LORA_ADAPTER_ID)
model = model.merge_and_unload()
model.eval()

print("Moving merged model to CUDA...")
model = model.to("cuda")

print("Loading processor...")
processor = AutoProcessor.from_pretrained(BASE_MODEL_ID)

print("Model loaded successfully!")


@spaces.GPU(duration=120)
def run_inference(video, question, fps=1.0, max_new_tokens=512):
    """Answer a spatial understanding question about an indoor video.

    Upload a video (or select an example) and ask a question about where a
    small object is located. The model will describe the object's spatial
    position relative to nearby reference objects.

    Args:
        video: Uploaded video file (mp4, etc.).
        question: A spatial question, e.g. "Where is the scissors?"
        fps: Frames per second to sample from the video.
        max_new_tokens: Maximum number of tokens to generate.
    """
    if video is None:
        return "Please upload a video or select an example."

    if not question or not question.strip():
        return "Please enter a question about the spatial location of an object."

    video_path = video

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "video",
                    "video": f"file://{video_path}",
                    "max_pixels": 360 * 420,
                    "fps": fps,
                },
                {"type": "text", "text": question},
            ],
        }
    ]

    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    images, videos = process_vision_info(messages)

    inputs = processor(
        text=[text],
        images=images,
        videos=videos,
        return_tensors="pt",
        padding=True,
    )
    inputs = inputs.to(model.device)

    with torch.inference_mode():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            use_cache=True,
        )

    generated_ids_trimmed = [
        out_ids[len(in_ids):]
        for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )[0]

    return output_text


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

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # PinpointQA: Spatial Grounding in Indoor Videos

        Ask questions about the spatial location of small objects in indoor video scenes.
        The model identifies objects and describes their position relative to nearby references.

        Built on [Qwen3-VL-8B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct) fine-tuned with
        [PinpointQA LoRA](https://huggingface.co/RainChow/PinpointQA-Qwen3-VL-8B-Instruct-LoRA) from the
        [PinpointQA benchmark](https://huggingface.co/papers/2604.08991).
        """
    )

    with gr.Row(elem_id="col-container"):
        with gr.Column(scale=2):
            video_input = gr.Video(label="Indoor Video")
            question_input = gr.Textbox(
                label="Spatial Question",
                placeholder="e.g. Where is the scissors? Describe its location relative to nearby objects.",
                lines=2,
            )
            run_btn = gr.Button("Run Spatial Grounding", variant="primary")
        with gr.Column(scale=3):
            output_text = gr.Textbox(
                label="Model Response",
                lines=12,
                interactive=False,
            )

    with gr.Accordion("Advanced Settings", open=False):
        fps_slider = gr.Slider(
            minimum=0.5, maximum=4.0, value=1.0, step=0.5,
            label="Sampling FPS",
            info="Frames per second to sample from the video",
        )
        max_tokens = gr.Slider(
            minimum=128, maximum=1024, value=512, step=64,
            label="Max New Tokens",
            info="Maximum number of tokens the model can generate",
        )

    examples_dir = Path("examples")
    gr.Examples(
        examples=[
            [
                str(examples_dir / "sample1.mp4"),
                "Where is the scissors? Describe its location relative to nearby objects.",
            ],
            [
                str(examples_dir / "sample2.mp4"),
                "Where is the pen holder? Describe its location relative to nearby objects.",
            ],
            [
                str(examples_dir / "sample3.mp4"),
                "Where is the mouse? Describe its location relative to nearby objects.",
            ],
        ],
        inputs=[video_input, question_input],
        outputs=output_text,
        fn=run_inference,
        cache_examples=True,
        cache_mode="lazy",
    )

    run_btn.click(
        fn=run_inference,
        inputs=[video_input, question_input, fps_slider, max_tokens],
        outputs=output_text,
        api_name="predict",
    )

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