# SPDX-License-Identifier: Apache-2.0
"""
VLM (Vision-Language Model) engine with continuous batching.

This engine extends BatchedEngine to support vision-language models via
mlx-vlm. It provides:

- Image input processing (URL, base64, local file)
- Multi-image chat support
- Pre-computed vision embeddings for efficient batched inference
- Full compatibility with oMLX's tiered KV cache and boundary snapshots

Architecture:
    1. Images are extracted from messages and loaded as PIL Images
    2. mlx-vlm's prepare_inputs() tokenizes text and preprocesses images
    3. model.get_input_embeddings() runs vision encoder + embedding merge
    4. VLMModelAdapter receives pre-computed embeddings for prefill injection
    5. After prefill, decode uses standard token IDs (vision context in KV cache)

Usage:
    Engine is automatically selected when model_discovery detects a VLM model
    (engine_type="vlm"). No changes needed for API callers — the OpenAI
    vision API format is transparently handled.
"""

import asyncio
import contextlib
import copy
import functools
import hashlib
import importlib
import inspect
import json
import logging
import os
import threading
from collections.abc import AsyncIterator
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import mlx.core as mx

from ..api.tool_calling import convert_tools_for_template
from ..api.utils import (
    clean_special_tokens,
    detect_and_strip_partial,
    remove_special_tokens_preserve_whitespace,
)
from ..cache.vision_feature_cache import VisionFeatureSSDCache
from ..exceptions import InvalidRequestError
from ..model_settings import ane_prefill_backend, ane_prefill_fraction
from ..models.vlm import VLMModelAdapter
from ..patches.gemma4_audio import apply_gemma4_audio_patch
from ..patches.mlx_vlm_pixtral_torch_free import apply_pixtral_torch_free_patch
from ..reasoning_effort import apply_chat_template_with_reasoning_effort_fallback
from ..utils.image import (
    compute_image_hash,
    compute_per_image_hashes,
    extract_images_from_messages,
    extract_media_from_messages,
)
from ..utils.video import (
    NATIVE_VIDEO_MODEL_TYPES,
    attach_native_video_processor,
    estimate_native_video_tokens,
    expand_video_parts,
    native_video_token_count,
    probe_video,
    write_video_data_uri,
)
from .base import (
    BaseEngine,
    GenerationOutput,
    _clear_teardown_references,
    _close_engine_core,
    _run_scheduler_preflight_with_cleanup_retry,
    _warn_scheduler_unreachable_once,
)

logger = logging.getLogger(__name__)

# OCR model types that require special handling.
# unlimited-ocr keeps its dashed config model_type (mlx-vlm resolves it to the
# unlimited_ocr package via MODEL_REMAPPING), so key it in the dashed form to
# match VLMBatchedEngine.model_type (== vlm_model.config.model_type).
OCR_MODEL_TYPES = {
    "deepseekocr",
    "deepseekocr_2",
    "unlimited-ocr",
    "dots_ocr",
    "glm_ocr",
}

# OCR model types and their default markdown conversion prompts.
# When an OCR model receives a generic user prompt with an image,
# the prompt is automatically adjusted for markdown output.
OCR_MODEL_PROMPTS: Dict[str, str] = {
    "deepseekocr": "Convert the document to markdown.",
    "deepseekocr_2": "Convert the document to markdown.",
    # baidu/Unlimited-OCR upstream-documented single-page baseline. Multi-page
    # / PDF workflows use "Multi page parsing." (pass it explicitly).
    "unlimited-ocr": "document parsing.",
    "dots_ocr": "Convert this page to clean Markdown while preserving reading order.",
    "glm_ocr": "Text Recognition:",
}

# Extra stop sequences for OCR models to prevent degeneration.
# Many OCR models lack proper EOS handling and generate chat-turn
# tokens (<|user|>, <|im_start|>, etc.) indefinitely after the OCR output.
OCR_EXTRA_STOP_SEQUENCES: List[str] = [
    "<|user|>",
    "<|im_start|>",
    "<|im_end|>",
    "<|endoftext|>",
    "<|endofassistant|>",
]

VLM_LANGUAGE_PROMPT_KWARGS = ("mm_token_type_ids", "token_type_ids")

COHERE2_MOE_MODEL_TYPE = "cohere2_moe"
QWEN4_EXP_MODEL_TYPE = "qwen4_exp"
MINIMAX_M3_VL_MODEL_TYPE = "minimax_m3_vl"
MINIMAX_M3_MODEL_TYPES = {"minimax_m3", MINIMAX_M3_VL_MODEL_TYPE}
# mlx-vlm's Gemma4Processor.apply_chat_template renders the messages again. It
# moves audio markers to the last user turn and prints list system content.
TOKENIZER_CHAT_TEMPLATE_MODEL_TYPES = {"gemma4", "gemma4_unified", "diffusion_gemma"}

DIFFUSION_PREFILL_STEP_SIZE = 2048

# Per-model OCR generation defaults from official configs.
# Applied automatically when no explicit user override is provided.
OCR_MODEL_GENERATION_DEFAULTS: Dict[str, Dict[str, Any]] = {
    "glm_ocr": {
        "temperature": 0.0,
        "repetition_penalty": 1.1,
        "max_tokens": 4096,
    },
    "deepseekocr": {
        "temperature": 0.0,
        "max_tokens": 8192,
    },
    "deepseekocr_2": {
        "temperature": 0.0,
        "max_tokens": 8192,
    },
    "unlimited-ocr": {
        "temperature": 0.0,
        "max_tokens": 8192,
    },
    "dots_ocr": {
        "temperature": 0.0,
        "max_tokens": 8192,
    },
}


def _read_config_model_type(model_path: str | Path) -> str | None:
    config_path = Path(model_path) / "config.json"
    if not config_path.exists():
        return None
    try:
        data = json.loads(config_path.read_text())
    except Exception:
        return None
    model_type = data.get("model_type")
    return model_type if isinstance(model_type, str) else None


def _is_missing_chat_template_error(exc: ValueError) -> bool:
    """True when apply_chat_template failed because no chat template is set.

    transformers raises ValueError both for a missing template and for real
    render errors (e.g. continue_final_message combined with
    add_generation_prompt). Only the missing-template case may fall back to
    mlx-vlm's plain rendering; anything else must propagate. Both the
    tokenizer and the processor spellings must match: the vision render
    prefers the processor's apply_chat_template, the counting render uses
    the tokenizer.
    """
    message = str(exc)
    return (
        # tokenizer wording (transformers tokenizers; the pair mlx-vlm's
        # prompt_utils._missing_template_error also matches)
        "chat_template is not set" in message
        or "no template argument was passed" in message
        # processor wording (transformers ProcessorMixin; mlx-vlm processors
        # that render their own template, e.g. phi3_v)
        or "does not have a chat template" in message
        or "No chat template found" in message
    )


def _capture_vlm_position_state(lm: Any, extra_kwargs: dict[str, Any]) -> None:
    """Capture returned position metadata before consulting model-owned state.

    Materialize it on the executor stream to prevent cross-stream ANE prefill deadlocks.
    """
    pid = extra_kwargs.get("position_ids", getattr(lm, "_position_ids", None))
    if pid is not None and "position_ids" not in extra_kwargs:
        extra_kwargs["position_ids"] = pid
    rd = extra_kwargs.get("rope_deltas", getattr(lm, "_rope_deltas", None))
    if rd is not None:
        extra_kwargs["_captured_rope_deltas"] = rd
    lazy_state = [v for v in (pid, rd) if isinstance(v, mx.array)]
    if lazy_state:
        mx.eval(*lazy_state)


def _apply_minimax_m3_thinking_mode(
    model_type: str | None,
    template_kwargs: dict[str, Any],
) -> None:
    """Map oMLX enable_thinking to MiniMax M3's thinking_mode template kwarg."""
    if model_type not in MINIMAX_M3_MODEL_TYPES:
        return
    enable_thinking = template_kwargs.pop("enable_thinking", None)
    if "thinking_mode" in template_kwargs:
        return

    if enable_thinking is True:
        template_kwargs["thinking_mode"] = "enabled"
    elif enable_thinking is False:
        template_kwargs["thinking_mode"] = "disabled"


def _attach_vlm_tokenizer_runtime(tokenizer: Any, model_path: Path, eos_token_id: Any):
    from mlx_vlm.tokenizer_utils import load_tokenizer
    from mlx_vlm.utils import StoppingCriteria

    if getattr(tokenizer, "pad_token", None) is None:
        tokenizer.pad_token = getattr(tokenizer, "eos_token", None)

    detokenizer_class = load_tokenizer(model_path, return_tokenizer=False)
    tokenizer.detokenizer = detokenizer_class(tokenizer)

    final_eos_token_ids = (
        eos_token_id
        or getattr(tokenizer, "eos_token_ids", None)
        or getattr(tokenizer, "eos_token_id", None)
    )
    tokenizer.stopping_criteria = StoppingCriteria(final_eos_token_ids, tokenizer)
    return tokenizer


def _load_cohere2_moe_text_model(
    model_name: str,
    *,
    trust_remote_code: bool = False,
):
    """Load Cohere2 MoE through mlx-vlm with a tokenizer-only fallback."""
    from mlx_vlm.utils import get_model_path, load_model, load_processor
    from transformers import AutoTokenizer

    model_path = get_model_path(model_name)
    model = load_model(
        model_path,
        lazy=False,
        strict=True,
        trust_remote_code=trust_remote_code,
    )

    eos_token_id = getattr(getattr(model, "config", None), "eos_token_id", None)
    try:
        processor = load_processor(
            model_path,
            True,
            eos_token_ids=eos_token_id,
            trust_remote_code=trust_remote_code,
        )
    except Exception as exc:
        logger.debug(
            "mlx-vlm processor load failed for Cohere2 MoE %s; "
            "falling back to AutoTokenizer: %s",
            model_name,
            exc,
        )
        tokenizer = AutoTokenizer.from_pretrained(
            model_path,
            trust_remote_code=trust_remote_code,
        )
        processor = _attach_vlm_tokenizer_runtime(tokenizer, model_path, eos_token_id)

    return model, processor


_video_processor_patched = False


def _patch_video_processor_bug():
    """Prevent video_processor from crashing processor loading.

    Two interrelated issues without torchvision:

    1. Gemma4's video_preprocessor_config.json triggers AutoVideoProcessor
       which requires torchvision. Removing ``video_processor`` from the
       MODALITY mapping prevents transformers from attempting to load it.

    2. When mlx-vlm's custom processor patch fails, it falls back to HF's
       ProcessorMixin which passes ``video_processor`` as a kwarg. HF's
       own ProcessorMixin.__init__ rejects unexpected kwargs, so it is
       patched to silently drop ``video_processor``.
    """
    global _video_processor_patched
    if _video_processor_patched:
        return

    try:
        from transformers.processing_utils import MODALITY_TO_AUTOPROCESSOR_MAPPING

        mapping = MODALITY_TO_AUTOPROCESSOR_MAPPING._MAPPING_NAMES
        if "video_processor" in mapping:
            del mapping["video_processor"]
            logger.debug(
                "Removed video_processor from MODALITY_TO_AUTOPROCESSOR_MAPPING"
            )
    except (ImportError, AttributeError):
        pass

    try:
        from transformers.processing_utils import ProcessorMixin

        _orig_pm_init = ProcessorMixin.__init__

        def _pm_init_drop_video(self, *args, **kwargs):
            kwargs.pop("video_processor", None)
            return _orig_pm_init(self, *args, **kwargs)

        ProcessorMixin.__init__ = _pm_init_drop_video
    except (ImportError, AttributeError):
        pass

    _video_processor_patched = True


_torch_free_ip_patched = False


def _patch_torch_free_image_processor():
    """Route mlx-vlm OCR processors around torch-gated AutoImageProcessor.

    transformers 5.5+ ships ``AutoImageProcessor`` as a ``DummyObject`` that
    raises ``ImportError`` on attribute access without torch+torchvision
    installed. mlx-vlm's ``GlmOcrProcessor.from_pretrained`` and
    ``DotsOcrProcessor.from_pretrained`` call ``AutoImageProcessor.from_pretrained``
    directly, so they raise on oMLX's torch-free env.
    ``install_auto_processor_patch`` then silently swallows the ``ImportError``
    and falls back to a ``TokenizersBackend`` with no ``image_processor`` —
    image content is dropped at ``prepare_inputs()``. See #1131, #1175.

    transformers ships torch-free PIL-backend variants of these image
    processors (e.g. ``Glm46VImageProcessorPil``, ``Qwen2VLImageProcessorPil``).
    This patch wraps the affected mlx-vlm processors' ``from_pretrained`` so
    they substitute the PIL class when ``AutoImageProcessor`` raises.
    """
    global _torch_free_ip_patched
    if _torch_free_ip_patched:
        return

    try:
        import transformers
    except ImportError:
        return

    if not getattr(transformers.AutoImageProcessor, "is_dummy", False):
        # torch+torchvision available, AutoImageProcessor works as-is.
        _torch_free_ip_patched = True
        return

    for module_path, cls_name in (
        ("mlx_vlm.models.glm_ocr.processing", "GlmOcrProcessor"),
        ("mlx_vlm.models.dots_ocr.processing_dots_ocr", "DotsVLProcessor"),
    ):
        try:
            mod = importlib.import_module(module_path)
            cls = getattr(mod, cls_name)
            _wrap_from_pretrained_with_pil_image_processor(cls)
            logger.debug("Wrapped %s.from_pretrained with PIL fallback", cls_name)
        except (ImportError, AttributeError) as exc:
            logger.debug(
                "Skipped torch-free image processor patch for %s: %s",
                cls_name,
                exc,
            )

    _torch_free_ip_patched = True


def _wrap_from_pretrained_with_pil_image_processor(cls):
    """Wrap a ProcessorMixin subclass's ``from_pretrained`` so an ``ImportError``
    from ``AutoImageProcessor`` triggers PIL fallback instantiation."""
    if getattr(cls.from_pretrained, "_omlx_torch_free_patched", False):
        return

    orig = cls.from_pretrained

    @classmethod
    def patched(cls_inner, path, **kwargs):
        try:
            return orig(path, **kwargs)
        except ImportError as exc:
            msg = str(exc)
            if "Torchvision" not in msg and "PyTorch" not in msg:
                raise
            logger.info(
                "AutoImageProcessor unavailable (torch-free env); routing %s "
                "to PIL image processor",
                cls_inner.__name__,
            )
            return _build_processor_via_pil_image_processor(cls_inner, path, **kwargs)

    patched.__func__._omlx_torch_free_patched = True
    cls.from_pretrained = patched


def _build_processor_via_pil_image_processor(cls, path, **kwargs):
    """Construct a ProcessorMixin instance using transformers' PIL-backend
    image processor (looked up via ``IMAGE_PROCESSOR_MAPPING_NAMES``) instead
    of the torch-gated ``AutoImageProcessor``."""
    from transformers import AutoTokenizer
    from transformers.models.auto.auto_mappings import IMAGE_PROCESSOR_MAPPING_NAMES

    trust = kwargs.pop("trust_remote_code", True)

    # Look up image_processor_type from processor_config.json (preferred,
    # nested under "image_processor") or preprocessor_config.json (legacy).
    p = Path(path)
    ip_type = None
    for fname in ("processor_config.json", "preprocessor_config.json"):
        cfg_path = p / fname
        if not cfg_path.exists():
            continue
        with open(cfg_path) as f:
            cfg = json.load(f)
        ip_type = cfg.get("image_processor", {}).get("image_processor_type") or cfg.get(
            "image_processor_type"
        )
        if ip_type:
            break

    if not ip_type:
        raise ImportError(
            f"Cannot determine image_processor_type for {path}; install "
            "torch+torchvision or upgrade mlx-vlm."
        )

    # Read feature_extractor config if present (needed for audio models like gemma4_unified)
    fe_config = {}
    fe_type = None
    for fname in ("processor_config.json", "preprocessor_config.json"):
        cfg_path = p / fname
        if not cfg_path.exists():
            continue
        with open(cfg_path) as f:
            cfg = json.load(f)
        fe_section = cfg.get("feature_extractor", {})
        if isinstance(fe_section, dict):
            fe_config = dict(fe_section)
            fe_type = fe_config.pop("feature_extractor_type", None)
            if fe_type:
                break

    feature_extractor = None
    if fe_type:
        # Dynamically import the feature extractor class
        fe_cls = _resolve_feature_extractor_class(fe_type)
        if fe_cls is not None:
            try:
                feature_extractor = fe_cls(**fe_config)
                logger.debug("Created feature_extractor %s from %s", fe_type, path)
            except Exception as e:
                logger.warning("Failed to create feature_extractor %s: %s", fe_type, e)
    pil_cls = _resolve_pil_image_processor_class(ip_type, IMAGE_PROCESSOR_MAPPING_NAMES)
    if pil_cls is None:
        raise ImportError(
            f"No torch-free PIL image processor for image_processor_type={ip_type}."
        )

    image_processor = pil_cls.from_pretrained(str(path), trust_remote_code=trust)
    tokenizer = AutoTokenizer.from_pretrained(
        str(path), trust_remote_code=trust, **kwargs
    )

    # mlx-vlm helper: load chat_template.jinja into tokenizer if present.
    try:
        from mlx_vlm.models.base import load_chat_template

        load_chat_template(tokenizer, str(path))
    except (ImportError, AttributeError):
        pass

    processor_kwargs = {"image_processor": image_processor, "tokenizer": tokenizer}
    if feature_extractor is not None:
        processor_kwargs["feature_extractor"] = feature_extractor

    return cls(**processor_kwargs)


def _resolve_pil_image_processor_class(ip_type, mapping_names):
    """Find a non-dummy PIL backend class matching ``ip_type`` via
    ``IMAGE_PROCESSOR_MAPPING_NAMES``."""
    for model_type, mapping in mapping_names.items():
        if mapping.get("torchvision") != ip_type and mapping.get("pil") != ip_type:
            continue
        pil_name = mapping.get("pil")
        if not pil_name:
            continue
        module_name = (
            f"transformers.models.{model_type}.image_processing_pil_{model_type}"
        )
        try:
            mod = importlib.import_module(module_name)
        except ImportError:
            continue
        candidate = getattr(mod, pil_name, None)
        if candidate is not None and not getattr(candidate, "is_dummy", False):
            return candidate
    return None


# Mapping from feature_extractor_type to (module, class) locations in mlx_vlm
_FEATURE_EXTRACTOR_MAP = {
    "Gemma4UnifiedAudioFeatureExtractor": (
        "mlx_vlm.models.gemma4_unified.processing_gemma4_unified",
        "Gemma4UnifiedAudioFeatureExtractor",
    ),
    "Gemma4AudioFeatureExtractor": (
        "mlx_vlm.models.gemma4.audio_feature_extractor",
        "Gemma4AudioFeatureExtractor",
    ),
}


def _resolve_feature_extractor_class(fe_type: str):
    """Resolve a feature extractor class by its ``feature_extractor_type`` string.

    Returns the class object, or None if not found.
    """
    import importlib

    if fe_type in _FEATURE_EXTRACTOR_MAP:
        mod_name, cls_name = _FEATURE_EXTRACTOR_MAP[fe_type]
        try:
            mod = importlib.import_module(mod_name)
            return getattr(mod, cls_name, None)
        except ImportError:
            return None

    return None


def _fix_processor_none_pixels(processor):
    """Set sensible defaults when preprocessor_config.json has null pixels.

    Some Qwen3-VL model configs ship ``"max_pixels": null`` which overrides
    the constructor default and causes ``int > NoneType`` comparison errors
    in ``_smart_resize_image``.
    """
    ip = getattr(processor, "image_processor", None)
    if ip is None:
        return
    if getattr(ip, "max_pixels", None) is None and hasattr(ip, "max_pixels"):
        ip.max_pixels = 14 * 14 * 4 * 1280
        logger.debug("Fixed image_processor.max_pixels: None → %d", ip.max_pixels)
    if getattr(ip, "min_pixels", None) is None and hasattr(ip, "min_pixels"):
        ip.min_pixels = 56 * 56
        logger.debug("Fixed image_processor.min_pixels: None → %d", ip.min_pixels)


# Config keys to strip when audio_tower weights are missing but config still
# advertises audio support. See `_strip_audio_config_if_orphaned`.
_AUDIO_CONFIG_KEYS = (
    "audio_config",
    "audio_token_id",
    "boa_token_id",
    "eoa_token_id",
    "eoa_token_index",
)


def _resolve_optiq_vision_sidecar(model_dir: Path) -> Path | None:
    """Resolve the OptiQ multimodal sidecar declared by ``config.json``."""
    config_path = model_dir / "config.json"
    try:
        config = json.loads(config_path.read_text())
    except Exception:
        return None

    optiq_vision = config.get("optiq_vision")
    if not isinstance(optiq_vision, dict):
        return None
    relative_path = optiq_vision.get("sidecar")
    if not isinstance(relative_path, str) or not relative_path:
        return None

    model_root = model_dir.resolve()
    sidecar = (model_root / relative_path).resolve()
    try:
        sidecar.relative_to(model_root)
    except ValueError as exc:
        raise ValueError(
            f"OptiQ vision sidecar must stay inside the model directory: "
            f"{relative_path}"
        ) from exc
    if sidecar.suffix != ".safetensors":
        raise ValueError(f"OptiQ vision sidecar must be a safetensors file: {sidecar}")
    if not sidecar.is_file():
        raise FileNotFoundError(f"OptiQ vision sidecar not found: {sidecar}")
    return sidecar


_AUDIO_WEIGHT_PREFIXES = ("audio_tower.", "embed_audio.", "audio_encoder.")


def _is_audio_weight_key(key: str) -> bool:
    """True for audio encoder weights under MLX or HF naming.

    HF-named checkpoints keep a leading `model.` (`model.audio_tower.*`) that
    mlx-vlm's sanitize strips only at load time.
    """
    if key.startswith("model."):
        key = key[len("model.") :]
    return key.startswith(_AUDIO_WEIGHT_PREFIXES)


def _has_audio_weights(model_dir: Path) -> bool:
    """Return True iff checkpoint shards or MiMo's sidecar contain audio weights."""
    import safetensors

    weight_files = list(model_dir.glob("*.safetensors"))
    sidecar = _resolve_optiq_vision_sidecar(model_dir)
    if sidecar is not None and all(sf.resolve() != sidecar for sf in weight_files):
        weight_files.append(sidecar)
    mimo_audio_sidecar = model_dir / "omnimodal" / "audio_encoder.safetensors"
    if mimo_audio_sidecar.is_file():
        weight_files.append(mimo_audio_sidecar)

    for sf in weight_files:
        try:
            with safetensors.safe_open(str(sf), framework="np") as f:
                # safetensors.safe_open exposes keys() but is not a Mapping.
                if any(_is_audio_weight_key(k) for k in f.keys()):  # noqa: SIM118
                    return True
        except Exception:
            # Corrupt or unreadable shard — treat as no audio info, let
            # downstream loader produce its own error.
            return False
    return False


# Text-only oQ checkpoints can retain embed_vision without a vision tower.
_VISION_TOWER_MARKER = "vision_tower"
_VISION_TENSOR_MARKERS = (_VISION_TOWER_MARKER, "embed_vision")
# Moondream stores its tower under `vision` (current, converted) or
# `vision_encoder` (2024 revisions) and never declares a vision_config.
_VISION_TOWER_PATH_MARKERS = (_VISION_TOWER_MARKER, "vision", "vision_encoder")


def _is_vision_tensor_key(key: str) -> bool:
    """True for parameter paths under `vision_tower` / `embed_vision`."""
    return any(marker in key.split(".") for marker in _VISION_TENSOR_MARKERS)


def _is_vision_tower_key(key: str) -> bool:
    """Exclude orphan projection weights when detecting a vision tower."""
    return any(marker in key.split(".") for marker in _VISION_TOWER_PATH_MARKERS)


def _has_vision_tower_weights(model_dir: Path) -> bool:
    """Return True iff any safetensors shard contains vision_tower weights."""
    import safetensors

    weight_files = list(model_dir.glob("*.safetensors"))
    sidecar = _resolve_optiq_vision_sidecar(model_dir)
    if sidecar is not None and all(sf.resolve() != sidecar for sf in weight_files):
        weight_files.append(sidecar)

    for sf in weight_files:
        with safetensors.safe_open(str(sf), framework="np") as f:
            # safetensors.safe_open exposes keys() but is not a Mapping.
            if any(_is_vision_tower_key(k) for k in f.keys()):  # noqa: SIM118
                return True
    return False


def _vision_config_is_orphaned(model_dir: Path) -> bool:
    """Require absent vision config and readable shards without a tower."""
    try:
        raw = json.loads((model_dir / "config.json").read_text())
    except Exception:
        return False
    if raw.get("vision_config"):
        return False
    try:
        return not _has_vision_tower_weights(model_dir)
    except Exception:
        return False


