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