From 23fbc2ec2f7219888bfc76cf987381f501151971 Mon Sep 17 00:00:00 2001 From: Lazurite <43494437+Rapisurazurite@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:54:09 +0800 Subject: [PATCH] [Model] Support precomputed image embeddings for LightOnOCR (#61585) Signed-off-by: Rapisurazurite <43494437+Rapisurazurite@users.noreply.github.com> --- docs/models/supported_models.md | 2 +- .../multimodal/processing/test_lightonocr.py | 50 +++++++++++++++++++ vllm/model_executor/models/lightonocr.py | 43 ++++++++++------ vllm/model_executor/models/mistral3.py | 21 ++++++-- 4 files changed, 96 insertions(+), 20 deletions(-) create mode 100644 tests/models/multimodal/processing/test_lightonocr.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 1b3c0d39bf..156210e434 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -581,7 +581,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I+ | `moonshotai/Kimi-K2.5` | | ✅︎ | | `KimiK3ForConditionalGeneration` | Kimi-K3 | T + I+ | `moonshotai/Kimi-K3` | | ✅︎ | | `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I+ | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ | -| `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I+ | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ | +| `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + IE+ | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ | | `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I+ | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ | | `Llama4ForConditionalGeneration` | Llama 4 | T + I+ | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. | ✅︎ | ✅︎ | | `Llama_Nemotron_Nano_VL` | Llama Nemotron Nano VL | T + IE+ | `nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1` | ✅︎ | ✅︎ | diff --git a/tests/models/multimodal/processing/test_lightonocr.py b/tests/models/multimodal/processing/test_lightonocr.py new file mode 100644 index 0000000000..3966bf8cb4 --- /dev/null +++ b/tests/models/multimodal/processing/test_lightonocr.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for LightOnOCR precomputed image embeddings.""" + +import pytest +import torch + +from vllm.multimodal import MULTIMODAL_REGISTRY + +from ...utils import build_model_context + +_LIGHTON_MODEL_ID = "lightonai/LightOnOCR-1B-1025" +_MODEL_CONFIG_KWARGS = {"config_format": "hf"} + + +@pytest.mark.usefixtures("default_vllm_config") +@pytest.mark.parametrize("tensor_input", [False, True]) +@pytest.mark.parametrize("prompt_as_tokens", [False, True]) +def test_lightonocr_image_embeddings_preserve_lengths_and_values( + tensor_input: bool, prompt_as_tokens: bool +): + """Reserve each image's feature count and pass projected embeddings through.""" + ctx = build_model_context( + _LIGHTON_MODEL_ID, + limit_mm_per_prompt={"image": 2}, + model_config_kwargs=_MODEL_CONFIG_KWARGS, + ) + ctx.model_config.get_multimodal_config().enable_mm_embeds = True + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + hf_processor = processor.info.get_hf_processor() + hidden_size = ctx.model_config.get_hidden_size() + feature_sizes = [3, 3] if tensor_input else [3, 7] + embeddings = [ + torch.full((size, hidden_size), float(i)) + for i, size in enumerate(feature_sizes) + ] + image_input = torch.stack(embeddings) if tensor_input else embeddings + prompt = " ".join([hf_processor.image_token] * len(embeddings)) + if prompt_as_tokens: + prompt = ctx.get_tokenizer().encode(prompt) + + items = processor.info.parse_mm_data({"image": image_input}) + processed = processor(prompt, mm_items=items) + + assert processed["prompt_token_ids"].count(hf_processor.image_token_id) == sum( + feature_sizes + ) + assert [p.length for p in processed["mm_placeholders"]["image"]] == feature_sizes + for item, expected in zip(processed["mm_kwargs"]["image"], embeddings): + torch.testing.assert_close(item.get_data()["image_embeds"], expected) diff --git a/vllm/model_executor/models/lightonocr.py b/vllm/model_executor/models/lightonocr.py index d89fa53299..f1e8eb5f57 100644 --- a/vllm/model_executor/models/lightonocr.py +++ b/vllm/model_executor/models/lightonocr.py @@ -4,15 +4,14 @@ from collections.abc import Iterable, Mapping, Sequence import torch import torch.nn as nn -from transformers import ( - BatchFeature, - PixtralVisionConfig, -) +from transformers import BatchFeature from vllm.config import VllmConfig from vllm.model_executor.models.mistral3 import ( Mistral3DummyInputsBuilder, Mistral3ForConditionalGeneration, + Mistral3ImageEmbeddingInputs, + Mistral3ImageInputs, Mistral3MultiModalProjector, Mistral3ProcessingInfo, init_vision_tower_for_mistral3, @@ -26,12 +25,15 @@ from vllm.model_executor.models.utils import ( ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems -from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems +from vllm.multimodal.parse import ( + ImageEmbeddingItems, + ImageProcessorItems, + MultiModalDataItems, +) from vllm.multimodal.processing import ( BaseMultiModalProcessor, PromptReplacement, PromptUpdate, - PromptUpdateDetails, ) @@ -102,18 +104,20 @@ class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[LightOnOCRProcessing hf_config = self.info.get_hf_config() image_token_id = hf_config.image_token_index - assert isinstance(hf_config.vision_config, PixtralVisionConfig) - encoder_info = PixtralHFEncoderInfo(hf_config) - def replace(item_idx: int): - images = mm_items.get_items("image", ImageProcessorItems) - size = images.get_image_size(item_idx) - ncols, nrows = encoder_info.get_patch_grid_size( - image_width=size.width, image_height=size.height + images = mm_items.get_items( + "image", (ImageEmbeddingItems, ImageProcessorItems) ) + if isinstance(images, ImageEmbeddingItems): + num_image_tokens = images.get_feature_size(item_idx) + else: + assert isinstance(images, ImageProcessorItems) + size = images.get_image_size(item_idx) + num_image_tokens = self.info.get_num_image_tokens( + image_width=size.width, image_height=size.height + ) # break/end tokens are not used in LightOnOCR - tokens = [image_token_id] * (ncols * nrows) - return PromptUpdateDetails.select_token_id(tokens, image_token_id) + return [image_token_id] * num_image_tokens return [ PromptReplacement( @@ -178,6 +182,15 @@ class LightOnOCRForConditionalGeneration(Mistral3ForConditionalGeneration): self.language_model.make_empty_intermediate_tensors ) + def _parse_and_validate_image_input( + self, **kwargs: object + ) -> Mistral3ImageInputs | None: + image_embeds = kwargs.get("image_embeds") + if image_embeds is not None: + return Mistral3ImageEmbeddingInputs(data=image_embeds) + + return super()._parse_and_validate_image_input(**kwargs) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/mistral3.py b/vllm/model_executor/models/mistral3.py index f5ffe32aad..93879bb76b 100644 --- a/vllm/model_executor/models/mistral3.py +++ b/vllm/model_executor/models/mistral3.py @@ -3,7 +3,7 @@ import math from collections.abc import Iterable, Mapping, Sequence -from typing import Annotated, Any, Literal, cast +from typing import Annotated, Any, Literal, TypeAlias, cast import torch import torch.nn as nn @@ -76,6 +76,19 @@ class Mistral3ImagePixelInputs(TensorSchema): ] +class Mistral3ImageEmbeddingInputs(TensorSchema): + """Projected features with shape (num_images, feature_size, hidden_size).""" + + type: Literal["image_embeds"] = "image_embeds" + data: Annotated[ + torch.Tensor | list[torch.Tensor], + TensorShape("bn", "ifs", "hs", dynamic_dims={"ifs"}), + ] + + +Mistral3ImageInputs: TypeAlias = Mistral3ImagePixelInputs | Mistral3ImageEmbeddingInputs + + class Mistral3PatchMerger(nn.Module): """Learned merging of spatial_merge_size ** 2 patches.""" @@ -494,7 +507,7 @@ class Mistral3ForConditionalGeneration( def _parse_and_validate_image_input( self, **kwargs: object - ) -> Mistral3ImagePixelInputs | None: + ) -> Mistral3ImageInputs | None: pixel_values = kwargs.pop("pixel_values", None) image_embeds = kwargs.pop("image_embeds", None) @@ -508,8 +521,8 @@ class Mistral3ForConditionalGeneration( def _process_image_input( self, - image_input: Mistral3ImagePixelInputs, - ) -> torch.Tensor | tuple[torch.Tensor, ...]: + image_input: Mistral3ImageInputs, + ) -> MultiModalEmbeddings: if image_input["type"] == "image_embeds": return image_input["data"]