app.py

"""sVLM-Council Space app — ask small local VLMs about an image, side by side.

Gradio app + MCP server with two tabs:
- Council: the council from `svlm_council.py` (the model registry is imported
  from there; the load/infer path here is ZeroGPU-specific).
- ExtractArena: targeted field extraction from `extract_arena.py` — its remote
  backends (HF Inference Providers + the moondream Space) run in the main
  process (unbilled); local council members reuse the resident GPU path.

All enabled members are loaded to CUDA at module level — on ZeroGPU that runs
under the CUDA emulation mode at startup (unbilled); only generation happens
inside the @spaces.GPU window. The four default members total ~37GB bf16, which
fits resident on the default `large` ZeroGPU slice (48GB).

Env knobs:
- MAX_COUNCIL_MODELS   max models per council request (default 3)
- DEFAULT_COUNCIL_MODEL model used when none selected (default qianfan-ocr)
- COUNCIL_MEMBERS      comma-separated subset to load at startup (default: all;
                       e.g. COUNCIL_MEMBERS=hunyuan-ocr for a light local run)
- MAX_EXTRACT_MODELS   max models per extract request (default 3)
- DEFAULT_EXTRACT_MODELS comma-separated models used when none selected
                       (default qwen3.5,muse-glimmer)
- HF_TOKEN             optional; without it the HF-Inference extract backends
                       fail per-model (logged) while everything else runs
"""

import spaces  # must be imported before any CUDA-touching torch usage on ZeroGPU
import os
import tempfile
import time
from pathlib import Path

import gradio as gr
import torch
from loguru import logger
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor

from svlm_council import COUNCIL

MODEL_RESEARCH_URLS = (
    "https://huggingface.co/spaces/GF-John/svlm-council/raw/main/MODEL_RESEARCH.md "
    "(mirrored at https://github.com/ohjho/hfs-svlm-council/blob/main/MODEL_RESEARCH.md )"
)
# Interactive walkthrough of one real ExtractArena request (hosting routes, prompt,
# normalization, payload); built from the surge-explainer skill, redeployable to the same URL.
EXTRACT_EXPLAINER_URL = "https://miroai-artifacts-extract-arena-explainer.surge.sh"


def _load_dotenv(path: Path = Path(".env")):
    """Minimal .env loader for local runs (no-op for keys already set)."""
    if not path.is_file():
        return
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, value = line.partition("=")
        os.environ.setdefault(key.strip(), value.strip().strip("'\""))


_load_dotenv()

# After _load_dotenv(): extract_arena reads HF_TOKEN at import time.
import extract_arena  # noqa: E402

MAX_COUNCIL_MODELS = int(os.environ.get("MAX_COUNCIL_MODELS", "3"))
MAX_EXTRACT_MODELS = int(os.environ.get("MAX_EXTRACT_MODELS", "3"))
MAX_IMAGE_PIXELS = int(os.environ.get("MAX_IMAGE_PIXELS", str(3_000_000)))
DEFAULT_MODEL = os.environ.get("DEFAULT_COUNCIL_MODEL", "qianfan-ocr")
MEMBERS = [
    name.strip()
    for name in os.environ.get("COUNCIL_MEMBERS", ",".join(COUNCIL)).split(",")
    if name.strip() in COUNCIL
]
if not MEMBERS:
    raise RuntimeError("COUNCIL_MEMBERS matched no registered council member")
if DEFAULT_MODEL not in MEMBERS:
    logger.warning(
        "Default model {} not loaded; falling back to {}", DEFAULT_MODEL, MEMBERS[0]
    )
    DEFAULT_MODEL = MEMBERS[0]

# ExtractArena roster: extract_arena's remote backends + whatever council members
# are actually loaded (local names in extract_arena.MODELS route through the
# resident GPU path here, never through its MPS single-slot loader).
EXTRACT_REMOTE = [n for n in extract_arena.MODELS if n not in COUNCIL]
EXTRACT_ROSTER = EXTRACT_REMOTE + MEMBERS
DEFAULT_EXTRACT_MODELS = [
    name.strip()
    for name in os.environ.get("DEFAULT_EXTRACT_MODELS", "qwen3.5,muse-glimmer").split(",")
    if name.strip() in EXTRACT_ROSTER
]
if not DEFAULT_EXTRACT_MODELS:
    logger.warning(
        "DEFAULT_EXTRACT_MODELS matched no roster entry; falling back to {}",
        EXTRACT_ROSTER[0],
    )
    DEFAULT_EXTRACT_MODELS = [EXTRACT_ROSTER[0]]