@contextlib.contextmanager
def _strip_audio_config_if_orphaned(model_dir: Path):
    """Drop `audio_config` from `mlx_vlm.utils.load_config` results when the
    safetensors shards lack audio_tower / embed_audio weights.

    Some quantization tooling (notably oMLX's pre-fix oQ pipeline) writes
    multimodal Gemma 4 checkpoints without audio weights but leaves
    `audio_config` in `config.json`. mlx-vlm then instantiates `AudioEncoder`
    and `model.load_weights(strict=True)` fails with "Missing 752 parameters".

    This wrap is scoped to a single `mlx_vlm.utils.load(...)` call: it swaps
    `load_config` on entry and restores it on exit. Other code paths that
    read config (model_discovery, admin UI) bypass mlx-vlm entirely so they
    are unaffected.
    """
    import mlx_vlm.utils as _vu

    original = _vu.load_config
    warned = set()

    def _patched(path, **kwargs):
        cfg = original(path, **kwargs)

        from ..utils.model_loading import expand_per_layer_quant_keys

        expand_per_layer_quant_keys(cfg)

        if cfg.get("audio_config") is None:
            return cfg
        try:
            p = Path(path) if not isinstance(path, Path) else path
            if not p.is_dir():
                return cfg
            if _has_audio_weights(p):
                return cfg
        except Exception:
            return cfg
        cfg = dict(cfg)
        # Explicit None instead of pop: mlx-vlm's load_model runs
        # `config.setdefault("audio_config", {})` which would otherwise
        # repopulate audio_config with `{}` and cause AudioEncoder to be
        # instantiated with default values.
        cfg["audio_config"] = None
        for k in _AUDIO_CONFIG_KEYS:
            if k != "audio_config":
                cfg.pop(k, None)
        if str(p) not in warned:
            warned.add(str(p))
            logger.warning(
                "audio_tower weights missing for %s; loading without audio support",
                p.name,
            )
        return cfg

    _vu.load_config = _patched
    try:
        yield
    finally:
        _vu.load_config = original


@contextlib.contextmanager
def _strip_vision_config_if_orphaned(model_dir: Path):
    """Suppress inferred vision modules and orphan weights for text-only loads."""
    if not _vision_config_is_orphaned(model_dir):
        yield
        return

    import mlx.nn as _nn
    import mlx_vlm.utils as _vu

    original_load_config = _vu.load_config
    original_update_module_configs = _vu.update_module_configs
    original_load_weights = _nn.Module.load_weights
    warned = False

    def _text_capable_load_config(path, **kwargs):
        cfg = original_load_config(path, **kwargs)
        # mlx-vlm's gemma4 builds a vision tower unconditionally; gemma4_unified
        # shares its language_model layout and runs without one.
        if cfg.get("model_type") == "gemma4":
            cfg = {**cfg, "model_type": "gemma4_unified"}
        return cfg

    def _patched_update_module_configs(model_config, model_class, config, modules):
        model_config = original_update_module_configs(
            model_config, model_class, config, modules
        )
        # Clear the deserialized config; the raw dict must stay valid for quantization.
        if hasattr(model_config, "vision_config") and not config.get(
            "vision_config"
        ):
            model_config.vision_config = None
        return model_config

    def _vision_filtering_load_weights(self, weights_items, *args, **kwargs):
        nonlocal warned
        if isinstance(weights_items, str):
            return original_load_weights(self, weights_items, *args, **kwargs)

        # MLX-format checkpoints skip upstream sanitize, leaving orphan projections.
        owned = {k for k, _ in _nn.utils.tree_flatten(self.parameters())}
        kept = []
        dropped = 0
        for key, value in weights_items:
            if _is_vision_tensor_key(key) and key not in owned:
                dropped += 1
                continue
            kept.append((key, value))
        if dropped and not warned:
            warned = True
            logger.warning(
                "vision_tower weights missing for %s; loading without "
                "vision support",
                model_dir.name,
            )
        return original_load_weights(self, kept, *args, **kwargs)

    _vu.load_config = _text_capable_load_config
    _vu.update_module_configs = _patched_update_module_configs
    _nn.Module.load_weights = _vision_filtering_load_weights
    try:
        yield
    finally:
        _vu.load_config = original_load_config
        _vu.update_module_configs = original_update_module_configs
        _nn.Module.load_weights = original_load_weights


@contextlib.contextmanager
def _load_optiq_vision_sidecar_on_load(model_dir: Path):
    """Include a config-declared OptiQ sidecar in mlx-vlm strict loading.

    Pinned mlx-vlm only globs ``*.safetensors`` in the model root, while
    current OptiQ VLM checkpoints keep their unquantized multimodal weights
    under ``optiq/``. Root-level legacy sidecars remain a no-op because the
    native glob already loads them.
    """
    sidecar = _resolve_optiq_vision_sidecar(model_dir)
    if sidecar is None or sidecar.parent == model_dir.resolve():
        yield
        return

    sidecar_weights = mx.load(str(sidecar))
    if not isinstance(sidecar_weights, dict):
        raise ValueError(f"OptiQ vision sidecar must contain named tensors: {sidecar}")

    import mlx.nn as _nn

    original_load_weights = _nn.Module.load_weights
    injected = False

    def _patched_load_weights(self, weights_items, *args, **kwargs):
        nonlocal injected
        if injected or isinstance(weights_items, str):
            return original_load_weights(self, weights_items, *args, **kwargs)

        model_weights = list(weights_items)
        model_keys = {
            item[0]
            for item in model_weights
            if isinstance(item, (tuple, list))
            and len(item) >= 2
            and isinstance(item[0], str)
        }
        duplicates = model_keys.intersection(sidecar_weights)
        if duplicates:
            sample = ", ".join(sorted(duplicates)[:3])
            raise ValueError(
                f"OptiQ vision sidecar duplicates model weights: {sample}"
            )

        injected = True
        result = original_load_weights(
            self,
            [*model_weights, *sidecar_weights.items()],
            *args,
            **kwargs,
        )
        logger.info(
            "Loaded %d OptiQ multimodal sidecar weights from %s",
            len(sidecar_weights),
            sidecar,
        )
        return result

    _nn.Module.load_weights = _patched_load_weights
    try:
        yield
    finally:
        _nn.Module.load_weights = original_load_weights


def _is_mlx_format_safetensors_dir(model_dir: Path) -> bool:
    """Return True when the first safetensors shard declares ``format=mlx``."""
    import safetensors

    try:
        weight_files = sorted(
            sf
            for sf in model_dir.glob("*.safetensors")
            if not sf.name.endswith("consolidated.safetensors")
        )
    except Exception:
        return False
    if not weight_files:
        return False

    try:
        with safetensors.safe_open(str(weight_files[0]), framework="np") as f:
            metadata = f.metadata()
    except Exception:
        return False
    return isinstance(metadata, dict) and metadata.get("format") == "mlx"


def _gemma4_global_kv_from_per_layer_config(config: dict) -> dict[str, int]:
    """Derive Gemma4's legacy full-attention head fields from ``per_layer_config``.

    Newer Gemma4 checkpoints (Transformers >= 5.15) record the ``head_dim`` /
    ``num_key_value_heads`` overrides of the full-attention layers under
    ``text_config.per_layer_config`` instead of the legacy global
    ``global_head_dim`` / ``num_global_key_value_heads`` fields. The pinned
    mlx-vlm Gemma4 loader reads only the legacy fields, so it sizes the
    full-attention K/V projections with the sliding-window head count and
    ``load_weights`` fails with a shape mismatch (#3537).

    Returns the legacy fields that are absent from ``text_config`` and can be
    derived unambiguously (every overridden full-attention layer agrees), or
    an empty dict.
    """
    text_config = config.get("text_config")
    if not isinstance(text_config, dict):
        return {}
    model_type = str(text_config.get("model_type") or config.get("model_type") or "")
    if not model_type.startswith("gemma4"):
        return {}
    per_layer = text_config.get("per_layer_config")
    if not isinstance(per_layer, dict) or not per_layer:
        return {}
    layer_types = text_config.get("layer_types")
    if not isinstance(layer_types, list):
        layer_types = None

    derived: dict[str, int] = {}
    for legacy_key, layer_key in (
        ("global_head_dim", "head_dim"),
        ("num_global_key_value_heads", "num_key_value_heads"),
    ):
        if text_config.get(legacy_key) is not None:
            continue
        values: set[int] = set()
        for layer_id, overrides in per_layer.items():
            if not isinstance(overrides, dict) or overrides.get(layer_key) is None:
                continue
            if layer_types is not None:
                try:
                    layer_idx = int(layer_id)
                except (TypeError, ValueError):
                    continue
                if not 0 <= layer_idx < len(layer_types):
                    continue
                if layer_types[layer_idx] != "full_attention":
                    continue
            try:
                values.add(int(overrides[layer_key]))
            except (TypeError, ValueError):
                continue
        if len(values) == 1:
            derived[legacy_key] = values.pop()
    return derived


@contextlib.contextmanager
def _derive_gemma4_global_kv_on_load(model_dir: Path):
    """Feed ``per_layer_config``-only Gemma4 head overrides to the mlx-vlm loader.

    Wraps ``mlx_vlm.utils.load_config`` for one ``vlm_load(...)`` so the
    config handed to the Gemma4 ``TextConfig`` carries ``global_head_dim`` /
    ``num_global_key_value_heads`` derived from ``per_layer_config`` when the
    checkpoint does not spell them out (#3537). Checkpoints that already
    carry the legacy fields, and non-Gemma4 models, are untouched.
    """
    config_path = model_dir / "config.json"
    try:
        config = json.loads(config_path.read_text())
    except Exception:
        yield
        return
    derived = (
        _gemma4_global_kv_from_per_layer_config(config)
        if isinstance(config, dict)
        else {}
    )
    if not derived:
        yield
        return

    import mlx_vlm.utils as _vu

    original_load_config = _vu.load_config

    def _patched_load_config(model_path, **kwargs):
        loaded = original_load_config(model_path, **kwargs)
        text_config = loaded.get("text_config") if isinstance(loaded, dict) else None
        if isinstance(text_config, dict):
            for key, value in derived.items():
                if text_config.get(key) is None:
                    text_config[key] = value
        return loaded

    logger.info(
        "derive_gemma4_global_kv_on_load: per_layer_config -> %s",
        ", ".join(f"{k}={v}" for k, v in derived.items()),
    )
    _vu.load_config = _patched_load_config
    try:
        yield
    finally:
        _vu.load_config = original_load_config


@contextlib.contextmanager
def _drop_gemma4_mlx_shared_kv_extras_on_load(model_dir: Path):
    """Drop Gemma4 shared-KV extra weights for MLX-format VLM checkpoints.

    mlx-vlm skips model sanitizers when safetensors metadata is ``format=mlx``.
    Gemma4 E2B/E4B MLX checkpoints still ship K/V tensors for shared-KV layers,
    but the mlx-vlm model tree intentionally omits those modules.  If the
    extras reach strict ``load_weights``, VLM loading fails and oMLX falls back
    to text-only LLM.  Scope the fix to Gemma4 MLX-format models whose config
    declares shared-KV layers; 26B/31B Gemma4 models have zero shared-KV layers
    and remain no-op.
    """
    config_path = model_dir / "config.json"
    try:
        config = json.loads(config_path.read_text())
    except Exception:
        yield
        return

    text_config = config.get("text_config")
    if not isinstance(text_config, dict):
        yield
        return

    if config.get("model_type") != "gemma4":
        yield
        return
    if text_config.get("model_type") != "gemma4_text":
        yield
        return
    if not _is_mlx_format_safetensors_dir(model_dir):
        yield
        return

    try:
        num_layers = int(text_config.get("num_hidden_layers") or 0)
        num_shared = int(text_config.get("num_kv_shared_layers") or 0)
    except (TypeError, ValueError):
        yield
        return
    if num_layers <= 0 or num_shared <= 0 or num_shared >= num_layers:
        yield
        return

    first_shared = num_layers - num_shared
    drop_modules = {"k_proj", "v_proj", "k_norm", "v_norm"}
    layer_prefix = "language_model.model.layers."

    def _is_shared_kv_extra(key: str) -> bool:
        if not key.startswith(layer_prefix):
            return False
        parts = key[len(layer_prefix) :].split(".")
        if len(parts) < 4 or parts[1] != "self_attn":
            return False
        try:
            layer_idx = int(parts[0])
        except ValueError:
            return False
        return first_shared <= layer_idx < num_layers and parts[2] in drop_modules

    import mlx.nn as _nn

    original_load_weights = _nn.Module.load_weights
    dropped = 0

    def _patched_load_weights(self, weights_items, *args, **kwargs):
        nonlocal dropped
        if isinstance(weights_items, str):
            return original_load_weights(self, weights_items, *args, **kwargs)

        filtered = []
        local_dropped = 0
        for item in weights_items:
            if (
                isinstance(item, (tuple, list))
                and len(item) >= 2
                and isinstance(item[0], str)
                and _is_shared_kv_extra(item[0])
            ):
                local_dropped += 1
                continue
            filtered.append(item)
        dropped += local_dropped
        return original_load_weights(self, filtered, *args, **kwargs)

    _nn.Module.load_weights = _patched_load_weights
    try:
        yield
    finally:
        _nn.Module.load_weights = original_load_weights
        if dropped:
            logger.info(
                "Dropped %d Gemma4 shared-KV extra weights for MLX-format "
                "checkpoint %s",
                dropped,
                model_dir.name,
            )


@contextlib.contextmanager
def _transpose_qwen35_mlx_vision_patch_embed_on_load(model_dir: Path):
    """Fix channels-first Qwen3.5 vision weights in MLX checkpoints.

    mlx-vlm skips model sanitizers when safetensors metadata declares
    ``format=mlx``. Some converted Qwen3.5/3.6 checkpoints retain the PyTorch
    Conv3d layout ``(out, in, time, height, width)`` for the vision patch
    embedding, while MLX expects ``(out, time, height, width, in)``. Correct
    only that unambiguously channels-first tensor during loading.
    """
    if _read_config_model_type(model_dir) not in {"qwen3_5", "qwen3_5_moe"}:
        yield
        return
    if not _is_mlx_format_safetensors_dir(model_dir):
        yield
        return

    import mlx_vlm.utils as _vu

    original_load_safetensors = _vu._load_safetensors
    transposed = 0

    def _patched_load_safetensors(path):
        nonlocal transposed
        weights = original_load_safetensors(path)
        key = "vision_tower.patch_embed.proj.weight"
        value = weights.get(key)
        if (
            value is not None
            and getattr(value, "ndim", None) == 5
            and value.shape[1] == 3
            and value.shape[-1] != 3
        ):
            weights[key] = value.transpose(0, 2, 3, 4, 1)
            transposed += 1
        return weights

    _vu._load_safetensors = _patched_load_safetensors
    try:
        yield
    finally:
        _vu._load_safetensors = original_load_safetensors
        if transposed:
            logger.info(
                "Transposed Qwen3.5 vision patch embedding to MLX Conv3d "
                "layout for %s",
                model_dir.name,
            )


_NESTED_VIS_PREFIX = "language_model.model.visual."
_VISION_TOWER_PREFIX = "vision_tower."


def _should_pack_minimax_m3_shared_expert(args: Any) -> bool:
    """Resolve the explicit MiniMax shared-expert layout override."""
    configured = getattr(args, "pack_shared_expert", None)
    if configured is not None:
        return bool(configured)
    return bool(
        args.n_shared_experts == 1
        and args.shared_intermediate_size == args.intermediate_size
    )


def _model_shard_matcher(model_dir: Path):
    """Return a predicate for safetensors shards directly under *model_dir*."""
    target_dir = model_dir.resolve()

    def matches(filename: object) -> bool:
        try:
            path = Path(filename)
            return path.suffix == ".safetensors" and path.parent.resolve() == target_dir
        except (TypeError, OSError, RuntimeError):
            return False

    return matches


@contextlib.contextmanager
def _force_minimax_m3_moe_sanitize_on_load(model_dir: Path):
    """Force mlx-vlm's MiniMax M3 MoE sanitize path for MLX-format checkpoints.

    mlx-vlm's MiniMax M3 loader can pack ``shared_experts`` into the routed
    ``switch_mlp`` when ``Model.sanitize`` runs. MLX-format checkpoints skip
    sanitize upstream, while MiniMax checkpoints can carry either packed or
    explicitly unpacked mixed-bit MoE weights. Hide only the safetensors
    ``format=mlx`` metadata during this load so the configured sanitize path
    runs before quantization and load_weights.
    """
    if _read_config_model_type(model_dir) != MINIMAX_M3_VL_MODEL_TYPE:
        yield
        return

    from ..patches.mlx_vlm_minimax_m3_compat import (
        apply_mlx_vlm_minimax_m3_compat_patch,
    )

    apply_mlx_vlm_minimax_m3_compat_patch()

    import safetensors
    from mlx_vlm.models.minimax_m3_vl import minimax_m3_vl as _minimax_m3_vl

    original_safe_open = safetensors.safe_open
    original_sanitize_moe_weights = _minimax_m3_vl._sanitize_moe_weights
    is_target_shard = _model_shard_matcher(model_dir)

    class _SafeOpenMetadataWrapper:
        def __init__(self, inner):
            self._inner = inner

        def __enter__(self):
            self._inner.__enter__()
            return self

        def __exit__(self, *args):
            return self._inner.__exit__(*args)

        def __getattr__(self, name):
            return getattr(self._inner, name)

        def metadata(self):
            metadata = self._inner.metadata()
            if isinstance(metadata, dict) and metadata.get("format") == "mlx":
                metadata = dict(metadata)
                metadata.pop("format", None)
            return metadata

    def _patched_safe_open(filename, *args, **kwargs):
        handle = original_safe_open(filename, *args, **kwargs)
        if is_target_shard(filename):
            return _SafeOpenMetadataWrapper(handle)
        return handle

    def _pack_mlx_unpacked_moe_weights(weights: dict, args: Any) -> int:
        if not _should_pack_minimax_m3_shared_expert(args):
            return 0

        packed = 0
        for layer_idx in range(args.num_hidden_layers):
            prefix = f"language_model.model.layers.{layer_idx}.block_sparse_moe"
            for suffix in ("weight", "scales", "biases", "bias"):
                gate_key = f"{prefix}.switch_mlp.gate_proj.{suffix}"
                up_key = f"{prefix}.switch_mlp.up_proj.{suffix}"
                shared_gate_key = f"{prefix}.shared_experts.gate_proj.{suffix}"
                shared_up_key = f"{prefix}.shared_experts.up_proj.{suffix}"
                gate_up_key = f"{prefix}.switch_mlp.gate_up_proj.{suffix}"

                if (
                    gate_up_key not in weights
                    and gate_key in weights
                    and up_key in weights
                    and shared_gate_key in weights
                    and shared_up_key in weights
                ):
                    gate = weights.pop(gate_key)
                    up = weights.pop(up_key)
                    shared_gate = weights.pop(shared_gate_key)
                    shared_up = weights.pop(shared_up_key)
                    routed_gate_up = mx.concatenate([gate, up], axis=1)
                    shared_gate_up = mx.expand_dims(
                        mx.concatenate([shared_gate, shared_up], axis=0), axis=0
                    )
                    weights[gate_up_key] = mx.concatenate(
                        [routed_gate_up, shared_gate_up], axis=0
                    )
                    packed += 1

                down_key = f"{prefix}.switch_mlp.down_proj.{suffix}"
                shared_down_key = f"{prefix}.shared_experts.down_proj.{suffix}"
                packed_down_key = f"{prefix}.switch_mlp.down_proj.{suffix}"
                if down_key in weights and shared_down_key in weights:
                    down = weights.pop(down_key)
                    shared_down = mx.expand_dims(weights.pop(shared_down_key), axis=0)
                    weights[packed_down_key] = mx.concatenate(
                        [down, shared_down], axis=0
                    )
                    packed += 1
        return packed

    def _patched_sanitize_moe_weights(weights: dict, args: Any) -> None:
        original_sanitize_moe_weights(weights, args)
        packed = _pack_mlx_unpacked_moe_weights(weights, args)
        if packed:
            logger.info(
                "MiniMax M3 MLX-format MoE sanitize packed %d tensor groups",
                packed,
            )

    safetensors.safe_open = _patched_safe_open
    _minimax_m3_vl._sanitize_moe_weights = _patched_sanitize_moe_weights
    try:
        logger.info(
            "MiniMax M3 MLX-format MoE sanitize patch active for %s",
            model_dir.name,
        )
        yield
    finally:
        safetensors.safe_open = original_safe_open
        _minimax_m3_vl._sanitize_moe_weights = original_sanitize_moe_weights


@contextlib.contextmanager
def _force_qwen4_exp_sanitize_on_load(model_dir: Path):
    """Run vendored-model key sanitization before quantization selection.

    Converted MLX checkpoints legitimately declare ``format=mlx``, but the
    published Qwen4 layout still uses ``model.language_model.*`` and
    ``model.visual.*`` names.  mlx-vlm normally skips ``Model.sanitize`` for
    MLX-format files, which makes its quantization predicate inspect the wrong
    names and leaves every ``scales``/``biases`` tensor unmatched.  Hide only
    the format marker for this model and this load so sanitization happens at
    the point expected by the upstream loader: before ``nn.quantize``.
    """
    model_type = _read_config_model_type(model_dir)
    if model_type not in ("qwen4_exp", "glm5_next"):
        yield
        return

    import safetensors

    from ..patches.mlx_vlm_qwen4_exp_compat.ple_load_resources import ple_load_resources

    original_safe_open = safetensors.safe_open
    is_target_shard = _model_shard_matcher(model_dir)

    class _SafeOpenMetadataWrapper:
        def __init__(self, inner):
            self._inner = inner

        def __enter__(self):
            self._inner.__enter__()
            return self

        def __exit__(self, *args):
            return self._inner.__exit__(*args)

        def __getattr__(self, name):
            return getattr(self._inner, name)

        def metadata(self):
            metadata = self._inner.metadata()
            if isinstance(metadata, dict) and metadata.get("format") == "mlx":
                metadata = dict(metadata)
                metadata.pop("format", None)
            return metadata

    def _patched_safe_open(filename, *args, **kwargs):
        handle = original_safe_open(filename, *args, **kwargs)
        if is_target_shard(filename):
            return _SafeOpenMetadataWrapper(handle)
        return handle

    safetensors.safe_open = _patched_safe_open
    try:
        logger.info(
            "%s pre-quantization sanitize active for %s",
            model_type,
            model_dir.name,
        )
        with ple_load_resources():
            yield
    finally:
        safetensors.safe_open = original_safe_open


@contextlib.contextmanager
def _remap_nested_visual_on_load(model_dir: Path):
    """Remap ``language_model.model.visual.*`` → ``vision_tower.*`` during
    ``load_model`` for MLX-format models where sanitize is skipped.

    mlx-vlm's ``load_model`` skips ``Model.sanitize`` when the safetensors
    metadata declares ``format=mlx``. oQ output is MLX-format, so the
    nested-visual key fixup that sanitize normally applies never fires.
    This context manager wraps ``load_model`` to intercept the weight dict
    and perform the remap before ``nn.Module.load_weights`` is called.

    Scoped to a single ``vlm_load(...)`` call.
    """
    import mlx_vlm.utils as _vu

    original_load_model = _vu.load_model

    def _patched_load_model(model_path, lazy=False, **kwargs):
        import mlx.nn as _nn

        orig_load_weights = _nn.Module.load_weights

        def _remapping_load_weights(self, weights_items, *args, **kw):
            if isinstance(weights_items, str):
                return orig_load_weights(self, weights_items, *args, **kw)
            remapped = []
            n = 0
            for k, v in weights_items:
                if k.startswith(_NESTED_VIS_PREFIX):
                    k = _VISION_TOWER_PREFIX + k[len(_NESTED_VIS_PREFIX) :]
                    n += 1
                remapped.append((k, v))
            if n:
                logger.info(
                    "remap_nested_visual_on_load: remapped %d keys "
                    "'language_model.model.visual.*' -> 'vision_tower.*'",
                    n,
                )
            return orig_load_weights(self, remapped, *args, **kw)

        _nn.Module.load_weights = _remapping_load_weights
        try:
            return original_load_model(model_path, lazy, **kwargs)
        finally:
            _nn.Module.load_weights = orig_load_weights

    _vu.load_model = _patched_load_model
    try:
        yield
    finally:
        _vu.load_model = original_load_model


# Models that only support a single image per request
SINGLE_IMAGE_ONLY_MODELS = {
    "llava_next",
    "llava-qwen2",
    "bunny-llama",
    "paligemma",
    "multi_modality",
    "mllama",
}


def _uses_mrope(vlm_model) -> bool:
    """Check if the VLM model uses multi-dimensional RoPE (mRoPE).

    mRoPE models use 3D position IDs (temporal/height/width) that are
    incompatible with the mlx-lm decode model's standard 1D RoPE.
    """
    config = getattr(vlm_model, "config", None)
    if config is None:
        return False
    text_config = getattr(config, "text_config", None)
    if text_config is None:
        return False
    rope_cfg = getattr(text_config, "rope_scaling", None) or getattr(
        text_config, "rope_parameters", None
    )
    if isinstance(rope_cfg, dict):
        return "mrope_section" in rope_cfg
    return False


