import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import base64 import copy import time from io import BytesIO import spaces import torch import gradio as gr import numpy as np from PIL import Image from huggingface_hub import snapshot_download from transformers import AutoConfig, AutoTokenizer, AutoProcessor from qwen_vl.data.utils import load_and_preprocess_images from qwen_vl.model.modeling_qwen2_5_vl import ( Qwen2_5_VLForConditionalGenerationWithGenerative, ) MODEL_ID = "H-EmbodVis/VEGA-3D-Spatial-Reasoning" WAN_REPO = "Wan-AI/Wan2.1-T2V-1.3B" MIN_PIXELS = 256 * 28 * 28 MAX_PIXELS = 1605632 MAX_NUM_FRAMES = 32 # --------------------------------------------------------------------------- # Download the frozen Wan2.1-T2V-1.3B generative encoder weights (VAE + DiT). # The VEGA-3D checkpoint stores intermediate spatiotemporal features from this # video diffusion model as its implicit 3D prior, so the weights are required # at inference time. We only pull the three files the encoder actually reads # (skipping the 11GB T5 text encoder — text conditioning uses a precomputed # prompt embedding shipped inside the qwen_vl package). # --------------------------------------------------------------------------- WAN_DIR = snapshot_download( WAN_REPO, allow_patterns=["config.json", "diffusion_pytorch_model.safetensors", "Wan2.1_VAE.pth"], ) print(f"Wan2.1-T2V-1.3B encoder weights at: {WAN_DIR}") # --------------------------------------------------------------------------- # Build the model. The config points the generative encoder at a non-existent # local path ("data/models/Wan2.1-T2V-1.3B"), so we override it to the real # downloaded directory before loading. # --------------------------------------------------------------------------- config = AutoConfig.from_pretrained(MODEL_ID) config.use_generative_encoder = True config.generative_encoder_path = WAN_DIR config.generative_vision_tower_checkpoint = WAN_DIR os.environ["WAN_T2V_CKPT_DIR"] = WAN_DIR model = Qwen2_5_VLForConditionalGenerationWithGenerative.from_pretrained( MODEL_ID, config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", ).eval().to("cuda") tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, padding_side="left") processor = AutoProcessor.from_pretrained( MODEL_ID, max_pixels=MAX_PIXELS, min_pixels=MIN_PIXELS, padding_side="left" ) print("Model loaded.") def _pil_to_data_uri(img: Image.Image) -> str: img = img.convert("RGB") buffer = BytesIO() img.save(buffer, format="JPEG") b64 = base64.b64encode(buffer.getvalue()).decode("utf-8") return f"data:image/jpeg;base64,{b64}" def _build_inputs(images, prompt, add_frame_index): """Replicates the official demo.ipynb `call_model` preprocessing.""" message = [{"role": "system", "content": "You are a helpful assistant."}] content = [] for i, img in enumerate(images): if add_frame_index: content.append({"type": "text", "text": "Frame-{}: ".format(i)}) content.append({"type": "image", "image": _pil_to_data_uri(img)}) content.append({"type": "text", "text": prompt}) message.append({"role": "user", "content": content}) messages = [message] text = processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) patch_size = processor.image_processor.patch_size merge_size = processor.image_processor.merge_size image_inputs = [] geometry_encoder_inputs = [] cur_geometry_encoder_inputs = [] for img in images: proc = load_and_preprocess_images([img.convert("RGB")])[0] cur_geometry_encoder_inputs.append(copy.deepcopy(proc)) _, height, width = proc.shape if (width // patch_size) % merge_size > 0: width = width - (width // patch_size) % merge_size * patch_size if (height // patch_size) % merge_size > 0: height = height - (height // patch_size) % merge_size * patch_size proc = proc[:, :height, :width] image_inputs.append(proc) geometry_encoder_inputs.append(torch.stack(cur_geometry_encoder_inputs)) inputs = processor( text=text, images=image_inputs, videos=None, padding=True, return_tensors="pt", do_rescale=False, ) return inputs, geometry_encoder_inputs def _sample_video_frames(video_path): import decord vr = decord.VideoReader(video_path) n = len(vr) if n <= MAX_NUM_FRAMES: idxs = np.arange(n) else: idxs = np.linspace(0, n - 1, MAX_NUM_FRAMES).astype(int) return [Image.fromarray(vr[i].asnumpy()).convert("RGB") for i in idxs] @spaces.GPU(duration=60) def _generate(inputs, geometry_encoder_inputs, max_new_tokens, temperature): device = model.device inputs["geometry_encoder_inputs"] = [f.to(device) for f in geometry_encoder_inputs] inputs = inputs.to(device) do_sample = temperature > 0 t0 = time.perf_counter() with torch.inference_mode(): cont = model.generate( **inputs, eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id, do_sample=do_sample, temperature=temperature if do_sample else None, top_p=None, num_beams=1, max_new_tokens=int(max_new_tokens), ) trimmed = [out[len(inp):] for inp, out in zip(inputs.input_ids, cont)] answer = processor.batch_decode( trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False )[0] dt = time.perf_counter() - t0 print(f"[generate] {dt:.1f}s, {int(max_new_tokens)} max tokens") return answer def infer(image, video, prompt, max_new_tokens=512, temperature=0.0): """Answer a spatial-reasoning question about an indoor scene. Args: image: A single scene image (PIL). Used when no video is supplied. video: Optional video of a scene; up to 32 frames are sampled from it. prompt: The question / instruction about the scene. max_new_tokens: Maximum number of tokens to generate. temperature: Sampling temperature; 0 for greedy decoding. """ if not prompt or not prompt.strip(): raise gr.Error("Please enter a question about the scene.") add_frame_index = False if video is not None: images = _sample_video_frames(video) add_frame_index = True elif image is not None: images = [image if isinstance(image, Image.Image) else Image.fromarray(image)] else: raise gr.Error("Please provide an image or a video of the scene.") inputs, geometry_encoder_inputs = _build_inputs(images, prompt, add_frame_index) return _generate(inputs, geometry_encoder_inputs, max_new_tokens, temperature) 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: with gr.Column(elem_id="col-container"): gr.Markdown( """ # VEGA-3D · Spatial Reasoning Ask 3D spatial-reasoning questions about an indoor scene. **VEGA-3D** augments a Qwen2.5-VL backbone with implicit 3D priors extracted from a frozen **Wan2.1-T2V** video diffusion model, giving it stronger geometric and spatial understanding. [Paper](https://huggingface.co/papers/2603.19235) · [Model](https://huggingface.co/H-EmbodVis/VEGA-3D-Spatial-Reasoning) · [Code](https://github.com/H-EmbodVis/VEGA-3D) """ ) with gr.Row(): with gr.Column(): image = gr.Image(label="Scene image", type="pil") video = gr.Video(label="…or a scene video (optional, samples 32 frames)") prompt = gr.Textbox( label="Question", placeholder="e.g. Which object is closest to the camera?", lines=2, ) run = gr.Button("Ask", variant="primary") with gr.Accordion("Advanced settings", open=False): max_new_tokens = gr.Slider( 16, 2048, value=512, step=16, label="Max new tokens" ) temperature = gr.Slider( 0.0, 1.0, value=0.0, step=0.05, label="Temperature (0 = greedy)", ) with gr.Column(): output = gr.Textbox(label="Answer", lines=12, show_copy_button=True) gr.Examples( examples=[ ["examples/living_room_minimalist.jpg", "Describe the spatial layout of this room. Which objects are on the left versus the right?"], ["examples/cafe_interior.jpg", "Which object is closest to the camera, and what is behind it?"], ["examples/library_interior.jpg", "Estimate the relative distances between the main objects in this scene."], ], inputs=[image, prompt], outputs=output, fn=lambda img, p: infer(img, None, p), cache_examples=True, cache_mode="lazy", ) run.click( fn=infer, inputs=[image, video, prompt, max_new_tokens, temperature], outputs=output, api_name="ask", ) if __name__ == "__main__": demo.launch(mcp_server=True)