import spaces  # MUST come before any torch/CUDA-touching import
import torch
import gradio as gr
from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info

MODEL_ID = "AvaXiao/ReToken-Qwen3VL-8B"
PROCESSOR_ID = "Qwen/Qwen3-VL-8B-Instruct"

model = Qwen3VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
)
model.to("cuda")
model.eval()

processor = AutoProcessor.from_pretrained(PROCESSOR_ID, use_fast=False)


@spaces.GPU(duration=60)
def answer_question(image, question, max_new_tokens=512, temperature=0.1, top_p=0.9,
                    progress=gr.Progress(track_tqdm=True)):
    """Answer a visual question about an image using ReToken-Qwen3VL-8B.

    Args:
        image: Input image to ask a question about.
        question: The question to ask about the image.
        max_new_tokens: Maximum number of new tokens to generate.
        temperature: Sampling temperature (lower = more deterministic).
        top_p: Nucleus sampling probability.
    """
    if image is None:
        return "Please upload an image first."
    if not question.strip():
        return "Please enter a question."

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": question},
            ],
        }
    ]

    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to("cuda")

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=temperature > 0,
            temperature=temperature if temperature > 0 else 1.0,
            top_p=top_p,
        )

    # Decode only the new tokens
    generated_ids = output_ids[:, inputs["input_ids"].shape[1]:]
    response = processor.batch_decode(
        generated_ids,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=True,
    )[0]

    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(
        """
        # ReToken: Visual Retrieval with Qwen3-VL-8B

        This demo showcases **ReToken-Qwen3VL-8B**, a vision-language model augmented with a learned retrieval token for improved visual retrieval. Ask questions about images — the model leverages a fine-tuned Qwen3-VL-8B backbone trained with a retrieval token that improves long-context visual understanding.

        [Paper](https://huggingface.co/papers/2607.28627) · [GitHub](https://github.com/avaxiao/ReToken) · [Model](https://huggingface.co/AvaXiao/ReToken-Qwen3VL-8B)
        """
    )

    with gr.Row(elem_id="col-container"):
        with gr.Column(scale=1):
            image_input = gr.Image(type="filepath", label="Input Image")
            question_input = gr.Textbox(
                label="Question",
                placeholder="Ask a question about the image…",
                lines=2,
            )
            run_btn = gr.Button("Answer", variant="primary")

            with gr.Accordion("Advanced settings", open=False):
                max_new_tokens = gr.Slider(
                    minimum=16, maximum=1024, value=512, step=16,
                    label="Max new tokens",
                )
                temperature = gr.Slider(
                    minimum=0.0, maximum=2.0, value=0.1, step=0.1,
                    label="Temperature (0 = greedy)",
                )
                top_p = gr.Slider(
                    minimum=0.1, maximum=1.0, value=0.9, step=0.05,
                    label="Top-p",
                )

        with gr.Column(scale=1):
            output_text = gr.Textbox(
                label="Answer",
                lines=12,
                interactive=False,
            )

    run_btn.click(
        fn=answer_question,
        inputs=[image_input, question_input, max_new_tokens, temperature, top_p],
        outputs=output_text,
        api_name="answer",
    )

    question_input.submit(
        fn=answer_question,
        inputs=[image_input, question_input, max_new_tokens, temperature, top_p],
        outputs=output_text,
        api_name="answer_submit",
    )

    gr.Examples(
        examples=[
            ["example1.jpg", "What breed is this dog and what is it doing?"],
            ["example2.jpg", "Describe this bird's colors and habitat."],
            ["example3.jpg", "What is this cat looking at?"],
        ],
        inputs=[image_input, question_input],
        outputs=output_text,
        fn=answer_question,
        cache_examples=True,
        cache_mode="lazy",
    )

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