iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Model] Support precomputed image embeddings for LightOnOCR (#61585)

Browse Source 
Signed-off-by: Rapisurazurite <43494437+Rapisurazurite@users.noreply.github.com>
This commit is contained in:
Lazurite
2026-10-14 10:54:09 +00:00
committed by GitHub
parent 81503550f3
commit 23fbc2ec2f
4 changed files with 96 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/models/supported_models.md
+1 -1
View File
@@ -581,7 +581,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I<sup>+</sup> | `moonshotai/Kimi-K2.5` | | ✅︎ |
| `KimiK3ForConditionalGeneration` | Kimi-K3 | T + I<sup>+</sup> | `moonshotai/Kimi-K3` | | ✅︎ |
| `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
| `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I<sup>+</sup> | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |
| `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I<sup>E+</sup> | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |
| `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |
| `Llama4ForConditionalGeneration` | Llama 4 | T + I<sup>+</sup> | `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 + I<sup>E+</sup> | `nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1` | ✅︎ | ✅︎ |
tests/models/multimodal/processing/test_lightonocr.py
+50
View File
@@ -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)
vllm/model_executor/models/lightonocr.py
+28 -15
View File
@@ -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)
vllm/model_executor/models/mistral3.py
+17 -4
View File
@@ -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"]