app.py

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)