How oMLX Handles Multi-Image Chat with Vision-Language Models (VLMs)

oMLX supports multi-image chat through the VLMBatchedEngine, which validates model capabilities against a whitelist, formats chat messages with image placeholder tokens, and injects vision embeddings directly into the prefill phase while maintaining per-image feature caches for efficient reuse.

The oMLX project (jundot/omlx) provides a high-performance inference engine for vision-language models on Apple Silicon. Its implementation of multi-image chat extends the standard batched generation pipeline with specialized preprocessing, model-specific validation, and adaptive embedding injection that enables complex conversations with multiple images per request.

Architecture Overview

The multi-image pipeline in omlx/engine/vlm.py orchestrates several distinct phases: capability validation, message formatting, vision preprocessing, and cache-aware decoding. The engine first checks if the requested model supports multiple images, then transforms user messages into a format with <|image|> placeholders, processes all images through the vision encoder, and finally injects the resulting embeddings into the language model's prefill phase via the VLMModelAdapter.

Model Capability Validation

Before processing begins, the engine validates whether the loaded model supports multiple images. A whitelist named SINGLE_IMAGE_ONLY_MODELS defined at lines 402–410 in omlx/engine/vlm.py enumerates architectures restricted to single-image inputs:

# omlx/engine/vlm.py – lines 402-410

SINGLE_IMAGE_ONLY_MODELS = {
    "llava_next",
    "llava-qwen2",
    "bunny-llama",
    "paligemma",
    "multi_modality",
    "mllama",
}

When a request contains multiple images, the _prepare_vision_inputs method (lines 1089–1095) explicitly checks against this set:

# omlx/engine/vlm.py – lines 1089-1095

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."
    )

This validation occurs early in the pipeline to fail fast before any expensive vision encoding begins.

Chat Formatting and Image Placeholder Injection

The engine converts user-provided messages into a template-compatible format using _format_messages_for_vlm_template (lines 838–885). This method counts explicit image parts per message turn and assigns them to the appropriate user role, handling both explicit image references and fallback attachments:

# omlx/engine/vlm.py – lines 838-885 (excerpt)

if role == "user":
    explicit_images = self._count_content_parts(raw_content, image_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
    ):
        # fallback – attach all remaining images to the first user turn

        msg_num_images = remaining_images
        remaining_images = 0
        assigned_fallback_images = True

The function returns image_message_ranges (e.g., [(2, 2), (5, 1)]), which maps each message index to its image count. These ranges enable precise cache key computation for each image-bearing turn in the conversation.

Vision Preprocessing and Embedding Injection

After template application, the engine calls mlx_vlm.utils.prepare_inputs to tokenize text and run the vision processor on all images simultaneously (lines 1179–1185):

# omlx/engine/vlm.py – lines 1179-1185

inputs = prepare_inputs(
    self._processor,
    images=images if images else None,
    prompts=[prompt] if isinstance(prompt, str) else prompt,
)

The resulting embeddings are passed to the VLMModelAdapter (defined in omlx/models/vlm.py) via set_pending_embeddings (lines 910–931):

# omlx/models/vlm.py – lines 910-931 (conceptual)

self._adapter = VLMModelAdapter(self._vlm_model)
self._adapter.set_pending_embeddings(inputs_embeds, extra_kwargs, start_offset=0)

During the prefill phase, the adapter slices these embeddings to match the chunk size requested by the BatchGenerator. After prefill completes, the adapter falls back to standard token-ID decoding, with the vision context preserved in the KV cache.

Cache-Aware Optimizations

oMLX implements two caching strategies to optimize multi-image workloads. First, an SSD-backed VisionFeatureSSDCache (defined in omlx/cache/vision_feature_cache.py) stores per-image vision features. Before encoding, the engine computes hashes for each image using compute_per_image_hashes from omlx/utils/image.py (lines 1165–1178):

# omlx/engine/vlm.py – lines 1165-1178

per_hashes = compute_per_image_hashes(images)
cached_per_image = [
    self._vision_cache.get(h, self._model_name) for h in per_hashes
]

Second, for each image-bearing turn, the engine builds a prefix-cache key that uniquely identifies the image sequence up to that turn (lines 1190–1244). This enables KV-cache reuse across requests sharing the same image history, significantly reducing latency for multi-turn conversations.

Decoding Pipeline

Once prefill completes with injected vision embeddings, the decoding process proceeds using standard token IDs. The vision context remains resident in the KV cache, eliminating the need for repeated image processing during generation. For architectures like Qwen-3.5 that require multi-RoPE (mRoPE) handling, the VLMModelAdapter maintains rope-delta bookkeeping to ensure positional encoding correctness across images.