DEVICE = (
    "cuda"
    if torch.cuda.is_available()  # true on ZeroGPU's main process (emulation mode)
    else "mps" if torch.backends.mps.is_available() else "cpu"
)
logger.info("Device: {}; loading council members: {}", DEVICE, MEMBERS)

MODELS: dict = {}
PROCESSORS: dict = {}
for _name in MEMBERS:
    _spec = COUNCIL[_name]
    logger.info("Loading {} ({})", _spec.name, _spec.model_id)
    PROCESSORS[_name] = AutoProcessor.from_pretrained(
        _spec.model_id, **_spec.processor_kwargs
    )
    # On CUDA, force sdpa over any spec'd eager: eager materializes the full
    # seq x seq attention matrix with a float32 softmax (~17GB transient at
    # hunyuan-ocr's 11.6k patches for a 3MP image), which OOMs the ~10GB of
    # VRAM left beside the resident council. sdpa's fused kernels don't.
    # Output verified identical to eager on the bib test image (2026-08-27).
    _attn = "sdpa" if DEVICE == "cuda" else _spec.attn_implementation
    MODELS[_name] = (
        AutoModelForImageTextToText.from_pretrained(
            _spec.model_id,
            dtype=_spec.dtype if _spec.dtype == "auto" else getattr(torch, _spec.dtype),
            attn_implementation=_attn,
        )
        .to(DEVICE)
        .eval()
    )
logger.info("Council loaded: {}", list(MODELS))


def _prepare_image(image_path: str) -> str:
    """Cap total input pixels before the processors see the image.

    Vision-tower cost scales quadratically with patch count and hunyuan-ocr's
    processor applies no resolution cap of its own: a ~3200x4800 upload becomes
    60k patches, whose 60k x 60k eager-attention matrix OOMs even a 48GB GPU
    (surfacing on ZeroGPU as an NVML internal assert in the caching allocator).
    """
    with Image.open(image_path) as im:
        w, h = im.size
        if w * h <= MAX_IMAGE_PIXELS:
            return image_path
        scale = (MAX_IMAGE_PIXELS / (w * h)) ** 0.5
        resized = im.convert("RGB").resize(
            (max(1, round(w * scale)), max(1, round(h * scale))), Image.LANCZOS
        )
    out = (
        Path(tempfile.mkdtemp(prefix="council_")) / "input.png"
    )  # lossless for OCR fidelity
    resized.save(out)
    logger.info(
        "Downscaled input {}x{} -> {}x{} (MAX_IMAGE_PIXELS={})",
        w,
        h,
        resized.width,
        resized.height,
        MAX_IMAGE_PIXELS,
    )
    return str(out)


def _generate_one(
    name: str,
    image_path: str,
    prompt: str,
    system_prompt: str | None,
    temperature: float,
    max_tokens: int,
) -> str:
    """Run one resident council member; mirrors svlm_council.query()'s generate path."""
    spec = COUNCIL[name]
    processor, vlm = PROCESSORS[name], MODELS[name]

    messages = []
    if system_prompt:
        messages.append(
            {"role": "system", "content": [{"type": "text", "text": system_prompt}]}
        )
    messages.append(
        {
            "role": "user",
            "content": [
                {"type": "image", "url": str(image_path)},
                {"type": "text", "text": prompt},
            ],
        }
    )
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(vlm.device)
    inputs.pop("token_type_ids", None)  # required for Qwen3-VL, harmless for the rest

    gen_kwargs: dict = {
        "max_new_tokens": max_tokens,
        "use_cache": True,  # Qianfan-OCR's config ships use_cache: false
        **spec.generate_kwargs,
    }
    if temperature > 0:
        gen_kwargs.update(do_sample=True, temperature=temperature)
    else:
        # Explicit Nones neutralize sampling defaults baked into some
        # generation_configs (Qwen3-VL ships do_sample=true, temperature=0.7).
        gen_kwargs.update(do_sample=False, temperature=None, top_p=None, top_k=None)

    with torch.inference_mode():
        generated = vlm.generate(**inputs, **gen_kwargs)
    return processor.batch_decode(
        generated[:, inputs["input_ids"].shape[1] :],
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,  # preserve whitespace fidelity in OCR-ish output
    )[0].strip()


