"""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)"


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