Spaces:
Running on Zero
Running on Zero
| import spaces # must be imported before torch / transformers | |
| import time | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| from wevisone_vl import WeVisOneVL | |
| from wv_render import ( | |
| PIXEL_TASKS, | |
| build_task_prompt, | |
| extract_count, | |
| parse_predictions_with_masks, | |
| render_from_cache, | |
| resolve_engine_task, | |
| ) | |
| MODEL_ID = "tencent/WeVisOne-VL-4B" | |
| # Released checkpoint (Qwen3-VL backbone, stored F32) is cast to bf16 and the | |
| # 9.6M-param pixel decoder (mask_decoder.pt) is loaded alongside it. | |
| engine = WeVisOneVL(MODEL_ID, backend="transformers", device="cuda", dtype="bfloat16") | |
| # UI task label -> (demo task, subtask) | |
| TASKS = { | |
| "Grounding": ("referring", ""), | |
| "Dense Detection": ("detection", ""), | |
| "Counting (boxes)": ("counting", "Box Counting"), | |
| "Counting (points)": ("counting", "Point Counting"), | |
| "Pointing": ("point", ""), | |
| "OCR (all text)": ("ocr", "All Text"), | |
| "OCR (locate text)": ("ocr", "Text Grounding"), | |
| "Caption / VQA": ("caption", "Free-form / VQA"), | |
| "Segmentation": ("segmentation", ""), | |
| "Matting": ("matting", ""), | |
| } | |
| TASK_HINTS = { | |
| "Grounding": "Referring expression, e.g. boys holding microphone", | |
| "Dense Detection": "Categories, one per line (sofa / pillow / lamp ...)", | |
| "Counting (boxes)": "What to count, e.g. white chess pieces", | |
| "Counting (points)": "What to count, e.g. white chess pieces", | |
| "Pointing": "What to point at, e.g. dogs", | |
| "OCR (all text)": "(no query needed: reads all text with boxes)", | |
| "OCR (locate text)": "Text to locate, e.g. COFFEE", | |
| "Caption / VQA": "Instruction or question, e.g. describe the image with bounding boxes in detail", | |
| "Segmentation": "Referring expression of the object to segment", | |
| "Matting": "What to matte, e.g. person", | |
| } | |
| DISPLAY_MODES = ["Auto", "Overlay", "Green Screen", "Mask / Alpha", "Cutout"] | |
| def _display_for(task: str, display: str) -> str: | |
| if display == "Auto": | |
| return "Green Screen" if task == "matting" else "Overlay" | |
| if display == "Mask / Alpha": | |
| return "Alpha Matte" if task == "matting" else "Binary Mask" | |
| return display | |
| def _duration(image, task="Grounding", query="", display="Auto", max_new_tokens=1024, *args, **kwargs): | |
| # Measured on ZeroGPU: ~50 tokens/s decode + a few seconds of prefill / mask decoding. | |
| return int(10 + int(max_new_tokens) * 0.02) | |
| def run( | |
| image: Image.Image, | |
| task: str = "Grounding", | |
| query: str = "", | |
| display: str = "Auto", | |
| max_new_tokens: int = 1024, | |
| temperature: float = 0.0, | |
| ): | |
| """Run WeVisOne-VL-4B perception on an image. | |
| Args: | |
| image: input image. | |
| task: one of Grounding, Dense Detection, Counting (boxes), Counting (points), | |
| Pointing, OCR (all text), OCR (locate text), Caption / VQA, Segmentation, Matting. | |
| query: target description / categories (one per line for Dense Detection) / question. | |
| display: how to render masks for Segmentation / Matting (Auto, Overlay, | |
| Green Screen, Mask / Alpha, Cutout). | |
| max_new_tokens: maximum number of generated tokens. | |
| temperature: sampling temperature (0 = greedy, as in the reference demo). | |
| Returns: | |
| Annotated image, a short summary, and the raw model response. | |
| """ | |
| if image is None: | |
| raise gr.Error("Please upload an image.") | |
| pil = image.convert("RGB") | |
| demo_task, subtask = TASKS.get(task, TASKS["Grounding"]) | |
| engine_task = resolve_engine_task(demo_task, subtask) | |
| is_pixel = demo_task in PIXEL_TASKS | |
| prompt = build_task_prompt(demo_task, subtask, query or "") | |
| t0 = time.perf_counter() | |
| result = engine.predict( | |
| pil, | |
| task=engine_task, | |
| query=query or "", | |
| prompt=prompt, | |
| decode_matting=is_pixel, | |
| max_new_tokens=int(max_new_tokens), | |
| temperature=float(temperature), | |
| top_p=0.95, | |
| top_k=50, | |
| repetition_penalty=1.05, | |
| ) | |
| latency = time.perf_counter() - t0 | |
| preds = parse_predictions_with_masks(result.response, pil.width, pil.height, result.masks) | |
| cache = {"pil": pil, "preds": preds, "masks": result.masks, "task": demo_task} | |
| vis = render_from_cache(cache, pixel_display=_display_for(demo_task, display)) or pil | |
| n_inst = sum(len(v) for v in preds.values()) | |
| parts = [f"**Task:** {task}", f"**Latency:** {latency:.1f}s"] | |
| if demo_task == "counting": | |
| count = extract_count(result.response, preds, counting=True) | |
| parts.insert(1, f"**Count:** {count if count is not None else 'N/A'}") | |
| if is_pixel: | |
| parts.append(f"**Masks:** {len(result.masks)}") | |
| if n_inst: | |
| parts.append(f"**Instances:** {n_inst}") | |
| parts.append( | |
| "**Found:** " + ", ".join(f"{k} ×{len(v)}" for k, v in preds.items()) | |
| ) | |
| elif result.response.strip() in ("", "None"): | |
| parts.append("**No matching target found.**") | |
| summary = " · ".join(parts) | |
| raw = f"Prompt: {result.prompt}\n{'-' * 40}\n{result.response}" | |
| return vis, summary, raw | |
| def on_task_change(task: str): | |
| hide = task == "OCR (all text)" | |
| return gr.update(placeholder=TASK_HINTS.get(task, ""), visible=not hide) | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="WeVisOne-VL-4B") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# WeVisOne-VL-4B\n" | |
| "Unified visual perception from objects to pixels: grounding, detection, counting, " | |
| "pointing, OCR, grounded captioning, segmentation and alpha matting with one model. \n" | |
| "[Model](https://huggingface.co/tencent/WeVisOne-VL-4B) · " | |
| "[GitHub](https://github.com/Tencent/WeVisOne-VL)" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_in = gr.Image(type="pil", label="Input image", height=420) | |
| task_in = gr.Dropdown(list(TASKS.keys()), value="Grounding", label="Task") | |
| query_in = gr.Textbox( | |
| label="Query", | |
| lines=2, | |
| placeholder=TASK_HINTS["Grounding"], | |
| ) | |
| run_btn = gr.Button("Run", variant="primary") | |
| with gr.Accordion("Advanced options", open=False): | |
| display_in = gr.Radio( | |
| DISPLAY_MODES, value="Auto", label="Mask display (Segmentation / Matting)" | |
| ) | |
| max_tokens_in = gr.Slider(64, 4096, value=1024, step=64, label="Max new tokens") | |
| temp_in = gr.Slider(0.0, 1.5, value=0.0, step=0.05, label="Temperature (0 = greedy)") | |
| with gr.Column(): | |
| image_out = gr.Image(type="pil", format="png", label="Result", height=420, interactive=False) | |
| summary_out = gr.Markdown() | |
| raw_out = gr.Textbox(label="Model output", lines=6, max_lines=16, interactive=False) | |
| task_in.change(on_task_change, inputs=task_in, outputs=query_in, api_visibility="private") | |
| inputs = [image_in, task_in, query_in, display_in, max_tokens_in, temp_in] | |
| outputs = [image_out, summary_out, raw_out] | |
| run_btn.click(run, inputs=inputs, outputs=outputs, api_name="run") | |
| query_in.submit(run, inputs=inputs, outputs=outputs, api_visibility="private") | |
| gr.Examples( | |
| examples=[ | |
| ["examples/boys.jpg", "Grounding", "boys holding microphone"], | |
| ["examples/livingroom.jpg", "Dense Detection", "sofa\npillow\ncoffee table\nbook\nlamp\nvase"], | |
| ["examples/chess.jpg", "Counting (boxes)", "white chess pieces"], | |
| ["examples/dogs.jpg", "Pointing", "dogs"], | |
| ["examples/coffee_sign.jpg", "OCR (all text)", ""], | |
| ["examples/livingroom.jpg", "Caption / VQA", "describe the image with bounding boxes in detail"], | |
| ["examples/parasol.jpg", "Segmentation", "parasol with alternating rings of light green and dark green with flowers"], | |
| ["examples/seabreeze.jpg", "Matting", "person"], | |
| ], | |
| inputs=[image_in, task_in, query_in], | |
| outputs=outputs, | |
| fn=run, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |