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

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

# SpatialCLI-8B is a fine-tune of Qwen3-VL-8B-Instruct. The published weights
# live at ZYT-MFM/SpatialCLI-8B; we try that first and fall back to the base
# model so the demo still runs if the fine-tuned checkpoint is unavailable.
SPATIALCLI_ID = "ZYT-MFM/SpatialCLI-8B"
BASE_ID = "Qwen/Qwen3-VL-8B-Instruct"

MODEL_ID = SPATIALCLI_ID
try:
    from huggingface_hub import list_repo_files
    sibs = list_repo_files(SPATIALCLI_ID, repo_type="model")
    has_weights = any(s.endswith(".safetensors") or s.endswith(".bin") for s in sibs)
    if not has_weights:
        print(f"[load] {SPATIALCLI_ID} has no weights yet; falling back to {BASE_ID}")
        MODEL_ID = BASE_ID
except Exception as e:
    print(f"[load] probe failed ({e!r}); falling back to {BASE_ID}")
    MODEL_ID = BASE_ID

print(f"[load] loading model + processor from {MODEL_ID}")
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = Qwen3VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval()
print("[load] model ready")


@spaces.GPU(duration=60)
def answer_image(
    image,
    question,
    max_new_tokens=1024,
    enable_thinking=False,
    temperature=0.7,
    top_p=0.8,
    top_k=20,
):
    """Answer a spatial-reasoning question about an image.

    Loads SpatialCLI-8B (a fine-tune of Qwen3-VL-8B-Instruct trained to
    internalize specialist spatial-tool capabilities for localization,
    segmentation, depth, and pose reasoning) and runs direct, tool-free
    inference: image + text question -> text answer.

    Args:
        image: the input image.
        question: the spatial-reasoning question to ask.
        max_new_tokens: maximum number of new tokens to generate.
        enable_thinking: enable the model's thinking/reasoning trace.
        temperature: sampling temperature.
        top_p: nucleus sampling probability.
        top_k: top-k sampling.
    """
    if image is None:
        return "Please provide an image."
    if not (question or "").strip():
        return "Please provide a question."

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": question},
            ],
        }
    ]
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to("cuda")

    gen_kwargs = dict(
        max_new_tokens=int(max_new_tokens),
        do_sample=True,
        temperature=float(temperature),
        top_p=float(top_p),
        top_k=int(top_k),
        repetition_penalty=1.0,
    )
    if enable_thinking:
        gen_kwargs["chat_template_kwargs"] = {"enable_thinking": True}

    with torch.inference_mode():
        out_ids = model.generate(**inputs, **gen_kwargs)
    in_ids = inputs["input_ids"]
    trimmed = [out[len(in_):] for in_, out in zip(in_ids, out_ids)]
    text = processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]
    return text.strip()


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

with gr.Blocks() as demo:
    gr.Markdown(
        "# SpatialCLI-8B — Spatial Reasoning VLM\n"
        "Ask spatial-reasoning questions (localization, depth, pose, "
        "spatial relationships) about an image. SpatialCLI-8B is trained "
        "to internalize specialist spatial-tool capabilities, performing "
        "tool-free spatial reasoning at inference time."
    )

    with gr.Row(elem_id="col-container"):
        with gr.Column(scale=1):
            image_in = gr.Image(type="filepath", label="Image", height=320)
            question_in = gr.Textbox(
                label="Question",
                placeholder="e.g. Which object is closer to the camera, the chair or the table?",
                lines=3,
            )
            run_btn = gr.Button("Run", variant="primary")
            with gr.Accordion("Advanced settings", open=False):
                max_tokens = gr.Slider(64, 8192, value=1024, step=64,
                                       label="Max new tokens")
                thinking = gr.Checkbox(value=False,
                                       label="Enable thinking trace")
                temp = gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="Temperature")
                top_p_s = gr.Slider(0.0, 1.0, value=0.8, step=0.05, label="Top-p")
                top_k_s = gr.Slider(1, 100, value=20, step=1, label="Top-k")
        with gr.Column(scale=1):
            out = gr.Textbox(label="Answer", lines=12)

    run_btn.click(
        answer_image,
        inputs=[image_in, question_in, max_tokens, thinking, temp, top_p_s, top_k_s],
        outputs=out,
        api_name="answer",
    )

    gr.Examples(
        examples=[
            ["cafe_interior.jpg",
             "Describe the spatial layout of this cafe. Which tables are closest to the camera, and which are furthest? How are the chairs arranged relative to the tables?"],
            ["city_skyline_night.jpg",
             "Estimate the relative depths of the buildings in this skyline. Which buildings appear closest to the camera, and which are furthest away?"],
            ["autumn_forest_path.jpg",
             "Describe the spatial structure of this path. Does it recede into the distance? Estimate which trees are nearest vs. furthest from the viewer."],
        ],
        inputs=[image_in, question_in],
        outputs=out,
        fn=answer_image,
        cache_examples=True,
        cache_mode="lazy",
    )

    gr.Markdown(
        "---\n"
        "**Model:** [ZYT-MFM/SpatialCLI-8B](https://huggingface.co/ZYT-MFM/SpatialCLI-8B) "
        "(fine-tune of [Qwen3-VL-8B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct)). "
        "**Paper:** [SpatialCLI: Learning to Reason With Spatial Tools, Then Without Them](https://huggingface.co/papers/2607.27703). "
        "**Code:** [IANNXANG/SpatialCLI](https://github.com/IANNXANG/SpatialCLI)."
    )

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