def _gpu_duration(
    image, prompt, models=None, system_prompt="", temperature=0.0, max_tokens=1024
):
    # Quota gate: visitors need this much remaining quota to be admitted at all,
    # so keep it proportional to the work actually requested.
    # qwen is the slowest ~30s per image, all others are 5-10s
    return 25 + 10 * max(1, len(models or []))


@spaces.GPU(duration=_gpu_duration)
def _gpu_run_council(
    image: str,
    prompt: str,
    models: list[str],
    system_prompt: str,
    temperature: float,
    max_tokens: int,
) -> list[dict]:
    results = []
    for name in models:
        logger.info("[{}] prompt: {!r}", name, prompt)
        t0 = time.perf_counter()
        try:
            text = _generate_one(
                name, image, prompt, system_prompt or None, temperature, max_tokens
            )
            result = {
                "model": name,
                "text": text,
                "latency_s": round(time.perf_counter() - t0, 2),
                "error": None,
            }
            logger.info("[{}] done in {}s", name, result["latency_s"])
        except Exception as exc:  # noqa: BLE001 — isolate failures per model
            logger.exception("[{}] failed", name)
            result = {
                "model": name,
                "text": "",
                "latency_s": round(time.perf_counter() - t0, 2),
                "error": f"{type(exc).__name__}: {exc}",
            }
        results.append(result)
        if torch.cuda.is_available():
            torch.cuda.empty_cache()  # release inter-model transients on multi-model calls
    return results


def council_query(
    image: str,
    prompt: str,
    models: list[str] | None = None,
    system_prompt: str = "",
    temperature: float = 0.0,
    max_tokens: int = 1024,
) -> list[dict]:
    # The real docstring (the MCP tool description) is assigned below — it needs
    # runtime values (roster, max) which a literal docstring can't interpolate.
    if not image:
        return [
            {"model": None, "text": "", "latency_s": 0.0, "error": "No image provided."}
        ]
    if not prompt or not prompt.strip():
        return [
            {
                "model": None,
                "text": "",
                "latency_s": 0.0,
                "error": "No prompt provided.",
            }
        ]

    if isinstance(models, str):  # tolerate a single name from API/MCP callers
        models = [models]
    models = [m for m in (models or []) if m] or [DEFAULT_MODEL]

    unknown = [m for m in models if m not in MODELS]
    if unknown:
        return [
            {
                "model": None,
                "text": "",
                "latency_s": 0.0,
                "error": f"Unknown or unloaded model(s): {', '.join(unknown)}. "
                f"Available: {', '.join(MODELS)}.",
            }
        ]
    if len(models) > MAX_COUNCIL_MODELS:
        return [
            {
                "model": None,
                "text": "",
                "latency_s": 0.0,
                "error": f"Too many models selected ({len(models)}); "
                f"the maximum per request is {MAX_COUNCIL_MODELS}.",
            }
        ]
    try:
        image = _prepare_image(image)  # main process — the downscale isn't billed
    except Exception as exc:  # noqa: BLE001 — unreadable/corrupt upload
        return [
            {
                "model": None,
                "text": "",
                "latency_s": 0.0,
                "error": f"Could not read the input image: {type(exc).__name__}: {exc}",
            }
        ]
    return _gpu_run_council(
        image, prompt, models, system_prompt, temperature, max_tokens
    )


