Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # MUST be the first import after the env var, before torch | |
| import re | |
| import threading | |
| import gradio as gr | |
| import torch | |
| from transformers import AutoProcessor, Qwen3VLForConditionalGeneration, TextIteratorStreamer | |
| MODEL_ID = "xinyili0624/GPD-4B" | |
| DEFAULT_SYSTEM_PROMPT = ( | |
| "You are a helpful assistant for 3D spatial reasoning. Think step by step using " | |
| "<think></think> tags, then give your final answer inside <answer> </answer> tags." | |
| ) | |
| processor = AutoProcessor.from_pretrained(MODEL_ID) | |
| model = ( | |
| Qwen3VLForConditionalGeneration.from_pretrained( | |
| MODEL_ID, | |
| dtype=torch.bfloat16, | |
| attn_implementation="sdpa", | |
| ) | |
| .to("cuda") | |
| .eval() | |
| ) | |
| ANSWER_RE = re.compile(r"<answer>\s*(.*?)\s*</answer>", re.DOTALL) | |
| def _extract_answer(text: str) -> str: | |
| """Return the content of the last <answer>...</answer> tag (or its best effort).""" | |
| matches = ANSWER_RE.findall(text) | |
| if matches: | |
| return matches[-1].strip() | |
| return "" | |
| def answer( | |
| image, | |
| question: str, | |
| system_prompt: str = DEFAULT_SYSTEM_PROMPT, | |
| max_new_tokens: int = 1024, | |
| temperature: float = 0.0, | |
| top_p: float = 0.95, | |
| ): | |
| """Answer a spatial-reasoning question about an image with GPD-4B. | |
| GPD (Geometry-Privileged Distillation) is an RGB-only vision-language model | |
| fine-tuned from Qwen3-VL-4B-Instruct for single-image 3D spatial reasoning | |
| (observer shifts, directional relations, depth/distance estimation). It | |
| thinks step by step between <think> and </think> tags and wraps its | |
| final answer in <answer> </answer> tags. | |
| Args: | |
| image: input image (PIL Image) or None for a text-only question. | |
| question: the spatial question to ask about the image. | |
| system_prompt: system prompt guiding the reasoning format. | |
| max_new_tokens: maximum number of generated tokens. | |
| temperature: sampling temperature; 0 = deterministic greedy decoding. | |
| top_p: nucleus sampling probability (used when temperature > 0). | |
| Yields: | |
| Tuple of (extracted_final_answer, full_generated_text), streamed live. | |
| """ | |
| question = (question or "").strip() | |
| if not question: | |
| yield "Please enter a question.", "" | |
| return | |
| content = [] | |
| if image is not None: | |
| content.append({"type": "image", "image": image}) | |
| content.append({"type": "text", "text": question}) | |
| messages = [ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": content}, | |
| ] | |
| prompt_text = processor.apply_chat_template( | |
| messages, tokenize=False, add_generation_prompt=True | |
| ) | |
| images = [image] if image is not None else None | |
| inputs = processor( | |
| text=[prompt_text], images=images, return_tensors="pt" | |
| ).to(model.device) | |
| gen_kwargs = { | |
| "max_new_tokens": int(max_new_tokens), | |
| "do_sample": float(temperature) > 0.0, | |
| } | |
| if gen_kwargs["do_sample"]: | |
| gen_kwargs["temperature"] = float(temperature) | |
| gen_kwargs["top_p"] = float(top_p) | |
| streamer = TextIteratorStreamer( | |
| processor.tokenizer, skip_prompt=True, skip_special_tokens=True | |
| ) | |
| gen_kwargs["streamer"] = streamer | |
| generation_thread = threading.Thread( | |
| target=lambda: model.generate(**inputs, **gen_kwargs), daemon=True | |
| ) | |
| generation_thread.start() | |
| full_text = "" | |
| for chunk in streamer: | |
| full_text += chunk | |
| yield _extract_answer(full_text), full_text | |
| generation_thread.join() | |
| final_answer = _extract_answer(full_text) | |
| if not final_answer: | |
| tail = full_text.strip().splitlines()[-1] if full_text.strip() else "" | |
| final_answer = tail.strip() | |
| yield final_answer, full_text | |
| EXAMPLES = [ | |
| [ | |
| "examples/ex_distance_to_chair.jpg", | |
| ( | |
| "From the person's position who took the photo, what is the distance " | |
| "to the center of chair (red point)? Express it in meters. Calculate " | |
| "or judge based on the 3D center points of these objects. Enter a " | |
| "single numeric response and no other details.\n" | |
| "Think step by step and provide your reasoning inside <think> </think> tags. " | |
| "Then give your final answer inside <answer> </answer> tags with a " | |
| "single number or value." | |
| ), | |
| ], | |
| [ | |
| "examples/ex_depth_difference.jpg", | |
| ( | |
| "The depth of pillow (red point) is 2.5. Use this information to " | |
| "find the absolute depth difference between plant (green point) and " | |
| "table (blue point). Answer in meters. Calculate or judge based on " | |
| "the 3D center points of these objects. Please give a response in " | |
| "the form of a single number.\n" | |
| "Think step by step and provide your reasoning inside <think> </think> tags. " | |
| "Then give your final answer inside <answer> </answer> tags with a " | |
| "single number or value." | |
| ), | |
| ], | |
| [ | |
| "examples/ex_closer_object.jpg", | |
| ( | |
| "From the red point laundry basket, which object, container (green " | |
| "point) or box (blue point), is closer? Calculate or judge based on " | |
| "the 3D center points of these objects. Choose the correct response " | |
| "from the given choices.\n" | |
| "A. box (blue point)\n" | |
| "B. container (green point)\n" | |
| "Your answer can only include one of options A, B.\n" | |
| "Think step by step and provide your reasoning inside <think> </think> tags. " | |
| "Then give your final answer inside <answer> </answer> tags. For " | |
| "example: <answer>A</answer>" | |
| ), | |
| ], | |
| [ | |
| "examples/ex_directional_relation.jpg", | |
| ( | |
| "Analyze the image and determine the directional relationship " | |
| "between the monitor (red bbox) and the whiteboard (blue bbox) as " | |
| "seen by the observer. Calculate or judge based on the 3D center " | |
| "points of these objects.\n" | |
| "The options describe the spatial relationship between two objects " | |
| "in terms of left-right (left, right, or empty if " | |
| "indistinguishable), above-below (above, below, or empty if " | |
| "indistinguishable), and farther-closer (farther, closer, or empty " | |
| "if indistinguishable).\n" | |
| "Select the right response from the available choices.\n" | |
| "A. left, , closer\n" | |
| "B. right, , farther\n" | |
| "C. , above, closer\n" | |
| "D. right, below, \n" | |
| "Your answer can only include one of options A, B, C or D.\n" | |
| "Think step by step and provide your reasoning inside <think> </think> tags. " | |
| "Then give your final answer inside <answer> </answer> tags. For " | |
| "example: <answer>A</answer>" | |
| ), | |
| ], | |
| [ | |
| "examples/ex_observer_shift.jpg", | |
| ( | |
| "From the observer's original viewpoint, where is dining table " | |
| "(red bbox)? How does this relationship shift once the observer " | |
| "moves to chair's 3D center (green bbox) and faces oven (blue " | |
| "bbox)? Calculate or judge based on the 3D center points of these " | |
| "objects. Please use the world coordinate system to determine the " | |
| "up and down position relationship of objects.\n" | |
| "For multiple-choice questions, consider only the state after the " | |
| "observer has moved.\n" | |
| "The options describe the spatial relationship between object and " | |
| "observer in terms of left-right (left, right, or empty if " | |
| "indistinguishable), above-below (above, below, or empty if " | |
| "indistinguishable), and front-behind (front, behind, or empty if " | |
| "indistinguishable).\n" | |
| "Choose the correct response from the given choices.\n" | |
| "A. right, below, front\n" | |
| "B. right, , behind\n" | |
| "C. left, above, behind\n" | |
| "D. right, below, \n" | |
| "Your answer can only include one of options A, B, C or D.\n" | |
| "Think step by step and provide your reasoning inside <think> </think> tags. " | |
| "Then give your final answer inside <answer> </answer> tags. For " | |
| "example: <answer>A</answer>" | |
| ), | |
| ], | |
| ] | |
| CSS = """ | |
| #col-container { max-width: 1100px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="GPD · Spatial Reasoning") as demo: | |
| gr.Markdown( | |
| """ | |
| # 🧭 GPD · Distilling Routed 3D Privilege for Spatial Reasoning | |
| **GPD-4B** is an RGB-only vision-language model (fine-tuned from | |
| [Qwen3-VL-4B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct)) that answers 3D | |
| spatial-reasoning questions about a single image — observer-shift, directional relations, | |
| and depth/distance estimation — thinking step by step and wrapping the final answer in | |
| `<answer>` tags. | |
| [Paper](https://huggingface.co/papers/2610.12355) · | |
| [Code](https://github.com/ZJU-REAL/GPD) · | |
| [GPD-4B](https://huggingface.co/xinyili0624/GPD-4B) · | |
| [GPD-2B](https://huggingface.co/xinyili0624/GPD-2B) · | |
| [GPD-15k Data](https://huggingface.co/datasets/xinyili0624/GPD-15k) | |
| """ | |
| ) | |
| with gr.Column(elem_id="col-container"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| image_input = gr.Image( | |
| type="pil", label="Image", height=420, sources=["upload", "clipboard"] | |
| ) | |
| with gr.Column(scale=1): | |
| question_input = gr.Textbox( | |
| label="Question", | |
| lines=10, | |
| placeholder=( | |
| "Ask a spatial question about the image, e.g. " | |
| "'What is closer to the camera, the chair or the table?'" | |
| ), | |
| ) | |
| submit_btn = gr.Button("Answer", variant="primary") | |
| answer_output = gr.Textbox( | |
| label="Final answer", | |
| lines=1, | |
| buttons=["copy"], | |
| ) | |
| reasoning_output = gr.Textbox( | |
| label="Full reasoning", | |
| lines=12, | |
| buttons=["copy"], | |
| ) | |
| with gr.Accordion("Advanced settings", open=False): | |
| system_prompt_input = gr.Textbox( | |
| label="System prompt", | |
| value=DEFAULT_SYSTEM_PROMPT, | |
| lines=2, | |
| ) | |
| with gr.Row(): | |
| max_new_tokens_input = gr.Slider( | |
| label="Max new tokens", minimum=64, maximum=3072, value=1024, step=64 | |
| ) | |
| temperature_input = gr.Slider( | |
| label="Temperature (0 = deterministic)", minimum=0.0, maximum=1.5, value=0.0, step=0.05 | |
| ) | |
| top_p_input = gr.Slider( | |
| label="Top-p (used when temperature > 0)", minimum=0.05, maximum=1.0, value=0.95, step=0.05 | |
| ) | |
| gr.Examples( | |
| examples=EXAMPLES, | |
| inputs=[image_input, question_input], | |
| outputs=[answer_output, reasoning_output], | |
| fn=answer, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Examples (from the authors' GPD-15k dataset — SPAR scenes)", | |
| ) | |
| submit_btn.click( | |
| answer, | |
| inputs=[image_input, question_input, system_prompt_input, max_new_tokens_input, temperature_input, top_p_input], | |
| outputs=[answer_output, reasoning_output], | |
| api_name="answer", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS) |