app.py

import gradio as gr
import spaces
import torch
from PIL import Image
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

MODEL_ID = "internlm/Spatial-SSRL-7B"
MAX_NEW_TOKENS = 2048

# Format prompt toggle content (auto-concatenated to user prompt when enabled)
FORMAT_PROMPT = (
    "You FIRST think about the reasoning process as an internal monologue and then provide the final answer. "
    "The reasoning process MUST BE enclosed within <think> </think> tags. "
    "The final answer MUST BE put in \\boxed{}."
)

# Example questions (base text only; format prompt appended automatically when toggle is on)
EXAMPLE_QUESTIONS = [
    "Consider the real-world 3D location of the objects. Which object is further away from the camera? A. boat B. fire hydrant.",
    "Consider the real-world 3D orientations of the objects. Are the kid and the teddy bear facing same or similar directions, or very different directions? A. very different directions B. same or similar directions.",
    "Consider the real-world 3D locations and orientations of the objects. If I stand at the recreational vehicle's position facing where it is facing, is the dog in front of me or behind me? A. behind B. in front of me."
]


def get_device() -> str:
    return "cuda" if torch.cuda.is_available() else "cpu"


def select_dtype(device: str):
    if device == "cuda":
        if torch.cuda.is_bf16_supported():
            return torch.bfloat16
        return torch.float16
    return torch.float32


def load_model():
    device = get_device()
    dtype = select_dtype(device)

    # Use device_map="auto" for proper GPU allocation with spaces.GPU decorator
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        MODEL_ID,
        dtype=dtype,
        device_map="auto",
        trust_remote_code=True,
    )

    processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
    return model, processor


MODEL, PROCESSOR = load_model()


@spaces.GPU
@torch.inference_mode()
def answer_question(image: Image.Image, question: str, use_format_prompt: bool):
    if image is None:
        return "Please upload an image.", 0

    if not question or question.strip() == "":
        return "Please enter a question.", 0

    try:
        # If the checkbox input wasn't provided (e.g., from examples with older gradio behavior), default to True
        if use_format_prompt is None:
            use_format_prompt = True
        # Concatenate the format prompt if the toggle is enabled
        if use_format_prompt:
            question = f"{question}{FORMAT_PROMPT}"
        # Validate image
        if not isinstance(image, Image.Image):
            return "Error: Invalid image format", 0

        # Check image size (warn if too large)
        max_size = 4096
        if image.width > max_size or image.height > max_size:
            # Resize if too large to prevent OOM
            image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)

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

        prompt_text = PROCESSOR.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

        inputs = PROCESSOR(
            text=[prompt_text],
            images=[image],
            return_tensors="pt",
        ).to(device)

        generated_ids = MODEL.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
        )

        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
        )
        answer = output_text[0].strip()

        input_ids = inputs.get("input_ids")
        input_length = input_ids.shape[-1] if input_ids is not None else 0
        total_length = generated_ids.shape[-1]
        num_generated_tokens = max(total_length - input_length, 0)

        return answer, int(num_generated_tokens)

    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        return "Error: Out of GPU memory. Please try with a smaller image.", 0
    except Exception as e:
        return f"Error generating answer: {str(e)}", 0


def load_example(example_idx):
    """Load example image and question based on index"""
    example_images = [
        "./examples/eg1.jpg",
        "./examples/eg2.jpg",
        "./examples/eg3.jpg"
    ]
    if 0 <= example_idx < len(EXAMPLE_QUESTIONS):
        return Image.open(example_images[example_idx]), EXAMPLE_QUESTIONS[example_idx]
    return None, ""


with gr.Blocks(title="Spatial-SSRL Spatial Reasoning") as demo:
    gr.Markdown("# 🌍 Spatial-SSRL: Spatial Reasoning with Vision-Language Models")
    gr.Markdown("### Understanding 3D Spatial Relationships from 2D Images")
    gr.Markdown("✨ Upload an image and ask questions about spatial relationships, locations, and orientations! ✨")

    gr.Markdown(
        """
📖 <a href="https://arxiv.org/abs/2510.27606">Paper</a> | 🏠 <a href="https://github.com/InternLM/Spatial-SSRL">Github</a> | 🤗 <a href="https://huggingface.co/internlm/Spatial-SSRL-7B">Spatial-SSRL-7B Model</a> | 🤗 <a href="https://huggingface.co/datasets/internlm/Spatial-SSRL-81k">Spatial-SSRL-81k | 📰 <a href="https://huggingface.co/papers/2510.27606">Daily Paper</a>
"""
    )

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Input Image")
            question_input = gr.Textbox(
                label="Question",
                placeholder="Ask a question about spatial relationships in the image...",
                lines=4
            )
            use_format_chk = gr.Checkbox(
                label="Apply format prompt (default on)",
                value=True,
                info="When enabled, the predefined format prompt is automatically concatenated to your question."
            )
            submit_button = gr.Button("Submit", variant="primary")

        with gr.Column():
            answer_output = gr.Textbox(label="Answer", lines=10)
            token_output = gr.Number(label="Generated Tokens", precision=0)

    submit_button.click(
        fn=answer_question,
        inputs=[image_input, question_input, use_format_chk],
        outputs=[answer_output, token_output],
        show_progress=True,
    )

    gr.Markdown("### 📸 Example Questions")
    gr.Markdown("Click on an example below to load it:")

    with gr.Row():
        example1_btn = gr.Button("Example 1: Boat vs Fire Hydrant")
        example2_btn = gr.Button("Example 2: Kid and Teddy Bear")
        example3_btn = gr.Button("Example 3: RV and Dog")

    example1_btn.click(
        fn=lambda: load_example(0),
        inputs=[],
        outputs=[image_input, question_input],
    )

    example2_btn.click(
        fn=lambda: load_example(1),
        inputs=[],
        outputs=[image_input, question_input],
    )

    example3_btn.click(
        fn=lambda: load_example(2),
        inputs=[],
        outputs=[image_input, question_input],
    )

    gr.Examples(
        examples=[
            ["./examples/eg1.jpg", EXAMPLE_QUESTIONS[0], True],
            ["./examples/eg2.jpg", EXAMPLE_QUESTIONS[1], True],
            ["./examples/eg3.jpg", EXAMPLE_QUESTIONS[2], True],
        ],
        inputs=[image_input, question_input, use_format_chk],
        outputs=[answer_output, token_output],
        fn=answer_question,
        cache_examples=True,
        label="Complete Examples"
    )

    gr.Markdown("### About")
    gr.Markdown(
        """
This demo showcases spatial reasoning capabilities of vision-language models. The model can:
- Understand 3D spatial relationships from 2D images
- Reason about object locations (near/far, front/behind)
- Analyze object orientations and facing directions
- Provide step-by-step reasoning before answering

"""
    )

    gr.Markdown("### Citation")
    gr.Markdown("If you find this project useful, please kindly cite:")

    citation_text = """@article{liu2025spatialssrl,
  title={{Spatial-SSRL}: Enhancing Spatial Understanding via Self-Supervised Reinforcement Learning}, 
  author={Liu, Yuhong and Zhang, Beichen and Zang, Yuhang and Cao, Yuhang and Xing, Long and Dong, Xiaoyi and Duan, Haodong and Lin, Dahua and Wang, Jiaqi},
  journal={arXiv preprint arXiv:2510.27606},
  year={2025}
}
}"""

    gr.Code(value=citation_text, language="markdown", label="BibTeX Citation")


demo.launch()