_ROSTER_DOCS = {
    "qianfan-ocr": "baidu/Qianfan-OCR, 4.7B — best key-information extraction (KIE) of the council; the default",
    "hunyuan-ocr": "tencent/HunyuanOCR, 1.1B — fastest; best scene-text under 3B, but prompt-sensitive on terse prompts",
    "granite-vision": "ibm-granite/granite-vision-4.1-4b, 4B — strong key-value pair extraction (VAREX 94.2% exact-match)",
    "qwen3-vl-8b": "Qwen/Qwen3-VL-8B-Instruct, 8.8B — best sub-10B OCRBench; robust to blur/tilt/low light; strongest general VQA",
}
council_query.__doc__ = f"""Ask a council of small vision-language models (VLMs) a question about an image and compare their answers.

Runs the same image + prompt through the selected council members — small (~1-9B) open VLMs
resident on this Space's GPU — and returns one result per model. Works for targeted data
extraction (e.g. "What is the bib number?") and vanilla visual question answering alike.
Available members: {"; ".join(f'"{m}" ({_ROSTER_DOCS.get(m, COUNCIL[m].model_id)})' for m in MEMBERS)}.
If "models" is omitted or empty, only the default member "{DEFAULT_MODEL}" runs; at most
{MAX_COUNCIL_MODELS} models may be selected per request. To decide which members fit your task,
read the maintained model research notes (benchmarks, strengths, quirks per member) at
{MODEL_RESEARCH_URLS}.
Inputs larger than {MAX_IMAGE_PIXELS / 1e6:.1f} megapixels are downscaled (aspect preserved)
before inference, so very small text in very large images may need pre-cropping by the caller.
The output is a JSON list with one object per requested model, in request order. Each object
has: "model" (the council member name, or null when the request itself was invalid); "text"
(the model's reply, "" on failure); "latency_s" (inference seconds, float); and "error" (null
on success, otherwise a message — per-model failures don't abort the other members).

Args:
    image: Filepath or URL of the input image to analyze.
    prompt: The question or instruction about the image.
    models: List of council member names to run (see the description for the roster); empty or omitted runs the default member.
    system_prompt: Optional system prompt prepended to the conversation; empty string means none.
    temperature: Sampling temperature; 0.0 (default) is deterministic greedy decoding, values above 0 enable sampling.
    max_tokens: Maximum number of new tokens each model may generate.
"""


def _extract_error(message: str) -> list[dict]:
    return [
        {
            "model": None,
            "value": None,
            "values": None,
            "raw": "",
            "latency_s": 0.0,
            "error": message,
        }
    ]


def _to_extract_result(
    name: str, raw: str, latency_s: float, multi: bool, error: str | None = None
) -> dict:
    """Shape one reply like extract_arena.ExtractResult (normalized + raw kept)."""
    if error:
        return {
            "model": name,
            "value": None,
            "values": None,
            "raw": "",
            "latency_s": latency_s,
            "error": error,
        }
    raw = raw.strip()
    if multi:
        parsed = extract_arena._parse_json_reply(raw)
        values = (
            {k: extract_arena._normalize_value(v) for k, v in parsed.items()}
            if parsed is not None
            else None
        )
        value = None
    else:
        values = None
        value = extract_arena._normalize_value(raw)
    return {
        "model": name,
        "value": value,
        "values": values,
        "raw": raw,
        "latency_s": latency_s,
        "error": None,
    }


