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) @spaces.GPU(duration=_duration) 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)