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 PIL import Image
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration

MODEL_ID = "Pokerme/view2space_4b"

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


def _load_images(file_objs):
    """Load uploaded file paths into PIL Images."""
    if not file_objs:
        return []
    images = []
    for f in file_objs:
        if isinstance(f, str):
            path = f
        else:
            path = f.name if hasattr(f, "name") else str(f)
        images.append(Image.open(path).convert("RGB"))
    return images


@spaces.GPU(duration=90)
def answer(
    images: list | None,
    prompt: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
) -> str:
    """Answer a visual-reasoning question about one or more images using VIEW2SPACE.

    Args:
        images: one or more input image files (multi-view observations supported).
        prompt: the question or instruction about the image(s).
        max_new_tokens: maximum number of tokens to generate.
        temperature: sampling temperature (0 = greedy, higher = more creative).
        top_p: nucleus-sampling probability mass.
    """
    pil_images = _load_images(images)
    if not pil_images:
        return "Please provide at least one image."

    content = []
    for _ in pil_images:
        content.append({"type": "image"})
    content.append({"type": "text", "text": prompt})

    messages = [{"role": "user", "content": content}]
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = processor(
        text=text,
        images=pil_images,
        return_tensors="pt",
    ).to("cuda")

    gen_kwargs = dict(
        **inputs,
        max_new_tokens=int(max_new_tokens),
    )
    if float(temperature) > 0:
        gen_kwargs["do_sample"] = True
        gen_kwargs["temperature"] = float(temperature)
        gen_kwargs["top_p"] = float(top_p)

    with torch.inference_mode():
        gen_ids = model.generate(**gen_kwargs)

    trimmed = gen_ids[0][inputs["input_ids"].shape[-1]:]
    result = processor.decode(trimmed, skip_special_tokens=True)
    return result


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

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown(
        """
        # VIEW2SPACE-4B: Multi-View Visual Reasoning
        A 4.4B-parameter vision-language model based on Qwen3-VL that reasons about
        scenes from sparse multi-view observations. Upload one or more images and ask
        a question — the model integrates information across views to answer.

        [Model card](https://huggingface.co/Pokerme/view2space_4b) |
        [Paper](http://arxiv.org/abs/2603.16506) |
        [Code](https://github.com/pokerme7777/VIEW2SPACE)
        """
    )

    with gr.Column(elem_id="col-container"):
        images_in = gr.File(
            label="Input image(s)",
            file_count="multiple",
            file_types=["image"],
            type="filepath",
        )
        prompt = gr.Textbox(
            label="Question / instruction",
            placeholder="e.g. What animal is on the candy?",
            lines=2,
        )
        run = gr.Button("Run", variant="primary")
        output = gr.Textbox(label="Model response", lines=6, interactive=False)

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

    gr.Examples(
        examples=[
            [[os.path.join("examples", "candy.jpg")], "What animal is on the candy?"],
            [[os.path.join("examples", "dog.jpg")], "Describe this image in one sentence."],
            [[os.path.join("examples", "rubber_ducks.jpg")], "How many rubber ducks are in the image?"],
            [[os.path.join("examples", "cake.jpg")], "What is the main subject of this image?"],
        ],
        inputs=[images_in, prompt],
        outputs=output,
        fn=answer,
        cache_examples=True,
        cache_mode="lazy",
    )

    run.click(
        fn=answer,
        inputs=[images_in, prompt, max_new_tokens, temperature, top_p],
        outputs=output,
        api_name="answer",
    )

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