def extract_fields(
    image: str,
    fields: list[str] | None = None,
    models: list[str] | None = None,
    answer_format: str = "",
) -> list[dict]:
    # The real docstring (the MCP tool description) is assigned below — it needs
    # runtime values (roster, defaults, max) a literal docstring can't interpolate.
    if not image:
        return _extract_error("No image provided.")
    if isinstance(fields, str):  # tolerate a single field from API/MCP callers
        fields = [fields]
    fields = [f.strip() for f in (fields or []) if f and f.strip()]
    if not fields:
        return _extract_error("No fields provided.")

    if isinstance(models, str):  # tolerate a single name from API/MCP callers
        models = [models]
    models = [m for m in (models or []) if m] or list(DEFAULT_EXTRACT_MODELS)
    models = list(dict.fromkeys(models))  # results are keyed by name below

    unknown = [m for m in models if m not in EXTRACT_ROSTER]
    if unknown:
        return _extract_error(
            f"Unknown or unloaded model(s): {', '.join(unknown)}. "
            f"Available: {', '.join(EXTRACT_ROSTER)}."
        )
    if len(models) > MAX_EXTRACT_MODELS:
        return _extract_error(
            f"Too many models selected ({len(models)}); "
            f"the maximum per request is {MAX_EXTRACT_MODELS}."
        )
    try:
        # One downscale for every backend: bounds remote upload size too, and
        # keeps all models comparing the same pixels.
        image = _prepare_image(image)
    except Exception as exc:  # noqa: BLE001 — unreadable/corrupt upload
        return _extract_error(
            f"Could not read the input image: {type(exc).__name__}: {exc}"
        )

    prompt = extract_arena.build_prompt(fields, answer_format.strip() or None)
    multi = len(fields) > 1
    image_path = Path(image)

    results: dict[str, dict] = {}
    local = [m for m in models if m in MODELS]
    for name in (m for m in models if m not in MODELS):  # remote — unbilled
        logger.info("[{}] prompt: {!r}", name, prompt)
        t0 = time.perf_counter()
        try:
            raw = extract_arena.MODELS[name].fn(image_path, prompt)
            results[name] = _to_extract_result(
                name, raw, round(time.perf_counter() - t0, 2), multi
            )
            logger.info("[{}] done in {}s", name, results[name]["latency_s"])
        except Exception as exc:  # noqa: BLE001 — isolate failures per model
            logger.exception("[{}] failed", name)  # incl. missing HF_TOKEN
            results[name] = _to_extract_result(
                name,
                "",
                round(time.perf_counter() - t0, 2),
                multi,
                error=f"{type(exc).__name__}: {exc}",
            )
    if local:  # resident members — one billed GPU window for the whole subset
        for r in _gpu_run_council(image, prompt, local, "", 0.0, 2048):
            results[r["model"]] = _to_extract_result(
                r["model"], r["text"], r["latency_s"], multi, error=r["error"]
            )
    return [results[name] for name in models]


_EXTRACT_ROSTER_DOCS = {
    "moondream3": "moondream/moondream3-preview via the GF-John/moondream-pointer Space — no token needed, but its ZeroGPU quota can exhaust",
    "qwen3.5": "Qwen/Qwen3.5-9B, hosted — most reliable remote entry; strong general VLM",
    "muse-glimmer": "meta-models/Muse-Glimmer-30B, hosted — 86.6% Roboflow Data Extraction, the best open <=40B on the closest proxy benchmark for this task",
    "gemma4-31b": "google/gemma-4-31B-it, hosted — fastest remote observed; best hosted availability; extraction quality unbenchmarked",
    "gemma4-26b-a4b": "google/gemma-4-26B-A4B-it, hosted — cheapest hosted option; 3.8B-active MoE sibling of the 31B",
    **_ROSTER_DOCS,
}
extract_fields.__doc__ = f"""Extract the value(s) of named fields from an image with several vision-language models (VLMs) and compare their answers side by side.

Builds one extraction prompt from the field names — with an abstention guardrail instructing
each model to answer exactly "N/A" rather than guess when a field is not visible — and runs it
through every selected model: hosted models via HF Inference Providers or a helper Space, local
council members on this Space's GPU. Available models:
{"; ".join(f'"{m}" ({_EXTRACT_ROSTER_DOCS.get(m, m)})' for m in EXTRACT_ROSTER)}.
If "models" is omitted or empty, the defaults ({", ".join(f'"{m}"' for m in DEFAULT_EXTRACT_MODELS)}) run;
at most {MAX_EXTRACT_MODELS} models may be selected per request. Hosted models need an HF_TOKEN
configured on the server; without it they fail per-model while the rest still run. To decide
which models fit your task, read the maintained model research notes (benchmarks, strengths,
quirks per member) at {MODEL_RESEARCH_URLS}.
Inputs larger than {MAX_IMAGE_PIXELS / 1e6:.1f} megapixels are downscaled (aspect preserved)
before inference, so very small text in very large images may need pre-cropping by the caller;
images sent to hosted backends are additionally recompressed as JPEG when needed to fit the HF
router's ~5MB request-body cap, so hosted models may see a slightly lossier image than local ones.
The output is a JSON list with one object per requested model, in request order. Each object
has: "model" (the model name, or null when the request itself was invalid); "value" (single-field
requests: the normalized answer, "" meaning the model abstained/field not found, null on
multi-field requests or failure); "values" (multi-field requests: an object mapping each field
name to its normalized value with "" for absent fields, or null if the reply didn't parse as
JSON — check "raw" then); "raw" (the model's verbatim reply, kept for debugging); "latency_s"
(seconds, float); and "error" (null on success, otherwise a message — per-model failures don't
abort the other models).

Args:
    image: Filepath or URL of the input image to extract from.
    fields: List of field names to extract, e.g. ["bib number"] or ["bib number", "race name"]; single-field requests are the most reliable regime.
    models: List of model names to run (see the description for the roster); empty or omitted runs the defaults.
    answer_format: Optional format directive appended verbatim to the prompt, e.g. "Format the answer as a continuous sequence of digits (e.g., 12345)."; with multiple fields, phrase it per value; empty string means none.
"""


