# SPDX-FileCopyrightText: The Docling Contributors # SPDX-License-Identifier: MIT """MLX-based VLM inference engine for Apple Silicon.""" from __future__ import annotations import logging import threading import time from contextlib import contextmanager from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, List, Union from PIL.Image import Image from docling.datamodel.vlm_engine_options import MlxVlmEngineOptions from docling.models.inference_engines.vlm._utils import ( check_min_engine_version, extract_generation_stoppers, preprocess_image_batch, ) from docling.models.inference_engines.vlm.base import ( BaseVlmEngine, VlmEngineInput, VlmEngineOutput, VlmEngineType, ) from docling.models.utils.generation_utils import GenerationStopper from docling.models.utils.hf_model_download import HuggingFaceModelDownloadMixin if TYPE_CHECKING: from docling.datamodel.stage_model_specs import EngineModelConfig _log = logging.getLogger(__name__) # Global lock for MLX model calls - MLX models are not thread-safe # All MLX models share this lock to prevent concurrent MLX operations _MLX_GLOBAL_LOCK = threading.Lock() # mlx-vlm injects enable_thinking=False into processor.apply_chat_template for any # processor whose signature accepts **kwargs. Models whose chat template does not # reference that variable (e.g. chandra-ocr-2) make transformers>=5.16 log a # warning about it. The kwarg only reaches the processor when tokenize=True, and # mlx-vlm calls with tokenize=False, so it is inert here -- drop just that record # rather than let it imply something went wrong. Remove once mlx-vlm stops # passing the kwarg unconditionally. _PROCESSOR_KWARGS_WARNING = "have to be in `processor_kwargs` dict" class _ProcessorKwargsWarningFilter(logging.Filter): def filter(self, record: logging.LogRecord) -> bool: return _PROCESSOR_KWARGS_WARNING not in record.getMessage() @contextmanager def _suppressed_processor_kwargs_warning(): logger = logging.getLogger("transformers.processing_utils") log_filter = _ProcessorKwargsWarningFilter() logger.addFilter(log_filter) try: yield finally: logger.removeFilter(log_filter) class MlxVlmEngine(BaseVlmEngine, HuggingFaceModelDownloadMixin): """MLX engine for VLM inference on Apple Silicon. This engine uses the mlx-vlm library to run vision-language models efficiently on Apple Silicon (M1/M2/M3) using the Metal Performance Shaders. Note: MLX models are not thread-safe and use a global lock. """ def __init__( self, options: MlxVlmEngineOptions, artifacts_path: Union[Path, str] | None, model_config: EngineModelConfig | None = None, ): """Initialize the MLX engine. Args: options: MLX-specific runtime options artifacts_path: Path to cached model artifacts model_config: Model configuration (repo_id, revision, extra_config) """ super().__init__(options, model_config=model_config) self.options: MlxVlmEngineOptions = options self.artifacts_path = ( artifacts_path if artifacts_path is None else Path(artifacts_path) ) # These will be set during initialization # MLX types are complex and external, using Any with type: ignore self.vlm_model: Any = None self.processor: Any = None self.config: Any = None self.apply_chat_template: Any = None self.stream_generate: Any = None # Initialize immediately if model_config is provided if self.model_config is not None: self.initialize() def initialize(self) -> None: """Initialize the MLX model and processor.""" if self._initialized: return _log.info("Initializing MLX VLM inference engine...") try: from mlx_vlm import load, stream_generate from mlx_vlm.prompt_utils import apply_chat_template from mlx_vlm.utils import load_config except ImportError: raise ImportError( "mlx-vlm is not installed. Please install it via `pip install mlx-vlm` " "to use MLX VLM models on Apple Silicon." ) check_min_engine_version( VlmEngineType.MLX, self.model_config.min_engine_version if self.model_config else None, ) self.apply_chat_template = apply_chat_template # type: ignore[assignment] self.stream_generate = stream_generate # type: ignore[assignment] # Load model if model_config is provided if self.model_config is not None and self.model_config.repo_id is not None: repo_id = self.model_config.repo_id revision = self.model_config.revision or "main" _log.info(f"Loading MLX model {repo_id} (revision: {revision})") self._load_model_for_repo(repo_id, revision=revision) self._initialized = True _log.info("MLX runtime initialized") def _load_model_for_repo(self, repo_id: str, revision: str = "main") -> None: """Load model and processor for a specific repository. Args: repo_id: HuggingFace repository ID revision: Model revision """ from mlx_vlm import load from mlx_vlm.utils import load_config # Download or locate model artifacts repo_cache_folder = repo_id.replace("/", "--") if self.artifacts_path is None: artifacts_path = self.download_models(repo_id, revision=revision) elif (self.artifacts_path / repo_cache_folder).exists(): artifacts_path = self.artifacts_path / repo_cache_folder else: # Model not found in artifacts_path - raise clear error available_models = [] if self.artifacts_path.exists(): available_models = [ p.name for p in self.artifacts_path.iterdir() if p.is_dir() ] raise FileNotFoundError( f"Model '{repo_id}' not found in artifacts_path.\n" f"Expected location: {self.artifacts_path / repo_cache_folder}\n" f"Available models in {self.artifacts_path}: " f"{', '.join(available_models) if available_models else 'none'}\n\n" f"To fix this issue:\n" f" 1. Download the model: docling-tools models download-hf-repo {repo_id}\n" f" 2. Or remove --artifacts-path to enable auto-download\n" f" 3. Or use a different model that exists in your artifacts_path" ) # Load the model start_time = time.monotonic() self.vlm_model, self.processor = load(artifacts_path) self.config = load_config(artifacts_path) load_time = time.monotonic() - start_time _log.info( f"Loaded MLX model {repo_id} (revision: {revision}) in {load_time:.2f} sec." ) def predict_batch(self, input_batch: List[VlmEngineInput]) -> List[VlmEngineOutput]: """Run inference on a batch of inputs. Note: MLX models are not thread-safe and use a global lock, so batch processing is done sequentially. This method is provided for API consistency but does not provide performance benefits over sequential processing. Args: input_batch: List of inputs to process Returns: List of outputs, one per input """ if not self._initialized: self.initialize() if not input_batch: return [] # Model should already be loaded via initialize() if self.vlm_model is None or self.processor is None or self.config is None: raise RuntimeError( "Model not loaded. Ensure EngineModelConfig was provided during initialization." ) _log.debug( f"MLX runtime processing batch of {len(input_batch)} images sequentially " "(MLX does not support batched inference)" ) outputs: List[VlmEngineOutput] = [] _log.info( "Running MLX inference on %s image(s) (max_new_tokens=%s)...", len(input_batch), input_batch[0].max_new_tokens, ) # MLX models are not thread-safe - use global lock to serialize access with _MLX_GLOBAL_LOCK: _log.debug("MLX model: Acquired global lock for thread safety") for input_data in input_batch: # Preprocess image images = preprocess_image_batch([input_data.image]) image = images[0] # Format prompt using MLX's chat template with _suppressed_processor_kwargs_warning(): formatted_prompt = self.apply_chat_template( self.processor, self.config, input_data.prompt, num_images=1 ) # Extract custom stopping criteria custom_stoppers = extract_generation_stoppers( input_data.extra_generation_config ) # Stream generate with stop strings and custom stopping criteria support start_time = time.time() _log.debug("Starting MLX generation...") output_text = "" num_tokens = 0 stop_reason = "unspecified" # Use stream_generate for proper stop string handling for token in self.stream_generate( self.vlm_model, self.processor, formatted_prompt, [image], # MLX stream_generate expects list of images max_tokens=input_data.max_new_tokens, verbose=False, temp=input_data.temperature, ): output_text += token.text num_tokens += 1 # Check for configured stop strings if input_data.stop_strings: if any( stop_str in output_text for stop_str in input_data.stop_strings ): _log.debug("Stopping generation due to stop string match") stop_reason = "stop_string" break # Check for custom stopping criteria if custom_stoppers: for stopper in custom_stoppers: # Determine the text window to check based on lookback_tokens lookback_tokens = stopper.lookback_tokens() text_to_check = ( output_text[-lookback_tokens:] if len(output_text) > lookback_tokens else output_text ) try: if stopper.should_stop(text_to_check): _log.info( f"Stopping generation due to GenerationStopper: {type(stopper).__name__}" ) stop_reason = "custom_criteria" break except Exception as e: _log.warning( f"Error in GenerationStopper.should_stop: {e}" ) continue else: # for-else: only executed if inner loop didn't break continue # Break outer loop if any stopper triggered break generation_time = time.time() - start_time _log.info( "MLX generated %s tokens in %.2f sec. (%.2f tok/s, " "stop_reason: %s)", num_tokens, generation_time, num_tokens / generation_time if generation_time > 0 else 0.0, stop_reason, ) # Create output outputs.append( VlmEngineOutput( text=output_text, stop_reason=stop_reason, metadata={ "generation_time": generation_time, "num_tokens": num_tokens, "model": self.model_config.repo_id if self.model_config else "unknown", }, ) ) _log.debug("MLX model: Released global lock") return outputs def cleanup(self) -> None: """Clean up model resources.""" if self.vlm_model is not None: del self.vlm_model self.vlm_model = None if self.processor is not None: del self.processor self.processor = None _log.info("MLX runtime cleaned up")