Spaces:
Running on Zero
Running on Zero
| """LOCUS-2B: referring-expression visual grounding with draft-and-verify boxes. | |
| Upload an image, type one referring expression per line, and the model draws each | |
| object's bounding box. A small box head drafts the box in one pass; the model's | |
| own coordinate tokens verify the draft in parallel (Jacobi iteration) until it | |
| settles. | |
| """ | |
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # noqa: E402 (must precede any CUDA-touching import) | |
| import gradio as gr # noqa: E402 | |
| import torch # noqa: E402 | |
| from PIL import Image # noqa: E402 | |
| from locus.hub import from_pretrained, ground # noqa: E402 | |
| MODEL_ID = "MTalha2001/LOCUS-2B" | |
| # Module-scope load: ZeroGPU intercepts .to("cuda"), packs weights to disk at | |
| # startup and streams them into VRAM on the first @spaces.GPU entry. | |
| model = from_pretrained(MODEL_ID) | |
| MODES = { | |
| "Verified, exact (greedy)": dict(accept=1.0, verify=True), | |
| "Verified, typical (α = 0.5)": dict(accept=0.5, verify=True), | |
| "Head only (one pass)": dict(verify=False), | |
| } | |
| DEFAULT_MODE = "Verified, exact (greedy)" | |
| def run_grounding( | |
| image: Image.Image | None, | |
| text: str, | |
| mode: str = DEFAULT_MODE, | |
| ): | |
| """Ground referring expressions in an image. | |
| Args: | |
| image: the input image. | |
| text: referring expressions, one per line. | |
| mode: output mode — verified exact (greedy), verified typical, or head-only draft. | |
| Returns: | |
| (annotated image with drawn boxes, results table) | |
| """ | |
| if image is None: | |
| raise gr.Error("Upload an image first.") | |
| exprs = [e.strip() for e in (text or "").splitlines() if e.strip()] | |
| if not exprs: | |
| raise gr.Error("Enter at least one referring expression (one per line).") | |
| if len(exprs) > 10: | |
| raise gr.Error("At most 10 expressions per run.") | |
| out = ground(model, image, exprs, **MODES[mode]) | |
| boxes = [(tuple(int(v) for v in r["box"]), r["expression"]) for r in out] | |
| rows = [ | |
| [r["expression"], [round(v) for v in r["box"]], round(r["confidence"], 3), r["passes"]] | |
| for r in out | |
| ] | |
| return (image, boxes), rows | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="LOCUS-2B · visual grounding") as demo: | |
| gr.Markdown( | |
| "# LOCUS-2B — Locate Once, Check Until Settled\n" | |
| "Ground a referring expression in an image: a box head drafts the box in one pass, " | |
| "then the model's own coordinate tokens verify the draft in parallel until it settles. " | |
| "Runs on [MTalha2001/LOCUS-2B](https://huggingface.co/MTalha2001/LOCUS-2B) " | |
| "(Qwen3-VL-2B backbone) · " | |
| "[Code](https://github.com/muhammad-talha-ad/LOCUS) · " | |
| "[Paper](https://arxiv.org/abs/2610.35053)" | |
| ) | |
| with gr.Column(elem_id="col-container"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| image = gr.Image(type="pil", label="Image", height=420) | |
| text = gr.Textbox( | |
| label="Expressions (one per line)", | |
| lines=3, | |
| placeholder="the cat on the left\nthe remote on the right", | |
| ) | |
| mode = gr.Radio( | |
| list(MODES), | |
| value=DEFAULT_MODE, | |
| label="Output", | |
| info=( | |
| "Exact: the model's greedy coordinate-token box (a few verifier passes). " | |
| "Typical: fewer passes, no longer exactly greedy. " | |
| "Head only: the one-pass draft, fastest." | |
| ), | |
| ) | |
| btn = gr.Button("Ground", variant="primary") | |
| with gr.Column(): | |
| result = gr.AnnotatedImage(label="Boxes", height=420) | |
| table = gr.Dataframe( | |
| headers=["expression", "box [x0, y0, x1, y1]", "confidence", "passes"], | |
| label="Results", | |
| interactive=False, | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["examples/elephant_asian.jpg", "the elephant", "Verified, exact (greedy)"], | |
| [ | |
| "examples/rubber_ducks.jpg", | |
| "the smallest duck\nthe largest duck", | |
| "Verified, exact (greedy)", | |
| ], | |
| ["examples/living_room_blue_couch.jpg", "the sofa", "Verified, typical (α = 0.5)"], | |
| ["examples/man_beach.jpg", "the backpack", "Head only (one pass)"], | |
| ], | |
| inputs=[image, text, mode], | |
| outputs=[result, table], | |
| fn=run_grounding, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| btn.click(run_grounding, [image, text, mode], [result, table], api_name="ground") | |
| demo.launch(mcp_server=True, css=CSS, theme=gr.themes.Citrus()) | |