council_iface = gr.Interface(
    fn=council_query,
    inputs=[
        gr.Image(type="filepath", label="Input Image"),
        gr.Textbox(
            label="Prompt",
            lines=2,
            value="Describe this image.",
            info="Question or instruction for the image — extraction or open-ended VQA",
        ),
        gr.CheckboxGroup(
            label="Council Members",
            choices=MEMBERS,
            value=[DEFAULT_MODEL],
            info=f"Select up to {MAX_COUNCIL_MODELS} models; see MODEL_RESEARCH.md for how to choose",
        ),
        gr.Textbox(
            label="System Prompt", lines=1, value="", info="Optional; empty = none"
        ),
        gr.Slider(label="Temperature", value=0.0, minimum=0.0, maximum=1.0, step=0.1),
        gr.Slider(label="Max Tokens", value=1024, minimum=32, maximum=4096, step=32),
    ],
    outputs=gr.JSON(label="Council Results"),
    description=(
        "Run the same image + prompt through small (~1-9B) open vision-language models and "
        "compare answers — data extraction or vanilla VQA. Model choices, benchmarks, and "
        "quirks are documented in "
        "[MODEL_RESEARCH.md](https://huggingface.co/spaces/GF-John/svlm-council/raw/main/MODEL_RESEARCH.md)."
    ),
    api_name="council_query",
    api_visibility="public",
)

extract_iface = gr.Interface(
    fn=extract_fields,
    inputs=[
        gr.Image(type="filepath", label="Input Image"),
        gr.Dropdown(
            label="Fields",
            multiselect=True,
            allow_custom_value=True,
            choices=["bib number", "race name"],
            value=["bib number"],
            info="Field name(s) to extract — type your own and press Enter; one field per request is the most reliable",
        ),
        gr.CheckboxGroup(
            label="Models",
            choices=EXTRACT_ROSTER,
            value=DEFAULT_EXTRACT_MODELS,
            info=f"Select up to {MAX_EXTRACT_MODELS} models; see MODEL_RESEARCH.md for how to choose",
        ),
        gr.Textbox(
            label="Format Directive",
            lines=1,
            value="",
            info='Optional; appended verbatim to the prompt, e.g. "Format the answer as a continuous sequence of digits (e.g., 12345)."',
        ),
    ],
    outputs=gr.JSON(label="Extraction Results"),
    description=(
        "Ask several VLMs — hosted and local — for the value of one or more named fields in an "
        'image, with an abstention guardrail (models answer "N/A", normalized to "", instead '
        "of guessing when a field is absent). Model choices, benchmarks, and quirks are "
        "documented in "
        "[MODEL_RESEARCH.md](https://huggingface.co/spaces/GF-John/svlm-council/raw/main/MODEL_RESEARCH.md). "
        f"New here? Read the [interactive explainer]({EXTRACT_EXPLAINER_URL}) — one real request "
        "traced through the three hosting routes (Inference Providers, a helper Space, this "
        "Space's ZeroGPU), the prompt, normalization, and the payload."
    ),
    api_name="extract_fields",
    api_visibility="public",
)

app = gr.TabbedInterface(
    [council_iface, extract_iface],
    ["Council", "ExtractArena"],
    title="sVLM Council",
)
app.launch(mcp_server=True, app_kwargs={"docs_url": "/docs"})