# Qwen-style VLMs: vision_tower takes (pixel_values, grid_thw).
_QWEN_VISION_MODELS = {
    "qwen3_5",
    "qwen3_5_moe",
    "prism_hadamard_qwen35",  # Ternary Bonsai 2 keeps the Qwen3.5 vision tower.
    "qwen3_vl",
    "qwen3_vl_moe",
    "qwen4_exp",  # Reuses the Qwen3.5/Qwen3-VL vision tower verbatim.
    "qwen2_vl",
    "qwen2_5_vl",
    "mimo_v2",
    "mimo_v2_flash",
}

# Grid-based VLMs whose flat vision features can be split with grid_thw.
_GRID_VISION_MODELS = _QWEN_VISION_MODELS | {"glm5_next"}

# Model types eligible for the processor-free cached-input path (see
# ``VLMBatchedEngine._try_build_cached_vision_inputs``). mimo variants are
# excluded: their audio-aware pipeline owns input assembly.
_CACHED_INPUT_FAST_PREPARE_MODEL_TYPES = _QWEN_VISION_MODELS - {
    "mimo_v2",
    "mimo_v2_flash",
}


def _grid_row(image_grid_thw: Any, i: int) -> Optional[List[int]]:
    """Row ``i`` of an ``image_grid_thw`` tensor as ``[t, h, w]``, or None."""
    try:
        row = [int(v) for v in image_grid_thw[i]]
    except (IndexError, TypeError, ValueError):
        return None
    return row if len(row) == 3 else None


