@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 808dd87b30
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 9, 03:45 UTC

[Model] Support DeepSeek-OCR-2 (#33165)

Signed-off-by: liuli <ll407707@alibaba-inc.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: liuli <ll407707@alibaba-inc.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

RED committed Feb 2, 2026 at 14:24 UTC 808dd87b305426c809a21b2a6a4fb3ea607bca3f
9 files changed +1099 -1
docs/models/supported_models.md
+1
@@ -672,6 +672,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
672 | `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ |
673 | `DeepseekVLV2ForCausalLM`<sup>^</sup> | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
674 | `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
675 +| `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ |
676 | `Eagle2_5_VLForConditionalGeneration` | Eagle2.5-VL | T + I<sup>E+</sup> | `nvidia/Eagle2.5-8B`, etc. | ✅︎ | ✅︎ |
677 | `Ernie4_5_VLMoeForConditionalGeneration` | Ernie4.5-VL | T + I<sup>+</sup>/ V<sup>+</sup> | `baidu/ERNIE-4.5-VL-28B-A3B-PT`, `baidu/ERNIE-4.5-VL-424B-A47B-PT` | | ✅︎ |
678 | `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ |
examples/offline_inference/vision_language.py
+44
@@ -270,6 +270,49 @@ def run_deepseek_ocr(questions: list[str], modality: str) -> ModelRequestData:
270 )
271
272
273 +def run_deepseek_ocr2(questions: list[str], modality: str) -> ModelRequestData:
274 + from vllm.model_executor.models.deepseek_ocr import NGramPerReqLogitsProcessor
275 +
276 + assert modality == "image"
277 +
278 + model_name = "deepseek-ai/DeepSeek-OCR-2"
279 +
280 + engine_args = EngineArgs(
281 + model=model_name,
282 + limit_mm_per_prompt={modality: 1},
283 + logits_processors=[NGramPerReqLogitsProcessor],
284 + )
285 +
286 + # deepseek-ocr use plain prompt template
287 + prompts = [f"<image>\n{question}" for question in questions]
288 +
289 + # The following sampling params config is taken from
290 + # the official Deepseek-OCR inference example.
291 + # (IMPORTANT) Use the custom logits processor and avoid skipping
292 + # special tokens for this model for the optimal OCR performance.
293 + sampling_params = [
294 + SamplingParams(
295 + temperature=0.0,
296 + max_tokens=8192,
297 + # ngram logit processor args
298 + extra_args=dict(
299 + ngram_size=30,
300 + window_size=90,
301 + # whitelist: <td>, </td>
302 + whitelist_token_ids={128821, 128822},
303 + ),
304 + skip_special_tokens=False,
305 + )
306 + for _ in questions
307 + ]
308 +
309 + return ModelRequestData(
310 + engine_args=engine_args,
311 + prompts=prompts,
312 + sampling_params=sampling_params,
313 + )
314 +
315 +
316 # Dots-OCR
317 def run_dots_ocr(questions: list[str], modality: str) -> ModelRequestData:
318 assert modality == "image"
@@ -2045,6 +2088,7 @@ model_example_map = {
2088 "command_a_vision": run_command_a_vision,
2089 "deepseek_vl_v2": run_deepseek_vl2,
2090 "deepseek_ocr": run_deepseek_ocr,
2091 + "deepseek_ocr2": run_deepseek_ocr2,
2092 "dots_ocr": run_dots_ocr,
2093 "eagle2_5": run_eagle2_5,
2094 "ernie45_vl": run_ernie45_vl,
tests/models/registry.py
+3
@@ -687,6 +687,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
687 "DeepseekOCRForCausalLM": _HfExamplesInfo(
688 "deepseek-ai/DeepSeek-OCR",
689 ),
690 + "DeepseekOCR2ForCausalLM": _HfExamplesInfo(
691 + "deepseek-ai/DeepSeek-OCR-2",
692 + ),
693 "DotsOCRForCausalLM": _HfExamplesInfo(
694 "rednote-hilab/dots.ocr", trust_remote_code=True
695 ),
vllm/model_executor/models/deepencoder.py
+2 -1
@@ -79,6 +79,7 @@ class ImageEncoderViT(nn.Module):
79 rel_pos_zero_init: bool = True,
80 window_size: int = 0,
81 global_attn_indexes: tuple[int, ...] = (),
82 + last_conv_output: int = 1024,
83 ) -> None:
84 """
85 Args:
@@ -155,7 +156,7 @@ class ImageEncoderViT(nn.Module):
156 256, 512, kernel_size=3, stride=2, padding=1, bias=False
157 )
158 self.net_3 = Conv2dLayer(
158 - 512, 1024, kernel_size=3, stride=2, padding=1, bias=False
159 + 512, last_conv_output, kernel_size=3, stride=2, padding=1, bias=False
160 )
161
162 def get_abs_pos(self, abs_pos: torch.Tensor, tgt_size: int):
vllm/model_executor/models/deepencoder2.py new
+283
@@ -0,0 +1,283 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# adapted from
5 +# https://github.com/deepseek-ai/DeepSeek-OCR-2/blob/main/DeepSeek-OCR2-master/DeepSeek-OCR2-vllm/deepencoderv2/qwen2_d2e.py
6 +
7 +# Copyright (c) Meta Platforms, Inc. and affiliates.
8 +# All rights reserved.
9 +
10 +# This source code is licensed under the license found in the
11 +# LICENSE file in the root directory of this source tree.
12 +
13 +import torch
14 +import torch.nn as nn
15 +import transformers
16 +
17 +
18 +class CustomQwen2Decoder(nn.Module):
19 + """
20 + Qwen2 visual encoder
21 + non-causal attention + causal attention
22 + token_type_ids ：0=non-causal, 1=causal
23 + """
24 +
25 + def __init__(
26 + self,
27 + decoder_layer: int = 24,
28 + max_position_embeddings: int = 131072,
29 + hidden_dimension: int = 896,
30 + num_attention_heads: int = 14,
31 + num_key_value_heads: int = 2,
32 + intermediate_size: int = 4864,
33 + vocab_size: int = 151936,
34 + attn_implementation: str = "sdpa", # ⭐
35 + rms_norm_eps: float = 1e-06,
36 + rope_theta: float = 1000000.0,
37 + attention_dropout: float = 0.0,
38 + hidden_act: str = "silu",
39 + initializer_range: float = 0.02,
40 + ):
41 + super().__init__()
42 +
43 + # load
44 + Qwen2Model = transformers.models.qwen2.modeling_qwen2.Qwen2Model
45 + Qwen2Config = transformers.Qwen2Config
46 +
47 + # config
48 + config = Qwen2Config(
49 + hidden_size=hidden_dimension,
50 + num_hidden_layers=decoder_layer,
51 + num_attention_heads=num_attention_heads,
52 + num_key_value_heads=num_key_value_heads,
53 + intermediate_size=intermediate_size,
54 + max_position_embeddings=max_position_embeddings,
55 + vocab_size=vocab_size,
56 + rms_norm_eps=rms_norm_eps,
57 + rope_theta=rope_theta,
58 + attention_dropout=attention_dropout,
59 + hidden_act=hidden_act,
60 + initializer_range=initializer_range,
61 + _attn_implementation=attn_implementation, # ⭐
62 + )
63 +
64 + #
65 + self.model = self._create_custom_model(Qwen2Model, config)
66 +
67 + del self.model.embed_tokens
68 +
69 + def _create_custom_model(self, Qwen2Model, config):
70 + """Qwen2Model"""
71 +
72 + class CustomQwen2ModelInner(Qwen2Model):
73 + def forward(
74 + self,
75 + input_ids=None,
76 + attention_mask=None,
77 + position_ids=None,
78 + past_key_values=None,
79 + inputs_embeds=None,
80 + token_type_ids=None, # ⭐
81 + use_cache=None,
82 + output_attentions=None,
83 + output_hidden_states=None,
84 + return_dict=None,
85 + cache_position=None,
86 + ):
87 + # token_type_ids
88 + self._current_token_type_ids = token_type_ids
89 + causal_mask_mapping = {
90 + "full_attention": self._update_causal_mask(
91 + attention_mask,
92 + inputs_embeds,
93 + cache_position,
94 + past_key_values,
95 + output_attentions,
96 + )
97 + }
98 + outputs = super().forward(
99 + input_ids=input_ids,
100 + attention_mask=causal_mask_mapping,
101 + position_ids=position_ids,
102 + past_key_values=past_key_values,
103 + inputs_embeds=inputs_embeds,
104 + use_cache=use_cache,
105 + output_attentions=output_attentions,
106 + output_hidden_states=output_hidden_states,
107 + return_dict=return_dict,
108 + cache_position=cache_position,
109 + )
110 +
111 + return outputs
112 +
113 + def _update_causal_mask(
114 + self,
115 + attention_mask,
116 + input_tensor,
117 + cache_position,
118 + past_key_values,
119 + output_attentions,
120 + ):
121 + dtype, device = input_tensor.dtype, input_tensor.device
122 + min_dtype = torch.finfo(dtype).min
123 + batch_size, sequence_length = (
124 + input_tensor.shape[0],
125 + input_tensor.shape[1],
126 + )
127 +
128 + token_type_ids = self._current_token_type_ids
129 +
130 + # attention mask
131 + causal_mask = self._create_custom_4d_mask(
132 + sequence_length=sequence_length,
133 + dtype=dtype,
134 + device=device,
135 + batch_size=batch_size,
136 + token_type_ids=token_type_ids,
137 + )
138 +
139 + # padding mask
140 + if attention_mask is not None and attention_mask.dim() == 2:
141 + padding_mask = attention_mask[:, None, None, :].to(dtype=dtype)
142 + padding_mask = (1.0 - padding_mask) * min_dtype
143 + causal_mask = causal_mask + padding_mask
144 +
145 + return causal_mask
146 +
147 + def _create_custom_4d_mask(
148 + self,
149 + sequence_length,
150 + dtype,
151 + device,
152 + batch_size,
153 + token_type_ids,
154 + ):
155 + min_dtype = torch.finfo(dtype).min
156 +
157 + masks = []
158 + for b in range(batch_size):
159 + mask = torch.full(
160 + (sequence_length, sequence_length),
161 + fill_value=min_dtype,
162 + dtype=dtype,
163 + device=device,
164 + )
165 +
166 + type_ids = token_type_ids[b]
167 +
168 + image_positions = (type_ids == 0).nonzero(as_tuple=True)[0]
169 + text_positions = (type_ids == 1).nonzero(as_tuple=True)[0]
170 +
171 + # non-casual
172 + if len(image_positions) > 0:
173 + mask[image_positions[:, None], image_positions] = 0.0
174 +
175 + # causal
176 + for i, text_pos in enumerate(text_positions):
177 + if len(image_positions) > 0:
178 + mask[text_pos, image_positions] = 0.0
179 + mask[text_pos, text_positions[: i + 1]] = 0.0
180 +
181 + masks.append(mask)
182 +
183 + mask = torch.stack(masks, dim=0).unsqueeze(1)
184 + return mask
185 +
186 + return CustomQwen2ModelInner(config)
187 +
188 + def forward(
189 + self,
190 + inputs_embeds: torch.Tensor,
191 + token_type_ids: torch.Tensor,
192 + attention_mask: torch.Tensor = None,
193 + **kwargs,
194 + ):
195 + """
196 + Args:
197 + inputs_embeds: [batch_size, seq_len, hidden_dim]
198 + token_type_ids: [batch_size, seq_len], 0=non-causal, 1=causal
199 + attention_mask: [batch_size, seq_len], optional
200 + """
201 + return self.model(
202 + inputs_embeds=inputs_embeds,
203 + token_type_ids=token_type_ids,
204 + attention_mask=attention_mask,
205 + **kwargs,
206 + )
207 +
208 +
209 +class Qwen2Decoder2Encoder(nn.Module):
210 + """
211 + Decoder based on Multilingual BART
212 + Set the initial weights and configuration with a pretrained multilingual BART model,
213 + and modify the detailed configurations as a Nougat decoder
214 + """
215 +
216 + def __init__(
217 + self,
218 + decoder_layer: int,
219 + hidden_dimension: int,
220 + num_attention_heads: int,
221 + num_key_value_heads: int,
222 + intermediate_size: int,
223 + ):
224 + super().__init__()
225 +
226 + self.model = CustomQwen2Decoder(
227 + decoder_layer=decoder_layer,
228 + hidden_dimension=hidden_dimension,
229 + num_attention_heads=num_attention_heads,
230 + num_key_value_heads=num_key_value_heads,
231 + intermediate_size=intermediate_size,
232 + attn_implementation="sdpa",
233 + )
234 + self.query_768 = nn.Embedding(144, hidden_dimension)
235 + self.query_1024 = nn.Embedding(256, hidden_dimension)
236 +
237 + def forward(self, x: torch.Tensor) -> torch.Tensor:
238 + x = x.flatten(2).transpose(1, 2)
239 +
240 + bs, n_query, _ = x.shape
241 +
242 + if n_query == 144:
243 + param_img = self.query_768.weight
244 + elif n_query == 256:
245 + param_img = self.query_1024.weight
246 +
247 + batch_query_imgs = param_img.unsqueeze(0).expand(
248 + bs, -1, -1
249 + ) # (batch_size, num_queries, hidden_size)
250 +
251 + x_combined = torch.cat([x, batch_query_imgs], dim=1)
252 +
253 + token_type_ids = torch.cat(
254 + [
255 + torch.zeros(bs, n_query, dtype=torch.long),
256 + torch.ones(bs, n_query, dtype=torch.long),
257 + ],
258 + dim=1,
259 + )
260 +
261 + y = self.model(x_combined, token_type_ids)[0]
262 +
263 + y = y[:, n_query:, :] # causal flow query
264 +
265 + return y
266 +
267 +
268 +def build_qwen2_decoder_as_encoder(
269 + decoder_layer=24,
270 + hidden_dimension=896,
271 + num_attention_heads=14,
272 + num_key_value_heads=2,
273 + intermediate_size=4864,
274 +):
275 + decoder_as_encoder = Qwen2Decoder2Encoder(
276 + decoder_layer=decoder_layer,
277 + hidden_dimension=hidden_dimension,
278 + num_attention_heads=num_attention_heads,
279 + num_key_value_heads=num_key_value_heads,
280 + intermediate_size=intermediate_size,
281 + )
282 +
283 + return decoder_as_encoder
vllm/model_executor/models/deepseek_ocr2.py new
+444
@@ -0,0 +1,444 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only Deepseek-OCR model compatible with HuggingFace weights."""
4 +
5 +import math
6 +from collections.abc import Iterable, Mapping, Sequence
7 +from functools import partial
8 +
9 +import torch
10 +import torch.nn as nn
11 +from transformers import BatchFeature
12 +
13 +from vllm.config import VllmConfig
14 +from vllm.config.multimodal import BaseDummyOptions
15 +from vllm.model_executor.models.interfaces import (
16 + MultiModalEmbeddings,
17 + SupportsLoRA,
18 + SupportsMultiModal,
19 + SupportsPP,
20 +)
21 +from vllm.model_executor.models.module_mapping import MultiModelKeys
22 +from vllm.model_executor.models.utils import (
23 + AutoWeightsLoader,
24 + WeightsMapper,
25 + init_vllm_registered_model,
26 + maybe_prefix,
27 +)
28 +from vllm.multimodal import MULTIMODAL_REGISTRY
29 +from vllm.multimodal.inputs import (
30 + MultiModalDataDict,
31 + MultiModalFieldConfig,
32 + MultiModalKwargsItems,
33 + NestedTensors,
34 +)
35 +from vllm.multimodal.parse import (
36 + ImageEmbeddingItems,
37 + ImageProcessorItems,
38 + ImageSize,
39 + MultiModalDataItems,
40 +)
41 +from vllm.multimodal.processing import (
42 + BaseDummyInputsBuilder,
43 + BaseMultiModalProcessor,
44 + BaseProcessingInfo,
45 + PromptReplacement,
46 + PromptUpdate,
47 +)
48 +from vllm.sequence import IntermediateTensors
49 +from vllm.tokenizers import cached_tokenizer_from_config
50 +from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config
51 +from vllm.transformers_utils.processors.deepseek_ocr2 import (
52 + BASE_SIZE,
53 + CROP_MODE,
54 + IMAGE_SIZE,
55 + DeepseekOCR2Processor,
56 +)
57 +
58 +from ...transformers_utils.processors.deepseek_ocr import count_tiles
59 +from .deepencoder import ImageEncoderViT
60 +from .deepencoder2 import build_qwen2_decoder_as_encoder
61 +from .deepseek_ocr import DeepseekOCRImagePixelInputs
62 +from .deepseek_vl2 import MlpProjector
63 +
64 +# The image token id may be various
65 +_IMAGE_TOKEN = "<image>"
66 +
67 +
68 +class DeepseekOCR2ProcessingInfo(BaseProcessingInfo):
69 + def get_hf_config(self):
70 + return self.ctx.get_hf_config(DeepseekVLV2Config)
71 +
72 + def get_hf_processor(self, **kwargs: object):
73 + return self.ctx.get_hf_processor(DeepseekOCR2Processor, **kwargs)
74 +
75 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
76 + return {"image": None}
77 +
78 + def get_num_image_tokens(
79 + self, *, image_width: int, image_height: int, cropping: bool = True
80 + ) -> int:
81 + image_size = IMAGE_SIZE
82 + base_size = BASE_SIZE
83 + patch_size = 16
84 + downsample_ratio = 4
85 +
86 + if CROP_MODE:
87 + if image_width <= 768 and image_height <= 768:
88 + crop_ratio = [1, 1]
89 + else:
90 + # find the closest aspect ratio to the target
91 + crop_ratio = count_tiles(
92 + image_width, image_height, image_size=IMAGE_SIZE
93 + )
94 +
95 + num_width_tiles, num_height_tiles = crop_ratio
96 + else:
97 + num_width_tiles = num_height_tiles = 1
98 +
99 + h = w = math.ceil((base_size // patch_size) / downsample_ratio)
100 +
101 + h2 = w2 = math.ceil((image_size // patch_size) / downsample_ratio)
102 +
103 + global_views_tokens = h * w
104 + if num_width_tiles > 1 or num_height_tiles > 1:
105 + local_views_tokens = (num_height_tiles * h2) * (num_width_tiles * w2)
106 + else:
107 + local_views_tokens = 0
108 +
109 + return global_views_tokens + local_views_tokens + 1
110 +
111 + def get_image_size_with_most_features(self) -> ImageSize:
112 + if IMAGE_SIZE == 1024 and BASE_SIZE == 1280:
113 + return ImageSize(width=1024 * 2, height=1024 * 2)
114 + return ImageSize(width=768 * 2, height=768 * 2)
115 +
116 +
117 +class DeepseekOCR2DummyInputsBuilder(
118 + BaseDummyInputsBuilder[DeepseekOCR2ProcessingInfo]
119 +):
120 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
121 + num_images = mm_counts.get("image", 0)
122 +
123 + processor = self.info.get_hf_processor()
124 + image_token = processor.image_token
125 +
126 + return image_token * num_images
127 +
128 + def get_dummy_mm_data(
129 + self,
130 + seq_len: int,
131 + mm_counts: Mapping[str, int],
132 + mm_options: Mapping[str, BaseDummyOptions] | None = None,
133 + ) -> MultiModalDataDict:
134 + num_images = mm_counts.get("image", 0)
135 +
136 + max_image_size = self.info.get_image_size_with_most_features()
137 +
138 + return {
139 + "image": self._get_dummy_images(
140 + width=max_image_size.width,
141 + height=max_image_size.height,
142 + num_images=num_images,
143 + )
144 + }
145 +
146 +
147 +class DeepseekOCR2MultiModalProcessor(
148 + BaseMultiModalProcessor[DeepseekOCR2ProcessingInfo]
149 +):
150 + def _call_hf_processor(
151 + self,
152 + prompt: str,
153 + mm_data: Mapping[str, object],
154 + mm_kwargs: Mapping[str, object],
155 + tok_kwargs: Mapping[str, object],
156 + ) -> BatchFeature:
157 + if mm_data:
158 + processed_outputs = self.info.ctx.call_hf_processor(
159 + self.info.get_hf_processor(**mm_kwargs),
160 + dict(prompt=prompt, **mm_data),
161 + mm_kwargs,
162 + )
163 +
164 + else:
165 + tokenizer = self.info.get_tokenizer()
166 + processed_outputs = tokenizer(
167 + prompt, add_special_tokens=True, return_tensors="pt"
168 + )
169 +
170 + return processed_outputs
171 +
172 + def _get_mm_fields_config(
173 + self,
174 + hf_inputs: BatchFeature,
175 + hf_processor_mm_kwargs: Mapping[str, object],
176 + ) -> Mapping[str, MultiModalFieldConfig]:
177 + images_spatial_crop = hf_inputs.get("images_spatial_crop", torch.empty((0, 2)))
178 + is_tiled = (images_spatial_crop[:, 0] > 1) | (images_spatial_crop[:, 1] > 1)
179 + patches_per_image = torch.where(is_tiled, images_spatial_crop.prod(dim=-1), 0)
180 + return dict(
181 + pixel_values=MultiModalFieldConfig.batched("image"),
182 + images_spatial_crop=MultiModalFieldConfig.batched("image"),
183 + images_crop=MultiModalFieldConfig.flat_from_sizes(
184 + "image", patches_per_image
185 + ),
186 + )
187 +
188 + def _get_prompt_updates(
189 + self,
190 + mm_items: MultiModalDataItems,
191 + hf_processor_mm_kwargs: Mapping[str, object],
192 + out_mm_kwargs: MultiModalKwargsItems,
193 + ) -> Sequence[PromptUpdate]:
194 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
195 +
196 + image_token_id = hf_processor.image_token_id
197 + assert isinstance(image_token_id, int)
198 +
199 + def get_replacement_deepseek_vl2(item_idx: int):
200 + images = mm_items.get_items(
201 + "image", (ImageEmbeddingItems, ImageProcessorItems)
202 + )
203 +
204 + if isinstance(images, ImageEmbeddingItems):
205 + num_image_tokens = images.get_feature_size(item_idx)
206 + else:
207 + size = images.get_image_size(item_idx)
208 +
209 + num_image_tokens = self.info.get_num_image_tokens(
210 + image_width=size.width,
211 + image_height=size.height,
212 + cropping=CROP_MODE,
213 + )
214 + return [image_token_id] * num_image_tokens
215 +
216 + return [
217 + PromptReplacement(
218 + modality="image",
219 + target=[image_token_id],
220 + replacement=get_replacement_deepseek_vl2,
221 + )
222 + ]
223 +
224 +
225 +@MULTIMODAL_REGISTRY.register_processor(
226 + DeepseekOCR2MultiModalProcessor,
227 + info=DeepseekOCR2ProcessingInfo,
228 + dummy_inputs=DeepseekOCR2DummyInputsBuilder,
229 +)
230 +class DeepseekOCR2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
231 + hf_to_vllm_mapper = WeightsMapper(
232 + orig_to_new_prefix={
233 + # map prefix for language backbone
234 + "model.embed_tokens.": "language_model.model.embed_tokens.",
235 + "model.layers.": "language_model.model.layers.",
236 + "model.norm.": "language_model.model.norm.",
237 + "lm_head.": "language_model.lm_head.",
238 + # remove "model." prefix for other components
239 + "model.": "",
240 + }
241 + )
242 +
243 + @classmethod
244 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
245 + if modality.startswith("image"):
246 + return "<image>"
247 +
248 + raise ValueError("Only image modality is supported")
249 +
250 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
251 + super().__init__()
252 +
253 + config: DeepseekVLV2Config = vllm_config.model_config.hf_config
254 + multimodal_config = vllm_config.model_config.multimodal_config
255 +
256 + self.config = config
257 + self.multimodal_config = multimodal_config
258 +
259 + self.vision_config = config.vision_config
260 + self.projector_config = config.projector_config
261 + self.text_config = config.text_config
262 + model_config = vllm_config.model_config
263 + tokenizer = cached_tokenizer_from_config(model_config)
264 + self.image_token_id = tokenizer.vocab[_IMAGE_TOKEN]
265 +
266 + with self._mark_tower_model(vllm_config, "image"):
267 + self.sam_model = ImageEncoderViT(
268 + depth=12,
269 + embed_dim=768,
270 + img_size=1024,
271 + mlp_ratio=4,
272 + norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
273 + num_heads=12,
274 + patch_size=16,
275 + qkv_bias=True,
276 + use_rel_pos=True,
277 + global_attn_indexes=[2, 5, 8, 11],
278 + window_size=14,
279 + out_chans=256,
280 + last_conv_output=896,
281 + )
282 + self.qwen2_model = build_qwen2_decoder_as_encoder()
283 +
284 + self.projector = MlpProjector(self.projector_config)
285 + self.tile_tag = config.tile_tag
286 + self.global_view_pos = config.global_view_pos
287 +
288 + # special token for image token sequence format
289 + n_embed = self.projector_config.n_embed
290 + embed_std = 1 / torch.sqrt(torch.tensor(n_embed, dtype=torch.float32))
291 + if self.tile_tag == "2D":
292 + # This is a typo in original implementation
293 + self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std)
294 + else:
295 + raise ValueError(
296 + f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"
297 + )
298 +
299 + with self._mark_language_model(vllm_config):
300 + self.language_model = init_vllm_registered_model(
301 + vllm_config=vllm_config,
302 + hf_config=self.text_config,
303 + prefix=maybe_prefix(prefix, "language_model"),
304 + )
305 +
306 + self.make_empty_intermediate_tensors = (
307 + self.language_model.make_empty_intermediate_tensors
308 + )
309 +
310 + def _parse_and_validate_image_input(
311 + self, **kwargs: object
312 + ) -> DeepseekOCRImagePixelInputs | None:
313 + pixel_values = kwargs.pop("pixel_values", None)
314 + images_spatial_crop = kwargs.pop("images_spatial_crop", None)
315 + images_crop = kwargs.pop("images_crop", None)
316 +
317 + if pixel_values is None or torch.sum(pixel_values).item() == 0:
318 + return None
319 +
320 + base_size = self.vision_config.image_size
321 + return DeepseekOCRImagePixelInputs(
322 + type="pixel_values",
323 + data=pixel_values,
324 + images_crop=images_crop,
325 + images_spatial_crop=images_spatial_crop,
326 + resolve_bindings={
327 + "base_size": base_size,
328 + },
329 + )
330 +
331 + def _encode_global_features(self, image_tensor: torch.Tensor) -> torch.Tensor:
332 + global_features_1 = self.sam_model(image_tensor)
333 + global_features_2 = self.qwen2_model(global_features_1)
334 +
335 + features = self.projector(global_features_2)
336 +
337 + _, hw, dim = features.shape
338 +
339 + return features.view(-1, dim)
340 +
341 + def _encode_local_features(self, patches: torch.Tensor) -> torch.Tensor | None:
342 + if torch.sum(patches).item() == 0:
343 + return None
344 +
345 + local_features = self.sam_model(patches)
346 + local_features = self.qwen2_model(local_features)
347 +
348 + features = self.projector(local_features)
349 +
350 + _, _, dim = features.shape
351 +
352 + return features.view(-1, dim)
353 +
354 + def _pixel_values_to_embedding(
355 + self,
356 + pixel_values: torch.Tensor,
357 + images_crop: torch.Tensor,
358 + images_spatial_crop: torch.Tensor,
359 + ) -> NestedTensors:
360 + images_in_this_batch = []
361 +
362 + is_tiled = (images_spatial_crop[:, 0] > 1) | (images_spatial_crop[:, 1] > 1)
363 + patches_per_image = torch.where(is_tiled, images_spatial_crop.prod(dim=-1), 0)
364 + images_crop = images_crop.split(patches_per_image.tolist())
365 + for jdx in range(images_spatial_crop.size(0)):
366 + patches = images_crop[jdx]
367 + image_ori = pixel_values[[jdx]]
368 +
369 + global_features = self._encode_global_features(image_ori)
370 + local_features = self._encode_local_features(patches)
371 +
372 + if local_features is not None:
373 + combined = torch.cat(
374 + [local_features, global_features, self.view_seperator[None, :]],
375 + dim=0,
376 + )
377 + else:
378 + combined = torch.cat(
379 + [global_features, self.view_seperator[None, :]], dim=0
380 + )
381 +
382 + images_in_this_batch.append(combined)
383 +
384 + return images_in_this_batch
385 +
386 + def _process_image_input(
387 + self, image_input: DeepseekOCRImagePixelInputs
388 + ) -> torch.Tensor:
389 + pixel_values = image_input.data
390 + images_crop = image_input.images_crop
391 + images_spatial_crop = image_input.images_spatial_crop.to(dtype=torch.long)
392 +
393 + vision_features = self._pixel_values_to_embedding(
394 + pixel_values=pixel_values,
395 + images_crop=images_crop,
396 + images_spatial_crop=images_spatial_crop,
397 + )
398 +
399 + return vision_features
400 +
401 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:
402 + image_input = self._parse_and_validate_image_input(**kwargs)
403 + if image_input is None:
404 + return None
405 + vision_embeddings = self._process_image_input(image_input)
406 + return vision_embeddings
407 +
408 + def forward(
409 + self,
410 + input_ids: torch.Tensor,
411 + positions: torch.Tensor,
412 + intermediate_tensors: IntermediateTensors | None = None,
413 + inputs_embeds: torch.Tensor | None = None,
414 + **kwargs: object,
415 + ):
416 + if intermediate_tensors is not None:
417 + inputs_embeds = None
418 +
419 + hidden_states = self.language_model(
420 + input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
421 + )
422 +
423 + return hidden_states
424 +
425 + def compute_logits(
426 + self,
427 + hidden_states: torch.Tensor,
428 + ) -> torch.Tensor | None:
429 + return self.language_model.compute_logits(hidden_states)
430 +
431 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
432 + loader = AutoWeightsLoader(self)
433 + autoloaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
434 + return autoloaded_weights
435 +
436 + def get_mm_mapping(self) -> MultiModelKeys:
437 + """
438 + Get the module prefix in multimodal models
439 + """
440 + return MultiModelKeys.from_string_field(
441 + language_model="language_model",
442 + connector="projector",
443 + tower_model=["sam_model", "qwen2_model"],
444 + )
vllm/model_executor/models/registry.py
+1
@@ -308,6 +308,7 @@ _MULTIMODAL_MODELS = {
308 ),
309 "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
310 "DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"),
311 + "DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"),
312 "DotsOCRForCausalLM": ("dots_ocr", "DotsOCRForCausalLM"),
313 "Eagle2_5_VLForConditionalGeneration": (
314 "eagle2_5_vl",
vllm/transformers_utils/chat_templates/registry.py
+1
@@ -34,6 +34,7 @@ _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK: dict[str, ChatTemplatePath] = {
34 "chameleon": CHAT_TEMPLATES_DIR / "template_basic.jinja",
35 "clip": CHAT_TEMPLATES_DIR / "template_basic.jinja",
36 "deepseek_ocr": CHAT_TEMPLATES_DIR / "template_deepseek_ocr.jinja",
37 + "deepseek_ocr2": CHAT_TEMPLATES_DIR / "template_deepseek_ocr.jinja",
38 "deepseek_vl_v2": CHAT_TEMPLATES_DIR / "template_deepseek_vl2.jinja",
39 "fuyu": CHAT_TEMPLATES_DIR / "template_fuyu.jinja",
40 "minicpmv": _get_minicpmv_chat_template_fallback,
vllm/transformers_utils/processors/deepseek_ocr2.py new
+320
@@ -0,0 +1,320 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# adapted from https://github.com/deepseek-ai/DeepSeek-OCR/blob/main/DeepSeek-OCR-master/DeepSeek-OCR-vllm/process/image_process.py
4 +import math
5 +
6 +import torch
7 +from PIL import Image, ImageOps
8 +from transformers import AutoProcessor, BatchFeature, LlamaTokenizerFast
9 +from transformers.processing_utils import ProcessorMixin
10 +
11 +from vllm.transformers_utils.processors.deepseek_ocr import (
12 + ImageTransform,
13 + dynamic_preprocess,
14 +)
15 +
16 +BASE_SIZE = 1024
17 +IMAGE_SIZE = 768
18 +CROP_MODE = True
19 +MIN_CROPS = 2
20 +MAX_CROPS = 6
21 +
22 +
23 +class DeepseekOCR2Processor(ProcessorMixin):
24 + tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast")
25 + attributes = ["tokenizer"]
26 +
27 + def __init__(
28 + self,
29 + tokenizer: LlamaTokenizerFast,
30 + patch_size: int = 16,
31 + downsample_ratio: int = 4,
32 + image_mean: tuple[float, float, float] = (0.5, 0.5, 0.5),
33 + image_std: tuple[float, float, float] = (0.5, 0.5, 0.5),
34 + normalize: bool = True,
35 + image_token: str = "<image>",
36 + pad_token: str = "<｜▁pad▁｜>",
37 + add_special_token: bool = False,
38 + sft_format: str = "deepseek",
39 + mask_prompt: bool = True,
40 + ignore_id: int = -100,
41 + **kwargs,
42 + ):
43 + self.image_size = IMAGE_SIZE
44 + self.base_size = BASE_SIZE
45 + self.patch_size = 16
46 + self.image_mean = image_mean
47 + self.image_std = image_std
48 + self.normalize = normalize
49 + self.downsample_ratio = 4
50 +
51 + self.image_transform = ImageTransform(
52 + mean=image_mean, std=image_std, normalize=normalize
53 + )
54 +
55 + self.tokenizer = tokenizer
56 + self.tokenizer.padding_side = "left" # must set this，padding side with make a difference in batch inference # noqa: E501
57 +
58 + # add the pad_token as special token to use 'tokenizer.pad_token'
59 + # and 'tokenizer.pad_token_id'
60 + if self.tokenizer.pad_token is None:
61 + self.tokenizer.add_special_tokens({"pad_token": pad_token})
62 +
63 + # add image token
64 + self.image_token_id = self.tokenizer.vocab.get(image_token)
65 + self.image_token = image_token
66 + self.pad_token = pad_token
67 + self.add_special_token = add_special_token
68 + self.sft_format = sft_format
69 + self.mask_prompt = mask_prompt
70 + self.ignore_id = ignore_id
71 +
72 + super().__init__(
73 + tokenizer,
74 + **kwargs,
75 + )
76 +
77 + @property
78 + def bos_id(self):
79 + return self.tokenizer.bos_token_id
80 +
81 + @property
82 + def eos_id(self):
83 + return self.tokenizer.eos_token_id
84 +
85 + @property
86 + def pad_id(self):
87 + return self.tokenizer.pad_token_id
88 +
89 + def encode(self, text: str, bos: bool = True, eos: bool = False):
90 + t = self.tokenizer.encode(text, add_special_tokens=False)
91 + if bos:
92 + t = [self.bos_id] + t
93 + if eos:
94 + t = t + [self.eos_id]
95 + return t
96 +
97 + def decode(self, t: list[int], **kwargs) -> str:
98 + return self.tokenizer.decode(t, **kwargs)
99 +
100 + def process_one(
101 + self,
102 + prompt: str,
103 + images: list[Image.Image],
104 + crop_mode: bool = CROP_MODE,
105 + ):
106 + """
107 +
108 + Args:
109 + prompt (str): the formatted prompt;
110 + images (List[ImageType]): the list of images;
111 + crop_mode (bool): if True, then crop the image;
112 +
113 + Returns:
114 + outputs (BaseProcessorOutput): the output of the processor,
115 + - input_ids (torch.LongTensor): [N + image tokens]
116 + - target_ids (torch.LongTensor): [N + image tokens]
117 + - pixel_values (torch.FloatTensor): [n_patches, 3, H, W]
118 + - image_id (int): the id of the image token
119 + - num_image_tokens (List[int]): the number of image tokens
120 + """
121 +
122 + assert prompt is not None and images is not None, (
123 + "prompt and images must be used at the same time."
124 + )
125 +
126 + sft_format = prompt
127 +
128 + (
129 + input_ids,
130 + pixel_values,
131 + images_crop,
132 + images_seq_mask,
133 + images_spatial_crop,
134 + num_image_tokens,
135 + _,
136 + ) = self.tokenize_with_images(
137 + conversation=sft_format,
138 + images=images,
139 + bos=True,
140 + eos=True,
141 + cropping=crop_mode,
142 + )
143 +
144 + prepare = BatchFeature(
145 + data=dict(
146 + input_ids=input_ids,
147 + pixel_values=pixel_values,
148 + images_crop=images_crop,
149 + images_seq_mask=images_seq_mask,
150 + images_spatial_crop=images_spatial_crop,
151 + num_image_tokens=num_image_tokens,
152 + ),
153 + tensor_type="pt",
154 + )
155 + return prepare
156 +
157 + def __call__(
158 + self,
159 + *,
160 + prompt: str,
161 + images: list[Image.Image],
162 + crop_mode: bool = CROP_MODE,
163 + **kwargs,
164 + ):
165 + prepare = self.process_one(
166 + prompt=prompt,
167 + images=images,
168 + crop_mode=crop_mode,
169 + )
170 +
171 + return prepare
172 +
173 + def tokenize_with_images(
174 + self,
175 + conversation: str,
176 + images: list[Image.Image],
177 + bos: bool = True,
178 + eos: bool = True,
179 + cropping: bool = True,
180 + ):
181 + """Tokenize text with <image> tags."""
182 +
183 + assert conversation.count(self.image_token) == len(images)
184 + text_splits = conversation.split(self.image_token)
185 + images_list, images_crop_list, images_seq_mask, images_spatial_crop = (
186 + [],
187 + [],
188 + [],
189 + [],
190 + )
191 + image_shapes = []
192 + num_image_tokens = []
193 + tokenized_str = []
194 + for text_sep, image in zip(text_splits, images):
195 + tokenized_sep = self.encode(text_sep, bos=False, eos=False)
196 + tokenized_str += tokenized_sep
197 + images_seq_mask += [False] * len(tokenized_sep)
198 +
199 + image_shapes.append(image.size)
200 +
201 + images_crop_raw = []
202 + if image.size[0] <= 768 and image.size[1] <= 768:
203 + crop_ratio = [1, 1]
204 + elif cropping:
205 + images_crop_raw, crop_ratio = dynamic_preprocess(
206 + image, image_size=IMAGE_SIZE
207 + )
208 + else:
209 + crop_ratio = [1, 1]
210 +
211 + if self.image_size <= 768 and not cropping:
212 + image = image.resize((self.image_size, self.image_size))
213 +
214 + global_view = ImageOps.pad(
215 + image,
216 + (self.base_size, self.base_size),
217 + color=tuple(int(x * 255) for x in self.image_transform.mean),
218 + )
219 + images_list.append(self.image_transform(global_view))
220 +
221 + num_width_tiles, num_height_tiles = crop_ratio
222 + images_spatial_crop.append([num_width_tiles, num_height_tiles])
223 +
224 + if num_width_tiles > 1 or num_height_tiles > 1:
225 + for cropped_image in images_crop_raw:
226 + images_crop_list.append(self.image_transform(cropped_image))
227 +
228 + num_queries = math.ceil(
229 + (self.image_size // self.patch_size) / self.downsample_ratio
230 + )
231 + num_queries_base = math.ceil(
232 + (self.base_size // self.patch_size) / self.downsample_ratio
233 + )
234 +
235 + tokenized_image = (
236 + [self.image_token_id] * num_queries_base
237 + ) * num_queries_base
238 + tokenized_image += [self.image_token_id]
239 + if num_width_tiles > 1 or num_height_tiles > 1:
240 + local_row = [self.image_token_id] * (num_queries * num_width_tiles)
241 + tokenized_image += local_row * (num_queries * num_height_tiles)
242 + tokenized_str += tokenized_image
243 + images_seq_mask += [True] * len(tokenized_image)
244 + num_image_tokens.append(len(tokenized_image))
245 +
246 + """process the last text split"""
247 + tokenized_sep = self.encode(text_splits[-1], bos=False, eos=False)
248 + tokenized_str += tokenized_sep
249 + images_seq_mask += [False] * len(tokenized_sep)
250 +
251 + """add the bos and eos tokens"""
252 + if bos:
253 + tokenized_str = [self.bos_id] + tokenized_str
254 + images_seq_mask = [False] + images_seq_mask
255 + if eos:
256 + tokenized_str = tokenized_str + [self.eos_id]
257 + images_seq_mask = images_seq_mask + [False]
258 +
259 + assert len(tokenized_str) == len(images_seq_mask), (
260 + f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} "
261 + f"is not equal to images_seq_mask's length {len(images_seq_mask)}."
262 + )
263 +
264 + masked_tokenized_str = []
265 + for token_index in tokenized_str:
266 + if token_index != self.image_token_id:
267 + masked_tokenized_str.append(token_index)
268 + else:
269 + masked_tokenized_str.append(self.ignore_id)
270 +
271 + assert (
272 + len(tokenized_str) == len(images_seq_mask) == len(masked_tokenized_str)
273 + ), (
274 + f"tokenized_str's length {len(tokenized_str)}, "
275 + f"input_ids' length {len(masked_tokenized_str)}, "
276 + f"images_seq_mask's length {len(images_seq_mask)}, are not equal."
277 + )
278 +
279 + input_ids = torch.LongTensor(tokenized_str)
280 + target_ids = torch.LongTensor(masked_tokenized_str)
281 + images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool)
282 +
283 + # set input_ids < 0 | input_ids == self.image_token_id as ignore_id
284 + target_ids[(input_ids < 0) | (input_ids == self.image_token_id)] = (
285 + self.ignore_id
286 + )
287 + input_ids[input_ids < 0] = self.pad_id
288 +
289 + # Remove the ending eos token
290 + assert input_ids[-1] == self.eos_id
291 + input_ids = input_ids[:-1]
292 + target_ids = target_ids[:-1]
293 + images_seq_mask = images_seq_mask[:-1]
294 +
295 + if len(images_list) == 0:
296 + pixel_values = torch.zeros((0, 3, self.base_size, self.base_size))
297 + images_spatial_crop = torch.zeros((0, 2), dtype=torch.long)
298 + images_crop = torch.zeros((0, 3, self.image_size, self.image_size))
299 + else:
300 + pixel_values = torch.stack(images_list, dim=0)
301 + images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long)
302 + if images_crop_list:
303 + images_crop = torch.stack(images_crop_list, dim=0)
304 + else:
305 + images_crop = torch.zeros((0, 3, self.image_size, self.image_size))
306 +
307 + input_ids = input_ids.unsqueeze(0)
308 +
309 + return (
310 + input_ids,
311 + pixel_values,
312 + images_crop,
313 + images_seq_mask,
314 + images_spatial_crop,
315 + num_image_tokens,
316 + image_shapes,
317 + )
318 +
319 +
320 +AutoProcessor.register("DeepseekOCR2Processor", DeepseekOCR2Processor)