qwen_test_runner/vision/runners.py

"""
runners.py — VLMRunner: load a Qwen VLM once, run image→JSON generation.

Promotes the working image→JSON patterns from colab/qwen_vit_json_test.py
(AutoProcessor + AutoModelFor{ImageTextToText,MultimodalLM}, image content
blocks, left-padding, batched generate with OOM-halving) into package code, and
reuses model_runner._XGrammarLogitsProcessor verbatim for constrained decoding.

This module imports torch/transformers at module load, so it is imported LAZILY
from qwen_test_runner.vision (the Phase-0 stub path never touches it).

xgrammar + images: the processor only inspects input_ids/scores, never image
tensors, so the grammar matcher works with image prompts PROVIDED the prompt_len
handed to it is the image-INCLUSIVE tokenized length and TokenizerInfo is built
from processor.tokenizer. Constrained mode runs at batch=1 (matcher is
per-sequence).
"""

from __future__ import annotations

import time
import warnings
from typing import Optional

import torch
from transformers import AutoProcessor, AutoModelForImageTextToText

from ..model_runner import _XGrammarLogitsProcessor, _HAS_XGRAMMAR
from .runner_types import VLMResult
from .tasks_vision import VisionTaskSpec, gbnf_for, resolved_system_prompt, tool_schema_for

try:  # newer transformers exposes a unified multimodal class (Qwen3.5)
    from transformers import AutoModelForMultimodalLM
    _HAS_MULTIMODAL = True
except ImportError:  # pragma: no cover
    AutoModelForMultimodalLM = None
    _HAS_MULTIMODAL = False

if _HAS_XGRAMMAR:  # pragma: no cover - GPU/optional path
    import xgrammar as xgr


_DTYPE = {"bf16": torch.bfloat16}


