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 torchvision.transforms as T from PIL import Image from torchvision.transforms.functional import InterpolationMode from transformers import AutoTokenizer # Register the custom SpatioLM model classes with AutoModel import spatiolm.models # noqa: F401 – registers InternVL3RChatModel, DA3Model, etc. from spatiolm.models import InternVL3RChatModel MODEL_ID = "xiaomi-research/SpatioLM-Understanding-InternVL3.5" # --------------------------------------------------------------------------- # Image preprocessing – replicated from lmms_eval.models.simple.internvl2 # (we avoid pulling in the full lmms-eval dependency for a single function) # --------------------------------------------------------------------------- IMAGENET_MEAN = (0.485, 0.456, 0.406) IMAGENET_STD = (0.229, 0.224, 0.225) INPUT_SIZE = 448 def _build_transform(input_size=INPUT_SIZE): transform = T.Compose([ T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img), T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC), T.ToTensor(), T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ]) return transform def _find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size): best_ratio_diff = float("inf") best_ratio = (1, 1) area = width * height for ratio in target_ratios: target_aspect_ratio = ratio[0] / ratio[1] ratio_diff = abs(aspect_ratio - target_aspect_ratio) if ratio_diff < best_ratio_diff: best_ratio_diff = ratio_diff best_ratio = ratio elif ratio_diff == best_ratio_diff: if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]: best_ratio = ratio return best_ratio def _dynamic_preprocess(image, min_num=1, max_num=6, image_size=INPUT_SIZE, use_thumbnail=False): orig_width, orig_height = image.size aspect_ratio = orig_width / orig_height target_ratios = set( (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if i * j <= max_num and i * j >= min_num ) target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1]) target_aspect_ratio = _find_closest_aspect_ratio( aspect_ratio, target_ratios, orig_width, orig_height, image_size ) target_width = image_size * target_aspect_ratio[0] target_height = image_size * target_aspect_ratio[1] blocks = target_aspect_ratio[0] * target_aspect_ratio[1] resized_img = image.resize((target_width, target_height)) processed_images = [] for i in range(blocks): box = ( (i % (target_width // image_size)) * image_size, (i // (target_width // image_size)) * image_size, ((i % (target_width // image_size)) + 1) * image_size, ((i // (target_width // image_size)) + 1) * image_size, ) split_img = resized_img.crop(box) processed_images.append(split_img) assert len(processed_images) == blocks if use_thumbnail and len(processed_images) != 1: thumbnail_img = image.resize((image_size, image_size)) processed_images.append(thumbnail_img) return processed_images def load_image(image, input_size=INPUT_SIZE, max_num=6): """Load and preprocess an image into a tensor of tiles. Args: image: a PIL.Image or a path to an image file. input_size: pixel size of each tile. max_num: maximum number of tiles. Returns: torch.Tensor of shape (num_tiles, 3, input_size, input_size). """ if isinstance(image, str): image = Image.open(image).convert("RGB") else: image = image.convert("RGB") transform = _build_transform(input_size=input_size) images = _dynamic_preprocess( image, image_size=input_size, use_thumbnail=True, max_num=max_num ) pixel_values = [transform(img) for img in images] pixel_values = torch.stack(pixel_values) return pixel_values # --------------------------------------------------------------------------- # Model loading (module scope, eager .to("cuda") per ZeroGPU rules) # --------------------------------------------------------------------------- tokenizer = AutoTokenizer.from_pretrained( MODEL_ID, trust_remote_code=True, use_fast=False, ) model = InternVL3RChatModel.from_pretrained( MODEL_ID, dtype=torch.bfloat16, low_cpu_mem_usage=True, ).eval().to("cuda") # --------------------------------------------------------------------------- # Inference # --------------------------------------------------------------------------- @spaces.GPU(duration=30) def answer( image: Image.Image, question: str, max_new_tokens: int = 512, temperature: float = 0.0, top_p: float = 1.0, ) -> str: """Answer a spatial-reasoning question about an image. Args: image: the input image to analyse. question: a natural-language question about the scene (e.g. "Which object is closer to the camera?"). max_new_tokens: maximum number of tokens to generate. temperature: sampling temperature; 0 for greedy. top_p: nucleus-sampling probability mass. Returns: The model's textual answer. """ if image is None: return "Please upload an image." if not question.strip(): return "Please enter a question." pixel_values = load_image(image).to(dtype=torch.bfloat16, device="cuda") num_patches_list = [pixel_values.shape[0]] gen_config = { "max_new_tokens": int(max_new_tokens), "do_sample": temperature > 0, "temperature": float(temperature) if temperature > 0 else 1.0, "top_p": float(top_p), } response = model.chat( tokenizer, pixel_values, question, gen_config, num_patches_list=num_patches_list, ) return response # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- import gradio as gr CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ EXAMPLE_IMAGES = [ "cafe_interior.jpg", "living_room_blue_couch.jpg", "library_interior.jpg", "city_skyline_toronto.jpg", "spiral_staircase.jpg", ] EXAMPLES = [ ["cafe_interior.jpg", "Which table is closer to the camera, the one on the left or the one on the right?"], ["living_room_blue_couch.jpg", "Describe the spatial arrangement of the furniture in this room."], ["library_interior.jpg", "How many shelves can you see, and which one is closest to the viewer?"], ["spiral_staircase.jpg", "Does this staircase spiral clockwise or counter-clockwise when viewed from above?"], ] with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): gr.Markdown( "# SpatioLM: Spatial Understanding with Vision-Language Models\n" "Ask spatial-reasoning questions about any image. SpatioLM enhances " "physical spatial intelligence in VLMs without requiring 3D input at " "inference time.\n\n" "[Paper](https://arxiv.org/abs/2608.01899) | " "[Code](https://github.com/xiaomi-research/spatio-lm) | " "[Model](https://huggingface.co/xiaomi-research/SpatioLM-Understanding-InternVL3.5)" ) with gr.Row(): with gr.Column(scale=1): image_input = gr.Image(type="pil", label="Input Image") question_input = gr.Textbox( label="Question", placeholder="e.g. Which object is closer to the camera?", lines=2, ) run_btn = gr.Button("Ask", variant="primary") with gr.Column(scale=1): output_text = gr.Textbox( label="Answer", lines=10, interactive=False, ) with gr.Accordion("Advanced settings", open=False): max_tokens = gr.Slider( minimum=32, maximum=2048, value=512, step=32, label="Max new tokens", ) temp_slider = gr.Slider( minimum=0.0, maximum=2.0, value=0.0, step=0.1, label="Temperature (0 = greedy)", ) top_p_slider = gr.Slider( minimum=0.1, maximum=1.0, value=1.0, step=0.05, label="Top-p", ) gr.Examples( examples=EXAMPLES, inputs=[image_input, question_input], outputs=output_text, fn=answer, cache_examples=True, cache_mode="lazy", ) run_btn.click( fn=answer, inputs=[image_input, question_input, max_tokens, temp_slider, top_p_slider], outputs=output_text, api_name="answer", ) if __name__ == "__main__": demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)