def _grid_image_token_starts(
    token_ids: list[int], image_grid_thw: Any, image_token_id: int, merge_size: int
) -> list[int]:
    """Locate each grid image in the final, expanded processor token sequence."""
    grid = (
        image_grid_thw.tolist()
        if hasattr(image_grid_thw, "tolist")
        else image_grid_thw
    )
    counts = []
    for t, h, w in grid:
        patches = int(t) * int(h) * int(w)
        if merge_size <= 0 or patches <= 0 or patches % (merge_size**2):
            raise ValueError("Invalid image grid for cache boundaries")
        counts.append(patches // (merge_size**2))
    positions = [i for i, token in enumerate(token_ids) if token == image_token_id]
    if sum(counts) != len(positions):
        raise ValueError("Image grids do not match the final image tokens")
    starts = []
    offset = 0
    for count in counts:
        start = positions[offset]
        if positions[offset + count - 1] != start + count - 1:
            raise ValueError("Image token span is not contiguous")
        starts.append(start)
        offset += count
    return starts


# Conservative fallback upper bound on image-placeholder tokens per image
# content part. Used by ``preflight_chat`` only when the actual
# ``max_pixels`` cannot be derived from the loaded processor config.
# Qwen-VL / Gemma-Vision typically expand each image to 256–1280 tokens
# at default settings, but a deployment that lifts ``max_pixels`` can
# legitimately exceed this — relying on a hard-coded 1280 in that case
# silently under-counts and re-opens the panic-prone MLX prefill path.
# Prefer ``_derive_image_token_upper_bound(processor)`` when the
# processor is loaded.
_IMAGE_TOKEN_UPPER_BOUND_FALLBACK = 1280


def _derive_image_token_upper_bound(processor: Any) -> int:
    """Derive the per-image token upper bound from the processor config.

    GLM processors expose the final ``max_image_tokens`` bound directly.
    Qwen-style image processors expose ``max_pixels`` (an *area*) and pack
    pixels into ``patch_size`` × ``patch_size`` patches, then merge
    ``merge_size`` × ``merge_size`` patches into one model token. Their
    per-image token bound is therefore::

        max_tokens = max_pixels / (patch_size**2 * merge_size**2)

    Falls back to the conservative module-level constant when the
    processor doesn't expose the expected attributes (other model
    families) so we never *under*-count.
    """
    if processor is None:
        return _IMAGE_TOKEN_UPPER_BOUND_FALLBACK
    ip = getattr(processor, "image_processor", None) or processor
    max_image_tokens = getattr(ip, "max_image_tokens", None)
    if isinstance(max_image_tokens, int) and max_image_tokens > 0:
        return max(max_image_tokens, _IMAGE_TOKEN_UPPER_BOUND_FALLBACK)
    max_pixels = getattr(ip, "max_pixels", None)
    patch_size = getattr(ip, "patch_size", None)
    merge_size = getattr(ip, "merge_size", None)
    if (
        isinstance(max_pixels, int)
        and max_pixels > 0
        and isinstance(patch_size, int)
        and patch_size > 0
        and isinstance(merge_size, int)
        and merge_size > 0
    ):
        derived = max_pixels // (patch_size * patch_size * merge_size * merge_size)
        # Never go *below* the conservative fallback — a model whose
        # processor reports a tiny max_pixels (e.g. test fixtures) should
        # not weaken the guard.
        return max(derived, _IMAGE_TOKEN_UPPER_BOUND_FALLBACK)
    return _IMAGE_TOKEN_UPPER_BOUND_FALLBACK


def _count_image_tokens(
    messages: list[dict[str, Any]],
    per_image_upper_bound: int = _IMAGE_TOKEN_UPPER_BOUND_FALLBACK,
) -> int:
    """Count image-bearing content parts in OpenAI-style messages and
    return the conservative token-budget contribution.

    Supports both the OpenAI ``image_url`` / ``image`` part types and the
    Anthropic ``image`` block shape that gets adapted into the same
    message-list before reaching the engine layer.
    """
    image_parts = 0
    for msg in messages:
        content = msg.get("content")
        if not isinstance(content, list):
            continue
        for part in content:
            if not isinstance(part, dict):
                continue
            ptype = part.get("type")
            if ptype in ("image_url", "image", "input_image"):
                image_parts += 1
    return image_parts * per_image_upper_bound


def _smart_resize_tokens(
    h: int, w: int, patch_size: int, merge_size: int,
    min_pixels: int, max_pixels: int,
) -> int:
    """Real merged-token count for one image of pixel size (h, w), mirroring
    the Qwen image processor's ``smart_resize`` -> grid_thw ->
    ``(t*h*w)//merge**2`` pipeline (t=1 for a still image). Pure arithmetic;
    no pixel decode. This is the *exact* count the real chat path produces, so
    it never under-counts the prefill-memory guard."""
    import math

    factor = patch_size * merge_size
    if h <= 0 or w <= 0:
        return 0
    h_bar = round(h / factor) * factor
    w_bar = round(w / factor) * factor
    if h_bar * w_bar > max_pixels:
        beta = math.sqrt((h * w) / max_pixels)
        h_bar = max(factor, math.floor(h / beta / factor) * factor)
        w_bar = max(factor, math.floor(w / beta / factor) * factor)
    elif h_bar * w_bar < min_pixels:
        beta = math.sqrt(min_pixels / (h * w))
        h_bar = math.ceil(h * beta / factor) * factor
        w_bar = math.ceil(w * beta / factor) * factor
    return (h_bar // patch_size) * (w_bar // patch_size) // (merge_size ** 2)


def _read_image_dims(part: dict) -> Optional[tuple]:
    """Best-effort, decode-free ``(width, height)`` for an OpenAI image part.

    Handles only ``data:`` base64 URIs. Returns ``None`` for anything else so
    callers fall back to the conservative per-image upper bound without opening
    request-supplied paths or fetching remote URLs."""
    import base64 as _b64
    import binascii
    import io as _io

    from PIL import Image as _Image

    obj = part.get("image_url")
    if obj is None:
        obj = part.get("input_image") or part.get("image")
    url = (
        obj
        if isinstance(obj, str)
        else (obj.get("url") if isinstance(obj, dict) else None)
    )
    if not isinstance(url, str) or not url:
        return None

    raw = None
    s = url.strip()
    if s.startswith("data:"):
        prefix, sep, encoded = s.partition(",")
        prefix_lower = prefix.lower()
        if (
            sep != ","
            or not prefix_lower.startswith("data:image/")
            or ";base64" not in prefix_lower
        ):
            return None
        try:
            raw = _b64.b64decode(encoded, validate=True)
        except (binascii.Error, ValueError):
            return None
    else:
        return None

    try:
        with _Image.open(_io.BytesIO(raw)) as im:
            return im.size
    except Exception:
        return None


def _count_image_tokens_real(
    messages: list[dict[str, Any]],
    processor: Any,
    *,
    upper_bound: int = _IMAGE_TOKEN_UPPER_BOUND_FALLBACK,
    images: list[Any] | None = None,
) -> int:
    """Sum per-image tokens using processed images when supplied, otherwise
    source dimensions, instead of charging every image the model's ``max_pixels``
    ceiling. Falls back to ``upper_bound`` per image when the dimensions can't
    be read decode-free or the processor isn't a Qwen-style one, so the guard
    still never under-counts."""
    ip = getattr(processor, "image_processor", None) or processor
    ps = getattr(ip, "patch_size", None)
    ms = getattr(ip, "merge_size", None)
    minp = getattr(ip, "min_pixels", None)
    maxp = getattr(ip, "max_pixels", None)
    qwen_ok = all(isinstance(x, int) and x > 0 for x in (ps, ms, minp, maxp))
    patch_counter = getattr(ip, "get_number_of_image_patches", None)
    glm_ok = (
        callable(patch_counter)
        and isinstance(ms, int)
        and ms > 0
        and isinstance(getattr(ip, "max_image_tokens", None), int)
    )

    if images is not None:
        dimensions = [image.size for image in images]
    else:
        dimensions = []
        for msg in messages:
            content = msg.get("content")
            if not isinstance(content, list):
                continue
            for part in content:
                if not isinstance(part, dict):
                    continue
                if part.get("type") not in ("image_url", "image", "input_image"):
                    continue
                dimensions.append(_read_image_dims(part))

    total = 0
    for wh in dimensions:
        if wh is None or not (qwen_ok or glm_ok):
            total += upper_bound
        elif glm_ok:
            try:
                total += int(patch_counter(wh[1], wh[0]) // (ms**2))
            except Exception:
                total += upper_bound
        else:
            total += _smart_resize_tokens(wh[1], wh[0], ps, ms, minp, maxp)
    return total


def _audio_feature_cache_key_ranges(
    token_ids: list[int],
    input_features: Any,
    input_features_mask: Any,
    audio_token_id: int | None,
    image_ranges: list[tuple[int, str]],
) -> list[tuple[int, str]]:
    """Extend image cache boundaries with cumulative audio-clip identities.

    Audio placeholder tokens only encode clip length, so without this two
    different clips of equal length share prefix-cache blocks. Each clip is
    hashed without its batch padding (clips are padded to the longest one in
    the request), so a clip keeps its key when a later turn adds a longer one.
    """
    import hashlib

    import numpy as np

    def _to_np(value):
        if isinstance(value, mx.array):
            if value.dtype == mx.bfloat16:
                value = value.astype(mx.float32)
            return np.array(value)
        return np.asarray(value)

    features = _to_np(input_features)
    mask = None if input_features_mask is None else _to_np(input_features_mask)

    runs = []
    position = 0
    while audio_token_id is not None and position < len(token_ids):
        if token_ids[position] != audio_token_id:
            position += 1
            continue
        runs.append(position)
        while position < len(token_ids) and token_ids[position] == audio_token_id:
            position += 1

    audio_hash = hashlib.sha256()
    audio_events = []
    if features.ndim >= 2 and len(runs) == features.shape[0]:
        for i, start in enumerate(runs):
            clip = features[i]
            if mask is not None and mask.shape[:2] == features.shape[:2]:
                valid = int(mask[i].sum())
                # Trim only right padding; any other mask layout keeps the
                # padded row, which can over-key but never collide.
                if mask[i][:valid].all():
                    clip = clip[:valid]
            audio_hash.update(str(clip.shape).encode())
            audio_hash.update(np.ascontiguousarray(clip).tobytes())
            audio_events.append((start, audio_hash.hexdigest()))
    else:
        # Clips can't be matched to token runs: key everything from the
        # first audio token on (or the whole request) on all audio input.
        audio_hash.update(str(features.shape).encode())
        audio_hash.update(np.ascontiguousarray(features).tobytes())
        audio_events.append((runs[0] if runs else 0, audio_hash.hexdigest()))

    events = [(start, key, None) for start, key in image_ranges]
    events += [(start, None, key) for start, key in audio_events]
    ranges = []
    image_key = ""
    audio_key = None
    for start, image_update, audio_update in sorted(events, key=lambda e: e[0]):
        if image_update is not None:
            image_key = image_update
        if audio_update is not None:
            audio_key = audio_update
        key = image_key
        if audio_key is not None:
            key = hashlib.sha256(f"audio:{image_key}:{audio_key}".encode()).hexdigest()
        if ranges and ranges[-1][0] == start:
            ranges[-1] = (start, key)
        else:
            ranges.append((start, key))
    return ranges


class VLMBatchedEngine(BaseEngine):
    """
    VLM engine with continuous batching, tiered KV cache, and boundary snapshots.

    Extends the standard batched engine approach with vision-language model
    support. Uses VLMModelAdapter to inject pre-computed vision embeddings
    during prefill while maintaining full BatchGenerator compatibility.
    """

    def __init__(
        self,
        model_name: str,
        trust_remote_code: bool = False,
        scheduler_config: Any | None = None,
        stream_interval: int = 1,
        enable_thinking: bool | None = None,
        model_settings: Any | None = None,
        prefill_eviction_callback: Any | None = None,
    ):
        self._model_name = model_name
        self._trust_remote_code = trust_remote_code
        self._scheduler_config = scheduler_config
        self._stream_interval = stream_interval
        self._enable_thinking = enable_thinking
        self._model_settings = model_settings
        self._prefill_eviction_callback = prefill_eviction_callback

        self._vlm_model = None
        self._processor = None
        # Set after load for checkpoints that take native video input.
        self._native_video = False
        self._tokenizer = None
        self._adapter = None
        self._engine = None
        self._loaded = False
        self._grammar_compiler = None
        self._grammar_compiler_init_attempted = False
        self._vision_cache = None
        self._vision_cache_enabled = True
        # Holds the loaded gemma4_assistant drafter when vlm_mtp_enabled.
        # Phase 2A: attached but not yet wired into the decode path.
        self._vlm_mtp_drafter: Any | None = None
        # Holds the DFlash block drafter attached to Lightning MTP when
        # dflash_enabled routes here instead of DFlashEngine.
        self._dflash_drafter: Any | None = None
        self._diffusion_family: str | None = None
        self._diffusion_lock = asyncio.Lock()
        self._diffusion_active_requests = 0
        self._diffusion_cancel_events: set[threading.Event] = set()

    async def _preflight_or_raise_with_eviction(
        self,
        scheduler: Any,
        *,
        num_prompt_tokens: int,
        request_id: str | None,
        text_only: bool = False,
    ) -> None:
        await _run_scheduler_preflight_with_cleanup_retry(
            scheduler,
            num_prompt_tokens=num_prompt_tokens,
            request_id=request_id,
            eviction_callback=self._prefill_eviction_callback,
            executor=getattr(
                getattr(getattr(self, "_engine", None), "engine", None),
                "_mlx_executor",
                None,
            ),
            text_only=text_only,
        )

    @property
    def model_name(self) -> str:
        return self._model_name

    @property
    def tokenizer(self) -> Any:
        return self._tokenizer

    @property
    def supports_early_tool_call_streaming(self) -> bool:
        """Opt in only when the local scheduler has no structured parser."""

        scheduler = getattr(
            getattr(getattr(self, "_engine", None), "engine", None),
            "scheduler",
            None,
        )
        return bool(
            scheduler is not None
            and hasattr(scheduler, "_output_parser_factory")
            and scheduler._output_parser_factory is None
        )

    @property
    def model_type(self) -> str | None:
        vlm_model = getattr(self, "_vlm_model", None)
        if vlm_model is not None and hasattr(vlm_model, "config"):
            config = vlm_model.config
            if hasattr(config, "model_type"):
                return config.model_type
        return None

    @property
    def message_extractor(self):
        """Return the model-specific message extractor function, or ``None``."""
        try:
            from ..adapter.output_parser import detect_message_extractor

            model_config = {"model_type": self.model_type} if self.model_type else None
            return detect_message_extractor(self._model_name, model_config)
        except Exception:
            return None

    @property
    def is_ocr_model(self) -> bool:
        return (self.model_type or "") in OCR_MODEL_TYPES

    @property
    def is_diffusion_model(self) -> bool:
        return getattr(self, "_diffusion_family", None) == "block"

    @property
    def supports_tool_calling(self) -> bool:
        """True when a tool parser was injected into the tokenizer.

        Tool calling is prompt-driven plus output parsing — it does not
        require grammar-constrained decoding, so it works on any lane
        (autoregressive or diffusion) whose chat template matched a
        parser in ``_inject_tool_calling``.
        """
        return bool(getattr(self._tokenizer, "has_tool_calling", False))

    @property
    def grammar_compiler(self):
        """Lazily create and return a GrammarCompiler for this VLM model."""
        if self.is_diffusion_model:
            # The diffusion lane denoises canvas positions in parallel —
            # there is no sequential logit stream to mask, so compiled
            # grammars cannot be enforced. Returning None routes
            # response_format through the existing prompt-injection
            # fallback (with the #1241 Warning header) instead of
            # compiling a grammar that the lane would have to reject.
            return None
        if self._grammar_compiler is not None:
            return self._grammar_compiler
        if self._grammar_compiler_init_attempted:
            return None
        self._grammar_compiler_init_attempted = True
        try:
            from ..api.grammar import create_grammar_compiler

            self._grammar_compiler = create_grammar_compiler(
                self._tokenizer, self._vlm_model
            )
            logger.info("GrammarCompiler initialized for %s", self._model_name)
        except Exception:
            from ..utils.install import get_install_method

            method = get_install_method()
            if method == "dmg":
                logger.warning(
                    "GrammarCompiler initialization failed for %s on the "
                    "DMG build. The bundle ships xgrammar against a torch "
                    "stub; this usually means the bundled xgrammar / tvm-"
                    "ffi version drifted past what the stub covers.",
                    self._model_name,
                )
            elif method == "homebrew":
                logger.info(
                    "Structured output requires xgrammar. "
                    "Reinstall with: brew reinstall omlx --with-grammar"
                )
            else:
                logger.info(
                    "Structured output requires xgrammar. "
                    "Install with: pip install 'omlx[grammar]'"
                )
        return self._grammar_compiler

    @property
    def prefix_cache_enabled(self) -> bool:
        """True when the scheduler has a BlockAwarePrefixCache wired up."""
        if self._engine is None:
            return False
        try:
            return self._engine.engine.scheduler.block_aware_cache is not None
        except AttributeError:
            return False

    def _detect_diffusion_family(self) -> str | None:
        """Route canvas diffusion models to the serial generation lane."""
        from mlx_vlm.generate.diffusion import is_diffusion_model

        config = getattr(self._vlm_model, "config", None)
        if (
            getattr(config, "canvas_length", None) is not None
            and is_diffusion_model(self._vlm_model)
        ):
            return "block"
        return None

    def _resolve_ocr_stop_token_ids(self) -> list[int]:
        """Convert OCR stop sequences to token IDs via the tokenizer.

        Caches the result after first call since the tokenizer doesn't change.
        """
        if hasattr(self, "_ocr_stop_ids_cache"):
            return self._ocr_stop_ids_cache

        ids: list[int] = []
        if self._tokenizer is None:
            return ids

        unk_id = getattr(self._tokenizer, "unk_token_id", None)
        for seq in OCR_EXTRA_STOP_SEQUENCES:
            try:
                token_id = self._tokenizer.convert_tokens_to_ids(seq)
                if token_id is not None and token_id != unk_id:
                    ids.append(token_id)
            except (AttributeError, KeyError, TypeError):
                pass

        self._ocr_stop_ids_cache = ids
        if ids:
            logger.debug(f"OCR stop token IDs resolved: {ids}")
        return ids

    async def start(self) -> None:
        """Load VLM model and processor via mlx-vlm, create engine with VLMModelAdapter."""
        if self._loaded:
            return

        from mlx_vlm.utils import load as vlm_load

        from ..engine_core import AsyncEngineCore, EngineConfig
        from ..scheduler import SchedulerConfig
        from ..utils.model_loading import maybe_load_custom_quantization

        # Apply pre-load patches (MTP runtime patch, etc.) before the model
        # is instantiated, so the patched ``__init__`` runs. ``maybe_apply``
        # is a no-op when the model is incompatible.
        try:
            from ..utils.model_loading import maybe_apply_pre_load_patches

            maybe_apply_pre_load_patches(
                self._model_name,
                model_settings=self._model_settings,
                for_vlm=True,
            )
        except ValueError:
            raise
        except Exception as e:
            logger.debug(f"pre-load patches skipped: {e}")

        # Load VLM model on the global MLX executor to avoid blocking the event loop
        # while ensuring no concurrent Metal operations. See issue #85.
        from ..engine_core import get_mlx_executor

        def _load_vlm_sync():
            _patch_video_processor_bug()
            _patch_torch_free_image_processor()
            apply_pixtral_torch_free_patch()
            apply_gemma4_audio_patch()
            with (
                _strip_audio_config_if_orphaned(Path(self._model_name)),
                _strip_vision_config_if_orphaned(Path(self._model_name)),
                _drop_gemma4_mlx_shared_kv_extras_on_load(Path(self._model_name)),
                _derive_gemma4_global_kv_on_load(Path(self._model_name)),
                _force_minimax_m3_moe_sanitize_on_load(Path(self._model_name)),
                _force_qwen4_exp_sanitize_on_load(Path(self._model_name)),
                _remap_nested_visual_on_load(Path(self._model_name)),
                _transpose_qwen35_mlx_vision_patch_embed_on_load(
                    Path(self._model_name)
                ),
            ):
                model_type = _read_config_model_type(self._model_name)
                if model_type in {"mimo_v2", "mimo_v2_flash"}:
                    from ..patches.mimo_v2.omnimodal import (
                        has_vision_sidecar,
                    )
                    from ..patches.mimo_v2.omnimodal import (
                        load as load_mimo_omnimodal,
                    )

                    if has_vision_sidecar(self._model_name):
                        return load_mimo_omnimodal(
                            self._model_name,
                            model_settings=self._model_settings,
                            trust_remote_code=self._trust_remote_code,
                        )

                custom_loaded = maybe_load_custom_quantization(
                    self._model_name,
                    is_vlm=True,
                )
                if custom_loaded is not None:
                    model, processor = custom_loaded
                    return model, processor

                if model_type == "deepseek_v41":
                    from ..patches.deepseek_v41.loading import load

                    return load(
                        self._model_name,
                        moe_expert_offload_resident_fraction=(
                            self._model_settings.moe_expert_offload_resident_fraction
                            if getattr(
                                self._model_settings,
                                "moe_expert_offload_enabled",
                                False,
                            )
                            and os.environ.get("OMLX_MOE_EXPERT_OFFLOAD", "1") != "0"
                            else None
                        ),
                        engram_ssd_offload=bool(
                            getattr(
                                self._model_settings,
                                "deepseek_v41_engram_ssd_offload",
                                False,
                            )
                        ),
                        ced_prefill=bool(
                            getattr(
                                self._model_settings,
                                "deepseek_v41_ced_prefill_enabled",
                                False,
                            )
                        ),
                    )
                if model_type == COHERE2_MOE_MODEL_TYPE:
                    return _load_cohere2_moe_text_model(
                        self._model_name,
                        trust_remote_code=self._trust_remote_code,
                    )
                with _load_optiq_vision_sidecar_on_load(
                    Path(self._model_name)
                ):
                    load_kwargs = {
                        "trust_remote_code": self._trust_remote_code,
                    }
                    if model_type == QWEN4_EXP_MODEL_TYPE:
                        load_kwargs["lazy"] = True
                    # Expert offload wraps BEFORE materialization so non-resident
                    # experts never load; keep the load lazy only when the feature
                    # is on. Threads into main's load_kwargs path (lazy is idempotent
                    # with the QWEN4_EXP case above).
                    if getattr(
                        self._model_settings, "moe_expert_offload_enabled", False
                    ):
                        load_kwargs["lazy"] = True
                    loaded = vlm_load(
                        self._model_name,
                        **load_kwargs,
                    )
                    return loaded

        loop = asyncio.get_running_loop()
        self._vlm_model, self._processor = await loop.run_in_executor(
            get_mlx_executor(), _load_vlm_sync
        )

        from ..models.vlm import restore_bonsai_quantized_modules

        restored = restore_bonsai_quantized_modules(self._vlm_model)
        if restored:
            logger.info("Restored oMLX Bonsai kernel paths for %d modules", restored)

        if self.model_type in NATIVE_VIDEO_MODEL_TYPES:
            self._native_video = attach_native_video_processor(
                self._processor, self._model_name
            )

        if self.model_type == "unlimited-ocr":
            from ..utils.tokenizer import (
                create_streaming_detokenizer,
                repair_misconverted_unlimited_ocr_tokenizer,
            )

            tokenizer_obj = getattr(self._processor, "tokenizer", self._processor)
            if repair_misconverted_unlimited_ocr_tokenizer(
                tokenizer_obj,
                model_path=self._model_name,
            ):
                # mlx-vlm also keeps a processor-level detokenizer for its own
                # generation helpers.  Replace that stale SPM instance even
                # though oMLX's scheduler creates fresh request-local copies.
                self._processor.detokenizer = create_streaming_detokenizer(
                    tokenizer_obj,
                    model_path=self._model_name,
                )
                logger.warning(
                    "Repaired misconverted Unlimited-OCR tokenizer metadata "
                    "in memory for %s",
                    self._model_name,
                )

        # MoE expert offload for the VLM path: Gemma 4 checkpoints are
        # detected as VLMs, so this — not BatchedEngine — is their default
        # engine. Same sequence as batched.py: wrap on the MLX executor
        # BEFORE materialize so non-resident experts never load.
        moe_offload_wrapped = 0
        if getattr(self._model_settings, "moe_expert_offload_enabled", False):
            from ..patches.moe_expert_offload import (
                apply_moe_expert_offload,
                materialize_offload_state,
            )

            fraction = float(
                getattr(
                    self._model_settings,
                    "moe_expert_offload_resident_fraction",
                    0.25,
                )
            )
            # glm5_next Lightning MTP: the draft head's experts stay resident
            # while the backbone streams (run_in_executor takes no kwargs,
            # so bind with partial).
            moe_offload_wrapped = await loop.run_in_executor(
                get_mlx_executor(),
                functools.partial(
                    apply_moe_expert_offload,
                    mtp_resident=bool(
                        getattr(self._model_settings, "mtp_enabled", False)
                    ),
                ),
                self._vlm_model,
                self._model_name,
                fraction,
            )
            if moe_offload_wrapped:
                # The caches' slot maps and resident slots live on plain
                # attributes outside the module tree, so the lazy-state
                # materialization below never reaches them; left lazy they
                # stay bound to the loader stream and the first request from
                # an inference thread dies with "There is no Stream(gpu, N)
                # in current thread". Same executor as the apply, so the
                # arrays realize on the stream that created them.
                await loop.run_in_executor(
                    get_mlx_executor(),
                    materialize_offload_state,
                    self._vlm_model,
                )
        self._moe_offload_wrapped = moe_offload_wrapped

        # Materialize lazy buffers (RoPE freqs, vision/audio towers) on the
        # loader thread so per-engine inference threads can read them (#1304).
        from ..utils.model_loading import materialize_lazy_state

        await loop.run_in_executor(
            get_mlx_executor(), materialize_lazy_state, self._vlm_model
        )

        # t5 ternary: free unused bias tensors to recover ~420 MB RAM.
        # The repacked safetensors carries 2-bit biases for format compat;
        # the t5 symmetric kernel (scale*(q-1)) never reads them.
        try:
            from ..patches.bonsai_t5_load import free_t5_biases

            freed = await loop.run_in_executor(
                get_mlx_executor(), free_t5_biases, self._vlm_model
            )
            if freed > 0:
                logger.info(
                    "t5 bias tensors freed: %.0f MB recovered", freed / 1e6
                )
        except Exception:
            logger.debug("t5 bias free skipped", exc_info=True)

        # Supported MoE gate+up regroup: concatenate the routed experts'
        # gate and up projections so decode runs 2 gather_qmm launches per
        # MoE layer instead of 3 (issue #2238). Bit-exact; also swaps the
        # mlx-vlm target-verify helper for a fused-aware version. Runs on
        # the MLX executor because it rewrites weights in place.
        if getattr(self, "_moe_offload_wrapped", 0):
            logger.info(
                "moe expert offload active (%d layers): skipping gate/up "
                "fusion on the VLM path",
                self._moe_offload_wrapped,
            )
        elif (
            getattr(self._model_settings, "moe_gate_up_fusion_enabled", True)
            is not False
        ):
            try:
                from ..patches.moe_gate_up_fusion import apply_moe_gate_up_fusion

                await loop.run_in_executor(
                    get_mlx_executor(),
                    apply_moe_gate_up_fusion,
                    self._vlm_model,
                )
            except Exception:
                logger.debug("MoE gate+up fusion not applied", exc_info=True)

        # Qwen ANE prefill compiles its slices from the stock QuantizedLinear
        # layout, so build them before the packed projections replace it. The
        # class hooks are installed later, outside the other Qwen class patches.
        ane_count = 0
        requested_ane_sequence_length = 2048
        if (
            getattr(self._model_settings, "qwen35_ane_prefill_enabled", False)
            and ane_prefill_backend(self.model_type) == "qwen"
        ):
            ane_fraction = ane_prefill_fraction(
                self._model_settings.qwen35_ane_prefill_fraction, self.model_type
            )
            try:
                from ..patches.qwen35_ane_prefill import enable_qwen35_ane_prefill

                requested_ane_sequence_length = int(
                    getattr(
                        self._model_settings,
                        "qwen35_ane_prefill_sequence_length",
                        2048,
                    )
                )

                def _enable_ane_prefill():
                    return enable_qwen35_ane_prefill(
                        self._vlm_model,
                        sequence_length=requested_ane_sequence_length,
                        tail_padding_min_tokens=int(
                            getattr(
                                self._model_settings,
                                "qwen35_ane_prefill_tail_padding_min_tokens",
                                0,
                            )
                            or 0
                        ),
                        fraction=ane_fraction,
                        max_layers=getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_max_layers",
                            64,
                        ),
                        gdn=getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_gdn",
                            True,
                        ),
                        gdn_fraction=getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_gdn_fraction",
                            0.50,
                        ),
                        gdn_max_layers=getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_gdn_max_layers",
                            48,
                        ),
                        dual_ane=getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_dual_ane",
                            True,
                        ),
                        ane_down_fraction=(
                            ane_fraction
                            if getattr(
                                self._model_settings,
                                "qwen35_ane_prefill_fused_down",
                                False,
                            )
                            else 0.0
                        ),
                        fused_down=getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_fused_down",
                            False,
                        ),
                        cpu_fraction=getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_cpu_fraction",
                            0.135,
                        )
                        if getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_cpu_enabled",
                            False,
                        )
                        else 0.0,
                        cpu_down_fraction=getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_cpu_down_fraction",
                            0.0,
                        )
                        if getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_cpu_enabled",
                            False,
                        )
                        else 0.0,
                        cpu_gdn_fraction=getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_cpu_gdn_fraction",
                            0.0,
                        )
                        if getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_cpu_enabled",
                            False,
                        )
                        else 0.0,
                        cpu_threads=getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_cpu_threads",
                            8,
                        ),
                        cpu_shared_resource=getattr(
                            self._model_settings,
                            "qwen35_ane_prefill_cpu_shared_resource",
                            True,
                        ),
                        install_dispatch=False,
                    )

                ane_count = await loop.run_in_executor(
                    get_mlx_executor(),
                    _enable_ane_prefill,
                )
            except Exception:
                logger.warning("Qwen ANE prefill not enabled", exc_info=True)

        # Dense Qwen3.5-family 4-bit projections -> tile-repacked layout for
        # the M5 tensor units. Replaces the layers in place, so resident
        # memory does not grow; runs on the MLX executor before any forward.
        try:
            from ..patches import qwen35_packed_linear

            if qwen35_packed_linear.enabled(self._vlm_model):
                packed = await loop.run_in_executor(
                    get_mlx_executor(),
                    qwen35_packed_linear.pack_model,
                    self._vlm_model,
                )
                if packed:
                    logger.info("Qwen packed 4-bit projections: %d layers", packed)
        except Exception:
            logger.warning("Qwen packed 4-bit projections not applied", exc_info=True)

        _fix_processor_none_pixels(self._processor)
        self._diffusion_family = self._detect_diffusion_family()
        if self.is_diffusion_model:
            logger.info(
                "Diffusion VLM detected; using serial diffusion lane for %s",
                self._model_name,
            )

        # Initialize vision feature cache
        vision_ssd_dir = None
        if not self.is_diffusion_model:
            if self._scheduler_config and getattr(
                self._scheduler_config, "paged_ssd_cache_dir", None
            ):
                vision_ssd_dir = (
                    Path(self._scheduler_config.paged_ssd_cache_dir) / "vision_features"
                )
            self._vision_cache = VisionFeatureSSDCache(
                cache_dir=vision_ssd_dir,
                # Agent sessions resend 20-90 screenshots; bound by bytes.
                max_memory_entries=4096,
                max_memory_bytes=1024**3,
            )
            logger.info(
                "Vision feature cache enabled (SSD: %s)",
                vision_ssd_dir or "disabled",
            )
        else:
            self._vision_cache = None
            self._vision_cache_enabled = False

        # Extract tokenizer from processor with deep-copy for thread safety.
        # The processor keeps the original tokenizer for executor-thread work
        # (_prepare_vision_inputs / prepare_inputs), while this deep copy is
        # used exclusively on the event loop (apply_chat_template, encode).
        # Without separate Rust tokenizer backends, concurrent access causes
        # "RuntimeError: Already borrowed".
        # See: https://github.com/huggingface/tokenizers/issues/537
        if hasattr(self._processor, "tokenizer"):
            self._tokenizer = copy.deepcopy(self._processor.tokenizer)
        else:
            self._tokenizer = copy.deepcopy(self._processor)
        if self._tokenizer is None or not callable(
            getattr(self._tokenizer, "encode", None)
        ):
            raise RuntimeError(
                f"VLM processor for {self._model_name} did not provide a usable tokenizer"
            )

        if self.is_diffusion_model:
            self._inject_tool_calling(self._tokenizer)
            self._loaded = True
            logger.info(f"VLMBatchedEngine loaded: {self._model_name}")
            return

        # Create VLM model adapter wrapping language_model.
        # mlx-vlm models now handle per-sequence mx.array offsets natively
        # and batched decode is fixed, so no separate mlx-lm decode model needed.
        self._adapter = VLMModelAdapter(self._vlm_model)

        # Create scheduler config
        scheduler_config = (
            copy.copy(self._scheduler_config)
            if self._scheduler_config
            else SchedulerConfig()
        )
        if (
            self._adapter.model_type == "deepseek_v41"
            and self._adapter.config.ced_prefill
            and scheduler_config.paged_ssd_cache_dir
        ):
            # Approximate decoder states must not become hits in full-prefill
            # mode (or vice versa) after reloading the model with new settings.
            scheduler_config.paged_ssd_cache_dir = str(
                Path(scheduler_config.paged_ssd_cache_dir) / "deepseek_v41_ced_v1"
            )

        engine_config = EngineConfig(
            model_name=self._model_name,
            scheduler_config=scheduler_config,
            stream_interval=self._stream_interval,
            prefill_eviction_callback=self._prefill_eviction_callback,
        )

        # Create engine with adapter as the "model"
        # The adapter exposes .layers, .make_cache() for cache infrastructure
        self._engine = AsyncEngineCore(
            model=self._adapter,
            tokenizer=self._tokenizer,
            config=engine_config,
        )

        await self._engine.engine.start()

        # TurboQuant KV cache
        scheduler = self._engine.engine.scheduler
        if self._model_settings is not None:
            tq_enabled = getattr(self._model_settings, "turboquant_kv_enabled", False)
            if tq_enabled and self.model_type == "glm5_next":
                logger.warning(
                    "TurboQuant KV cache is not supported for GLM-5.3-Flash's "
                    "composite latent/indexer cache; using the native cache layout"
                )
                tq_enabled = False
            if tq_enabled:
                from ..patches.turboquant_attention import (
                    apply_turboquant_attention_patch,
                )

                apply_turboquant_attention_patch()
                tq_bits = float(getattr(self._model_settings, "turboquant_kv_bits", 4))
                scheduler._turboquant_kv_bits = tq_bits
                scheduler._turboquant_skip_last = getattr(
                    self._model_settings, "turboquant_skip_last", True
                )
                scheduler._set_model_info_for_monitor()
                logger.info(f"TurboQuant KV cache enabled for VLM: {tq_bits} bits")

        # head_dim=256 long-context prefill -> O(L) tiled SDPA kernel. See
        # batched.py for rationale. Passthrough-safe; strictly gated route.
        if getattr(self._model_settings, "sdpa256_prefill_enabled", True) is not False:
            try:
                from ..patches.sdpa256_attention import (
                    apply_sdpa256_attention_patch,
                )

                apply_sdpa256_attention_patch()
            except Exception:
                logger.debug("sdpa256 attention patch not applied", exc_info=True)

        # Qwen3.5/3.6 head_dim=256 causal prefill -> native steel FA kernel.
        # Installed after sdpa256 so matched Qwen dense attention takes the
        # simdgroup-MMA path, while unsupported cases fall through unchanged.
        if (
            getattr(self._model_settings, "fa256_steel_prefill_enabled", True)
            is not False
        ):
            try:
                from ..patches.qwen35_fa256_attention import (
                    apply_qwen35_fa256_attention_patch,
                )

                apply_qwen35_fa256_attention_patch()
            except Exception:
                logger.debug("Qwen FA-256 steel patch not applied", exc_info=True)

        # Qwen3.5/3.6 verify-width GDN prework -> one fused Metal launch
        # (conv+SiLU+split+RMS+scale+conv-state), bit-exact to the chain.
        try:
            from ..patches.qwen35_gdn_prework import (
                apply_qwen35_gdn_prework_patch,
                configure_qwen4_decode,
            )

            configure_qwen4_decode(
                self._vlm_model,
                wide_projections=scheduler_config.qwen4_gdn_decode_wide_proj,
            )
            apply_qwen35_gdn_prework_patch()
        except Exception:
            logger.debug("Qwen GDN prework patch not applied", exc_info=True)

        # Qwen3.5/3.6 verify-width (MTP target-verify) attention -> chunked
        # causal vector-kernel calls instead of the per-row SDPA loop.
        try:
            from ..patches.qwen35_verify_sdpa_split import (
                apply_qwen35_verify_sdpa_split_patch,
            )

            apply_qwen35_verify_sdpa_split_patch()
        except Exception:
            logger.debug(
                "Qwen verify-split attention patch not applied", exc_info=True
            )

        # Qwen3.5/3.6 Gated DeltaNet prefill -> optimized Metal kernel.
        # Decode and masked paths keep the original mlx-vlm kernel.
        gdn_prefill_enabled = getattr(
            self._model_settings,
            "gdn_prefill_enabled",
            getattr(self._model_settings, "gdn_chunked_prefill_enabled", True),
        )
        if gdn_prefill_enabled is not False:
            try:
                from ..patches.qwen35_gdn_chunked import (
                    apply_qwen35_gdn_prefill_patch,
                )

                apply_qwen35_gdn_prefill_patch()
            except Exception:
                logger.debug("GDN prefill patch not applied", exc_info=True)

        # Qwen3.5/3.6 q4 MLP prefill -> native qmm tile tuned for long batches.
        # Decode and target-verify paths keep the original QuantizedLinear.
        if (
            getattr(self._model_settings, "qwen35_q4_mlp_prefill_enabled", True)
            is not False
        ):
            try:
                from ..patches.qwen35_q4_mlp import (
                    apply_muse_glimmer_q4_prefill_patch,
                    apply_qwen35_q4_mlp_patch,
                    apply_qwen35_q4_prefill_linear_patch,
                )

                apply_qwen35_q4_mlp_patch()
                apply_qwen35_q4_prefill_linear_patch(self._vlm_model)
                # Muse Glimmer rides the same native qmm tile (MLP plus the
                # q/gate/o attention projections); no-op unless the muse
                # compat patch installed the vendored module.
                apply_muse_glimmer_q4_prefill_patch()
            except Exception:
                logger.debug("Qwen q4 MLP prefill patch not applied", exc_info=True)

        # Qwen MoE decode router: fuse the top-k select + renormalize chain
        # into one launch for short rows (decode + MTP verify widths).
        try:
            from ..patches.qwen35_moe_router import (
                apply_qwen35_moe_router_patch,
            )

            apply_qwen35_moe_router_patch()
        except Exception:
            logger.debug("Qwen MoE router patch not applied", exc_info=True)

        # One-token routed experts: gate+up/SwiGLU and down/weighted-sum in
        # two launches instead of five (bit-identical; needs the router patch).
        try:
            from ..patches.qwen35_moe_routed_decode import (
                apply_qwen35_moe_routed_decode_patch,
            )

            apply_qwen35_moe_routed_decode_patch()
        except Exception:
            logger.debug("Qwen MoE routed decode patch not applied", exc_info=True)

        # oQ mixed-bit QxA8 prefill kernels. Gated on the per-model setting
        # because it quantizes activations to INT8, which changes numerics;
        # the patch itself falls through for anything it cannot route.
        if getattr(self._model_settings, "qwen35_oq_a8_enabled", False):
            try:
                from ..patches.qwen35_oq_a8 import apply_qwen35_oq_a8_patch

                # The model itself is what gets opted in: the patch tags its
                # modules, so a model loaded with the setting off is never
                # routed even though the class wrapper is process-wide.
                apply_qwen35_oq_a8_patch(
                    self._vlm_model,
                    min_tokens=int(
                        getattr(self._model_settings, "qwen35_oq_a8_min_tokens", 128)
                    ),
                )
            except Exception:
                logger.debug("oQ A8 prefill patch not applied", exc_info=True)

        if ane_count or getattr(self._vlm_model, "_omlx_ane_gdn_prefill_count", 0):
            try:
                from ..patches.qwen35_ane_prefill import (
                    ane_prefill_transient_bytes,
                    configure_qwen35_ane_prefill_scheduler,
                    install_qwen35_ane_prefill_dispatch,
                )

                if install_qwen35_ane_prefill_dispatch():
                    configure_qwen35_ane_prefill_scheduler(
                        scheduler,
                        requested_ane_sequence_length,
                    )
                    # Price the compiled I/O surfaces off the full model so
                    # admission charges them while the banks are resident.
                    monitor = getattr(scheduler, "memory_monitor", None)
                    if monitor is not None:
                        monitor.set_ane_prefill_transient_bytes(
                            ane_prefill_transient_bytes(self._vlm_model)
                        )
            except Exception:
                logger.warning("Qwen ANE prefill not enabled", exc_info=True)

        # Qwen3.5/3.6 sparse MoE prefill -> native weighted-sum after sorted
        # SwitchGLU. Decode and target-verify keep the original path.
        if (
            getattr(self._model_settings, "qwen35_moe_weighted_sum_enabled", True)
            is not False
        ):
            try:
                from ..patches.qwen35_moe_weighted_sum import (
                    apply_qwen35_moe_weighted_sum_patch,
                )

                apply_qwen35_moe_weighted_sum_patch()
            except Exception:
                logger.debug(
                    "Qwen MoE weighted-sum patch not applied", exc_info=True
                )

        if (
            getattr(self._model_settings, "qwen35_ragged_decode_fallback_enabled", True)
            is not False
        ):
            try:
                from ..patches.qwen35_ragged_decode import (
                    apply_qwen35_ragged_decode_patch,
                )

                apply_qwen35_ragged_decode_patch()
            except Exception:
                logger.debug("qwen3_5 ragged decode patch not applied", exc_info=True)
        scheduler.refresh_ssd_layer_signature()

        # SpecPrefill: load draft model and pass to scheduler
        if self._model_settings is not None:
            specprefill_draft = getattr(
                self._model_settings, "specprefill_draft_model", None
            )
            specprefill_enabled = getattr(
                self._model_settings, "specprefill_enabled", False
            )
            if specprefill_enabled and specprefill_draft:
                try:
                    from ..utils.model_loading import (
                        lm_load_compat as mlx_lm_load,
                    )
                    from ..utils.model_loading import (
                        maybe_load_custom_quantization,
                    )
                    from ..utils.tokenizer import get_tokenizer_config

                    def _load_draft():
                        from ..patches.mlx_lm_mtp import set_mtp_active
                        from ..utils.model_loading import materialize_lazy_state

                        was_mtp = False
                        try:
                            from ..patches.mlx_lm_mtp import is_mtp_active

                            was_mtp = is_mtp_active()
                        except Exception:
                            pass
                        set_mtp_active(False)
                        try:
                            custom_loaded = maybe_load_custom_quantization(
                                specprefill_draft,
                                is_vlm=False,
                            )
                            if custom_loaded is not None:
                                draft_model, _ = custom_loaded
                            else:
                                draft_tokenizer_config = get_tokenizer_config(
                                    specprefill_draft,
                                    trust_remote_code=self._trust_remote_code,
                                )
                                draft_model, _ = mlx_lm_load(
                                    specprefill_draft,
                                    tokenizer_config=draft_tokenizer_config,
                                    trust_remote_code=self._trust_remote_code,
                                )
                            # Materialize frozen buffers (RoPE freqs, etc.)
                            # on the loader thread. mlx_lm.load only does
                            # mx.eval(model.parameters()) and leaves siblings
                            # lazy bound to this thread's stream. Without
                            # this, the first score_tokens() call from
                            # Scheduler.step on the per-engine executor
                            # thread raises "no Stream(gpu, X) in current
                            # thread". Same root cause and fix as e93c408
                            # for the VLM MTP drafter.
                            materialize_lazy_state(draft_model)
                            return draft_model
                        finally:
                            set_mtp_active(was_mtp)

                    draft_model = await loop.run_in_executor(
                        get_mlx_executor(), _load_draft
                    )
                    self._engine.engine.scheduler.set_specprefill_draft_model(
                        draft_model, draft_model_name=specprefill_draft
                    )
                    logger.info(
                        f"SpecPrefill: draft model loaded ({specprefill_draft})"
                    )
                except Exception as e:
                    logger.error(f"SpecPrefill: draft model load failed: {e}")

        # Inject mlx-lm tool calling support into VLM tokenizer
        self._inject_tool_calling(self._tokenizer)

        self._loaded = True
        logger.info(f"VLMBatchedEngine loaded: {self._model_name}")

    def set_vlm_mtp_drafter(self, drafter: Any) -> None:
        """Attach a loaded MTP drafter for VLM MTP decoding.

        Passes the drafter (and the configured draft-block size) down to
        the scheduler so eligible requests get routed to mlx-vlm's MTP
        round loop at decode time.  Supports gemma4_assistant and
        qwen3_5_mtp drafter types.
        """
        self._vlm_mtp_drafter = drafter
        block_size = None
        if self._model_settings is not None:
            block_size = getattr(self._model_settings, "vlm_mtp_draft_block_size", None)
        scheduler = None
        if self._engine is not None and hasattr(self._engine, "engine"):
            scheduler = getattr(self._engine.engine, "scheduler", None)
        if scheduler is not None and hasattr(scheduler, "set_vlm_mtp_drafter"):
            scheduler.set_vlm_mtp_drafter(drafter, draft_block_size=block_size)
        logger.info(
            "VLM MTP drafter attached to engine: %s (block_size=%s)",
            self._model_name,
            block_size,
        )

    @property
    def vlm_mtp_drafter(self) -> Any | None:
        return self._vlm_mtp_drafter

    @property
    def vlm_model(self) -> Any | None:
        """The loaded mlx-vlm model (target for external drafters)."""
        return self._vlm_model

    def set_dflash_drafter(self, drafter: Any) -> None:
        """Attach a DFlash block drafter to the Lightning MTP decode path."""
        from ..speculative.dflash_drafter import attach_drafter

        language_model = getattr(self._adapter, "_language_model", None)
        if language_model is None:
            raise RuntimeError("VLM engine has no language model for the drafter")
        attach_drafter(language_model, drafter)
        self._dflash_drafter = drafter
        logger.info(
            "DFlash drafter attached to engine: %s (kind=%s, block=%d, layers=%s)",
            self._model_name,
            drafter.kind,
            drafter.block_size,
            drafter.target_layer_ids,
        )

    @property
    def dflash_drafter(self) -> Any | None:
        return self._dflash_drafter

    async def stop(self) -> None:
        """Stop the engine and cleanup resources."""
        cancelled = False
        engine = self._engine

        for cancel_event in getattr(self, "_diffusion_cancel_events", ()):
            cancel_event.set()

        if engine:
            await engine.stop()

        if self._vision_cache is not None:
            try:
                self._vision_cache.close()
            except Exception:
                logger.warning("Error closing vision feature cache", exc_info=True)
            self._vision_cache = None

        # Drop wrapper-side references before EngineCore.close() performs its
        # final worker-thread MLX reclaim. Otherwise the VLM wrapper can keep
        # model weights or cached feature arrays alive until after the reclaim
        # pass has already run. Mark unloaded first: the scheduler is no longer
        # reachable once _engine is cleared.
        self._loaded = False
        _clear_teardown_references(
            self,
            none_attrs=(
                "_engine",
                "_vlm_model",
                "_processor",
                "_adapter",
                "_tokenizer",
                "_grammar_compiler",
                "_vlm_mtp_drafter",
                "_dflash_drafter",
                "_diffusion_family",
            ),
            false_attrs=("_grammar_compiler_init_attempted",),
        )

        if engine:
            if hasattr(engine, "engine") and engine.engine is not None:
                try:
                    cancelled = await _close_engine_core(engine.engine)
                except Exception as e:
                    logger.warning(f"Error closing engine: {e}")
        self._diffusion_cancel_events = set()
        self._diffusion_active_requests = 0
        logger.info("VLMBatchedEngine stopped")
        if cancelled:
            raise asyncio.CancelledError

    def _inject_tool_calling(self, tokenizer) -> None:
        """Inject tool calling attributes into VLM tokenizer.

        mlx-vlm's TokenizerWrapper lacks tool calling support (has_tool_calling,
        tool_parser, etc). We prefer mlx_vlm.tools.parsers which is a superset of
        mlx_lm's — it recognises additional markers such as Gemma4's <|tool_call>
        and loads the correct per-model parser.  Falls back to mlx_lm if the
        mlx_vlm.tools.parsers package is not present.
        """
        chat_template = getattr(tokenizer, "chat_template", None)
        if not chat_template:
            return

        # MiniMax M3's template contains generic <tool_call> text in comments
        # and examples, so upstream inference incorrectly selects json_tools.
        # Use the same oMLX protocol adapter as distributed ranks.
        from ..adapter.output_parser import (
            install_minimax_m3_tokenizer_protocol,
        )

        if install_minimax_m3_tokenizer_protocol(
            tokenizer,
            self._model_name,
            {"model_type": self.model_type} if self.model_type else None,
        ):
            logger.info("VLM tool calling enabled: parser=minimax_m3")
            return

        # Prefer mlx_vlm.tools.parsers (superset; knows about Gemma4 etc.)
        try:
            from mlx_vlm.tools.registry import (
                _infer_tool_parser,
                load_tool_module,
            )

            tool_parser_type = _infer_tool_parser(chat_template)
            if tool_parser_type is None:
                return
            try:
                tool_module = load_tool_module(tool_parser_type)
            except ImportError:
                logger.warning(f"VLM tool parser module not found: {tool_parser_type}")
                return
        except ImportError:
            # Fallback: mlx_lm only (no Gemma4 support)
            try:
                import importlib

                from mlx_lm.tokenizer_utils import (
                    _infer_tool_parser as _mlx_lm_infer,
                )
            except ImportError:
                return
            tool_parser_type = _mlx_lm_infer(tokenizer)
            if tool_parser_type is None:
                return
            try:
                tool_module = importlib.import_module(
                    f"mlx_lm.tool_parsers.{tool_parser_type}"
                )
            except ImportError:
                logger.warning(f"VLM tool parser module not found: {tool_parser_type}")
                return

        tool_call_start = tool_module.tool_call_start
        tool_call_end = tool_module.tool_call_end

        # Validate tokens exist in vocab (same check as mlx-lm)
        vocab = tokenizer.get_vocab()
        if (tool_call_start and tool_call_start not in vocab) or (
            tool_call_end and tool_call_end not in vocab
        ):
            return

        # Set instance attributes on the mlx-vlm TokenizerWrapper.
        # Python's __getattr__ is only called when normal lookup fails,
        # so instance attributes take precedence over delegation to HF tokenizer.
        tokenizer.has_tool_calling = True
        tokenizer.tool_call_start = tool_call_start
        tokenizer.tool_call_end = tool_call_end
        tokenizer.tool_parser = tool_module.parse_tool_call

        logger.info(f"VLM tool calling enabled: parser={tool_parser_type}")

    def _chat_template_target(self, model_type: str) -> Any:
        """Return the object that renders messages already formatted by oMLX."""
        if model_type in TOKENIZER_CHAT_TEMPLATE_MODEL_TYPES or not hasattr(
            self._processor, "apply_chat_template"
        ):
            return getattr(self._processor, "tokenizer", self._processor)
        return self._processor

    def _video_cache_identity(self, video_hashes: list[str]) -> str:
        """Prefix-cache identity for a request's clips.

        Sampling settings decide which frames reach the vision tower, so they
        are part of the identity along with the clip contents.
        """
        processor = getattr(self._processor, "video_processor", None)
        sampling = ":".join(
            str(getattr(processor, name, None))
            for name in ("fps", "min_frames", "max_frames", "min_pixels", "max_pixels")
        )
        return hashlib.sha256(
            f"video:{sampling}:{','.join(video_hashes)}".encode()
        ).hexdigest()

    def _video_cache_key_ranges(
        self, token_ids: list[int], video_grid_thw: Any, video_hashes: list[str]
    ) -> list[tuple[int, str]]:
        """Per-clip prefix-cache ranges, like the per-turn ranges for images.

        Each clip's ``video_pad`` runs start where its key takes over, and the
        key covers that clip and every earlier one. Text before the first clip
        stays shareable, and a follow-up turn that adds a clip keeps the cached
        prefix of the earlier ones. Returns ``[]`` (whole-request keying) when
        the tokens do not line up with the grids.
        """
        config = getattr(self._vlm_model, "config", None)
        video_token_id = getattr(config, "video_token_id", None)
        if video_token_id is None or video_grid_thw is None:
            return []
        try:
            grids = [int(row[0]) for row in video_grid_thw.tolist()]
        except (AttributeError, TypeError, ValueError, IndexError):
            return []
        if len(grids) != len(video_hashes):
            return []
        # The processor emits one run of video_pad tokens per temporal patch.
        run_starts = [
            index
            for index, token in enumerate(token_ids)
            if token == video_token_id
            and (index == 0 or token_ids[index - 1] != video_token_id)
        ]
        if len(run_starts) != sum(grids):
            return []
        ranges = []
        run = 0
        for clip, temporal_patches in enumerate(grids):
            ranges.append(
                (run_starts[run], self._video_cache_identity(video_hashes[: clip + 1]))
            )
            run += temporal_patches
        return ranges

    @staticmethod
    def _count_content_parts(content: Any, part_types: set[str]) -> int:
        """Count multimodal parts in list content by type."""
        if not isinstance(content, list):
            return 0

        count = 0
        for item in content:
            if isinstance(item, dict):
                item_type = item.get("type", "")
            else:
                item_type = getattr(item, "type", "")
            if item_type in part_types:
                count += 1
        return count

    def _format_messages_for_vlm_template(
        self,
        messages: list[dict[str, Any]],
        num_images: int,
        num_audios: int = 0,
        num_videos: int = 0,
    ) -> tuple[list[dict[str, Any]], list[tuple[int, int]]]:
        """Format VLM messages with image/audio/video tokens on media-bearing user turns."""
        from mlx_vlm.prompt_utils import extract_text_from_content, get_message_json

        model_type = self.model_type or getattr(
            self._vlm_model.config, "model_type", ""
        )
        if not model_type:
            raise ValueError("Missing VLM model_type for chat template formatting")

        if model_type == "deepseek_v41":
            if num_audios:
                raise ValueError("DeepSeek V4.1 supports text and images, not audio")
            from ..patches.deepseek_v41.processing import format_messages

            return format_messages(messages, num_images)

        image_part_types = {"image", "image_url", "input_image"}
        audio_part_types = {"input_audio"}
        video_part_types = {"video", "video_url", "input_video"}
        remaining_videos = num_videos
        has_explicit_images = any(
            isinstance(msg, dict)
            and self._count_content_parts(msg.get("content"), image_part_types) > 0
            for msg in messages
        )

        has_explicit_audio = any(
            isinstance(msg, dict)
            and self._count_content_parts(msg.get("content"), audio_part_types) > 0
            for msg in messages
        )

        remaining_images = num_images
        remaining_audios = num_audios
        assigned_fallback_images = False
        assigned_fallback_audios = False
        formatted_messages: list[dict[str, Any]] = []
        image_message_ranges: list[tuple[int, int]] = []

        for idx, msg in enumerate(messages):
            if not isinstance(msg, dict):
                msg = {"role": "user", "content": str(msg)}

            role = msg.get("role", "user")
            raw_content = msg.get("content")
            content = extract_text_from_content(raw_content)

            msg_num_images = 0
            msg_num_audios = 0
            msg_num_videos = 0
            if role == "user" and remaining_videos > 0:
                msg_num_videos = min(
                    self._count_content_parts(raw_content, video_part_types),
                    remaining_videos,
                )
                remaining_videos -= msg_num_videos
            if role == "user":
                explicit_images = self._count_content_parts(
                    raw_content, image_part_types
                )
                explicit_audios = self._count_content_parts(
                    raw_content, audio_part_types
                )
                if explicit_images > 0 and remaining_images > 0:
                    msg_num_images = min(explicit_images, remaining_images)
                    remaining_images -= msg_num_images
                elif (
                    not has_explicit_images
                    and remaining_images > 0
                    and not assigned_fallback_images
                ):
                    msg_num_images = remaining_images
                    remaining_images = 0
                    assigned_fallback_images = True

                if explicit_audios > 0 and remaining_audios > 0:
                    msg_num_audios = min(explicit_audios, remaining_audios)
                    remaining_audios -= msg_num_audios
                elif (
                    not has_explicit_audio
                    and remaining_audios > 0
                    and not assigned_fallback_audios
                ):
                    msg_num_audios = remaining_audios
                    remaining_audios = 0
                    assigned_fallback_audios = True

            if msg_num_images > 0:
                image_message_ranges.append((idx, msg_num_images))

            # Preserve tool-related messages and reasoning_content verbatim
            # so the chat template receives tool_calls, tool_call_id,
            # tool_responses, and reasoning_content fields. get_message_json()
            # only handles (content, role) and strips every other top-level
            # key, which would make tool results and Qwen 3.6+ reasoning
            # blocks invisible to the model.
            if role == "tool" or (
                role == "assistant"
                and (
                    msg.get("tool_calls")
                    or msg.get("tool_responses")
                    or msg.get("reasoning_content")
                )
            ):
                formatted_messages.append(msg)
            elif model_type == "glm5_next" and (
                msg_num_images > 0 or msg_num_audios > 0
            ):
                # mlx-vlm does not register GLM-5 in MODEL_CONFIG. Preserve
                # media parts and their relative order for its native template.
                glm_content: list[Any] = []
                inserted_images = 0
                inserted_audios = 0
                if isinstance(raw_content, list):
                    for item in raw_content:
                        if isinstance(item, dict):
                            item_type = item.get("type", "")
                            item_text = item.get("text", "")
                        else:
                            item_type = getattr(item, "type", "")
                            item_text = getattr(item, "text", "")

                        if item_type in image_part_types:
                            if inserted_images < msg_num_images:
                                glm_content.append({"type": "image"})
                                inserted_images += 1
                        elif item_type in audio_part_types:
                            if inserted_audios < msg_num_audios:
                                glm_content.append({"type": "audio"})
                                inserted_audios += 1
                        elif item_type == "text":
                            glm_content.append({"type": "text", "text": item_text})
                        elif isinstance(item, str):
                            glm_content.append(item)

                missing_media = [
                    *(
                        {"type": "image"}
                        for _ in range(msg_num_images - inserted_images)
                    ),
                    *(
                        {"type": "audio"}
                        for _ in range(msg_num_audios - inserted_audios)
                    ),
                ]
                if missing_media:
                    glm_content[:0] = missing_media
                if not any(
                    isinstance(item, str)
                    or (isinstance(item, dict) and item.get("type") == "text")
                    for item in glm_content
                ):
                    glm_content.append({"type": "text", "text": content})

                formatted_messages.append({"role": role, "content": glm_content})
            elif model_type in {"mimo_v2", "mimo_v2_flash"}:
                # mlx-vlm's get_message_json does not support MiMo, even for
                # text-only history. Format every turn here so one earlier text
                # message cannot trigger the fallback that drops image markers.
                # Render media tokens before templates that drop media dicts.
                mimo_content: list[str] = []
                inserted_images = 0
                inserted_audios = 0
                inserted_text = False
                if isinstance(raw_content, list):
                    for item in raw_content:
                        if isinstance(item, dict):
                            item_type = item.get("type", "")
                            item_text = item.get("text", "")
                        else:
                            item_type = getattr(item, "type", "")
                            item_text = getattr(item, "text", "")

                        if item_type in image_part_types:
                            if inserted_images < msg_num_images:
                                mimo_content.append(
                                    "<|vision_start|><|image_pad|><|vision_end|>"
                                )
                                inserted_images += 1
                        elif item_type in audio_part_types:
                            if inserted_audios < msg_num_audios:
                                mimo_content.append(
                                    "<|mimo_audio_start|><|audio_pad|>"
                                    "<|mimo_audio_end|>"
                                )
                                inserted_audios += 1
                        elif item_type == "text":
                            mimo_content.append(str(item_text))
                            inserted_text = True
                        elif isinstance(item, str):
                            mimo_content.append(item)
                            inserted_text = True

                missing_media = [
                    *(
                        "<|vision_start|><|image_pad|><|vision_end|>"
                        for _ in range(msg_num_images - inserted_images)
                    ),
                    *(
                        "<|mimo_audio_start|><|audio_pad|><|mimo_audio_end|>"
                        for _ in range(msg_num_audios - inserted_audios)
                    ),
                ]
                if missing_media:
                    mimo_content[:0] = missing_media
                if not inserted_text:
                    mimo_content.append(content)

                formatted_messages.append(
                    {"role": role, "content": "".join(mimo_content)}
                )
            elif msg_num_videos > 0:
                # Qwen3-VL chat templates render {"type": "video"} as the
                # vision_start/video_pad/vision_end block the processor expands
                # into timestamped temporal patches. Keep the clip where the
                # user placed it relative to the text.
                video_content: list[dict[str, Any]] = []
                inserted_videos = 0
                for item in raw_content:
                    if isinstance(item, dict):
                        item_type = item.get("type", "")
                        item_text = item.get("text", "")
                    else:
                        item_type = getattr(item, "type", "")
                        item_text = getattr(item, "text", "")
                    if item_type in video_part_types:
                        if inserted_videos < msg_num_videos:
                            video_content.append({"type": "video"})
                            inserted_videos += 1
                    elif item_type == "text" and item_text:
                        video_content.append({"type": "text", "text": item_text})
                formatted_messages.append({"role": role, "content": video_content})
            else:
                formatted = get_message_json(
                    model_type,
                    content,
                    role,
                    skip_image_token=role != "user" or msg_num_images == 0,
                    skip_audio_token=role != "user" or msg_num_audios == 0,
                    num_images=msg_num_images,
                    num_audios=msg_num_audios,
                )
                # Collapse text-only list content to plain string so that
                # simplified chat templates (without render_content macro)
                # can handle it. Image/audio parts stay as a list.
                fc = formatted.get("content")
                if isinstance(fc, list) and all(
                    isinstance(p, dict) and p.get("type") == "text" for p in fc
                ):
                    formatted["content"] = "\n".join(p.get("text", "") for p in fc)
                formatted_messages.append(formatted)

        return formatted_messages, image_message_ranges

    def _compute_vision_features(
        self, pixel_values: Any, extra_model_inputs: dict
    ) -> Optional[mx.array]:
        """Compute vision features for caching.

        Tries multiple strategies based on model architecture:
        1. model.encode_image() — upstream mlx-vlm API (e.g. gemma4)
        2. Direct vision_tower call for qwen-style models
        3. Direct vision_tower + projector for llava-style models
        4. Returns None for unsupported models

        Args:
            pixel_values: Preprocessed image tensors from prepare_inputs().
            extra_model_inputs: Additional model-specific inputs (e.g. image_grid_thw).

        Returns:
            Computed vision features (mx.array), or None if unsupported.
        """
        model = self._vlm_model
        model_type = self.model_type or ""

        if model_type == "deepseek_v4":
            features = model.encode_images(pixel_values, **extra_model_inputs)
            return mx.concatenate(features, axis=0)

        # Strategy 1: upstream encode_image (gemma4 and future models)
        if hasattr(model, "encode_image"):
            image_grid_thw = extra_model_inputs.get("image_grid_thw")
            image_position_ids = extra_model_inputs.get("image_position_ids")
            if image_grid_thw is not None or image_position_ids is not None:
                try:
                    signature = inspect.signature(model.encode_image)
                except (TypeError, ValueError):
                    signature = None

                if signature is None:
                    try:
                        if image_grid_thw is not None:
                            return model.encode_image(
                                pixel_values, image_grid_thw=image_grid_thw
                            )
                        return model.encode_image(
                            pixel_values, image_position_ids=image_position_ids
                        )
                    except TypeError:
                        logger.debug(
                            "encode_image rejected image metadata; "
                            "retrying without it",
                            exc_info=True,
                        )
                else:
                    parameters = signature.parameters
                    accepts_kwargs = any(
                        p.kind == inspect.Parameter.VAR_KEYWORD
                        for p in parameters.values()
                    )
                    if image_grid_thw is not None and (
                        "image_grid_thw" in parameters or accepts_kwargs
                    ):
                        return model.encode_image(
                            pixel_values, image_grid_thw=image_grid_thw
                        )

                    if image_position_ids is not None and (
                        "image_position_ids" in parameters or accepts_kwargs
                    ):
                        return model.encode_image(
                            pixel_values, image_position_ids=image_position_ids
                        )

                    positional_parameters = [
                        p
                        for p in parameters.values()
                        if p.kind
                        in (
                            inspect.Parameter.POSITIONAL_ONLY,
                            inspect.Parameter.POSITIONAL_OR_KEYWORD,
                        )
                    ]
                    if image_position_ids is not None and len(positional_parameters) >= 2:
                        return model.encode_image(pixel_values, image_position_ids)

            return model.encode_image(pixel_values)

        # Strategy 2: qwen-style (vision_tower + grid_thw)
        if model_type in _QWEN_VISION_MODELS:
            grid_thw = extra_model_inputs.get("image_grid_thw")
            if grid_thw is None:
                return None
            dtype = model.vision_tower.patch_embed.proj.weight.dtype
            pv = (
                mx.array(pixel_values)
                if not isinstance(pixel_values, mx.array)
                else pixel_values
            )
            pv = pv.astype(dtype)
            result = model.vision_tower(pv, grid_thw)
            # qwen3_5 returns (hidden_states, _), qwen2_vl returns hidden_states
            if isinstance(result, tuple):
                return result[0]
            return result

        # Strategy 3: llava-style (vision_tower → layer select → projector)
        if model_type == "llava":
            pv = pixel_values
            if not isinstance(pv, mx.array):
                pv = mx.array(pv)
            *_, hidden_states = model.vision_tower(
                pv.transpose(0, 2, 3, 1), output_hidden_states=True
            )
            selected = hidden_states[model.vision_feature_layer]
            if isinstance(model.vision_feature_layer, int):
                if (
                    getattr(model, "vision_feature_select_strategy", "default")
                    == "default"
                ):
                    selected = selected[:, 1:]
            else:
                hs_pool = [hidden_states[idx] for idx in model.vision_feature_layer]
                if (
                    getattr(model, "vision_feature_select_strategy", "default")
                    == "default"
                ):
                    hs_pool = [hs[:, 1:] for hs in hs_pool]
                selected = mx.concatenate(hs_pool, axis=-1)
            return model.multi_modal_projector(selected)

        # Unsupported model: skip caching
        return None

    def _encode_missing_vision_features(
        self,
        pixel_values: Any,
        extra_model_inputs: dict,
        cached_per_image: List[Any],
        per_hashes: List[str],
        image_token_count: Optional[int],
    ) -> Optional[mx.array]:
        """Encode only uncached images and combine with cached ones in order.

        Qwen-style towers attend within each image, so a subset encodes
        independently. Returns None when the request cannot be split safely.
        """
        model = self._vlm_model
        model_type = self.model_type or ""
        grid_thw = extra_model_inputs.get("image_grid_thw")
        if (
            model_type not in _QWEN_VISION_MODELS
            or hasattr(model, "encode_image")
            or grid_thw is None
            or pixel_values is None
            or not hasattr(pixel_values, "shape")
            or pixel_values.ndim != 2
        ):
            return None

        num_images = len(cached_per_image)
        miss_idx = [i for i, f in enumerate(cached_per_image) if f is None]
        if not miss_idx or len(miss_idx) == num_images:
            return None

        grids = [_grid_row(grid_thw, i) for i in range(num_images)]
        if any(g is None for g in grids):
            return None

        rows = [t * h * w for t, h, w in grids]
        if sum(rows) != pixel_values.shape[0]:
            return None

        vision_tower = getattr(model, "vision_tower", None)
        merge_sq = getattr(vision_tower, "spatial_merge_size", 2) ** 2
        # A token count mismatch means a different resize regime.
        for f, (t, h, w) in zip(cached_per_image, grids):
            if f is not None and f.shape[0] != (t * h * w) // merge_sq:
                return None

        try:
            offsets = [0]
            for r in rows:
                offsets.append(offsets[-1] + r)
            pv_miss = mx.concatenate(
                [pixel_values[offsets[i] : offsets[i + 1]] for i in miss_idx],
                axis=0,
            )
            grid_miss = mx.array([grids[i] for i in miss_idx])
            miss_inputs = dict(extra_model_inputs)
            miss_inputs["image_grid_thw"] = grid_miss
            features_miss = self._compute_vision_features(pv_miss, miss_inputs)
            if features_miss is None:
                return None
            mx.eval(features_miss)
            split_miss = self._split_vision_features(
                features_miss, len(miss_idx), miss_inputs
            )
            if split_miss is None or len(split_miss) != len(miss_idx):
                return None

            full = list(cached_per_image)
            for i, f in zip(miss_idx, split_miss):
                full[i] = f
                self._vision_cache.put(
                    per_hashes[i], self._model_name, f, grid=grids[i]
                )
            combined = mx.concatenate(full, axis=0)
            if not self._vision_features_match_image_tokens(
                combined, image_token_count
            ):
                return None
            logger.debug(
                "Vision feature cache partial hit: encoded %d of %d images",
                len(miss_idx),
                num_images,
            )
            return combined
        except Exception:
            logger.debug(
                "Partial vision encoding failed, recomputing all images",
                exc_info=True,
            )
            return None

    def _try_build_cached_vision_inputs(
        self,
        prompt: str,
        images: List[Any],
    ) -> Optional[Dict[str, Any]]:
        """Build ``prepare_inputs``-style inputs, preprocessing only cache misses.

        Splits the prompt on the vision marker, tokenizes the text chunks, and
        expands each marker from the cached or fresh patch grid. Adds
        ``cached_image_features``. Returns None to use the full path.
        """
        model_type = self.model_type or ""
        if (
            model_type not in _CACHED_INPUT_FAST_PREPARE_MODEL_TYPES
            or self._vision_cache is None
            or not self._vision_cache_enabled
            or not images
            or not isinstance(prompt, str)
            or hasattr(self._vlm_model, "encode_image")
        ):
            return None
        try:
            processor = self._processor
            tokenizer = getattr(processor, "tokenizer", None)
            image_processor = getattr(processor, "image_processor", None)
            config = self._vlm_model.config
            vs_id = getattr(config, "vision_start_token_id", None)
            ve_id = getattr(config, "vision_end_token_id", None)
            it_id = getattr(config, "image_token_id", None)
            if (
                tokenizer is None
                or image_processor is None
                or vs_id is None
                or ve_id is None
                or it_id is None
            ):
                return None
            merge_sq = int(getattr(image_processor, "merge_size", 2)) ** 2
            vs_tok = tokenizer.convert_ids_to_tokens(vs_id)
            it_tok = tokenizer.convert_ids_to_tokens(it_id)
            ve_tok = tokenizer.convert_ids_to_tokens(ve_id)
            if not all(isinstance(t, str) for t in (vs_tok, it_tok, ve_tok)):
                return None
            marker = vs_tok + it_tok + ve_tok
            # Manual expansion is only equivalent to the processor when every
            # placeholder in the prompt is the canonical marker triple.
            if prompt.count(marker) != len(images):
                return None

            per_hashes = compute_per_image_hashes(images)
            feats: List[Optional[mx.array]] = []
            grids: List[Optional[List[int]]] = []
            for h in per_hashes:
                feat = self._vision_cache.get(h, self._model_name)
                grid = self._vision_cache.get_grid(h, self._model_name)
                # A row count mismatch means a different resize regime.
                if (
                    feat is not None
                    and grid is not None
                    and (grid[0] * grid[1] * grid[2]) % merge_sq == 0
                    and feat.shape[0] == (grid[0] * grid[1] * grid[2]) // merge_sq
                ):
                    feats.append(feat)
                    grids.append(grid)
                else:
                    feats.append(None)
                    grids.append(None)
            miss_idx = [i for i, f in enumerate(feats) if f is None]
            if len(miss_idx) == len(images):
                return None  # Nothing cached; the full path costs the same.

            if miss_idx:
                # Marker-only prompt; the real text is tokenized below.
                from mlx_vlm.utils import prepare_inputs

                miss_images = [images[i] for i in miss_idx]
                miss_inputs = prepare_inputs(
                    processor,
                    images=miss_images,
                    prompts=[marker * len(miss_idx)],
                )
                miss_pv = miss_inputs.get("pixel_values")
                miss_grid = miss_inputs.get("image_grid_thw")
                if miss_pv is None or miss_grid is None:
                    return None
                miss_grids = [
                    _grid_row(miss_grid, k) for k in range(len(miss_idx))
                ]
                if any(g is None for g in miss_grids):
                    return None
                miss_kw = {"image_grid_thw": mx.array(miss_grids)}
                features_miss = self._compute_vision_features(miss_pv, miss_kw)
                if features_miss is None:
                    return None
                mx.eval(features_miss)
                split_miss = self._split_vision_features(
                    features_miss, len(miss_idx), miss_kw
                )
                if split_miss is None or len(split_miss) != len(miss_idx):
                    return None
                for k, i in enumerate(miss_idx):
                    feats[i] = split_miss[k]
                    grids[i] = miss_grids[k]
                    self._vision_cache.put(
                        per_hashes[i],
                        self._model_name,
                        split_miss[k],
                        grid=miss_grids[k],
                    )

            token_ids: List[int] = []
            chunks = prompt.split(marker)
            for k, chunk in enumerate(chunks):
                if chunk:
                    enc = tokenizer(chunk, add_special_tokens=False)
                    token_ids.extend(enc["input_ids"])
                if k < len(chunks) - 1:
                    t, h, w = grids[k]
                    token_ids.append(vs_id)
                    token_ids.extend([it_id] * ((t * h * w) // merge_sq))
                    token_ids.append(ve_id)

            combined = mx.concatenate(feats, axis=0)
            pad_total = sum(1 for t in token_ids if t == it_id)
            if pad_total != combined.shape[0]:
                return None

            pixel_values = (
                miss_pv if miss_idx else mx.zeros((0, 1), dtype=mx.float32)
            )
            return {
                "input_ids": mx.array([token_ids]),
                "attention_mask": mx.ones((1, len(token_ids)), dtype=mx.int32),
                # Non-None empty tensor keeps the model's multimodal branch
                # taken; with cached_image_features the tower skips it anyway.
                "pixel_values": pixel_values,
                "image_grid_thw": mx.array(grids),
                "mm_token_type_ids": mx.array(
                    [[1 if t == it_id else 0 for t in token_ids]]
                ),
                "cached_image_features": combined,
            }
        except Exception:
            logger.debug(
                "Cached-input fast path failed; falling back to full preprocessing",
                exc_info=True,
            )
            return None

    def _split_vision_features(
        self,
        features: mx.array,
        num_images: int,
        extra_model_inputs: dict,
    ) -> Optional[List[mx.array]]:
        """Split batched vision features into per-image tensors for caching.

        Returns a list of per-image feature tensors, or None if the model
        architecture does not support splitting.
        """
        if self.model_type == "deepseek_v4":
            grid = extra_model_inputs["image_grid_hw"].tolist()
            ratio = self._vlm_model.config.vision_downsample_ratio
            counts = [
                ((h + ratio - 1) // ratio) * ((w + ratio - 1) // ratio)
                for h, w in grid
            ]
            if len(counts) != num_images or sum(counts) != features.shape[0]:
                raise ValueError("DeepSeek V4 cached features do not match image grids")
            offsets = [sum(counts[:i + 1]) for i in range(len(counts))]
            return list(mx.split(features, offsets[:-1], axis=0))

        if num_images <= 1:
            return [features]

        model_type = self.model_type or ""

        # Gemma4 / LLaVA: batch dimension = number of images
        if features.ndim >= 3 and features.shape[0] == num_images:
            return [features[i : i + 1] for i in range(num_images)]

        # Some mlx-vlm models, including Gemma4 unified, return compacted flat
        # features after applying per-image position IDs.
        if features.ndim == 2:
            soft_tokens = self._as_int_list(
                extra_model_inputs.get("num_soft_tokens_per_image")
            )
            if soft_tokens is not None:
                if len(soft_tokens) != num_images:
                    logger.debug(
                        "Per-image soft token count mismatch: expected %d entries, got %d",
                        num_images,
                        len(soft_tokens),
                    )
                    return None
                if sum(soft_tokens) != features.shape[0]:
                    logger.debug(
                        "Per-image soft token total mismatch: expected %d, got %d",
                        sum(soft_tokens),
                        features.shape[0],
                    )
                    return None
                result = []
                offset = 0
                for count in soft_tokens:
                    result.append(features[offset : offset + count])
                    offset += count
                return result

        # Qwen: flat (total_merged_tokens, dim) → split using grid_thw
        if model_type in _GRID_VISION_MODELS and features.ndim == 2:
            grid_thw = extra_model_inputs.get("image_grid_thw")
            if grid_thw is None:
                return None
            vision_tower = getattr(self._vlm_model, "vision_tower", None)
            if vision_tower is None:
                vision_tower = getattr(self._vlm_model, "vision_model", None)
            spatial_merge_size = getattr(vision_tower, "spatial_merge_size", 2)
            merge_sq = spatial_merge_size**2
            per_image_tokens = []
            for i in range(num_images):
                t, h, w = int(grid_thw[i, 0]), int(grid_thw[i, 1]), int(grid_thw[i, 2])
                per_image_tokens.append((t * h * w) // merge_sq)
            if sum(per_image_tokens) != features.shape[0]:
                logger.debug(
                    "Per-image token count mismatch: expected %d, got %d",
                    sum(per_image_tokens),
                    features.shape[0],
                )
                return None
            result = []
            offset = 0
            for count in per_image_tokens:
                result.append(features[offset : offset + count])
                offset += count
            return result

        return None

    @staticmethod
    def _as_int_list(value: Any) -> Optional[List[int]]:
        if value is None:
            return None
        if hasattr(value, "tolist"):
            value = value.tolist()
        if isinstance(value, (int, float)):
            return [int(value)]
        if not isinstance(value, (list, tuple)):
            return None

        result: List[int] = []
        for item in value:
            if hasattr(item, "tolist"):
                item = item.tolist()
            if isinstance(item, (list, tuple)):
                if len(item) != 1:
                    return None
                item = item[0]
            try:
                result.append(int(item))
            except (TypeError, ValueError):
                return None
        return result

    @staticmethod
    def _vision_feature_token_count(features: Any) -> Optional[int]:
        if isinstance(features, (list, tuple)):
            total = 0
            for feature in features:
                count = VLMBatchedEngine._vision_feature_token_count(feature)
                if count is None:
                    return None
                total += count
            return total

        shape = getattr(features, "shape", None)
        if not shape:
            return None
        if len(shape) == 1:
            return 1

        count = 1
        for dim in shape[:-1]:
            count *= int(dim)
        return count

    def _image_token_count(
        self, input_ids: Any, token_id_attr: str = "image_token_id"
    ) -> Optional[int]:
        config = getattr(self._vlm_model, "config", None)
        image_token_id = getattr(config, token_id_attr, None)
        if image_token_id is None:
            return None

        try:
            ids = input_ids if isinstance(input_ids, mx.array) else mx.array(input_ids)
            return int(mx.sum(ids == int(image_token_id)).item())
        except Exception:
            logger.debug("Failed to count VLM image tokens", exc_info=True)
            return None

    def _cached_video_features(
        self, video_identity: str, input_ids: Any, extra_model_inputs: dict
    ) -> mx.array | None:
        """Vision features of a request's clips, through the vision feature cache.

        Follow-up turns resend the same clip. Without the cache the vision
        tower re-encodes every sampled frame on each turn, even when the prefix
        cache already covers the whole prompt.
        """
        if self._vision_cache is None or not self._vision_cache_enabled:
            return None
        token_count = self._image_token_count(input_ids, "video_token_id")
        features = self._vision_cache.get(video_identity, self._model_name)
        if features is not None and self._vision_features_match_image_tokens(
            features, token_count
        ):
            return features
        try:
            # The vision tower encodes a video grid exactly like an image grid.
            features = self._compute_vision_features(
                extra_model_inputs["pixel_values_videos"],
                {"image_grid_thw": extra_model_inputs.get("video_grid_thw")},
            )
        except Exception:
            logger.debug(
                "Video feature computation failed, using full pipeline", exc_info=True
            )
            return None
        if features is None or not self._vision_features_match_image_tokens(
            features, token_count
        ):
            return None
        mx.eval(features)
        self._vision_cache.put(video_identity, self._model_name, features)
        return features

    def _vision_features_match_image_tokens(
        self, features: Any, image_token_count: Optional[int]
    ) -> bool:
        if image_token_count is None:
            return True

        feature_token_count = self._vision_feature_token_count(features)
        if feature_token_count is None:
            return True

        if feature_token_count == image_token_count:
            return True

        logger.debug(
            "Ignoring cached vision features: feature_tokens=%d, image_tokens=%d",
            feature_token_count,
            image_token_count,
        )
        return False

    @staticmethod
    def _language_prompt_kwargs(extra_model_inputs: dict[str, Any]) -> dict[str, Any]:
        """Return processor kwargs that must survive into language prefill."""
        return {
            key: extra_model_inputs[key]
            for key in VLM_LANGUAGE_PROMPT_KWARGS
            if extra_model_inputs.get(key) is not None
        }

    def _prepare_vision_inputs(
        self,
        messages: list[dict[str, Any]],
        images: list[Any],
        audio: list | None = None,
        chat_template_kwargs: dict[str, Any] | None = None,
        tools: list[dict] | None = None,
        is_partial: bool | None = None,
        videos: list[str] | None = None,
        video_hashes: list[str] | None = None,
    ) -> Tuple[
        List[int],
        Optional[mx.array],
        Optional[Dict[str, Any]],
        Optional[str],
        int,
        List[Tuple[int, str]],
    ]:
        """
        Run the full VLM preprocessing pipeline:
        1. Apply chat template with image placeholders
        2. Tokenize and preprocess images via processor
        3. Run vision encoder to produce merged embeddings
        4. Compute image hash for prefix cache

        Args:
            messages: Chat messages (text-only, media already extracted)
            images: List of PIL Image objects
            audio: List of audio data (BytesIO buffers, tuples, or numpy arrays)
            videos: Video file paths for checkpoints with native video input
            video_hashes: Content hashes of ``videos``, for prefix-cache identity
            is_partial: Explicit partial-mode signal from the API server.
                ``True``/``False`` — the server has already decided.  ``None``
                (default) — auto-detect from ``messages`` for direct engine
                callers.

        Returns:
            Tuple of (
                token_ids,
                inputs_embeds,
                extra_kwargs,
                image_hash,
                image_cache_key_start,
                image_cache_key_ranges,
            ):
            - token_ids: List of token IDs for BatchGenerator
            - inputs_embeds: Merged vision+text embeddings (or None if text-only)
            - extra_kwargs: Model-specific kwargs for language model
            - image_hash: Image/audio identity for prefix cache
            - image_cache_key_start: Token index where media-aware keying begins
            - image_cache_key_ranges: Media boundaries with cumulative hashes
        """
        from mlx_vlm.prompt_utils import apply_chat_template, get_chat_template
        from mlx_vlm.utils import load_audio as _load_audio
        from mlx_vlm.utils import prepare_inputs

        num_images = len(images)
        num_audios = len(audio) if audio else 0
        num_videos = len(videos) if videos else 0

        model_type = self.model_type or _read_config_model_type(self._model_name) or ""
        if model_type == COHERE2_MOE_MODEL_TYPE and (
            num_images > 0 or num_audios > 0
        ):
            raise InvalidRequestError(
                "Cohere2 MoE is a text-only model and does not support "
                "image or audio input.",
                field="messages",
            )
        if model_type == QWEN4_EXP_MODEL_TYPE and num_audios > 0:
            raise InvalidRequestError(
                "Qwen4-Exp supports text and image input but not audio.",
                field="messages",
            )

        # Normalize audio to numpy float32 arrays expected by processor.
        # Request-facing string paths are rejected before this point; remaining
        # sources are inline buffers, arrays, or (array, sample_rate) tuples.
        if audio:
            if any(not isinstance(a, tuple) for a in audio):
                from ..patches.mlx_audio_compat import (
                    ensure_mlx_audio_resample_export,
                )

                ensure_mlx_audio_resample_export()
            feature_extractor = getattr(self._processor, "feature_extractor", None)
            target_sample_rate = (
                getattr(feature_extractor, "sampling_rate", 16000)
                if feature_extractor is not None
                else 16000
            )
            if not isinstance(target_sample_rate, (int, float)):
                target_sample_rate = 16000
            audio = [
                _load_audio(a, target_sample_rate) if not isinstance(a, tuple) else a
                for a in audio
            ]
        # Validate multi-image support
        if num_images > 1 and model_type in SINGLE_IMAGE_ONLY_MODELS:
            raise ValueError(
                f"Model {model_type} does not support multi-image chat. "
                f"Please use only 1 image."
            )

        # Apply VLM-specific chat template with image placeholders.
        # Build per-message placeholders in oMLX so image-bearing turns always
        # receive image tokens, regardless of conversation history shape.
        try:
            formatted_messages, image_message_ranges = (
                self._format_messages_for_vlm_template(
                    messages,
                    num_images=num_images,
                    num_audios=num_audios,
                    num_videos=num_videos,
                )
            )
        except Exception as e:
            # The upstream formatter has no video placeholders; falling back
            # would send frames the prompt never references.
            if num_videos:
                raise
            logger.debug(
                "Falling back to mlx-vlm apply_chat_template for VLM formatting: %s",
                e,
            )
            # Fallback to upstream formatter for unknown model/format edge cases.
            formatted_messages = apply_chat_template(
                self._processor,
                self._vlm_model.config,
                messages,
                num_images=num_images,
                num_audios=num_audios,
                return_messages=True,
            )
            image_message_ranges = []
            for idx, msg in enumerate(messages):
                if not isinstance(msg, dict):
                    continue
                image_count = self._count_content_parts(
                    msg.get("content"), {"image", "image_url", "input_image"}
                )
                if image_count > 0:
                    image_message_ranges.append((idx, image_count))

        if is_partial is None:
            # get_message_json() keeps only (role, content), so the flag has to
            # be read from the pre-format messages for direct engine callers.
            is_partial = detect_and_strip_partial(messages)
        # Tool and reasoning_content turns are appended verbatim by
        # _format_messages_for_vlm_template(), so a residual key can survive
        # formatting; the chat template must never see the non-standard field.
        detect_and_strip_partial(formatted_messages)
        template_kwargs = {
            "tokenize": False,
            "add_generation_prompt": not is_partial,
        }
        if is_partial:
            template_kwargs["continue_final_message"] = True
        if self._enable_thinking is not None:
            template_kwargs["enable_thinking"] = self._enable_thinking
        # Per-model/request kwargs override global defaults (e.g. enable_thinking,
        # reasoning_effort).  This mirrors the text-only _apply_chat_template().
        if tools:
            template_kwargs["tools"] = tools
        if chat_template_kwargs:
            template_kwargs.update(chat_template_kwargs)
        _apply_minimax_m3_thinking_mode(model_type, template_kwargs)

        template_target = self._chat_template_target(model_type)
        try:
            prompt = apply_chat_template_with_reasoning_effort_fallback(
                template_target,
                formatted_messages,
                template_kwargs,
                is_harmony=model_type == "gpt_oss",
            )
        except TypeError:
            # Fallback: template doesn't support some kwargs
            if chat_template_kwargs:
                for key in chat_template_kwargs:
                    template_kwargs.pop(key, None)
            template_kwargs.pop("enable_thinking", None)
            prompt = template_target.apply_chat_template(
                formatted_messages, **template_kwargs
            )
        except ValueError as exc:
            if not _is_missing_chat_template_error(exc):
                raise
            # Processor/tokenizer has apply_chat_template but no chat_template
            # set. Some OCR checkpoints (e.g. raw baidu/Unlimited-OCR) ship no
            # chat template at all. mlx-vlm's get_chat_template handles this by
            # rendering the messages into a plain prompt (the <image> tokens are
            # already in the message content from get_message_json), preferring
            # processor.chat_template -> processor.tokenizer.chat_template ->
            # plain rendering, so it subsumes the tokenizer fallback too.
            template_kwargs.pop("tokenize", None)
            template_kwargs.pop("add_generation_prompt", None)
            # get_chat_template() has no continue_final_message equivalent, so
            # partial mode cannot be honoured on this fallback. Drop the kwarg
            # and say so rather than passing it through as an unknown argument.
            template_kwargs.pop("continue_final_message", None)
            if is_partial:
                logger.warning(
                    "Partial mode requested but %s exposes no chat template; "
                    "mlx-vlm plain rendering always starts a new assistant "
                    "turn, so the final message will not be continued.",
                    self._model_name,
                )
            prompt = get_chat_template(
                self._processor,
                formatted_messages,
                add_generation_prompt=True,
                **template_kwargs,
            )

        # Images with cached features and grids skip the image processor.
        fast_cached_features = None
        inputs = None
        if num_audios == 0 and num_videos == 0:
            fast = self._try_build_cached_vision_inputs(prompt, images)
            if fast is not None:
                fast_cached_features = fast.pop("cached_image_features", None)
                inputs = fast
        if inputs is None:
            # Tokenize text and preprocess images and audio
            inputs = prepare_inputs(
                self._processor,
                images=images if images else None,
                audio=audio if audio else None,
                videos=videos if videos else None,
                prompts=[prompt] if isinstance(prompt, str) else prompt,
            )

        input_ids = inputs["input_ids"]
        pixel_values = inputs.get("pixel_values")
        attention_mask = inputs.get("attention_mask")

        token_ids = input_ids[0].tolist() if input_ids.ndim > 1 else input_ids.tolist()
        image_cache_key_start = 0
        image_cache_key_ranges: list[Tuple[int, str]] = []
        if image_message_ranges:
            try:
                image_starts = None
                if (
                    model_type in _GRID_VISION_MODELS
                    and inputs.get("image_grid_thw") is not None
                ):
                    image_starts = _grid_image_token_starts(
                        token_ids,
                        inputs["image_grid_thw"],
                        self._vlm_model.config.image_token_id,
                        self._processor.image_processor.merge_size,
                    )
                    if (
                        len(image_starts) != num_images
                        or sum(count for _, count in image_message_ranges) != num_images
                    ):
                        raise ValueError("Image boundary count does not match images")
                prefix_template_kwargs = {
                    "tokenize": False,
                    "add_generation_prompt": False,
                }
                if self._enable_thinking is not None:
                    prefix_template_kwargs["enable_thinking"] = self._enable_thinking
                if tools:
                    prefix_template_kwargs["tools"] = tools
                if chat_template_kwargs:
                    prefix_template_kwargs.update(chat_template_kwargs)
                _apply_minimax_m3_thinking_mode(model_type, prefix_template_kwargs)

                images_consumed = 0
                for msg_idx, msg_num_images in image_message_ranges:
                    prefix_messages = formatted_messages[:msg_idx]
                    boundary_tokens = 0
                    if image_starts is not None:
                        boundary_tokens = image_starts[images_consumed]
                    elif prefix_messages:
                        try:
                            prefix_prompt = (
                                apply_chat_template_with_reasoning_effort_fallback(
                                    template_target,
                                    prefix_messages,
                                    prefix_template_kwargs,
                                    is_harmony=model_type == "gpt_oss",
                                )
                            )
                        except TypeError:
                            local_kwargs = dict(prefix_template_kwargs)
                            if chat_template_kwargs:
                                for key in chat_template_kwargs:
                                    local_kwargs.pop(key, None)
                            local_kwargs.pop("enable_thinking", None)
                            prefix_prompt = template_target.apply_chat_template(
                                prefix_messages, **local_kwargs
                            )
                        prefix_inputs = prepare_inputs(
                            self._processor,
                            images=(
                                images[:images_consumed]
                                if images_consumed > 0
                                else None
                            ),
                            prompts=(
                                [prefix_prompt]
                                if isinstance(prefix_prompt, str)
                                else prefix_prompt
                            ),
                        )
                        prefix_ids = prefix_inputs["input_ids"]
                        prefix_tokens = (
                            prefix_ids[0].tolist()
                            if prefix_ids.ndim > 1
                            else prefix_ids.tolist()
                        )
                        # Rendering a shorter conversation can retain reasoning
                        # that the full template removes. Only a matching token
                        # prefix is a valid position in the final model input.
                        for actual, prefix in zip(token_ids, prefix_tokens):
                            if actual != prefix:
                                break
                            boundary_tokens += 1

                    images_consumed += msg_num_images
                    cumulative_hash = compute_image_hash(images[:images_consumed])
                    image_cache_key_ranges.append((boundary_tokens, cumulative_hash))

                # A later image's prefix can diverge earlier. Its cumulative
                # hash must apply there, including all preceding images.
                for i in range(len(image_cache_key_ranges) - 2, -1, -1):
                    start, image_key = image_cache_key_ranges[i]
                    next_start = image_cache_key_ranges[i + 1][0]
                    image_cache_key_ranges[i] = (min(start, next_start), image_key)
                image_cache_key_start = image_cache_key_ranges[0][0]
            except Exception:
                logger.debug(
                    "Failed to compute segmented VLM cache boundaries, "
                    "falling back to whole-request keying",
                )
                image_cache_key_start = 0
                image_cache_key_ranges = []

        # Extract additional model-specific inputs (filter None values
        # since prepare_inputs may include them after mlx-vlm 348466f)
        extra_model_inputs = {
            k: v
            for k, v in inputs.items()
            if k not in (
                "input_ids",
                "attention_mask",
                "pixel_values",
            )
            and v is not None
        }

        # Check for any multimodal inputs. Most processors expose audio as
        # ``input_features``; MiMo's tokenizer produces discrete ``audio_codes``.
        has_audio = any(
            key in extra_model_inputs for key in ("input_features", "audio_codes")
        )
        has_video = extra_model_inputs.get("pixel_values_videos") is not None
        has_multimodal = (
            (pixel_values is not None and num_images > 0) or has_audio or has_video
        )

        if has_multimodal:
            # Build call kwargs from extra_model_inputs (includes input_features
            # for audio, image_grid_thw, etc.)
            call_kwargs = dict(extra_model_inputs)
            if fast_cached_features is not None:
                call_kwargs["cached_image_features"] = fast_cached_features

            # Image-specific: compute hash and try vision feature cache
            image_hash = None
            image_token_count = None
            if num_images > 0:
                image_hash = compute_image_hash(images)
                image_token_count = self._image_token_count(input_ids)
            elif has_video:
                image_cache_key_ranges = self._video_cache_key_ranges(
                    token_ids,
                    extra_model_inputs.get("video_grid_thw"),
                    video_hashes or [],
                )
                if image_cache_key_ranges:
                    image_cache_key_start = image_cache_key_ranges[0][0]
                    image_hash = image_cache_key_ranges[-1][1]
                else:
                    image_cache_key_start = 0
                    image_hash = self._video_cache_identity(video_hashes or [])
                video_features = self._cached_video_features(
                    image_hash, input_ids, extra_model_inputs
                )
                if video_features is not None:
                    call_kwargs["cached_image_features"] = video_features

            if (
                num_images > 0
                and self._vision_cache is not None
                and self._vision_cache_enabled
                # Fast path already assembled the combined features.
                and fast_cached_features is None
            ):
                per_hashes = compute_per_image_hashes(images)
                cached_per_image = [
                    self._vision_cache.get(h, self._model_name) for h in per_hashes
                ]

                # Per-image entries are keyed by the image alone, but the
                # number of soft tokens an image encodes to depends on the
                # resize regime, which depends on the *other* images in the
                # request (Gemma 4 per-image resize: 1024x1024 -> 256 tokens,
                # 1536x640 -> 250). Entries cached from separate single-image
                # requests can therefore disagree, and mx.concatenate raises
                # before _vision_features_match_image_tokens below ever gets to
                # reject them. Check the shapes agree first and fall through to
                # the whole-request entry (and then a recompute) when they do
                # not.
                per_image_usable = (
                    all(f is not None for f in cached_per_image)
                    and len({f.shape[1:] for f in cached_per_image}) == 1
                )

                cached_whole = None
                if not per_image_usable:
                    # Fallback: whole-request entry (stored when per-image split
                    # is unsupported, e.g. Gemma 4 multi-image with per-image
                    # resize). Mirrors the store-side branch below.
                    cached_whole = self._vision_cache.get(image_hash, self._model_name)

                used_cached_features = False
                if per_image_usable:
                    # All images cached individually — combine and use
                    combined = mx.concatenate(cached_per_image, axis=0)
                    if self._vision_features_match_image_tokens(
                        combined, image_token_count
                    ):
                        call_kwargs["cached_image_features"] = combined
                        used_cached_features = True
                        logger.debug(
                            "Vision feature cache hit (per-image): all %d images cached",
                            num_images,
                        )
                elif cached_whole is not None:
                    if self._vision_features_match_image_tokens(
                        cached_whole, image_token_count
                    ):
                        call_kwargs["cached_image_features"] = cached_whole
                        used_cached_features = True
                        logger.debug(
                            "Vision feature cache hit (whole-request): %s",
                            image_hash[:16],
                        )

                if not used_cached_features:
                    # Partial hit: encode only the uncached images when the
                    # vision tower supports per-image slicing.
                    partial = self._encode_missing_vision_features(
                        pixel_values,
                        extra_model_inputs,
                        cached_per_image,
                        per_hashes,
                        image_token_count,
                    )
                    if partial is not None:
                        call_kwargs["cached_image_features"] = partial
                        used_cached_features = True

                if not used_cached_features:
                    # Some or all uncached — compute all, then cache per-image
                    try:
                        features = self._compute_vision_features(
                            pixel_values, extra_model_inputs
                        )
                        if (
                            features is not None
                            and self._vision_features_match_image_tokens(
                                features, image_token_count
                            )
                        ):
                            mx.eval(features)
                            call_kwargs["cached_image_features"] = features
                            # Split and cache each image individually
                            per_features = self._split_vision_features(
                                features, num_images, extra_model_inputs
                            )
                            if per_features is not None:
                                grid_thw = extra_model_inputs.get("image_grid_thw")
                                for j, (h, f) in enumerate(
                                    zip(per_hashes, per_features)
                                ):
                                    grid = _grid_row(grid_thw, j)
                                    self._vision_cache.put(
                                        h, self._model_name, f, grid=grid
                                    )
                                logger.debug(
                                    "Vision feature cache miss, stored %d per-image entries",
                                    len(per_features),
                                )
                            else:
                                # Split unsupported for this model — store whole-request
                                self._vision_cache.put(
                                    image_hash, self._model_name, features
                                )
                                logger.debug(
                                    "Vision feature cache miss, stored whole-request: %s",
                                    image_hash[:16],
                                )
                    except Exception:
                        logger.debug(
                            "Vision feature computation failed, using full pipeline",
                            exc_info=True,
                        )

            if (
                self.model_type == "deepseek_v4"
                and "cached_image_features" in call_kwargs
            ):
                call_kwargs["cached_image_features"] = self._split_vision_features(
                    call_kwargs["cached_image_features"], num_images, extra_model_inputs
                )

            # Run vision encoder + embedding merge.
            # Pass attention_mask as 'mask' — mlx-vlm models (e.g. Gemma 3)
            # expect it as a positional/keyword arg named 'mask'.
            try:
                embed_features = self._vlm_model.get_input_embeddings(
                    input_ids, pixel_values, mask=attention_mask, **call_kwargs
                )
            except TypeError:
                # cached_image_features kwarg not supported — disable and retry
                if "cached_image_features" in call_kwargs:
                    logger.warning(
                        "cached_image_features not supported by %s, "
                        "disabling vision feature cache",
                        self.model_type,
                    )
                    self._vision_cache_enabled = False
                    call_kwargs.pop("cached_image_features")
                    embed_features = self._vlm_model.get_input_embeddings(
                        input_ids, pixel_values, mask=attention_mask, **call_kwargs
                    )
                else:
                    raise
            mx.eval(embed_features.inputs_embeds)

            # Convert InputEmbeddingsFeatures to dict for extra kwargs
            extra_kwargs = {}
            if hasattr(embed_features, "to_dict"):
                feat_dict = embed_features.to_dict()
                for k, v in feat_dict.items():
                    if k != "inputs_embeds" and v is not None:
                        extra_kwargs[k] = v
            for k, v in self._language_prompt_kwargs(extra_model_inputs).items():
                extra_kwargs.setdefault(k, v)

            # Capture per-request mRoPE state set by
            # get_input_embeddings(). The language model stores these as
            # global state that gets overwritten by subsequent calls.
            # Storing per-request ensures correct position computation
            # when multiple VLM requests are batched.
            _capture_vlm_position_state(
                getattr(self._vlm_model, "language_model", None), extra_kwargs
            )

            if model_type in {"mimo_v2", "mimo_v2_flash"} and has_audio:
                from ..patches.mimo_v2.audio import audio_cache_key_ranges

                # Keep vision-feature identities independent of audio inputs.
                image_ranges = image_cache_key_ranges
                if image_hash is not None and not image_ranges:
                    image_ranges = [(0, image_hash)]
                image_cache_key_ranges = audio_cache_key_ranges(
                    token_ids,
                    extra_model_inputs["audio_codes"],
                    self._vlm_model.config.audio_token_id,
                    image_ranges,
                )
                image_cache_key_start = image_cache_key_ranges[0][0]
                image_hash = image_cache_key_ranges[-1][1]
            elif has_audio and "input_features" in extra_model_inputs:
                image_ranges = image_cache_key_ranges
                if image_hash is not None and not image_ranges:
                    image_ranges = [(0, image_hash)]
                config = self._vlm_model.config
                audio_token_id = getattr(config, "audio_token_id", None)
                if audio_token_id is None:
                    audio_token_id = getattr(config, "audio_token_index", None)
                image_cache_key_ranges = _audio_feature_cache_key_ranges(
                    token_ids,
                    extra_model_inputs["input_features"],
                    extra_model_inputs.get("input_features_mask"),
                    audio_token_id,
                    image_ranges,
                )
                image_cache_key_start = image_cache_key_ranges[0][0]
                image_hash = image_cache_key_ranges[-1][1]

            return (
                token_ids,
                embed_features.inputs_embeds,
                extra_kwargs,
                image_hash,
                image_cache_key_start,
                image_cache_key_ranges,
            )
        else:
            # Text-only (no images in this message)
            return token_ids, None, None, None, 0, []

    def _apply_chat_template(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict] | None = None,
        chat_template_kwargs: dict[str, Any] | None = None,
        is_partial: bool | None = None,
        add_generation_prompt: bool | None = None,
    ) -> str:
        """Apply chat template for text-only messages (no images).

        Args:
            is_partial: Explicit partial-mode signal from the API server.
                ``True``/``False`` — the server has already decided; the
                ``partial`` key is cleaned from message dicts but no detection
                is performed.  ``None`` (default) — auto-detect from messages
                for direct engine callers.
            add_generation_prompt: Overrides the partial-derived default, used
                to render the same messages without the generation prompt.
        """
        if hasattr(self._tokenizer, "apply_chat_template"):
            if is_partial is None:
                is_partial = detect_and_strip_partial(messages)
            else:
                # Server already resolved partial; just clean residual keys
                # so the chat template never sees the non-standard field.
                for msg in messages:
                    msg.pop("partial", None)
            if add_generation_prompt is None:
                add_generation_prompt = not is_partial
            template_kwargs = {
                "tokenize": False,
                "add_generation_prompt": add_generation_prompt,
            }
            if is_partial:
                template_kwargs["continue_final_message"] = True
            if tools:
                template_kwargs["tools"] = tools
            if self._enable_thinking is not None:
                template_kwargs["enable_thinking"] = self._enable_thinking
            if chat_template_kwargs:
                template_kwargs.update(chat_template_kwargs)
            _apply_minimax_m3_thinking_mode(self.model_type, template_kwargs)

            try:
                return apply_chat_template_with_reasoning_effort_fallback(
                    self._tokenizer,
                    messages,
                    template_kwargs,
                    is_harmony=self.model_type == "gpt_oss",
                )
            except TypeError:
                if chat_template_kwargs:
                    for key in chat_template_kwargs:
                        template_kwargs.pop(key, None)
                template_kwargs.pop("tools", None)
                template_kwargs.pop("enable_thinking", None)
                return self._tokenizer.apply_chat_template(messages, **template_kwargs)
            except ValueError as exc:
                if not _is_missing_chat_template_error(exc):
                    raise
                # Tokenizer exposes apply_chat_template but has no chat_template
                # set (e.g. raw baidu/Unlimited-OCR ships none). Fall back to
                # mlx-vlm's plain-message rendering, matching the vision path.
                from mlx_vlm.prompt_utils import get_chat_template

                if is_partial:
                    logger.warning(
                        "Partial mode requested but %s exposes no chat "
                        "template; mlx-vlm plain rendering always starts a "
                        "new assistant turn, so the final message will not "
                        "be continued.",
                        self._model_name,
                    )
                return get_chat_template(
                    self._processor,
                    messages,
                    add_generation_prompt=add_generation_prompt,
                )
        else:
            prompt = "\n".join(f"{m['role']}: {m['content']}" for m in messages)
            return prompt + "\nassistant:"

    @staticmethod
    def _pop_specprefill_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
        """Pop SpecPrefill per-request overrides out of ``kwargs``.

        The engine's ``add_request`` accepts these as dedicated arguments, so
        they must be forwarded explicitly rather than left in ``**kwargs``.
        Shared by ``generate`` and ``stream_generate`` so both request paths
        honour SpecPrefill overrides identically.
        """
        specprefill_kwargs: dict[str, Any] = {}
        for key in (
            "specprefill",
            "specprefill_keep_pct",
            "specprefill_threshold",
            "specprefill_system_end",
            "generation_prompt_text",
            "generation_prompt_persists",
        ):
            if kwargs.get(key) is not None:
                specprefill_kwargs[key] = kwargs.pop(key)
        return specprefill_kwargs

    def _inject_specprefill_system_end(
        self,
        messages: list[dict[str, Any]],
        prompt: str | list[int],
        kwargs: dict[str, Any],
    ) -> None:
        """Compute the system-prompt token boundary and add it to ``kwargs``.

        SpecPrefill protects the system-prompt region from token dropping. The
        boundary is derived by subtracting the non-system prompt token count
        from the full prompt token count (system-only messages usually can't be
        templated on their own). Shared by ``chat`` and ``stream_chat`` so the
        non-streaming path protects the system prompt identically. No-op unless
        the model has SpecPrefill enabled and the request has a system prompt.

        ``prompt`` is the already-tokenized VLM prompt (a list of token IDs,
        per ``_process_chat_messages``), so the full-prompt count is just its
        length rather than a re-encode.
        """
        specprefill_model_enabled = (
            getattr(self._model_settings, "specprefill_enabled", False)
            if self._model_settings
            else False
        )
        if not (specprefill_model_enabled and kwargs.get("specprefill") is not False):
            return
        non_system = [
            m for m in messages if m.get("role") not in ("system", "developer")
        ]
        if len(non_system) < len(messages) and non_system:
            try:
                non_system_prompt = self._tokenizer.apply_chat_template(
                    non_system,
                    tokenize=False,
                    add_generation_prompt=True,
                )
                full_tokens = len(prompt)
                non_system_tokens = len(self._tokenizer.encode(non_system_prompt))
                system_end = full_tokens - non_system_tokens
                if system_end > 0:
                    kwargs["specprefill_system_end"] = system_end
            except Exception as e:
                logger.debug(f"SpecPrefill: system_end calc failed: {e}")

    async def generate(
        self,
        prompt: str | list[int],
        max_tokens: int = 256,
        temperature: float = 0.7,
        top_p: float = 0.9,
        top_k: int = 0,
        min_p: float = 0.0,
        repetition_penalty: float = 1.0,
        presence_penalty: float = 0.0,
        stop: list[str] | None = None,
        vlm_inputs_embeds: Any = None,
        vlm_extra_kwargs: dict[str, Any] | None = None,
        vlm_image_hash: str | None = None,
        vlm_cache_key_start: int = 0,
        vlm_cache_key_ranges: Optional[List[Tuple[int, str]]] = None,
        **kwargs,
    ) -> GenerationOutput:
        """Generate a complete response (non-streaming)."""
        if not self._loaded:
            await self.start()

        if self.is_diffusion_model:
            full_text = ""
            last_output: GenerationOutput | None = None
            async for output in self.stream_generate(
                prompt,
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                min_p=min_p,
                repetition_penalty=repetition_penalty,
                presence_penalty=presence_penalty,
                stop=stop,
                **kwargs,
            ):
                full_text += output.new_text
                last_output = output
            if last_output is None:
                return GenerationOutput(text="", prompt_tokens=0, completion_tokens=0)
            return GenerationOutput(
                text=full_text,
                prompt_tokens=last_output.prompt_tokens,
                completion_tokens=last_output.completion_tokens,
                finish_reason=last_output.finish_reason,
                cached_tokens=0,
            )

        # OCR models: add extra stop token IDs to prevent degeneration.
        # Sampling params (temperature, repetition_penalty, max_tokens) are
        # resolved by get_sampling_params() with OCR defaults as a fallback
        # layer, so admin/API overrides are respected.
        extra_stop_ids: list[int] = []
        if self.is_ocr_model:
            extra_stop_ids = self._resolve_ocr_stop_token_ids()

        from ..request import SamplingParams

        sampling_params = SamplingParams(
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            min_p=min_p,
            xtc_probability=kwargs.get("xtc_probability", 0.0),
            xtc_threshold=kwargs.get("xtc_threshold", 0.1),
            repetition_penalty=repetition_penalty,
            repetition_context_size=kwargs.get("repetition_context_size"),
            presence_penalty=presence_penalty,
            frequency_penalty=kwargs.get("frequency_penalty", 0.0),
            stop=stop or [],
            stop_token_ids=extra_stop_ids or None,
            thinking_budget=kwargs.get("thinking_budget", None),
            compiled_grammar=kwargs.get("compiled_grammar", None),
            seed=kwargs.get("seed", None),
        )

        # SpecPrefill: forward per-request overrides to the engine, mirroring
        # stream_generate so the non-streaming path is not silently ignored.
        specprefill_kwargs = self._pop_specprefill_kwargs(kwargs)
        tools = kwargs.pop("tools", None)

        output = await self._engine.generate(
            prompt=prompt,
            sampling_params=sampling_params,
            vlm_inputs_embeds=vlm_inputs_embeds,
            vlm_extra_kwargs=vlm_extra_kwargs,
            vlm_image_hash=vlm_image_hash,
            vlm_cache_key_start=vlm_cache_key_start,
            vlm_cache_key_ranges=vlm_cache_key_ranges,
            tools=tools,
            preserve_reasoning=bool(kwargs.get("preserve_reasoning", False)),
            **specprefill_kwargs,
        )

        text = clean_special_tokens(output.output_text)

        return GenerationOutput(
            text=text,
            prompt_tokens=output.prompt_tokens,
            completion_tokens=output.completion_tokens,
            finish_reason=output.finish_reason,
            tool_calls=output.tool_calls,
            cached_tokens=output.cached_tokens,
            first_token_at=output.first_token_at,
        )

    async def stream_generate(
        self,
        prompt: str | list[int],
        max_tokens: int = 256,
        temperature: float = 0.7,
        top_p: float = 0.9,
        top_k: int = 0,
        min_p: float = 0.0,
        repetition_penalty: float = 1.0,
        presence_penalty: float = 0.0,
        stop: list[str] | None = None,
        vlm_inputs_embeds: Any = None,
        vlm_extra_kwargs: dict[str, Any] | None = None,
        vlm_image_hash: str | None = None,
        vlm_cache_key_start: int = 0,
        vlm_cache_key_ranges: Optional[List[Tuple[int, str]]] = None,
        **kwargs,
    ) -> AsyncIterator[GenerationOutput]:
        """Stream generation token by token."""
        if not self._loaded:
            await self.start()

        if self.is_diffusion_model:
            if (
                vlm_inputs_embeds is not None
                or vlm_extra_kwargs is not None
                or vlm_image_hash is not None
                or vlm_cache_key_ranges is not None
                or vlm_cache_key_start
            ):
                raise InvalidRequestError(
                    "Precomputed VLM embeddings and cache metadata are not "
                    "supported with diffusion models."
                )
            self._validate_diffusion_request(
                stop=stop,
                kwargs=kwargs,
            )
            loop = asyncio.get_running_loop()
            from ..engine_core import get_mlx_executor

            diffusion_inputs = await loop.run_in_executor(
                get_mlx_executor(),
                self._prepare_diffusion_inputs_from_prompt,
                prompt,
            )
            async for output in self._stream_diffusion_inputs(
                diffusion_inputs,
                max_tokens=max_tokens,
                temperature=temperature,
                seed=kwargs.get("seed"),
            ):
                yield output
            return

        # OCR models: add extra stop token IDs to prevent degeneration.
        # Sampling params (temperature, repetition_penalty, max_tokens) are
        # resolved by get_sampling_params() with OCR defaults as a fallback
        # layer, so admin/API overrides are respected.
        extra_stop_ids: list[int] = []
        if self.is_ocr_model:
            extra_stop_ids = self._resolve_ocr_stop_token_ids()

        from ..request import SamplingParams

        sampling_params = SamplingParams(
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            min_p=min_p,
            xtc_probability=kwargs.get("xtc_probability", 0.0),
            xtc_threshold=kwargs.get("xtc_threshold", 0.1),
            repetition_penalty=repetition_penalty,
            repetition_context_size=kwargs.get("repetition_context_size"),
            presence_penalty=presence_penalty,
            frequency_penalty=kwargs.get("frequency_penalty", 0.0),
            stop=stop or [],
            stop_token_ids=extra_stop_ids or None,
            thinking_budget=kwargs.get("thinking_budget", None),
            compiled_grammar=kwargs.get("compiled_grammar", None),
            seed=kwargs.get("seed", None),
        )

        # SpecPrefill: pass per-request overrides
        specprefill_kwargs = self._pop_specprefill_kwargs(kwargs)
        tools = kwargs.pop("tools", None)

        engine = self._engine
        request_id = await engine.add_request(
            prompt=prompt,
            sampling_params=sampling_params,
            vlm_inputs_embeds=vlm_inputs_embeds,
            vlm_extra_kwargs=vlm_extra_kwargs,
            vlm_image_hash=vlm_image_hash,
            vlm_cache_key_start=vlm_cache_key_start,
            vlm_cache_key_ranges=vlm_cache_key_ranges,
            skip_cache_store=bool(kwargs.get("skip_cache_store", False)),
            preserve_reasoning=bool(kwargs.get("preserve_reasoning", False)),
            benchmark_trace=bool(kwargs.get("benchmark_trace", False)),
            benchmark_ane_sequence_length=int(
                kwargs.get("benchmark_ane_sequence_length", 0) or 0
            ),
            tools=tools,
            **specprefill_kwargs,
        )

        finished_normally = False
        try:
            async for output in engine.stream_outputs(request_id):
                text = clean_special_tokens(output.output_text)

                if output.finished:
                    finished_normally = True

                yield GenerationOutput(
                    text=text,
                    new_text=output.new_text,
                    prompt_tokens=output.prompt_tokens,
                    completion_tokens=output.completion_tokens,
                    finished=output.finished,
                    finish_reason=output.finish_reason,
                    tool_calls=output.tool_calls,
                    cached_tokens=output.cached_tokens,
                    generated_at=getattr(output, "generated_at", None),
                    generated_until=getattr(output, "generated_until", None),
                    first_token_at=getattr(output, "first_token_at", None),
                    benchmark_prefill_chunks=(
                        list(chunks)
                        if (chunks := getattr(output, "benchmark_prefill_chunks", []))
                        else []
                    ),
                    benchmark_requested_steps=(
                        list(steps)
                        if (steps := getattr(output, "benchmark_requested_steps", []))
                        else []
                    ),
                    benchmark_boundary_enabled=bool(
                        getattr(output, "benchmark_boundary_enabled", False)
                    ),
                    benchmark_cache_block_size=int(
                        getattr(output, "benchmark_cache_block_size", 0) or 0
                    ),
                )
        except GeneratorExit:
            logger.info(f"[vlm_stream_generate] GeneratorExit for request {request_id}")
        finally:
            if not finished_normally:
                logger.info(f"[vlm_stream_generate] Aborting request {request_id}")
                await engine.abort_request(request_id)

    async def chat(
        self,
        messages: list[dict[str, Any]],
        max_tokens: int = 256,
        temperature: float = 0.7,
        top_p: float = 0.9,
        top_k: int = 0,
        min_p: float = 0.0,
        repetition_penalty: float = 1.0,
        presence_penalty: float = 0.0,
        tools: list[dict] | None = None,
        **kwargs,
    ) -> GenerationOutput:
        """Chat completion with vision support (non-streaming)."""
        if not self._loaded:
            await self.start()

        if self.is_diffusion_model:
            full_text = ""
            last_output: GenerationOutput | None = None
            async for output in self.stream_chat(
                messages,
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                min_p=min_p,
                repetition_penalty=repetition_penalty,
                presence_penalty=presence_penalty,
                tools=tools,
                **kwargs,
            ):
                full_text += output.new_text
                last_output = output
            if last_output is None:
                return GenerationOutput(text="", prompt_tokens=0, completion_tokens=0)
            return GenerationOutput(
                text=full_text,
                prompt_tokens=last_output.prompt_tokens,
                completion_tokens=last_output.completion_tokens,
                finish_reason=last_output.finish_reason,
                cached_tokens=0,
            )

        loop = asyncio.get_running_loop()
        # _process_chat_messages pops these; the tail marker needs them too.
        ct_kwargs = kwargs.get("chat_template_kwargs")
        partial = kwargs.get("is_partial")
        (
            prompt,
            vlm_embeds,
            vlm_kwargs,
            image_hash,
            image_cache_key_start,
            image_cache_key_ranges,
        ) = await loop.run_in_executor(
            self._engine._mlx_executor,
            self._process_chat_messages,
            messages,
            tools,
            kwargs,
        )

        # SpecPrefill: protect the system-prompt region, mirroring stream_chat.
        self._inject_specprefill_system_end(messages, prompt, kwargs)
        generation_prompt, persists = self._generation_prompt_text(ct_kwargs, partial)
        if generation_prompt:
            kwargs["generation_prompt_text"] = generation_prompt
            kwargs["generation_prompt_persists"] = persists

        return await self.generate(
            prompt=prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            min_p=min_p,
            repetition_penalty=repetition_penalty,
            presence_penalty=presence_penalty,
            vlm_inputs_embeds=vlm_embeds,
            vlm_extra_kwargs=vlm_kwargs,
            vlm_image_hash=image_hash,
            vlm_cache_key_start=image_cache_key_start,
            vlm_cache_key_ranges=image_cache_key_ranges,
            tools=tools,
            **kwargs,
        )

    async def preflight_chat(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict] | None = None,
        request_id: str | None = None,
        **kwargs,
    ) -> None:
        """Early prefill memory check for chat completions (VLM path).

        The actual VLM prompt is built by ``_process_chat_messages`` →
        ``_prepare_vision_inputs``, which expands each image content-part
        into 256–1280 model-specific image-placeholder tokens before the
        chat template runs. Doing that work here would require image
        decoding + the heavy preprocessor pipeline; for preflight we only
        need a conservative upper bound on the prompt size, so we instead:

          1. Apply the *text-only* chat template (cheap).
          2. Count its tokens.
          3. Add a per-image upper-bound budget (``_IMAGE_TOKEN_UPPER_BOUND``)
             for each image-bearing content part — over-counts somewhat
             on small images (false-positive 400s for borderline-and-image
             cases) but never under-counts, which is the property the
             guard needs to stay safe against the Apple IOGPUFamily
             panic path.

        Tools (when supplied as Pydantic ``ToolDefinition`` objects by
        direct API callers) must be converted to dict form for the
        template — ``BatchedEngine.preflight_chat`` does this and we
        mirror it here. Without conversion the template's ``TypeError``
        retry path silently drops tools entirely, which not only
        miscalibrates the token count but also bypasses the actual
        tool-prompt rendering on the real chat path.

        Raises ``PrefillMemoryExceededError`` if the conservative estimate
        would exceed the configured memory ceiling. See
        ``BatchedEngine.preflight_chat`` for the upstream rationale
        (avoiding the ``StreamingResponse`` 200 commit so HTTP 400
        actually reaches the client).
        """
        if not self._loaded:
            await self.start()
        if self.is_diffusion_model:
            _, _, audio = extract_images_from_messages(messages)
            self._validate_diffusion_request(
                tools=tools,
                audio=audio if audio else None,
                stop=kwargs.get("stop"),
                kwargs=kwargs,
            )
            return
        template_tools = convert_tools_for_template(tools) if tools else None
        ct_kwargs = kwargs.get("chat_template_kwargs")
        partial = kwargs.get("is_partial")
        # Strip image content-parts BEFORE templating. Modern HF chat
        # templates (Qwen2.5-VL, Gemma-Vision, Llama-3.2-Vision) render
        # ``image_url`` / ``image`` content parts as literal placeholder
        # strings inline with the text; if we leave them in, the
        # tokenized prompt already contains some image-placeholder
        # tokens AND we then add the per-image budget on top — a double
        # count that rejects borderline image-bearing prompts the real
        # chat path would have handled. The real ``chat`` flow itself
        # strips images first via ``extract_images_from_messages`` (see
        # ``_process_chat_messages``), so mirroring that here keeps
        # preflight and execution on the same template input.
        _, text_messages, images, _, videos = self._extract_request_media(messages)
        if (
            images
            and self.model_type in {"gemma4", "gemma4_unified"}
            and self._vlm_model.config.vision_config is None
        ):
            raise InvalidRequestError(
                "This text-only Gemma 4 model does not support image input.",
                field="messages",
            )
        prompt = self._apply_chat_template(
            text_messages,
            template_tools,
            chat_template_kwargs=ct_kwargs,
            is_partial=partial,
        )
        # Tokenizer errors propagate as 500 today regardless of where they
        # fire; the real chat path's add_request → tokenize call has no
        # path-specific 400 handler. Don't introduce a NEW failure mode
        # in preflight: skip the memory check on tokenizer error and let
        # the real chat path surface the same error through the existing
        # handler chain.
        try:
            num_tokens = len(self._tokenizer.encode(prompt))
        except Exception as e:
            logger.warning(
                "VLMBatchedEngine.preflight_chat: tokenizer.encode raised "
                "%s; skipping prefill memory check, real chat path will "
                "surface the error",
                type(e).__name__,
            )
            return
        # Use the decoded dimensions, including resizing and EXIF orientation,
        # so the memory estimate matches the images passed to the processor.
        image_tokens = _count_image_tokens_real(
            messages,
            getattr(self, "_processor", None),
            images=images,
            upper_bound=_derive_image_token_upper_bound(
                getattr(self, "_processor", None)
            ),
        )
        num_tokens += image_tokens
        video_tokens = 0
        if videos:
            video_processor = getattr(self._processor, "video_processor", None)
            video_tokens = sum(
                estimate_native_video_tokens(uri, video_processor) for uri in videos
            )
        num_tokens += video_tokens
        scheduler = getattr(getattr(self._engine, "engine", None), "scheduler", None)
        if scheduler is None:
            _warn_scheduler_unreachable_once(self, "preflight_chat")
            return
        await self._preflight_or_raise_with_eviction(
            scheduler,
            num_prompt_tokens=num_tokens,
            request_id=request_id,
            text_only=image_tokens == 0 and video_tokens == 0,
        )

    async def preflight_completion(
        self,
        prompt: str,
        request_id: str | None = None,
        **kwargs,
    ) -> None:
        """Early prefill memory check for plain /v1/completions calls (VLM)."""
        if not self._loaded:
            await self.start()
        if self.is_diffusion_model:
            self._validate_diffusion_request(
                stop=kwargs.get("stop"),
                kwargs=kwargs,
            )
            return
        try:
            num_tokens = len(self._tokenizer.encode(prompt))
        except Exception as e:
            logger.warning(
                "VLMBatchedEngine.preflight_completion: tokenizer.encode "
                "raised %s; skipping prefill memory check, real completion "
                "path will surface the error",
                type(e).__name__,
            )
            return
        scheduler = getattr(getattr(self._engine, "engine", None), "scheduler", None)
        if scheduler is None:
            _warn_scheduler_unreachable_once(self, "preflight_completion")
            return
        await self._preflight_or_raise_with_eviction(
            scheduler,
            num_prompt_tokens=num_tokens,
            request_id=request_id,
            text_only=True,
        )

    async def stream_chat(
        self,
        messages: list[dict[str, Any]],
        max_tokens: int = 256,
        temperature: float = 0.7,
        top_p: float = 0.9,
        top_k: int = 0,
        min_p: float = 0.0,
        repetition_penalty: float = 1.0,
        presence_penalty: float = 0.0,
        tools: list[dict] | None = None,
        **kwargs,
    ) -> AsyncIterator[GenerationOutput]:
        """Stream chat completion with vision support."""
        if not self._loaded:
            await self.start()

        if self.is_diffusion_model:
            self._validate_diffusion_request(
                tools=tools,
                stop=kwargs.get("stop"),
                kwargs=kwargs,
            )
            loop = asyncio.get_running_loop()
            from ..engine_core import get_mlx_executor

            diffusion_inputs = await loop.run_in_executor(
                get_mlx_executor(),
                self._process_diffusion_chat_messages,
                messages,
                tools,
                dict(kwargs),
            )
            async for output in self._stream_diffusion_inputs(
                diffusion_inputs,
                max_tokens=max_tokens,
                temperature=temperature,
                seed=kwargs.get("seed"),
            ):
                yield output
            return

        # Run vision encoding on the MLX executor thread to avoid blocking
        # the event loop.  Blocking here (synchronous mx.eval) prevents
        # uvicorn from managing HTTP keep-alive connections, causing
        # TransferEncodingError on the next request (issue #80).
        loop = asyncio.get_running_loop()
        # _process_chat_messages pops these; the tail marker needs them too.
        ct_kwargs = kwargs.get("chat_template_kwargs")
        partial = kwargs.get("is_partial")
        (
            prompt,
            vlm_embeds,
            vlm_kwargs,
            image_hash,
            image_cache_key_start,
            image_cache_key_ranges,
        ) = await loop.run_in_executor(
            self._engine._mlx_executor,
            self._process_chat_messages,
            messages,
            tools,
            kwargs,
        )

        # SpecPrefill: protect the system-prompt region from token dropping.
        self._inject_specprefill_system_end(messages, prompt, kwargs)
        generation_prompt, persists = self._generation_prompt_text(ct_kwargs, partial)
        if generation_prompt:
            kwargs["generation_prompt_text"] = generation_prompt
            kwargs["generation_prompt_persists"] = persists

        async for output in self.stream_generate(
            prompt=prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            min_p=min_p,
            repetition_penalty=repetition_penalty,
            presence_penalty=presence_penalty,
            vlm_inputs_embeds=vlm_embeds,
            vlm_extra_kwargs=vlm_kwargs,
            vlm_image_hash=image_hash,
            vlm_cache_key_start=image_cache_key_start,
            vlm_cache_key_ranges=image_cache_key_ranges,
            tools=tools,
            **kwargs,
        ):
            yield output

    def _apply_ocr_prompt(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Apply a default OCR prompt only when the user sends no text.

        OCR models (DeepSeek-OCR, GLM-OCR, DOTS-OCR) work best with specific
        prompt formats. When the user sends an image without any text, this
        injects the model's default OCR prompt. If the user provides their own
        text, it is preserved as-is so they can use custom prompts (e.g.
        structured extraction with JSON schema).

        Only activates when:
        - The model_type is in OCR_MODEL_PROMPTS
        - The last user message contains image content
        - The last user message has no meaningful text
        """
        model_type = self.model_type or ""
        if model_type not in OCR_MODEL_PROMPTS:
            return messages

        ocr_prompt = OCR_MODEL_PROMPTS[model_type]
        messages = copy.deepcopy(messages)

        # Find last user message
        for msg in reversed(messages):
            if msg.get("role") != "user":
                continue
            content = msg.get("content")
            if isinstance(content, list):
                # Multi-part content: check if it has images
                has_image = any(
                    isinstance(p, dict) and p.get("type") == "image_url"
                    for p in content
                )
                if not has_image:
                    break
                # Check if user provided meaningful text
                user_text = " ".join(
                    p.get("text", "")
                    for p in content
                    if isinstance(p, dict) and p.get("type") == "text"
                ).strip()
                if user_text:
                    # User provided their own prompt, keep it
                    break
                # No user text — inject default OCR prompt
                new_content = [{"type": "text", "text": ocr_prompt}]
                new_content.extend(
                    p
                    for p in content
                    if not (isinstance(p, dict) and p.get("type") == "text")
                )
                msg["content"] = new_content
            break

        return messages

    def _extract_request_media(
        self, messages: list[dict[str, Any]]
    ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list, list, list[str]]:
        """Split a request into text and media, honoring each model's video path.

        Returns ``(media_messages, text_messages, images, audio, videos)``.
        MiMo consumes video as a bounded, chronological sequence of frames
        through the same vision tower used for still images. Checkpoints with
        native video input keep clips as inline URIs for the video processor.
        Every other model rejects video parts.
        """
        media_messages = messages
        model_type = self.model_type or _read_config_model_type(self._model_name)
        if model_type in {"mimo_v2", "mimo_v2_flash"}:
            media_messages = expand_video_parts(messages)
        if not getattr(self, "_native_video", False):
            text_messages, images, audio = extract_images_from_messages(media_messages)
            return media_messages, text_messages, images, audio, []

        text_messages, images, audio, videos = extract_media_from_messages(
            media_messages
        )
        # mlx-vlm's Qwen3-VL embedding path takes the grid of one modality per
        # call, so a request that mixes them would silently lose one.
        if videos and (images or audio):
            raise InvalidRequestError(
                "Video input cannot be combined with image or audio input "
                "in the same request.",
                field="messages",
            )
        return media_messages, text_messages, images, audio, videos

    def _process_chat_messages(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict] | None,
        kwargs: dict,
    ) -> Tuple[
        str | list[int], Any, dict | None, str | None, int, List[Tuple[int, str]]
    ]:
        """
        Process chat messages, extracting images and preparing VLM inputs.

        Returns:
            Tuple of (prompt_or_token_ids, vlm_embeds, vlm_kwargs, image_hash)
        """
        media_messages, text_messages, images, audio, videos = (
            self._extract_request_media(messages)
        )
        model_type = self.model_type or _read_config_model_type(self._model_name)

        ct_kwargs = kwargs.pop("chat_template_kwargs", None)
        partial = kwargs.pop("is_partial", None)

        # Keep VLM-capable models on one prompt-rendering path, even before the
        # first image arrives. Otherwise the conversation switches prompt families
        # on the first image-bearing turn and invalidates early prefix blocks.
        if images:
            vlm_messages = self._apply_ocr_prompt(media_messages)
        elif videos or (audio and model_type in {"mimo_v2", "mimo_v2_flash"}):
            vlm_messages = media_messages
        else:
            vlm_messages = text_messages
        template_tools = convert_tools_for_template(tools) if tools else None
        # mlx-vlm samples frames from a file, so inline clips are written to
        # temporary files that live only until preprocessing is done.
        video_files: list[tuple[Path, str]] = []
        try:
            video_processor = getattr(self._processor, "video_processor", None)
            for uri in videos:
                video_files.append(write_video_data_uri(uri))
                # Rejects undecodable or degenerate clips with a 400 before
                # mlx-vlm samples them.
                native_video_token_count(
                    probe_video(video_files[-1][0]), video_processor
                )
            video_kwargs = (
                {
                    "videos": [str(path) for path, _ in video_files],
                    "video_hashes": [digest for _, digest in video_files],
                }
                if video_files
                else {}
            )
            (
                token_ids,
                vlm_embeds,
                vlm_kwargs,
                image_hash,
                image_cache_key_start,
                image_cache_key_ranges,
            ) = self._prepare_vision_inputs(
                vlm_messages,
                images,
                audio=audio if audio else None,
                chat_template_kwargs=ct_kwargs,
                tools=template_tools,
                is_partial=partial,
                **video_kwargs,
            )
        finally:
            for path, _ in video_files:
                path.unlink(missing_ok=True)

        if images or videos:
            # Free Metal intermediates from vision encoding.
            mx.synchronize()
            mx.clear_cache()

        return (
            token_ids,
            vlm_embeds,
            vlm_kwargs,
            image_hash,
            image_cache_key_start,
            image_cache_key_ranges,
        )

    def _validate_diffusion_request(
        self,
        *,
        tools: list[dict] | None = None,
        audio: list | None = None,
        stop: list[str] | None = None,
        kwargs: dict[str, Any] | None = None,
    ) -> None:
        if not self.is_diffusion_model:
            return
        kwargs = kwargs or {}
        if tools and not self.supports_tool_calling:
            raise InvalidRequestError(
                "Tool calling is not supported for this diffusion model "
                "(no tool parser matched its chat template).",
                field="tools",
            )
        if audio:
            raise InvalidRequestError(
                "Audio input is not supported with diffusion models.",
                field="messages",
            )
        if stop:
            raise InvalidRequestError(
                "Custom stop sequences are not supported with diffusion models.",
                field="stop",
            )
        if kwargs.get("compiled_grammar") is not None:
            raise InvalidRequestError(
                "Structured response_format is not supported with diffusion models.",
                field="response_format",
            )
        if kwargs.get("specprefill") is True:
            raise InvalidRequestError(
                "SpecPrefill is not supported with diffusion models.",
                field="specprefill",
            )

    def _diffusion_apply_chat_template(
        self,
        messages: list[dict[str, Any]],
        *,
        images: list[Any],
        chat_template_kwargs: dict[str, Any] | None = None,
        tools: list[dict] | None = None,
    ) -> str | list[int]:
        from mlx_vlm.prompt_utils import apply_chat_template

        num_images = len(images)
        model_type = self.model_type or ""
        if num_images > 1 and model_type in SINGLE_IMAGE_ONLY_MODELS:
            raise ValueError(
                f"Model {model_type} does not support multi-image chat. "
                f"Please use only 1 image."
            )

        try:
            formatted_messages, _ = self._format_messages_for_vlm_template(
                messages, num_images=num_images, num_audios=0
            )
        except Exception as e:
            logger.debug(
                "Falling back to mlx-vlm apply_chat_template for diffusion: %s",
                e,
            )
            formatted_messages = apply_chat_template(
                self._processor,
                self._vlm_model.config,
                messages,
                num_images=num_images,
                num_audios=0,
                return_messages=True,
            )

        detect_and_strip_partial(formatted_messages)
        template_kwargs = {
            "tokenize": False,
            "add_generation_prompt": True,
        }
        if self._enable_thinking is not None:
            template_kwargs["enable_thinking"] = self._enable_thinking
        if tools:
            template_kwargs["tools"] = tools
        if chat_template_kwargs:
            template_kwargs.update(chat_template_kwargs)
        _apply_minimax_m3_thinking_mode(model_type, template_kwargs)

        template_target = self._chat_template_target(model_type)
        try:
            return template_target.apply_chat_template(
                formatted_messages, **template_kwargs
            )
        except TypeError:
            if chat_template_kwargs:
                for key in chat_template_kwargs:
                    template_kwargs.pop(key, None)
            template_kwargs.pop("tools", None)
            template_kwargs.pop("enable_thinking", None)
            return template_target.apply_chat_template(
                formatted_messages, **template_kwargs
            )
        except ValueError:
            fallback = getattr(self._processor, "tokenizer", None)
            if fallback is None or fallback is template_target:
                raise
            try:
                return fallback.apply_chat_template(
                    formatted_messages, **template_kwargs
                )
            except TypeError:
                if chat_template_kwargs:
                    for key in chat_template_kwargs:
                        template_kwargs.pop(key, None)
                template_kwargs.pop("tools", None)
                template_kwargs.pop("enable_thinking", None)
                return fallback.apply_chat_template(
                    formatted_messages, **template_kwargs
                )

    def _prepare_diffusion_inputs_from_prompt(
        self,
        prompt: str | list[int],
        *,
        images: list[Any] | None = None,
    ) -> dict[str, Any]:
        from mlx_vlm.utils import prepare_inputs

        images = images or []
        if isinstance(prompt, list):
            input_ids = mx.array([prompt])
            return {
                "input_ids": input_ids,
                "pixel_values": None,
                "attention_mask": None,
                "mm_token_type_ids": None,
                "prompt_tokens": int(input_ids.size),
            }

        inputs = prepare_inputs(
            self._processor,
            images=images if images else None,
            prompts=[prompt],
        )
        input_ids = inputs["input_ids"]
        return {
            "input_ids": input_ids,
            "pixel_values": inputs.get("pixel_values"),
            "attention_mask": inputs.get("attention_mask"),
            "mm_token_type_ids": inputs.get("mm_token_type_ids"),
            "prompt_tokens": int(input_ids.size),
        }

    def _process_diffusion_chat_messages(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict] | None,
        kwargs: dict[str, Any],
    ) -> dict[str, Any]:
        text_messages, images, audio = extract_images_from_messages(messages)
        self._validate_diffusion_request(
            tools=tools,
            audio=audio if audio else None,
            stop=kwargs.get("stop"),
            kwargs=kwargs,
        )
        chat_template_kwargs = kwargs.pop("chat_template_kwargs", None)
        diffusion_messages = messages if images else text_messages
        prompt = self._diffusion_apply_chat_template(
            diffusion_messages,
            images=images,
            chat_template_kwargs=chat_template_kwargs,
            tools=tools,
        )
        return self._prepare_diffusion_inputs_from_prompt(prompt, images=images)

    def _iter_diffusion_outputs_sync(
        self,
        diffusion_inputs: dict[str, Any],
        *,
        max_tokens: int,
        temperature: float,
        seed: int | None = None,
        cancel_event: threading.Event | None = None,
    ):
        from mlx_vlm.generate.diffusion import stream_diffusion_generate

        try:
            from mlx_vlm.generate.common import generation_stream, wired_limit

            limit_ctx = wired_limit(self._vlm_model, [generation_stream])
        except Exception:
            limit_ctx = contextlib.nullcontext()

        if seed is not None:
            mx.random.seed(seed)

        tokenizer = self._tokenizer
        if hasattr(tokenizer, "stopping_criteria"):
            tokenizer.stopping_criteria.reset(
                getattr(self._vlm_model.config, "eos_token_id", None)
            )

        prompt_tokens = int(diffusion_inputs.get("prompt_tokens") or 0)
        results = None
        full_text = ""
        block_text: list[str] = []
        emitted_tokens = 0
        last_stream_segment = ""

        # Special tokens are stripped from the stream, EXCEPT protocol
        # markers the model's output parser needs to see in the text:
        # tool-call markers (e.g. Gemma's <|tool_call> / <tool_call|>)
        # for the tool parser, and channel/turn markers for the output
        # parser session (thought-channel → <think> conversion). They
        # are removed downstream (parser session / parse_tool_calls /
        # ToolCallStreamFilter) so they never leak to clients.
        skip_special_ids = set(getattr(tokenizer, "all_special_ids", None) or [])
        preserved_marker_texts: list[str] = []
        if getattr(tokenizer, "has_tool_calling", False):
            preserved_marker_texts.extend(
                m
                for m in (
                    getattr(tokenizer, "tool_call_start", None),
                    getattr(tokenizer, "tool_call_end", None),
                )
                if m
            )

        # Detect a protocol output parser (e.g. gemma4 channel markers).
        # The diffusion lane emits detokenized text segments, so only
        # sessions exposing ``process_text`` can be used here.
        parser_session = None
        try:
            from ..adapter.output_parser import detect_output_parser

            model_config = {"model_type": self.model_type} if self.model_type else None
            factory = detect_output_parser(self._model_name, tokenizer, model_config)
            if factory is not None:
                session = factory.create_session(tokenizer)
                if hasattr(session, "process_text"):
                    parser_session = session
                    preserved_marker_texts.extend(factory.protocol_marker_texts)
        except Exception as e:
            logger.debug("Diffusion output parser unavailable: %s", e)
            parser_session = None

        for marker in preserved_marker_texts:
            try:
                marker_id = tokenizer.convert_tokens_to_ids(marker)
            except Exception:
                marker_id = None
            if marker_id is not None:
                skip_special_ids.discard(marker_id)

        def _parse_block(text: str, *, final: bool = False) -> str:
            if parser_session is None:
                return text
            parsed = parser_session.process_text(text).visible_text
            if final:
                parsed += parser_session.finalize().visible_text
            return parsed

        try:
            with limit_ctx:
                results = stream_diffusion_generate(
                    self._vlm_model,
                    self._processor,
                    tokenizer,
                    diffusion_inputs["input_ids"],
                    diffusion_inputs.get("pixel_values"),
                    diffusion_inputs.get("attention_mask"),
                    max_tokens=max_tokens,
                    temperature=temperature,
                    skip_special_token_ids=skip_special_ids,
                    mm_token_type_ids=diffusion_inputs.get("mm_token_type_ids"),
                    prefill_step_size=DIFFUSION_PREFILL_STEP_SIZE,
                )
                for result in results:
                    if cancel_event is not None and cancel_event.is_set():
                        break
                    if getattr(result, "is_draft", False):
                        continue
                    result_tokens = getattr(result, "generation_tokens", None)
                    finish_reason = getattr(result, "finish_reason", None)
                    result_text = result.text or ""
                    if result_text:
                        has_token_progress = (
                            result_tokens is None or int(result_tokens) > emitted_tokens
                        )
                        has_final_flush = (
                            finish_reason is not None
                            and result_text != last_stream_segment
                        )
                        if has_token_progress or has_final_flush:
                            block_text.append(result_text)
                            last_stream_segment = result_text
                    is_boundary = bool(
                        getattr(result, "diffusion_block_complete", False)
                    )
                    if not is_boundary and not finish_reason:
                        continue

                    new_text = remove_special_tokens_preserve_whitespace(
                        _parse_block(
                            "".join(block_text),
                            final=finish_reason is not None,
                        )
                    )
                    full_text += new_text
                    completion_tokens = int(result_tokens or emitted_tokens)
                    emitted_tokens = max(emitted_tokens, completion_tokens)
                    if new_text or finish_reason:
                        yield GenerationOutput(
                            text=full_text,
                            new_text=new_text,
                            prompt_tokens=int(
                                getattr(result, "prompt_tokens", prompt_tokens)
                                or prompt_tokens
                            ),
                            completion_tokens=emitted_tokens,
                            finished=finish_reason is not None,
                            finish_reason=finish_reason,
                            cached_tokens=0,
                            prompt_tps=float(getattr(result, "prompt_tps", 0.0) or 0.0),
                            generation_tps=float(
                                getattr(result, "generation_tps", 0.0) or 0.0
                            ),
                            diffusion_canvas_tokens=int(
                                getattr(result, "diffusion_canvas_tokens", 0) or 0
                            ),
                            diffusion_denoising_steps=int(
                                getattr(result, "diffusion_denoising_steps", 0) or 0
                            ),
                            diffusion_work_tokens=int(
                                getattr(result, "diffusion_work_tokens", 0) or 0
                            ),
                            diffusion_canvas_tps=float(
                                getattr(result, "diffusion_canvas_tps", 0.0) or 0.0
                            ),
                            diffusion_work_tps=float(
                                getattr(result, "diffusion_work_tps", 0.0) or 0.0
                            ),
                        )
                    block_text = []
                    if finish_reason:
                        break
        finally:
            if results is not None and callable(getattr(results, "close", None)):
                results.close()
            mx.synchronize()
            mx.clear_cache()

    async def _stream_diffusion_inputs(
        self,
        diffusion_inputs: dict[str, Any],
        *,
        max_tokens: int,
        temperature: float,
        seed: int | None = None,
    ) -> AsyncIterator[GenerationOutput]:
        from ..engine_core import get_mlx_executor

        async with self._diffusion_lock:
            self._diffusion_active_requests += 1
            queue: asyncio.Queue[Any] = asyncio.Queue()
            cancel_event = threading.Event()
            self._diffusion_cancel_events.add(cancel_event)
            loop = asyncio.get_running_loop()

            def _put(item: Any) -> None:
                loop.call_soon_threadsafe(queue.put_nowait, item)

            def _worker() -> None:
                try:
                    for item in self._iter_diffusion_outputs_sync(
                        diffusion_inputs,
                        max_tokens=max_tokens,
                        temperature=temperature,
                        seed=seed,
                        cancel_event=cancel_event,
                    ):
                        _put(item)
                        if cancel_event.is_set():
                            break
                except BaseException as e:
                    _put(e)
                finally:
                    _put(None)

            future = loop.run_in_executor(get_mlx_executor(), _worker)
            try:
                while True:
                    item = await queue.get()
                    if item is None:
                        break
                    if isinstance(item, BaseException):
                        raise item
                    yield item
            finally:
                cancel_event.set()
                await future
                self._diffusion_cancel_events.discard(cancel_event)
                self._diffusion_active_requests -= 1

    def count_chat_tokens(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict] | None = None,
        chat_template_kwargs: dict[str, Any] | None = None,
        is_partial: bool | None = None,
    ) -> int:
        """Count prompt tokens for chat messages (text-only approximation).

        For VLM messages with images, this counts only the text tokens.
        Image tokens are added during vision encoding and vary by model.
        """
        # Extract text-only version for token counting through the same media
        # split as the real request, so video parts are not rejected here
        # before they reach the multimodal path.
        _, text_messages, _, _, _ = self._extract_request_media(messages)

        template_tools = convert_tools_for_template(tools) if tools else None
        prompt = self._apply_chat_template(
            text_messages,
            template_tools,
            chat_template_kwargs=chat_template_kwargs,
            is_partial=is_partial,
        )
        return len(self._tokenizer.encode(prompt))

    def has_active_requests(self) -> bool:
        """Check if the engine has active in-flight requests."""
        if self.is_diffusion_model:
            return getattr(self, "_diffusion_active_requests", 0) > 0
        engine_core = getattr(self, "_engine", None)
        if engine_core is not None:
            inner = getattr(engine_core, "engine", None)
            if inner is not None:
                collectors = getattr(inner, "_output_collectors", {})
                return len(collectors) > 0
        return False

    def get_stats(self) -> dict[str, Any]:
        """Get engine statistics."""
        stats = {
            "engine_type": "vlm",
            "model_name": self._model_name,
            "loaded": self._loaded,
            "stream_interval": self._stream_interval,
        }
        if self._diffusion_family is not None:
            stats["diffusion_family"] = self._diffusion_family
            stats["active_requests"] = self._diffusion_active_requests
        if self._engine:
            stats.update(self._engine.get_stats())
        return stats

    def get_cache_stats(self) -> dict[str, Any] | None:
        """Get cache statistics."""
        if self._engine:
            return self._engine.get_cache_stats()
        return None

    async def abort_all_requests(
        self,
        *,
        reason: str | None = None,
        error_code: str | None = None,
    ) -> int:
        """Abort all active requests."""
        if self.is_diffusion_model:
            cancel_events = list(getattr(self, "_diffusion_cancel_events", ()))
            for cancel_event in cancel_events:
                cancel_event.set()
            return len(cancel_events)
        if self._engine and self._engine.engine:
            return await self._engine.engine.abort_all_requests(
                reason=reason,
                error_code=error_code,
            )
        return 0