class VLMRunner:
    """Loads one Qwen VLM checkpoint and runs the four generation modes."""

    def __init__(
        self,
        model_id: str,
        loader_kind: str = "image_text_to_text",
        precision: str = "bf16",
        device: Optional[str] = None,
        device_map: str = "cuda",
        enable_thinking: bool = False,
        trust_remote_code: bool = True,
    ):
        self.model_id = model_id
        self.precision = precision
        self.loader_kind = loader_kind
        self.enable_thinking = enable_thinking
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        print(f"[VLMRunner] loading {model_id} ({loader_kind}, {precision}) on {device_map}")
        self.processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=trust_remote_code)
        # left-pad for correct batched decoding; ensure a pad token exists — Llama-based
        # checkpoints (e.g. JoyCaption) ship without one, which breaks processor padding.
        if getattr(self.processor, "tokenizer", None) is not None:
            self.processor.tokenizer.padding_side = "left"
            if self.processor.tokenizer.pad_token_id is None:
                self.processor.tokenizer.pad_token = self.processor.tokenizer.eos_token

        load_cls = AutoModelForImageTextToText
        if loader_kind == "multimodal_lm" and _HAS_MULTIMODAL:
            load_cls = AutoModelForMultimodalLM
        elif loader_kind == "llava_conditional":  # JoyCaption (SigLIP + Llama 3.1)
            from transformers import LlavaForConditionalGeneration
            load_cls = LlavaForConditionalGeneration

        load_kwargs = dict(device_map=device_map, trust_remote_code=trust_remote_code)
        if precision in _DTYPE:
            load_kwargs["dtype"] = _DTYPE[precision]  # quant repos carry their own config
        self.model = load_cls.from_pretrained(model_id, **load_kwargs)
        self.model.eval()

        tok = getattr(self.processor, "tokenizer", self.processor)
        self._pad_id = getattr(tok, "pad_token_id", None) or getattr(tok, "eos_token_id", None)

        # xgrammar compiler — reusable; per-category grammars compiled on demand.
        # Detect xgrammar LAZILY here (not just at module import) so installing it
        # AFTER the package was first imported still enables constrained mode — and
        # flip model_runner's captured flag so _XGrammarLogitsProcessor works too.
        self._xgr = None
        self._xgr_compiler = None
        self._xgr_tokenizer_info = None
        self._compiled: dict[str, object] = {}
        try:
            import xgrammar as _xgr_mod
            self._xgr = _xgr_mod
        except ImportError:
            pass
        if self._xgr is not None:
            import qwen_test_runner.model_runner as _mr
            if not _mr._HAS_XGRAMMAR:           # _XGrammarLogitsProcessor reads these
                _mr.xgr = self._xgr
                _mr._HAS_XGRAMMAR = True
            try:
                self._xgr_tokenizer_info = self._xgr.TokenizerInfo.from_huggingface(tok)
                self._xgr_compiler = self._xgr.GrammarCompiler(self._xgr_tokenizer_info)
            except Exception as e:  # pragma: no cover
                warnings.warn(f"xgrammar init failed: {e}; constrained mode unavailable")
        print(f"[VLMRunner] ready. xgrammar={self._xgr_compiler is not None}")

    def close(self) -> None:
        """Free the model from VRAM (used between models in a sweep)."""
        try:
            del self.model
        except AttributeError:
            pass
        import gc
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    # ── message construction ─────────────────────────────────────────────────

    def _messages(self, spec: VisionTaskSpec, image, user_prompt=None) -> list:
        system = resolved_system_prompt(spec)
        user_content = []
        if image is not None:
            user_content.append({"type": "image", "image": image})
        # per-sample prompt override (e.g. the question for VQA), else the task default
        user_content.append({"type": "text", "text": user_prompt or spec.user_prompt})
        return [
            {"role": "system", "content": system},
            {"role": "user", "content": user_content},
        ]

    def _encode(self, messages_list: list, tools=None):
        kw = dict(add_generation_prompt=True, tokenize=True, return_dict=True,
                  return_tensors="pt", padding=True)
        if tools is not None:                       # don't pass tools=None (some templates warn)
            kw["tools"] = tools
        # `enable_thinking` is a Qwen3.5 (multimodal_lm) toggle; Qwen3-VL's processor rejects it.
        if self.loader_kind == "multimodal_lm":
            kw["enable_thinking"] = self.enable_thinking
        return self.processor.apply_chat_template(messages_list, **kw).to(self.model.device)

    def _encode_llava(self, messages_list: list):
        """LLaVA/JoyCaption encode. Its chat template expects STRING `content` (it does
        string ops like `.replace` on it) and prepends the `<image>` token itself, so we
        cannot pass Qwen's structured content-parts list — we rebuild a string-content
        conversation. Render to TEXT (no `enable_thinking` no-op kwarg; no `tools`, since
        the Llama template would render a tools block), then let the processor attach the
        PIL image and expand `<image>` into feature tokens in `input_ids` — keeping the
        prompt length image-inclusive for xgrammar."""
        prompts, images = [], []
        for messages in messages_list:
            system_txt, user_txt, img = None, [], None
            for msg in messages:
                content = msg.get("content")
                if msg.get("role") == "system":
                    system_txt = content if isinstance(content, str) else None
                    continue
                if isinstance(content, list):
                    for part in content:
                        if not isinstance(part, dict):
                            continue
                        if part.get("type") == "image":
                            img = part.get("image")
                        elif part.get("type") == "text":
                            user_txt.append(part.get("text", ""))
                elif isinstance(content, str):
                    user_txt.append(content)
            convo = []
            if system_txt:
                convo.append({"role": "system", "content": system_txt})
            convo.append({"role": "user", "content": " ".join(t for t in user_txt if t)})
            prompts.append(self.processor.apply_chat_template(
                convo, add_generation_prompt=True, tokenize=False))
            images.append(img)
        inputs = self.processor(
            images=images, text=prompts, return_tensors="pt", padding=True,
        ).to(self.model.device)
        # The SigLIP vision tower is bf16 but the processor emits float32 pixel_values —
        # cast to the model dtype or the vision tower raises a dtype mismatch (the JoyCaption
        # model card does exactly this).
        if "pixel_values" in inputs:
            mdtype = getattr(self.model, "dtype", None)
            if mdtype is not None and inputs["pixel_values"].dtype != mdtype:
                inputs["pixel_values"] = inputs["pixel_values"].to(mdtype)
        return inputs

    # ── modes ─────────────────────────────────────────────────────────────────

    @torch.no_grad()
    def generate(self, spec: VisionTaskSpec, image, mode: str, *,
                 image_id: str = "", image_size=None, gt=None, user_prompt=None) -> VLMResult:
        if mode == "constrained":
            return self._generate_constrained(spec, image, image_id, user_prompt)
        tools = [self._tool_def(spec)] if mode == "tool_use" else None
        return self._generate_free(spec, image, mode, image_id, tools, user_prompt)

    @torch.no_grad()
    def _generate_free(self, spec, image, mode, image_id, tools, user_prompt=None) -> VLMResult:
        msgs = [self._messages(spec, image, user_prompt)]
        inputs = (self._encode_llava(msgs) if self.loader_kind == "llava_conditional"
                  else self._encode(msgs, tools=tools))
        n_in = inputs["input_ids"].shape[1]
        t0 = time.perf_counter()
        out = self.model.generate(**inputs, max_new_tokens=spec.max_new_tokens,
                                  do_sample=False, pad_token_id=self._pad_id)
        dt = time.perf_counter() - t0
        cont = out[0, n_in:]
        text = self.processor.decode(cont, skip_special_tokens=True)
        return VLMResult(mode, text, "transformers", int(n_in), int(cont.shape[0]), dt, image_id)

    @torch.no_grad()
    def _generate_constrained(self, spec, image, image_id, user_prompt=None) -> VLMResult:
        if self._xgr_compiler is None:
            warnings.warn("xgrammar unavailable; constrained falling back to json_mode")
            return self._generate_free(spec, image, "json_mode", image_id, None, user_prompt)
        grammar = gbnf_for(spec)
        compiled = self._compiled.get(spec.category)
        if compiled is None:
            compiled = self._xgr_compiler.compile_grammar(grammar)
            self._compiled[spec.category] = compiled

        msgs = [self._messages(spec, image, user_prompt)]
        inputs = (self._encode_llava(msgs) if self.loader_kind == "llava_conditional"
                  else self._encode(msgs))
        n_in = inputs["input_ids"].shape[1]  # image-INCLUSIVE length — critical
        lp = _XGrammarLogitsProcessor(
            compiled_grammar=compiled,
            vocab_size=self._xgr_tokenizer_info.vocab_size,
            prompt_len=n_in,
        )
        t0 = time.perf_counter()
        out = self.model.generate(**inputs, max_new_tokens=spec.max_new_tokens,
                                  do_sample=False, pad_token_id=self._pad_id,
                                  logits_processor=[lp])
        dt = time.perf_counter() - t0
        cont = out[0, n_in:]
        text = self.processor.decode(cont, skip_special_tokens=True)
        return VLMResult("constrained", text, "xgrammar", int(n_in), int(cont.shape[0]),
                         dt, image_id, grammar_conformant=True)

    def _tool_def(self, spec: VisionTaskSpec) -> dict:
        return {
            "type": "function",
            "function": {
                "name": "emit_" + spec.category,
                "description": spec.probes,
                "parameters": tool_schema_for(spec),
            },
        }

    # ── batched json_mode with OOM-halving (throughput path) ───────────────────

    @torch.no_grad()
    def generate_batch(self, spec: VisionTaskSpec, images: list, mode: str = "json_mode",
                       image_ids: Optional[list] = None) -> list[VLMResult]:
        image_ids = image_ids or [""] * len(images)
        return self._batch_with_fallback(spec, images, image_ids, mode)

    def _batch_with_fallback(self, spec, images, ids, mode) -> list[VLMResult]:
        if not images:
            return []
        try:
            return self._batch(spec, images, ids, mode)
        except torch.cuda.OutOfMemoryError:
            torch.cuda.empty_cache()
            if len(images) == 1:
                return [VLMResult(mode, "", "transformers", 0, 0, 0.0, ids[0])]
            half = max(1, len(images) // 2)
            return (self._batch_with_fallback(spec, images[:half], ids[:half], mode)
                    + self._batch_with_fallback(spec, images[half:], ids[half:], mode))
        except Exception:
            if len(images) == 1:
                return [VLMResult(mode, "", "transformers", 0, 0, 0.0, ids[0])]
            return [self._batch_with_fallback(spec, [im], [i], mode)[0]
                    for im, i in zip(images, ids)]

    @torch.no_grad()
    def _batch(self, spec, images, ids, mode) -> list[VLMResult]:
        tools = [self._tool_def(spec)] if mode == "tool_use" else None
        msgs = [self._messages(spec, im) for im in images]
        inputs = (self._encode_llava(msgs) if self.loader_kind == "llava_conditional"
                  else self._encode(msgs, tools=tools))
        n_in = inputs["input_ids"].shape[1]
        t0 = time.perf_counter()
        out = self.model.generate(**inputs, max_new_tokens=spec.max_new_tokens,
                                  do_sample=False, pad_token_id=self._pad_id)
        dt = time.perf_counter() - t0
        per = dt / len(images)
        results = []
        for row, iid in zip(out[:, n_in:], ids):
            text = self.processor.decode(row, skip_special_tokens=True)
            results.append(VLMResult(mode, text, "transformers", int(n_in), int(row.shape[0]), per, iid))
        return results