Practical Code Examples

Example 1: Multi-Image Chat with the Python Client

from omlx import Client  # oMLX high-level client

client = Client(model="qwen2_vl")          # a VLM that supports multi-image

messages = [
    {"role": "user", "content": [
        {"type": "text", "text": "Describe both pictures."},
        {"type": "image_url", "image_url": {"url": "https://example.com/cat.jpg"}},
        {"type": "image_url", "image_url": {"url": "https://example.com/dog.jpg"}},
    ]}
]

response = client.chat(messages)
print(response["choices"][0]["message"]["content"])

The engine assigns both images to the single user turn, inserts two <|image|> tokens, computes vision embeddings once, and returns a combined description.

Example 2: Handling Single-Image-Only Models

client = Client(model="paligemma")   # single-image only model

# This will raise a clear error

try:
    client.chat(messages)   # messages contain 2 images

except ValueError as err:
    print(err)   # → "Model paligemma does not support multi-image chat..."

Example 3: Low-Level Engine Access for Custom Pipelines

from omlx.engine.vlm import VLMBatchedEngine
import asyncio

async def run():
    engine = VLMBatchedEngine(
        model_name="qwen3_5_vl",
        trust_remote_code=True,
    )
    await engine.start()

    # Prepare a request with three images

    msgs = [{"role": "user", "content": "What are the objects?"}]
    images = ["cat.png", "dog.png", "bird.png"]   # paths or PIL.Image objects

    token_ids, inputs_embeds, extra_kwargs, _, _, _ = await engine._prepare_vision_inputs(
        messages=msgs,
        images=images,
    )

    # Inject embeddings and run the first prefill chunk

    engine._adapter.set_pending_embeddings(inputs_embeds, extra_kwargs)
    logits = await engine._engine.engine.run_step(token_ids)   # simplified

    print(logits.shape)

    await engine.stop()

asyncio.run(run())

This snippet demonstrates manual use of _prepare_vision_inputs and the adapter's set_pending_embeddings call for custom inference pipelines.

Summary

• Model Validation: The SINGLE_IMAGE_ONLY_MODELS whitelist in omlx/engine/vlm.py prevents incompatible multi-image requests before processing begins.
• Message Formatting: The _format_messages_for_vlm_template function handles explicit image placement and fallback attachment, generating image_message_ranges for cache management.
• Embedding Injection: Vision features are computed once via prepare_inputs and injected into the prefill phase through VLMModelAdapter.set_pending_embeddings.
• Caching: Per-image SSD caching (VisionFeatureSSDCache) and prefix-cache key generation enable efficient reuse of vision features and KV cache across requests.
• Zero-Overhead Decoding: After prefill, the pipeline uses standard token-ID decoding with vision context preserved in the KV cache.

Frequently Asked Questions

Which VLM architectures support multi-image chat in oMLX?

According to the source code in omlx/engine/vlm.py, models not listed in SINGLE_IMAGE_ONLY_MODELS support multi-image inputs. Supported architectures include qwen2_vl and qwen3_5_vl, while restricted models include paligemma, llava_next, llava-qwen2, bunny-llama, multi_modality, and mllama. Attempting to pass multiple images to restricted models raises a ValueError during the validation phase.

How does oMLX handle images that aren't explicitly referenced in message content?

The _format_messages_for_vlm_template method implements a fallback mechanism. When no message contains explicit image parts but images are provided in the request, the engine attaches all remaining images to the first user turn automatically. This ensures compatibility with simple text prompts that accompany image inputs without requiring strict message formatting.

What is the role of the VLMModelAdapter in the multi-image pipeline?

The VLMModelAdapter (defined in omlx/models/vlm.py) serves as the bridge between the vision encoder and the language model. It receives pre-computed vision embeddings via set_pending_embeddings and manages their injection into the prefill phase. The adapter handles chunking of embeddings to match the batch generator's requested sizes and manages positional encoding adjustments (including multi-RoPE) required by specific model architectures.

How does vision feature caching improve multi-image performance?

oMLX uses the VisionFeatureSSDCache class combined with image hashing utilities in omlx/utils/image.py to cache vision encoder outputs on disk. When processing a request, the engine computes hashes for each image and checks the cache before running the vision model. If all images are cached, the combined feature tensor is passed directly to the model without recomputation, significantly reducing latency for repeated images or multi-turn conversations with the same visual context.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}