# Copyright 2022-2026 Xinference Holdings Pte. Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import concurrent.futures
import importlib
import inspect
import logging
import pathlib
import platform
import sys
import threading
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import (
    Any,
    AsyncGenerator,
    Callable,
    Dict,
    Iterator,
    List,
    Optional,
    Tuple,
    TypedDict,
    Union,
)

import xoscar as xo

from ....constants import XINFERENCE_MAX_TOKENS, XINFERENCE_TRUST_REMOTE_CODE
from ....fields import max_tokens_field
from ....types import (
    ChatCompletion,
    ChatCompletionChunk,
    Completion,
    CompletionChunk,
    CompletionUsage,
    LoRA,
)
from ...utils import check_dependency_available
from ..core import LLM, chat_context_var
from ..llm_family import LLMFamilyV2, LLMSpecV1
from ..utils import (
    DEEPSEEK_TOOL_CALL_FAMILY,
    GEMMA_TOOL_CALL_FAMILY,
    GLM5_TOOL_CALL_FAMILY,
    QWEN_TOOL_CALL_FAMILY,
    ChatModelMixin,
    generate_completion_chunk,
    get_context_length_from_config,
)

logger = logging.getLogger(__name__)

_mlx_vlm_stream_lock = threading.Lock()
_mlx_executor_lock = threading.Lock()


def _ensure_mlx_vlm_thread_local_stream() -> Any:
    """Use an MLX generation stream that is safe across worker threads."""
    import mlx.core as mx

    mlx_vlm_generate = importlib.import_module("mlx_vlm.generate")
    thread_local_stream_type = getattr(mx, "ThreadLocalStream", None)
    new_thread_local_stream = getattr(mx, "new_thread_local_stream", None)
    if thread_local_stream_type is None or new_thread_local_stream is None:
        return mlx_vlm_generate

    generation_stream = getattr(mlx_vlm_generate, "generation_stream", None)
    if not isinstance(generation_stream, thread_local_stream_type):
        with _mlx_vlm_stream_lock:
            generation_stream = getattr(mlx_vlm_generate, "generation_stream", None)
            if not isinstance(generation_stream, thread_local_stream_type):
                setattr(
                    mlx_vlm_generate,
                    "generation_stream",
                    new_thread_local_stream(mx.default_device()),
                )
                logger.debug(
                    "Replaced mlx-vlm generation stream with a thread-local stream"
                )

    return mlx_vlm_generate


class MLXBatchModel:
    """Wrapper around MLX-LM BatchGenerator for continuous batching."""

    # Class-level storage for multiple batch generators keyed by (temperature, top_p)
    _batch_generators: Dict[Tuple[float, float], Dict[str, Any]] = (
        {}
    )  # (temp, top_p) -> {'generator': BatchGenerator, 'queues': dict, 'pending': dict, 'active': set, 'task': task}
    _model_ref = None
    _tokenizer_ref = None
    _batch_size = 4
    _max_context_length = 2048
    _stop_tokens: set = set()
    _lock: Optional[asyncio.Lock] = None  # Will be initialized lazily
    _mlx_lm_version: Optional[str] = None  # Cache mlx-lm version

    def __init__(
        self, model, tokenizer, batch_size: int = 4, max_context_length: int = 2048
    ):
        # Store references for creating new generators on demand
        MLXBatchModel._model_ref = model
        MLXBatchModel._tokenizer_ref = tokenizer
        MLXBatchModel._batch_size = batch_size
        MLXBatchModel._max_context_length = max_context_length
        eos_token_ids = tokenizer.eos_token_ids
        if isinstance(eos_token_ids, int):
            eos_token_ids = [eos_token_ids]
        MLXBatchModel._stop_tokens = set(eos_token_ids or [])

    @staticmethod
    def _get_lock() -> asyncio.Lock:
        """Get or create the async lock."""
        if MLXBatchModel._lock is None:
            MLXBatchModel._lock = asyncio.Lock()
        return MLXBatchModel._lock

    @staticmethod
    def _get_mlx_lm_version() -> str:
        """Get mlx-lm version and cache it."""
        if MLXBatchModel._mlx_lm_version is None:
            try:
                import mlx_lm

                MLXBatchModel._mlx_lm_version = mlx_lm.__version__
            except (ImportError, AttributeError):
                MLXBatchModel._mlx_lm_version = "0.0.0"
        return MLXBatchModel._mlx_lm_version

    @staticmethod
    def _is_new_mlx_lm() -> bool:
        """Check if mlx-lm version is >= 0.31.2 (new API)."""
        version = MLXBatchModel._get_mlx_lm_version()
        try:
            from packaging import version as pkg_version

            return pkg_version.parse(version) >= pkg_version.parse("0.31.2")
        except Exception:
            # Fallback: check if new API exists
            try:
                from mlx_lm.generate import BatchGenerator

                # Check if next_generated method exists (new API)
                return hasattr(BatchGenerator, "next_generated")
            except ImportError:
                return False

    def _get_or_create_generator(self, temperature: float, top_p: float):
        """Get or create a BatchGenerator for the given sampling parameters."""
        key = (round(temperature, 6), round(top_p, 6))

        if key not in MLXBatchModel._batch_generators:
            logger.info(
                f"Creating new BatchGenerator for temperature={temperature}, top_p={top_p}"
            )
            from mlx_lm.generate import BatchGenerator
            from mlx_lm.sample_utils import make_sampler

            # Create sampler with specific settings
            sampler = make_sampler(temp=temperature, top_p=top_p)

            # Create batch generator
            batch_generator = BatchGenerator(
                model=MLXBatchModel._model_ref,
                max_tokens=MLXBatchModel._max_context_length,
                stop_tokens=MLXBatchModel._stop_tokens,
                sampler=sampler,
                completion_batch_size=MLXBatchModel._batch_size,
                prefill_batch_size=1,  # Use 1 for lowest latency
            )

            MLXBatchModel._batch_generators[key] = {
                "generator": batch_generator,
                "queues": {},  # uid -> asyncio.Queue
                "pending": {},  # uid -> list of results
                "active": set(),  # active uids
                "task": None,
            }

        return MLXBatchModel._batch_generators[key]

    def _ensure_background_worker(self, gen_dict):
        """Ensure background worker is running for this generator."""
        if gen_dict["task"] is None:
            loop = asyncio.get_event_loop()
            gen_dict["task"] = loop.create_task(self._background_worker(gen_dict))

    async def _background_worker(self, gen_dict):
        """Background worker that continuously calls next() and distributes results."""
        key = f"temp={gen_dict['generator'].sampler}"
        logger.info(f"Starting BatchGenerator background worker for {key}")
        batch_generator = gen_dict["generator"]

        while True:
            try:
                # Get next batch of results for ALL active requests
                # Use different API based on mlx-lm version
                if MLXBatchModel._is_new_mlx_lm():
                    # New API (mlx-lm >= 0.31.2): use next_generated()
                    batch_results = batch_generator.next_generated()
                else:
                    # Old API (mlx-lm < 0.31.2): use next() and unpack
                    batch_results = batch_generator.next()

                if not batch_results:
                    # No active requests, sleep briefly
                    await asyncio.sleep(0.001)
                    continue

                # Distribute results to respective request queues
                async with MLXBatchModel._get_lock():
                    for result in batch_results:
                        if result.uid in gen_dict["queues"]:
                            queue = gen_dict["queues"][result.uid]
                            # Put result in queue
                            try:
                                queue.put_nowait(result)
                            except asyncio.QueueFull:
                                logger.warning(f"Queue full for uid {result.uid}")
                        else:
                            # Queue not ready yet, cache the result
                            if result.uid not in gen_dict["pending"]:
                                gen_dict["pending"][result.uid] = []
                            gen_dict["pending"][result.uid].append(result)
                            logger.debug(
                                f"Cached result for uid {result.uid}, queue not ready yet"
                            )

                        # Remove from active requests if finished
                        if result.finish_reason is not None:
                            if result.uid in gen_dict["active"]:
                                gen_dict["active"].remove(result.uid)
                                logger.debug(
                                    f"Request {result.uid} finished, removed from active set"
                                )

                # Small delay to prevent busy waiting
                await asyncio.sleep(0.0001)

            except Exception as e:
                logger.error(f"Error in background worker: {e}", exc_info=True)
                await asyncio.sleep(0.01)

    async def generate_stream(
        self,
        prompt: str,
        max_tokens: int,
        temperature: float = 0.0,
        top_p: float = 1.0,
        request_id: Optional[str] = None,
        skip_special_tokens: bool = True,
    ) -> AsyncGenerator[CompletionChunk, None]:
        """Async stream generate method using BatchGenerator."""
        # Get or create generator for this sampling configuration
        gen_dict = self._get_or_create_generator(temperature, top_p)
        batch_generator = gen_dict["generator"]

        # Ensure background worker is running for this generator
        self._ensure_background_worker(gen_dict)

        # Prepare request
        assert MLXBatchModel._tokenizer_ref is not None
        prompt_tokens = MLXBatchModel._tokenizer_ref.encode(prompt)
        input_echo_len = len(prompt_tokens)

        # Create queue first
        queue: asyncio.Queue = asyncio.Queue()

        # Insert prompt into batch to get the real uid
        # Use different API based on mlx-lm version
        if MLXBatchModel._is_new_mlx_lm():
            # New API (mlx-lm >= 0.31.2): max_tokens should be a list
            request_ids = batch_generator.insert(
                [prompt_tokens], max_tokens=[max_tokens]
            )
        else:
            # Old API (mlx-lm < 0.31.2): max_tokens is a single value
            request_ids = batch_generator.insert([prompt_tokens], max_tokens=max_tokens)
        inserted_uid = request_ids[0]

        logger.debug(
            f"Inserted request {inserted_uid} into batch (temp={temperature}, top_p={top_p})"
        )

        # Add to active requests set
        async with MLXBatchModel._get_lock():
            gen_dict["active"].add(inserted_uid)
            gen_dict["queues"][inserted_uid] = queue

            # Check if there are any pending results (early results that arrived before queue was ready)
            if inserted_uid in gen_dict["pending"]:
                logger.debug(
                    f"Found {len(gen_dict['pending'][inserted_uid])} pending results for uid {inserted_uid}"
                )
                for result in gen_dict["pending"][inserted_uid]:
                    try:
                        queue.put_nowait(result)
                    except asyncio.QueueFull:
                        logger.warning(
                            f"Queue full for uid {inserted_uid} while adding pending results"
                        )
                del gen_dict["pending"][inserted_uid]

        try:
            # Track generated text
            generated_tokens = []
            chunk_id = str(uuid.uuid4())
            token_count = 0

            # Wait for results from background worker
            while True:
                try:
                    # Wait for result with short timeout (so we can check if request is still active)
                    result = await asyncio.wait_for(queue.get(), timeout=0.5)

                    if result.token is not None:
                        generated_tokens.append(result.token)
                        token_count += 1

                        # Decode current token
                        assert MLXBatchModel._tokenizer_ref is not None
                        token_text = MLXBatchModel._tokenizer_ref.decode(
                            [result.token], skip_special_tokens=skip_special_tokens
                        )
                        token_text = token_text.strip("�")

                        # Generate completion chunk
                        yield generate_completion_chunk(
                            chunk_text=token_text,
                            finish_reason=None,
                            chunk_id=chunk_id,
                            model_uid=None,
                            prompt_tokens=input_echo_len,
                            completion_tokens=token_count,
                            total_tokens=input_echo_len + token_count,
                        )

                    # Check if generation is finished
                    if result.finish_reason is not None:
                        # Generate final chunk with finish reason
                        finish_reason = (
                            "length" if result.finish_reason == "length" else "stop"
                        )
                        yield generate_completion_chunk(
                            chunk_text="",
                            finish_reason=finish_reason,
                            chunk_id=chunk_id,
                            model_uid=None,
                            prompt_tokens=input_echo_len,
                            completion_tokens=token_count,
                            total_tokens=input_echo_len + token_count,
                        )
                        logger.debug(
                            f"Request {inserted_uid} finished with {finish_reason}"
                        )
                        return

                except asyncio.TimeoutError:
                    # Check if request is still in active set
                    async with MLXBatchModel._get_lock():
                        is_active = inserted_uid in gen_dict["active"]

                    if not is_active:
                        # Request has finished (removed from active set by background worker)
                        logger.debug(
                            f"Request {inserted_uid} is no longer active, finishing"
                        )
                        return
                    # Otherwise, continue waiting (request is still being processed)

        except Exception as e:
            logger.error(f"Error in generate_stream for request: {e}", exc_info=True)
            # Generate error chunk
            yield generate_completion_chunk(
                chunk_text="",
                finish_reason="error",
                chunk_id=chunk_id,
                model_uid=None,
                prompt_tokens=input_echo_len,
                completion_tokens=token_count,
                total_tokens=input_echo_len + token_count,
            )
        finally:
            # Clean up queue using the correct uid
            async with MLXBatchModel._get_lock():
                if inserted_uid in gen_dict["queues"]:
                    del gen_dict["queues"][inserted_uid]
                    logger.debug(f"Cleaned up queue for uid {inserted_uid}")
                # Also clean up any pending results
                if inserted_uid in gen_dict["pending"]:
                    del gen_dict["pending"][inserted_uid]
                # And remove from active set (in case it's still there)
                if inserted_uid in gen_dict["active"]:
                    gen_dict["active"].remove(inserted_uid)

    async def generate(
        self,
        prompt: str,
        max_tokens: int,
        temperature: float = 0.0,
        top_p: float = 1.0,
        stream: bool = False,
        skip_special_tokens: bool = True,
    ) -> str:
        """Async generate method using BatchGenerator."""
        # Non-streaming: collect all tokens
        result_text = ""
        async for chunk in self.generate_stream(
            prompt=prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            request_id=None,
            skip_special_tokens=skip_special_tokens,
        ):
            if chunk.get("choices") and len(chunk["choices"]) > 0:
                result_text += chunk["choices"][0].get("text", "")
        return result_text


class MLXModelConfig(TypedDict, total=False):
    revision: Optional[str]
    max_gpu_memory: str
    trust_remote_code: bool
    reasoning_content: bool
    enable_thinking: bool
    # distributed
    address: Optional[str]
    shard: Optional[int]
    n_worker: Optional[int]
    # batch configuration
    allow_batch: bool
    batch_size: int
    # speculative decoding, currently only the MLX vision engine consumes them.
    # ``draft_model_path`` is resolved by ``create_llm_model_instance`` from the
    # spec's ``draft_model_id`` when ``enable_mtp`` is passed at launch time.
    draft_model_path: Optional[str]
    num_speculative_tokens: Optional[int]


class MLXGenerateConfig(TypedDict, total=False):
    max_tokens: int
    temperature: float
    repetition_penalty: Optional[float]
    repetition_context_size: Optional[float]
    top_p: float
    logit_bias: Optional[Dict[int, float]]
    stop: Optional[Union[str, List[str]]]
    stop_token_ids: Optional[Union[int, List[int]]]
    stream: bool
    stream_options: Optional[Union[dict, None]]
    tools: Optional[List[Dict]]
    lora_name: Optional[str]


@dataclass
class PromptCache:
    cache: List[Any] = field(default_factory=list)
    model_key: Tuple[str, Optional[str]] = ("", None)
    tokens: List[int] = field(default_factory=list)


class MLXModel(LLM, ChatModelMixin):
    _rank_to_addresses: Optional[Dict[int, str]]
    allow_batch: bool = False

    def __init__(
        self,
        model_uid: str,
        model_family: "LLMFamilyV2",
        model_path: str,
        model_config: Optional[MLXModelConfig] = None,
        peft_model: Optional[List[LoRA]] = None,
    ):
        LLM.__init__(self, model_uid, model_family, model_path)
        ChatModelMixin.__init__(self)
        self._use_fast_tokenizer = True
        self._model_config: MLXModelConfig = self._sanitize_model_config(model_config)
        self._context_length: Optional[int] = None
        # for distributed
        assert model_config is not None
        self._address = model_config.pop("address", None)
        self._n_worker = model_config.pop("n_worker", 1)
        self._shard = model_config.pop("shard", 0)
        self._driver_info = model_config.pop("driver_info", None)  # type: ignore
        self._rank_to_addresses = None
        self._loading_thread = None
        self._loading_error = None
        self._all_worker_started = asyncio.Event()
        self._max_kv_size = None
        self._prompt_cache = None
        if peft_model is not None:
            raise ValueError("MLX engine has not supported lora yet")
        # used to call async
        self._loop = None
        self._batch_model = None

    def set_loop(self, loop: asyncio.AbstractEventLoop):
        # loop will be passed into ModelWrapper,
        # to call aynsc method with asyncio.run_coroutine_threadsafe
        self._loop = loop  # type: ignore

    def _cleanup_memory(self):
        import gc

        import mlx.core as mx

        # mandatory recycling
        gc.collect()
        # clear the MLX cache
        mx.clear_cache()

    @property
    def driver_info(self) -> Optional[dict]:
        return self._driver_info

    def set_shard_info(self, shard: int, address: str):
        # set shard info to rank 0
        if self._rank_to_addresses is None:
            self._rank_to_addresses = {}
        self._rank_to_addresses[shard] = address
        if len(self._rank_to_addresses) == self._n_worker:
            self._all_worker_started.set()

    async def get_rank_addresses(self) -> Optional[Dict[int, str]]:
        await self._all_worker_started.wait()
        return self._rank_to_addresses

    def _sanitize_model_config(
        self, model_config: Optional[MLXModelConfig]
    ) -> MLXModelConfig:
        if model_config is None:
            model_config = MLXModelConfig()
        model_config.setdefault("revision", self.model_spec.model_revision)
        # Respect the XINFERENCE_TRUST_REMOTE_CODE setting.
        model_config["trust_remote_code"] = (
            bool(model_config.get("trust_remote_code", XINFERENCE_TRUST_REMOTE_CODE))
            and XINFERENCE_TRUST_REMOTE_CODE
        )
        model_config.setdefault("reasoning_content", False)
        model_config.setdefault("enable_thinking", False)
        return model_config

    def _sanitize_generate_config(
        self,
        generate_config: Optional[MLXGenerateConfig],
    ) -> MLXGenerateConfig:
        if generate_config is None:
            generate_config = MLXGenerateConfig()

        # default config is adapted from
        # https://github.com/ml-explore/mlx-examples/blob/f212b770d8b5143e23102eda20400ae43340f844/llms/mlx_lm/utils.py#L129
        generate_config.setdefault("temperature", 0.0)
        generate_config.setdefault("logit_bias", None)
        generate_config.setdefault("repetition_penalty", None)
        generate_config.setdefault("repetition_context_size", 20)
        generate_config.setdefault("top_p", 1.0)

        max_tokens = max_tokens_field.default or XINFERENCE_MAX_TOKENS
        if not generate_config.get("max_tokens") and max_tokens:
            generate_config["max_tokens"] = max_tokens  # type: ignore
        return generate_config

    def _load_model(self, **kwargs):
        if self._model_config.pop("draft_model_path", None):
            self._model_config.pop("num_speculative_tokens", None)
            raise ValueError(
                "Speculative decoding is only supported by the MLX vision engine "
                "for now, e.g. gemma-4 with its `*-it-assistant` drafter."
            )
        try:
            import mlx.core as mx
            from mlx_lm import load
        except ImportError:
            error_message = "Failed to import module 'mlx_lm'"
            installation_guide = [
                "Please make sure 'mlx_lm' is installed. ",
                "You can install it by `pip install mlx_lm`\n",
            ]

            raise ImportError(f"{error_message}\n\n{''.join(installation_guide)}")

        tokenizer_config = dict(
            use_fast=self._use_fast_tokenizer,
            trust_remote_code=kwargs["trust_remote_code"],
            revision=kwargs["revision"],
        )
        logger.debug(
            "loading model with tokenizer config: %s, model config: %s",
            tokenizer_config,
            self._model_config,
        )

        cache_limit_gb = kwargs.get("cache_limit_gb", None)
        if cache_limit_gb:
            logger.debug(f"Setting cache limit to {cache_limit_gb} GB")
            mx.metal.set_cache_limit(cache_limit_gb * 1024 * 1024 * 1024)

        self._max_kv_size = kwargs.get("max_kv_size", None)
        self._prompt_cache = PromptCache()

        model, tokenizer = load(
            self.model_path,
            tokenizer_config=tokenizer_config,
            model_config=self._model_config,
        )
        if stop_token_ids := self.model_family.stop_token_ids:
            for stop_token_id in stop_token_ids:
                # mlx_lm's TokenizerWrapper used to expose ``add_eos_token``
                # as a method; newer versions removed it and HF tokenizers in
                # transformers>=5 expose ``add_eos_token`` as a bool attribute.
                # Fall back to extending the wrapper's eos token id set.
                add_eos = getattr(tokenizer, "add_eos_token", None)
                if callable(add_eos):
                    add_eos(stop_token_id)
                else:
                    eos_ids = getattr(tokenizer, "_eos_token_ids", None)
                    if eos_ids is None:
                        eos_ids = set()
                        try:
                            setattr(tokenizer, "_eos_token_ids", eos_ids)
                        except Exception:
                            pass
                    try:
                        if isinstance(eos_ids, set):
                            eos_ids.add(stop_token_id)
                        elif isinstance(eos_ids, list):
                            if stop_token_id not in eos_ids:
                                eos_ids.append(stop_token_id)
                        else:
                            logger.warning(
                                "Unsupported type %s for tokenizer._eos_token_ids; "
                                "stop token %s was not registered.",
                                type(eos_ids).__name__,
                                stop_token_id,
                            )
                    except Exception:
                        logger.warning(
                            "Failed to register stop token %s on tokenizer.",
                            stop_token_id,
                            exc_info=True,
                        )

        # Store model and tokenizer for batch model creation later
        self._model = model
        self._tokenizer = tokenizer

        return model, tokenizer

    def _load_model_shard(self, **kwargs):
        try:
            import mlx.core as mx
            from mlx_lm.utils import load_model, load_tokenizer
        except ImportError:
            error_message = "Failed to import module 'mlx_lm'"
            installation_guide = [
                "Please make sure 'mlx_lm' is installed. ",
                "You can install it by `pip install mlx_lm`\n",
            ]

            raise ImportError(f"{error_message}\n\n{''.join(installation_guide)}")

        # Ensure some attributes correctly inited by model actor
        assert (
            self._loop is not None and self._rank_to_addresses is not None
        ), "Service not started correctly"

        tokenizer_config = dict(
            use_fast=self._use_fast_tokenizer,
            trust_remote_code=kwargs["trust_remote_code"],
            revision=kwargs["revision"],
        )
        logger.debug(
            "loading model with tokenizer config: %s, model config: %s, shard: %d, n_worker: %d",
            tokenizer_config,
            self._model_config,
            self._shard,
            self._n_worker,
        )

        cache_limit_gb = kwargs.get("cache_limit_gb", None)
        if cache_limit_gb:
            logger.debug(f"Setting cache limit to {cache_limit_gb} GB")
            mx.metal.set_cache_limit(cache_limit_gb * 1024 * 1024 * 1024)

        self._max_kv_size = kwargs.get("max_kv_size", None)
        self._prompt_cache = PromptCache()

        self._model, config = load_model(
            pathlib.Path(self.model_path),
            lazy=True,
            get_model_classes=self._get_classes,
        )
        model = self._model.model
        model.rank = self._shard
        model.world_size = self._n_worker
        model.model_uid = self.model_uid
        model.loop = self._loop
        model.address = self._address
        model.rank_to_addresses = self._rank_to_addresses

        # create actors and so forth
        model.prepare()
        # real load the partial weights
        model.pipeline()
        mx.eval(model.parameters())

        self._tokenizer = load_tokenizer(
            pathlib.Path(self.model_path),
            tokenizer_config,
            eos_token_ids=config.get("eos_token_id", None),
        )

    @staticmethod
    def _get_classes(config: dict):
        """
        Retrieve the model and model args classes based on the configuration
        that supported distributed inference.

        Args:
            config (dict): The model configuration.

        Returns:
            A tuple containing the Model class and the ModelArgs class.
        """
        from mlx_lm.utils import MODEL_REMAPPING

        model_type = config["model_type"]
        model_type = MODEL_REMAPPING.get(model_type, model_type)
        try:
            arch = importlib.import_module(
                f"xinference.model.llm.mlx.distributed_models.{model_type}"
            )
        except ImportError:
            msg = f"Model type {model_type} not supported for distributed inference."
            logger.error(msg)
            raise ValueError(msg)

        return arch.Model, arch.ModelArgs

    def load(self):
        reasoning_content = self._model_config.pop("reasoning_content")
        enable_thinking = self._model_config.pop("enable_thinking", True)
        self.prepare_parse_reasoning_content(
            reasoning_content, enable_thinking=enable_thinking
        )
        self.prepare_parse_tool_calls()

        kwargs = {}
        kwargs["revision"] = self._model_config.get(
            "revision", self.model_spec.model_revision
        )
        kwargs["trust_remote_code"] = self._model_config.get("trust_remote_code")
        kwargs["cache_limit_gb"] = self._model_config.pop("cache_limit_gb", None)

        if self._n_worker <= 1:
            self._model, self._tokenizer = self._load_model(**kwargs)
        else:

            def _load():
                try:
                    if self._shard == 0:
                        self._driver_info = {"address": self._address}
                        self.set_shard_info(0, self._address)
                    else:
                        assert self._driver_info is not None
                        driver_address = self._driver_info["address"]

                        async def wait_for_all_shards():
                            model_ref = await xo.actor_ref(
                                address=driver_address, uid=self.raw_model_uid
                            )
                            # set shard info
                            await model_ref.set_shard_info(self._shard, self._address)
                            # wait for all shards
                            self._rank_to_addresses = (
                                await model_ref.get_rank_addresses()
                            )

                        asyncio.run_coroutine_threadsafe(
                            wait_for_all_shards(), self._loop
                        ).result()

                    self._load_model_shard(**kwargs)
                except Exception:
                    logger.exception("Loading mlx shard model failed")
                    self._loading_error = sys.exc_info()

            # distributed inference
            self._loading_thread = threading.Thread(target=_load)
            self._loading_thread.start()

    def wait_for_load(self):
        from mlx_lm.utils import load_config

        if self._loading_thread:
            self._loading_thread.join()
            if self._loading_error:
                _, err, tb = self._loading_error
                raise err.with_traceback(tb)

        # get context length
        config = load_config(Path(self.model_path))
        config.update(self._model_config)
        self._context_length = get_context_length_from_config(config)

        # Update allow_batch based on distributed inference
        # Only enable continuous batching for non-distributed inference (single worker)
        n_worker = self._n_worker if self._n_worker is not None else 1
        if self.__class__.allow_batch and n_worker > 1:
            # Distributed inference: disable continuous batching
            self.allow_batch = False

        # Create MLXBatchModel for continuous batching after context length is available
        # Only enable continuous batching for non-distributed inference (single worker)
        if (
            self._batch_model is None
            and hasattr(self, "_model")
            and hasattr(self, "_tokenizer")
            and self.allow_batch  # Check instance-level allow_batch
        ):
            batch_size = self._model_config.get("batch_size", 4)
            self._batch_model = MLXBatchModel(
                model=self._model,
                tokenizer=self._tokenizer,
                batch_size=batch_size,
                max_context_length=self._context_length,
            )

    @classmethod
    def check_lib(cls) -> Union[bool, Tuple[bool, str]]:
        dep_check = check_dependency_available("mlx_lm", "mlx_lm")
        if dep_check != True:
            return dep_check
        return True

    @classmethod
    def match_json(
        cls, llm_family: "LLMFamilyV2", llm_spec: "LLMSpecV1", quantization: str
    ) -> Union[bool, Tuple[bool, str]]:
        if llm_spec.model_format not in ["mlx"]:
            return False, "MLX base engine only supports mlx format"
        if sys.platform != "darwin" or platform.processor() != "arm":
            return False, "MLX base engine only works on Apple silicon Macs"
        if "generate" not in llm_family.model_ability:
            return False, "MLX base engine requires generate ability"
        if "chat" in llm_family.model_ability or "vision" in llm_family.model_ability:
            return False, "MLX base engine does not handle chat or vision models"
        return True

    def _get_prompt_cache(
        self, prompt, lora_name: Optional[str] = None, model: Any = None
    ):
        from mlx_lm.models.cache import make_prompt_cache

        assert self._prompt_cache is not None
        cache_len = len(self._prompt_cache.tokens)
        model_key = (self.model_path, lora_name)
        if (
            self._prompt_cache.model_key != model_key
            or cache_len >= len(prompt)
            or self._prompt_cache.tokens != prompt[:cache_len]
        ):
            self._prompt_cache.model_key = model_key
            self._prompt_cache.cache = make_prompt_cache(
                model or self._model, self._max_kv_size
            )
            self._prompt_cache.tokens = []
            logger.debug("Making new prompt cache for %s", self.model_uid)
        else:
            prompt = prompt[cache_len:]
            logger.debug("Cache hit for %s", self.model_uid)
        self._prompt_cache.tokens.extend(prompt)
        return prompt

    def _generate_stream_inner(self, **kwargs):
        try:
            from mlx_lm.utils import (
                make_logits_processors,
                make_sampler,
                stream_generate,
            )
        except ImportError:
            # for mlx-lm >= 0.22.3
            from mlx_lm.generate import stream_generate
            from mlx_lm.sample_utils import make_logits_processors, make_sampler

        sampler = make_sampler(
            temp=kwargs.pop("temperature"), top_p=kwargs.pop("top_p")
        )
        prompt_token_ids = kwargs.pop("prompt_token_ids")
        logits_processors = make_logits_processors(
            logit_bias=kwargs.pop("logits_bias", None),
            repetition_penalty=kwargs.pop("repetition_penalty"),
            repetition_context_size=kwargs.pop("repetition_context_size"),
        )
        try:
            yield from stream_generate(
                self._model,
                self._tokenizer,
                prompt_token_ids,
                sampler=sampler,
                logits_processors=logits_processors,
                **kwargs,
            )
        finally:
            # after completing the inference, clear the memory.
            self._cleanup_memory()

    def _prepare_inputs(
        self, prompt: Union[str, Dict[str, Any]], kwargs
    ) -> Tuple[Any, int]:
        prompt_token_ids = self._tokenizer.encode(prompt)
        prompt_token_ids = self._get_prompt_cache(
            prompt_token_ids, kwargs.get("lora_name")
        )
        return prompt_token_ids, len(prompt_token_ids)

    def _generate_stream(
        self, prompt: Union[str, Dict[str, Any]], kwargs: MLXGenerateConfig
    ):
        model_uid = self.model_uid
        tokenizer = self._tokenizer
        max_tokens = kwargs["max_tokens"]
        chunk_id = str(uuid.uuid4())
        stop_token_ids = kwargs.get("stop_token_ids", [])
        stream = kwargs.get("stream", False)
        stream_options = kwargs.pop("stream_options", None)
        include_usage = (
            stream_options["include_usage"]
            if isinstance(stream_options, dict)
            else False
        )

        prompt_token_ids, input_echo_len = self._prepare_inputs(prompt, kwargs)

        if max_tokens is None:
            # not set max_tokens
            max_tokens = self._context_length - input_echo_len
            logger.debug("No max_tokens set, setting to: %s", max_tokens)

        i = 0
        start = time.time()
        tokens = []
        # For VLM models, let generate_step handle cache creation internally
        # to avoid incompatible cache types between mlx_lm and mlx_vlm
        stream_kwargs = {
            "prompt_token_ids": prompt_token_ids,
            "max_tokens": max_tokens,
            "temperature": kwargs["temperature"],
            "top_p": kwargs["top_p"],
            "repetition_penalty": kwargs["repetition_penalty"],
            "repetition_context_size": kwargs["repetition_context_size"],
            "stop_token_ids": stop_token_ids,
            # Don't pass prompt_cache for VLM models to let mlx_vlm handle it
        }

        for i, chunk_resp in enumerate(self._generate_stream_inner(**stream_kwargs)):
            token = chunk_resp.token
            tokens.append(token)

            out = chunk_resp.text
            if stream:
                # this special character is mainly for qwen
                out = out.strip("")

            # MINIMAL FIX: Always yield the delta (out) rather than an accumulated output.
            # This prevents the calling generate() method from double-accumulating text.

            completion_usage = CompletionUsage(
                prompt_tokens=input_echo_len,
                completion_tokens=i,
                total_tokens=(input_echo_len + i),
            )

            yield generate_completion_chunk(
                chunk_text=out,
                finish_reason=None,
                chunk_id=chunk_id,
                model_uid=model_uid,
                prompt_tokens=input_echo_len,
                completion_tokens=i,
                total_tokens=(input_echo_len + i),
            ), completion_usage

            if token == tokenizer.eos_token_id or token in stop_token_ids:  # type: ignore
                break

        logger.info(
            f"Average generation speed: {i / (time.time() - start):.2f} tokens/s."  # noqa: E231
        )

        if self._prompt_cache:
            self._prompt_cache.tokens.extend(tokens)  # type: ignore

        if i == max_tokens - 1:
            finish_reason = "length"
        else:
            finish_reason = "stop"

        completion_usage = CompletionUsage(
            prompt_tokens=input_echo_len,
            completion_tokens=i,
            total_tokens=(input_echo_len + i),
        )
        yield generate_completion_chunk(
            "",
            finish_reason=finish_reason,
            chunk_id=chunk_id,
            model_uid=model_uid,
            prompt_tokens=input_echo_len,
            completion_tokens=i,
            total_tokens=(input_echo_len + i),
        ), completion_usage

        if include_usage:
            completion_chunk = CompletionChunk(
                id=chunk_id,
                object="text_completion",
                created=int(time.time()),
                model=model_uid,
                choices=[],
            )
            yield completion_chunk, completion_usage

    def _run_non_drivers(
        self, method: str, stream: bool, *args, **kwargs
    ) -> Optional[concurrent.futures.Future]:
        assert self._n_worker is not None and self._shard is not None
        if self._n_worker == 1 or self._shard > 0:
            # only run for distributed driver
            return None

        async def run_other_shard(shard: int):
            assert self._rank_to_addresses is not None
            address = self._rank_to_addresses[shard]
            model_actor_ref = await xo.actor_ref(
                address=address, uid=self.raw_model_uid
            )
            # we don't actually need to get the result from shard >= 1
            if stream:
                async for _ in await getattr(model_actor_ref, method)(*args, **kwargs):
                    pass
            else:
                await getattr(model_actor_ref, method)(*args, **kwargs)

        async def run_non_driver_shards():
            logger.debug("Start to run non driver %s", method)
            coros = []
            for rank in range(1, self._n_worker):
                coros.append(run_other_shard(rank))
            await asyncio.gather(*coros)

        assert self._loop is not None
        return asyncio.run_coroutine_threadsafe(run_non_driver_shards(), self._loop)

    async def async_generate(
        self,
        prompt: Union[str, Dict[str, Any]],
        generate_config: Optional[MLXGenerateConfig] = None,
        request_id: Optional[str] = None,
    ) -> Union[Completion, AsyncGenerator[CompletionChunk, None]]:
        logger.debug(
            "Enter async_generate, prompt: %s, generate config: %s",
            prompt,
            generate_config,
        )

        generate_config = self._sanitize_generate_config(generate_config)

        assert self._model is not None
        assert self._tokenizer is not None

        stream = generate_config.get("stream", False)

        # If batch model is available (non-distributed inference), use continuous batching
        if self._batch_model is not None:
            # Extract prompt text
            if isinstance(prompt, dict):
                prompt_text = prompt.get("prompt", "")
            else:
                prompt_text = prompt

            # Handle max_tokens - use auto-calculation if not set
            max_tokens = generate_config.get("max_tokens")
            if max_tokens is None:
                # Calculate max_tokens automatically based on context length and prompt length
                prompt_tokens = self._tokenizer.encode(prompt_text)
                input_echo_len = len(prompt_tokens)
                max_tokens = self._context_length - input_echo_len
                logger.debug(
                    f"Auto-calculated max_tokens: {max_tokens} (context_length: {self._context_length}, input_echo_len: {input_echo_len})"
                )

            if stream:
                # Return async generator for streaming
                async def stream_generator():
                    async for chunk in self._batch_model.generate_stream(
                        prompt=prompt_text,
                        max_tokens=max_tokens,
                        temperature=generate_config.get("temperature", 1.0),
                        top_p=generate_config.get("top_p", 1.0),
                        request_id=request_id or str(uuid.uuid4()),
                        skip_special_tokens=generate_config.get(
                            "skip_special_tokens", True
                        ),
                    ):
                        # Set model_uid for each chunk
                        if hasattr(chunk, "get"):
                            chunk["model"] = self.model_uid
                            if chunk.get("choices") and len(chunk["choices"]) > 0:
                                chunk["choices"][0]["index"] = 0
                        yield chunk

                return stream_generator()
            else:
                # Non-streaming: get full result and return completion
                result = await self._batch_model.generate(
                    prompt=prompt_text,
                    max_tokens=max_tokens,
                    temperature=generate_config.get("temperature", 1.0),
                    top_p=generate_config.get("top_p", 1.0),
                    stream=False,
                    skip_special_tokens=generate_config.get(
                        "skip_special_tokens", True
                    ),
                )

                # Return completion object
                completion = Completion(
                    id=str(uuid.uuid4()),
                    object="text_completion",
                    created=int(time.time()),
                    model=self.model_uid,
                    choices=[
                        {
                            "text": str(result),
                            "index": 0,
                            "logprobs": None,
                            "finish_reason": "stop",
                        }
                    ],
                    usage=CompletionUsage(
                        prompt_tokens=len(prompt_text),
                        completion_tokens=len(str(result)),
                        total_tokens=len(prompt_text) + len(str(result)),
                    ),
                )
                return completion

        # Fallback to original implementation for distributed inference
        # For distributed inference, use _generate_stream which doesn't use BatchGenerator
        if stream:

            async def stream_generator():
                loop = asyncio.get_event_loop()
                chunk_iterator = await loop.run_in_executor(
                    None, self._generate_stream, prompt, generate_config
                )
                for chunk, usage in chunk_iterator:
                    chunk["model"] = self.model_uid
                    yield chunk

            return stream_generator()
        else:
            # Non-streaming: collect all chunks and return completion
            loop = asyncio.get_event_loop()
            chunk_iterator = await loop.run_in_executor(
                None, self._generate_stream, prompt, generate_config
            )

            text = ""
            finish_reason = None
            usage: Optional[CompletionUsage] = None
            for chunk, chunk_usage in chunk_iterator:
                if chunk.get("choices") and len(chunk["choices"]) > 0:
                    text += chunk["choices"][0].get("text", "")
                    finish_reason = chunk["choices"][0].get("finish_reason")
                usage = chunk_usage

            return Completion(
                id=str(uuid.uuid4()),
                object="text_completion",
                created=int(time.time()),
                model=self.model_uid,
                choices=[
                    {
                        "text": text,
                        "index": 0,
                        "logprobs": None,
                        "finish_reason": finish_reason,
                    }
                ],
                usage=(
                    usage
                    if usage is not None
                    else CompletionUsage(
                        prompt_tokens=0, completion_tokens=0, total_tokens=0
                    )
                ),
            )


class MLXChatModel(MLXModel, ChatModelMixin):
    allow_batch: bool = True

    def _sanitize_generate_config(
        self,
        generate_config: Optional[MLXGenerateConfig],
    ) -> MLXGenerateConfig:
        generate_config = super()._sanitize_generate_config(generate_config)
        if (not generate_config.get("stop")) and self.model_family.stop:
            generate_config["stop"] = self.model_family.stop.copy()
        if (
            generate_config.get("stop_token_ids", None) is None
            and self.model_family.stop_token_ids
        ):
            generate_config["stop_token_ids"] = self.model_family.stop_token_ids.copy()

        return generate_config

    @classmethod
    def match_json(
        cls, llm_family: "LLMFamilyV2", llm_spec: "LLMSpecV1", quantization: str
    ) -> Union[bool, Tuple[bool, str]]:
        if llm_spec.model_format not in ["mlx"]:
            return False, "MLX chat engine only supports mlx format"
        if sys.platform != "darwin" or platform.processor() != "arm":
            return False, "MLX chat engine only works on Apple silicon Macs"
        if "chat" not in llm_family.model_ability:
            return False, "MLX chat engine requires chat ability"
        if "vision" in llm_family.model_ability:
            return False, "MLX chat engine does not support vision models"
        return True

    async def async_chat(
        self,
        messages: List[Dict],
        generate_config: Optional[MLXGenerateConfig] = None,
        request_id: Optional[str] = None,
    ) -> Union[ChatCompletion, AsyncGenerator[ChatCompletionChunk, None]]:
        model_family = self.model_family.model_family or self.model_family.model_name
        tools = generate_config.pop("tools", []) if generate_config else None
        chat_template_kwargs = (
            self._get_chat_template_kwargs_from_generate_config(
                generate_config, self.reasoning_parser
            )
            or {}
        )
        chat_context_var.set(chat_template_kwargs)
        full_context_kwargs = chat_template_kwargs.copy()
        if tools:
            if (
                model_family in QWEN_TOOL_CALL_FAMILY
                or model_family in GEMMA_TOOL_CALL_FAMILY
                or model_family in DEEPSEEK_TOOL_CALL_FAMILY
                or model_family in GLM5_TOOL_CALL_FAMILY
            ):
                full_context_kwargs["tools"] = tools
        chat_template = self.model_family.chat_template
        tokenizer = None
        if not chat_template:
            tokenizer = self._tokenizer
            if tokenizer is not None:
                chat_template = getattr(tokenizer, "chat_template", None)
        if not chat_template:
            raise ValueError(
                f"chat_template is required for model {self.model_uid}, but none was provided."
            )
        full_prompt = self.get_full_context(
            messages, chat_template, tokenizer=tokenizer, **full_context_kwargs
        )

        generate_config = self._sanitize_generate_config(generate_config)

        stream = generate_config.get("stream", False)
        if stream:
            # Use async_generate for streaming
            chunks = await self.async_generate(full_prompt, generate_config)
            assert hasattr(
                chunks, "__aiter__"
            ), "async_generate should return AsyncGenerator for streaming"

            async def _log_streaming_chunks():
                full_text = ""
                full_reasoning = ""
                async for chunk in chunks:  # type: ignore[arg-type]
                    choices = chunk.get("choices")
                    if choices:
                        first = choices[0]
                        delta = first.get("delta")
                        if isinstance(delta, dict):
                            delta_reasoning = delta.get("reasoning_content")
                            if isinstance(delta_reasoning, str):
                                full_reasoning += delta_reasoning
                            delta_text = delta.get("content")
                            if delta_text:
                                full_text += delta_text
                        else:
                            text = first.get("text")
                            if isinstance(text, str):
                                full_text += text
                    yield chunk
                logger.debug(
                    "[MLX] Full accumulated output: reasoning=%r, content=%r",
                    full_reasoning,
                    full_text,
                )

            return self._async_to_chat_completion_chunks(
                _log_streaming_chunks(),
                self.reasoning_parser,
                chat_template_kwargs,
            )
        else:
            # Use async_generate for non-streaming
            c = await self.async_generate(full_prompt, generate_config)
            assert not hasattr(
                c, "__aiter__"
            ), "async_generate should return Completion for non-streaming"
            if tools:
                return self._post_process_completion(
                    self.model_family, self.model_uid, c
                )
            return self._to_chat_completion(c, self.reasoning_parser)


class MLXVisionModel(MLXModel, ChatModelMixin):
    allow_batch: bool = False
    support_draft_model: bool = True

    def __init__(
        self,
        model_uid: str,
        model_family: "LLMFamilyV2",
        model_path: str,
        model_config: Optional[MLXModelConfig] = None,
        peft_model: Optional[List[LoRA]] = None,
    ):
        super().__init__(
            model_uid,
            model_family,
            model_path,
            model_config,
            peft_model,
        )
        self._mlx_executor: Optional[concurrent.futures.ThreadPoolExecutor] = None
        self._draft_model: Optional[Any] = None
        self._draft_kind: Optional[str] = None
        self._draft_block_size: Optional[int] = None

    def _run_on_mlx_thread(
        self, fn: Callable[..., Any], *args: Any, **kwargs: Any
    ) -> Any:
        executor = self._mlx_executor
        if executor is None:
            with _mlx_executor_lock:
                if self._mlx_executor is None:
                    self._mlx_executor = concurrent.futures.ThreadPoolExecutor(
                        max_workers=1,
                        thread_name_prefix=f"mlx-{self.model_uid}",
                    )
                executor = self._mlx_executor
        return executor.submit(fn, *args, **kwargs).result()

    def stop(self):
        with _mlx_executor_lock:
            executor = self._mlx_executor
            self._mlx_executor = None
        if executor is not None:
            executor.shutdown(wait=False)

    def _iterate_on_mlx_thread(
        self, iterator: Iterator[CompletionChunk]
    ) -> Iterator[CompletionChunk]:
        def _next_or_done():
            try:
                return True, next(iterator)
            except StopIteration:
                return False, None

        try:
            while True:
                has_item, item = self._run_on_mlx_thread(_next_or_done)
                if not has_item:
                    break
                yield item
        finally:
            close = getattr(iterator, "close", None)
            if close is not None:
                self._run_on_mlx_thread(close)

    @classmethod
    def check_lib(cls) -> Union[bool, Tuple[bool, str]]:
        dep_check = check_dependency_available("mlx_vlm", "mlx_vlm")
        if dep_check != True:
            return dep_check
        return True

    @classmethod
    def match_json(
        cls, llm_family: "LLMFamilyV2", llm_spec: "LLMSpecV1", quantization: str
    ) -> Union[bool, Tuple[bool, str]]:
        if llm_spec.model_format not in ["mlx"]:
            return False, "MLX vision engine only supports mlx format"
        if sys.platform != "darwin" or platform.processor() != "arm":
            return False, "MLX vision engine only works on Apple silicon Macs"
        if "vision" not in llm_family.model_ability:
            return False, "MLX vision engine requires vision ability"
        return True

    def _sanitize_generate_config(
        self,
        generate_config: Optional[MLXGenerateConfig],
    ) -> MLXGenerateConfig:
        generate_config = super()._sanitize_generate_config(generate_config)
        if (not generate_config.get("stop")) and self.model_family.stop:
            generate_config["stop"] = self.model_family.stop.copy()
        if (
            generate_config.get("stop_token_ids", None) is None
            and self.model_family.stop_token_ids
        ):
            generate_config["stop_token_ids"] = self.model_family.stop_token_ids.copy()

        return generate_config

    def generate(
        self,
        prompt: Union[str, Dict[str, Any]],
        generate_config: Optional[MLXGenerateConfig] = None,
        from_chat: bool = False,
    ) -> Union[Completion, Iterator[CompletionChunk]]:
        stream = bool(generate_config and generate_config.get("stream", False))
        result = self._run_on_mlx_thread(
            self._generate, prompt, generate_config, from_chat
        )
        if stream:
            assert isinstance(result, Iterator)
            return self._iterate_on_mlx_thread(result)
        assert not isinstance(result, Iterator)
        return result

    def _generate(
        self,
        prompt: Union[str, Dict[str, Any]],
        generate_config: Optional[MLXGenerateConfig] = None,
        from_chat: bool = False,
    ) -> Union[Completion, Iterator[CompletionChunk]]:
        """Generate method for vision models (not using continuous batching)."""
        generate_config = self._sanitize_generate_config(generate_config)
        logger.debug(f"[MLXVisionModel] generation params: {generate_config}")

        assert self._model is not None
        assert self._tokenizer is not None

        stream = generate_config.get("stream", False)

        if stream:
            # _generate_stream yields (chunk, usage) tuples; unwrap for the caller
            # so that _to_chat_completion_chunks receives plain CompletionChunk objects.
            def _unwrap():
                for chunk, _usage in self._generate_stream(prompt, generate_config):
                    yield chunk

            return _unwrap()
        else:
            # Non-streaming: collect all chunks and return completion
            text = ""
            finish_reason = None
            usage: Optional[CompletionUsage] = None
            for chunk, chunk_usage in self._generate_stream(prompt, generate_config):
                if chunk.get("choices") and len(chunk["choices"]) > 0:
                    text += chunk["choices"][0].get("text", "")
                    finish_reason = chunk["choices"][0].get("finish_reason")
                usage = chunk_usage

            return Completion(
                id=str(uuid.uuid4()),
                object="text_completion",
                created=int(time.time()),
                model=self.model_uid,
                choices=[
                    {
                        "text": text,
                        "index": 0,
                        "logprobs": None,
                        "finish_reason": finish_reason,
                    }
                ],
                usage=(
                    usage
                    if usage is not None
                    else CompletionUsage(
                        prompt_tokens=0, completion_tokens=0, total_tokens=0
                    )
                ),
            )

    def wait_for_load(self):
        """Override parent wait_for_load to skip MLXBatchModel creation."""
        if self._loading_thread:
            self._loading_thread.join()
            if self._loading_error:
                _, err, tb = self._loading_error
                raise err.with_traceback(tb)

    def _load_model(self, **kwargs):
        try:
            from mlx_vlm import load
        except ImportError:
            error_message = "Failed to import module 'mlx_vlm'"
            installation_guide = [
                "Please make sure 'mlx_vlm' is installed. ",
                "You can install it by `pip install mlx_vlm`\n",
            ]

            raise ImportError(f"{error_message}\n\n{''.join(installation_guide)}")

        self._prompt_cache = PromptCache()

        return load(self.model_path)

    def _load_draft_model(self, draft_model_path: str) -> Tuple[Any, str]:
        """Load the speculative drafter, e.g. a Gemma 4 ``*-it-assistant``
        checkpoint. Must run on the dedicated MLX thread, same as the target
        model, so that both share one thread-local generation stream."""
        try:
            from mlx_vlm.speculative.drafters import load_drafter
        except ImportError:
            raise ImportError(
                "Speculative decoding requires mlx-vlm>=0.5.0, "
                "you can upgrade it by `pip install -U mlx-vlm`"
            )

        generate_step = _ensure_mlx_vlm_thread_local_stream().generate_step
        if "draft_model" not in inspect.signature(generate_step).parameters:
            raise RuntimeError(
                "The installed mlx-vlm does not support speculative decoding, "
                "you can upgrade it by `pip install -U mlx-vlm`"
            )

        # ``kind`` is auto-detected from the drafter's ``model_type``, so a
        # gemma4_assistant checkpoint resolves to "mtp" on its own.
        draft_model, draft_kind = load_drafter(draft_model_path)
        try:
            from mlx_vlm.speculative.drafters import validate_drafter_compatibility
        except ImportError:
            # not available on older mlx-vlm, generate_step validates as well
            pass
        else:
            # fail at load time instead of on the first request
            validate_drafter_compatibility(self._model, draft_model, draft_kind)
        logger.info(
            "Loaded drafter %s for %s, kind: %s",
            draft_model_path,
            self.model_uid,
            draft_kind,
        )
        return draft_model, draft_kind

    def load(self):
        self._run_on_mlx_thread(self._load)

    def _load(self):
        from mlx_lm.utils import load_config

        if self._n_worker > 1:
            raise NotImplementedError(
                "Distributed inference is not supported for vision models"
            )

        reasoning_content = self._model_config.pop("reasoning_content")
        enable_thinking = self._model_config.pop("enable_thinking", True)
        self.prepare_parse_reasoning_content(
            reasoning_content, enable_thinking=enable_thinking
        )
        self.prepare_parse_tool_calls()

        kwargs = {}
        kwargs["revision"] = self._model_config.get(
            "revision", self.model_spec.model_revision
        )
        kwargs["trust_remote_code"] = self._model_config.get("trust_remote_code")
        kwargs["cache_limit_gb"] = self._model_config.pop("cache_limit_gb", None)

        draft_model_path = self._model_config.pop("draft_model_path", None)
        from ..core import parse_num_speculative_tokens

        # None keeps the depth the drafter was trained for
        self._draft_block_size = parse_num_speculative_tokens(
            self._model_config.pop("num_speculative_tokens", None)
        )

        self._model, self._processor = self._load_model(**kwargs)
        self._tokenizer = self._processor.tokenizer

        if draft_model_path:
            self._draft_model, self._draft_kind = self._load_draft_model(
                draft_model_path
            )

        # get context length
        config = load_config(Path(self.model_path))
        config.update(self._model_config)
        self._context_length = get_context_length_from_config(config)

    def _draft_generate_kwargs(self) -> Dict[str, Any]:
        """Extra ``generate_step`` kwargs enabling speculative decoding."""
        if self._draft_model is None:
            return {}
        draft_kwargs: Dict[str, Any] = {
            "draft_model": self._draft_model,
            "draft_kind": self._draft_kind,
        }
        if self._draft_block_size:
            # otherwise fall back to the block size configured by the drafter
            draft_kwargs["draft_block_size"] = self._draft_block_size
        return draft_kwargs

    def _generate_stream_inner(self, **kwargs):
        import mlx.core as mx

        try:
            from mlx_lm.utils import GenerationResponse
        except ImportError:
            # for mlx-lm >= 0.22.3
            from mlx_lm.generate import GenerationResponse
        mlx_vlm_generate = _ensure_mlx_vlm_thread_local_stream()
        generate_step = mlx_vlm_generate.generate_step

        inputs = kwargs.pop("prompt_token_ids")
        stop_token_ids = kwargs.pop("stop_token_ids", [])

        extra_kwargs = kwargs.copy()
        input_ids, pixel_values, mask, kwargs = inputs
        kwargs.update(extra_kwargs)
        kwargs.update(self._draft_generate_kwargs())

        tokenizer = self._processor.tokenizer
        detokenizer = self._processor.detokenizer

        detokenizer.reset()
        tic = time.perf_counter()
        try:
            for n, (token, logprobs) in enumerate(
                generate_step(input_ids, self._model, pixel_values, mask, **kwargs),
            ):
                if n == 0:
                    prompt_time = time.perf_counter() - tic
                    prompt_tps = len(input_ids) / prompt_time
                    tic = time.perf_counter()
                if token == tokenizer.eos_token_id or token in stop_token_ids:
                    break
                detokenizer.add_token(token)

                # Yield the last segment if streaming
                yield GenerationResponse(
                    text=detokenizer.last_segment,
                    token=token,
                    logprobs=logprobs,
                    from_draft=False,
                    prompt_tokens=len(input_ids),
                    prompt_tps=prompt_tps,
                    generation_tokens=n + 1,
                    generation_tps=(n + 1) / (time.perf_counter() - tic),
                    peak_memory=mx.metal.get_peak_memory() / 1e9,
                )

            detokenizer.finalize()
            yield GenerationResponse(
                text=detokenizer.last_segment,
                token=token,
                logprobs=logprobs,
                from_draft=False,
                prompt_tokens=len(input_ids),
                prompt_tps=prompt_tps,
                generation_tokens=n + 1,
                generation_tps=(n + 1) / (time.perf_counter() - tic),
                peak_memory=mx.metal.get_peak_memory() / 1e9,
            )
        finally:
            # after completing the inference, clear the memory
            self._cleanup_memory()

    def _prepare_inputs(
        self, prompt: Union[str, Dict[str, Any]], kwargs
    ) -> Tuple[Any, int]:
        import mlx.core as mx
        from mlx_vlm import prepare_inputs

        prompt_str = prompt.get("prompt")  # type: ignore
        images = prompt.get("multi_modal_data", {}).get("image")  # type: ignore
        if images and not isinstance(images, list):
            images = [images]
        resize_shape = kwargs.pop("resize_shape", None)
        image_token_index = getattr(self._model.config, "image_token_index", None)

        processor = self._processor
        tokenizer = processor if hasattr(processor, "encode") else processor.tokenizer
        prompt_tokens = mx.array(tokenizer.encode(prompt_str))

        if not images:
            input_ids = prompt_tokens[None, :]
            pixel_values = mask = None
            kwargs = {}
            input_token_len = input_ids.size
        else:
            # Check if processor supports audio to avoid feature_extractor errors
            supports_audio = hasattr(processor, "feature_extractor") or (
                hasattr(processor, "audio_tokenizer")
                and processor.audio_tokenizer is not None
            )

            # For processors that don't support audio (like Qwen3VLProcessor),
            # explicitly set audio=None to prevent mlx-vlm from trying to process audio
            if not supports_audio:
                inputs = prepare_inputs(
                    processor=processor,
                    images=images,
                    audio=None,
                    prompts=prompt_str,
                    image_token_index=image_token_index,
                    resize_shape=resize_shape,
                )
            else:
                inputs = prepare_inputs(
                    processor=processor,
                    images=images,
                    prompts=prompt_str,
                    image_token_index=image_token_index,
                    resize_shape=resize_shape,
                )
            input_ids = inputs["input_ids"]
            pixel_values = inputs["pixel_values"]
            mask = inputs["attention_mask"]
            kwargs = {
                k: v
                for k, v in inputs.items()
                if k not in ["input_ids", "pixel_values", "attention_mask"]
            }
            input_token_len = int(mask.sum())
        return (input_ids, pixel_values, mask, kwargs), input_token_len

    def chat(
        self,
        messages: List[Dict],
        generate_config: Optional[MLXGenerateConfig] = None,
    ) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
        messages = self._transform_messages(messages)  # type: ignore
        tools = generate_config.pop("tools", []) if generate_config else None

        model_family = self.model_family.model_family or self.model_family.model_name

        if "internvl2" not in model_family.lower():
            from qwen_vl_utils import process_vision_info

            chat_template_kwargs = (
                self._get_chat_template_kwargs_from_generate_config(
                    generate_config, self.reasoning_parser
                )
                or {}
            )
            chat_context_var.set(chat_template_kwargs)
            full_context_kwargs = chat_template_kwargs.copy()
            if tools and (
                model_family in QWEN_TOOL_CALL_FAMILY
                or model_family in GEMMA_TOOL_CALL_FAMILY
                or model_family in GLM5_TOOL_CALL_FAMILY
            ):
                full_context_kwargs["tools"] = tools
            chat_template = self.model_family.chat_template
            tokenizer = None
            if not chat_template:
                tokenizer = self._tokenizer
                if tokenizer is not None:
                    chat_template = getattr(tokenizer, "chat_template", None)
            if not chat_template:
                raise ValueError(
                    f"chat_template is required for model {self.model_uid}, but none was provided."
                )
            prompt = self.get_full_context(
                messages,
                chat_template,
                tokenizer=tokenizer,
                **full_context_kwargs,
            )
            images, video_inputs = process_vision_info(messages)
            if video_inputs:
                raise ValueError("Not support video input now.")
        else:
            prompt, images = self.get_specific_prompt(model_family, messages)  # type: ignore

        if not images:
            inputs = {
                "prompt": prompt,
            }
        elif len(images) == 1:
            inputs = {
                "prompt": prompt,
                "multi_modal_data": {"image": images[-1]},  # type: ignore
            }
        else:
            inputs = {
                "prompt": prompt,
                "multi_modal_data": {"image": images},  # type: ignore
            }
        generate_config = self._sanitize_generate_config(generate_config)

        stream = generate_config.get("stream", False)
        if stream:
            it = self.generate(inputs, generate_config)
            assert isinstance(it, Iterator)

            def _log_streaming_chunks():
                full_text = ""
                full_reasoning = ""
                for chunk in it:
                    choices = chunk.get("choices")
                    if choices:
                        first = choices[0]
                        delta = first.get("delta")
                        if isinstance(delta, dict):
                            delta_reasoning = delta.get("reasoning_content")
                            if isinstance(delta_reasoning, str):
                                full_reasoning += delta_reasoning
                            delta_text = delta.get("content")
                            if isinstance(delta_text, str):
                                full_text += delta_text
                        elif first.get("text"):
                            text = first["text"]
                            if text:
                                full_text += text  # type: ignore[arg-type]
                    yield chunk
                logger.debug(
                    "[MLX] Full accumulated output: reasoning=%r, content=%r",
                    full_reasoning,
                    full_text,
                )

            return self._to_chat_completion_chunks(
                _log_streaming_chunks(), self.reasoning_parser
            )
        else:
            c = self.generate(inputs, generate_config)
            assert not isinstance(c, Iterator)
            if tools:
                return self._post_process_completion(
                    self.model_family, self.model_uid, c
                )
            return self._to_chat_completion(c, self.reasoning_parser)
