@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 40b69e33e7
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 13, 20:45 UTC

[Model] Add PaddleOCR-VL Model Support (#27758)

Signed-off-by: zhangyue <zhangyue66@baidu.com> Signed-off-by: Roger Wang <hey@rogerw.io> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: zhangyue66 <zhangyue66@baidu.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

zhang-prog committed Nov 3, 2025 at 19:04 UTC 40b69e33e796efdc75e774a1c38cc73397ea6e17
7 files changed +1475
docs/models/supported_models.md
+1
@@ -675,6 +675,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
675 | `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
676 | `Ovis` | Ovis2, Ovis1.6 | T + I<sup>+</sup> | `AIDC-AI/Ovis2-1B`, `AIDC-AI/Ovis1.6-Llama3.2-3B`, etc. | | ✅︎ |
677 | `Ovis2_5` | Ovis2.5 | T + I<sup>+</sup> + V | `AIDC-AI/Ovis2.5-9B`, etc. | | |
678 +| `PaddleOCRVLForConditionalGeneration` | Paddle-OCR | T + I<sup>+</sup> | `PaddlePaddle/PaddleOCR-VL`, etc. | | |
679 | `PaliGemmaForConditionalGeneration` | PaliGemma, PaliGemma 2 | T + I<sup>E</sup> | `google/paligemma-3b-pt-224`, `google/paligemma-3b-mix-224`, `google/paligemma2-3b-ft-docci-448`, etc. | | ✅︎ |
680 | `Phi3VForCausalLM` | Phi-3-Vision, Phi-3.5-Vision | T + I<sup>E+</sup> | `microsoft/Phi-3-vision-128k-instruct`, `microsoft/Phi-3.5-vision-instruct`, etc. | | ✅︎ |
681 | `Phi4MMForCausalLM` | Phi-4-multimodal | T + I<sup>+</sup> / T + A<sup>+</sup> / I<sup>+</sup> + A<sup>+</sup> | `microsoft/Phi-4-multimodal-instruct`, etc. | ✅︎ | ✅︎ |
examples/offline_inference/vision_language.py
+27
@@ -1242,6 +1242,32 @@ def run_ovis2_5(questions: list[str], modality: str) -> ModelRequestData:
1242 )
1243
1244
1245 +# PaddleOCR-VL
1246 +def run_paddleocr_vl(questions: list[str], modality: str) -> ModelRequestData:
1247 + assert modality == "image"
1248 +
1249 + model_name = "PaddlePaddle/PaddleOCR-VL"
1250 +
1251 + engine_args = EngineArgs(
1252 + model=model_name,
1253 + max_model_len=4096,
1254 + max_num_seqs=2,
1255 + limit_mm_per_prompt={modality: 1},
1256 + trust_remote_code=True,
1257 + )
1258 +
1259 + placeholder = "<|IMAGE_START|><|IMAGE_PLACEHOLDER|><|IMAGE_END|>"
1260 + prompts = [
1261 + (f"<|begin_of_sentence|>User: {question}{placeholder}\nAssistant: ")
1262 + for question in questions
1263 + ]
1264 +
1265 + return ModelRequestData(
1266 + engine_args=engine_args,
1267 + prompts=prompts,
1268 + )
1269 +
1270 +
1271 # PaliGemma
1272 def run_paligemma(questions: list[str], modality: str) -> ModelRequestData:
1273 assert modality == "image"
@@ -1817,6 +1843,7 @@ model_example_map = {
1843 "NVLM_D": run_nvlm_d,
1844 "ovis": run_ovis,
1845 "ovis2_5": run_ovis2_5,
1846 + "paddleocr_vl": run_paddleocr_vl,
1847 "paligemma": run_paligemma,
1848 "paligemma2": run_paligemma2,
1849 "phi3_v": run_phi3v,
examples/offline_inference/vision_language_multi_image.py
+22
@@ -801,6 +801,27 @@ def load_ovis2_5(question: str, image_urls: list[str]) -> ModelRequestData:
801 )
802
803
804 +def load_paddleocr_vl(question: str, image_urls: list[str]) -> ModelRequestData:
805 + model_name = "PaddlePaddle/PaddleOCR-VL"
806 +
807 + engine_args = EngineArgs(
808 + model=model_name,
809 + trust_remote_code=True,
810 + max_model_len=8192,
811 + max_num_seqs=2,
812 + limit_mm_per_prompt={"image": len(image_urls)},
813 + )
814 +
815 + placeholders = "<|IMAGE_START|><|IMAGE_PLACEHOLDER|><|IMAGE_END|>" * len(image_urls)
816 + prompt = f"<|begin_of_sentence|>User: {question}{placeholders}\nAssistant: "
817 +
818 + return ModelRequestData(
819 + engine_args=engine_args,
820 + prompt=prompt,
821 + image_data=[fetch_image(url) for url in image_urls],
822 + )
823 +
824 +
825 def load_pixtral_hf(question: str, image_urls: list[str]) -> ModelRequestData:
826 model_name = "mistral-community/pixtral-12b"
827
@@ -1312,6 +1333,7 @@ model_example_map = {
1333 "NVLM_D": load_nvlm_d,
1334 "ovis": load_ovis,
1335 "ovis2_5": load_ovis2_5,
1336 + "paddleocr_vl": load_paddleocr_vl,
1337 "phi3_v": load_phi3v,
1338 "phi4_mm": load_phi4mm,
1339 "phi4_multimodal": load_phi4_multimodal,
tests/models/registry.py
+4
@@ -712,6 +712,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
712 },
713 ),
714 "Ovis2_5": _HfExamplesInfo("AIDC-AI/Ovis2.5-2B", trust_remote_code=True),
715 + "PaddleOCRVLForConditionalGeneration": _HfExamplesInfo(
716 + "PaddlePaddle/PaddleOCR-VL",
717 + trust_remote_code=True,
718 + ),
719 "PaliGemmaForConditionalGeneration": _HfExamplesInfo(
720 "google/paligemma-3b-mix-224",
721 extras={"v2": "google/paligemma2-3b-ft-docci-448"},
vllm/model_executor/models/ernie45.py
+10
@@ -23,12 +23,22 @@
23 # limitations under the License.
24 """Inference-only Erine model compatible with HuggingFace weights."""
25
26 +from vllm.compilation.decorators import support_torch_compile
27 from vllm.config import VllmConfig
28 from vllm.model_executor.models.llama import LlamaForCausalLM
29
30 from .utils import PPMissingLayer
31
32
33 +@support_torch_compile(
34 + # set dynamic_arg_dims to support mrope
35 + dynamic_arg_dims={
36 + "input_ids": 0,
37 + "positions": -1,
38 + "intermediate_tensors": 0,
39 + "inputs_embeds": 0,
40 + }
41 +)
42 class Ernie4_5ForCausalLM(LlamaForCausalLM):
43 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
44 super().__init__(vllm_config=vllm_config, prefix=prefix)
vllm/model_executor/models/paddleocr_vl.py new
+1407
@@ -0,0 +1,1407 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
4 +#
5 +# Licensed under the Apache License, Version 2.0 (the "License");
6 +# you may not use this file except in compliance with the License.
7 +# You may obtain a copy of the License at
8 +#
9 +# http://www.apache.org/licenses/LICENSE-2.0
10 +#
11 +# Unless required by applicable law or agreed to in writing, software
12 +# distributed under the License is distributed on an "AS IS" BASIS,
13 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 +# See the License for the specific language governing permissions and
15 +# limitations under the License.
16 +
17 +import math
18 +from collections.abc import Iterable, Mapping, Sequence
19 +from functools import partial
20 +from typing import Annotated, Literal
21 +
22 +import numpy as np
23 +import torch
24 +import torch.nn as nn
25 +import torch.nn.functional as F
26 +from einops import rearrange, repeat
27 +from transformers import BatchFeature, PretrainedConfig
28 +from transformers.activations import GELUActivation
29 +from transformers.modeling_outputs import (
30 + BaseModelOutputWithPooling,
31 +)
32 +from transformers.utils import torch_int
33 +
34 +from vllm.attention.backends.registry import _Backend
35 +from vllm.attention.layer import (
36 + check_upstream_fa_availability,
37 + maybe_get_vit_flash_attn_backend,
38 +)
39 +from vllm.attention.ops.vit_attn_wrappers import (
40 + vit_flash_attn_wrapper,
41 + vit_xformers_attn_wrapper,
42 +)
43 +from vllm.config import VllmConfig
44 +from vllm.config.multimodal import BaseDummyOptions
45 +from vllm.distributed import parallel_state
46 +from vllm.distributed import utils as dist_utils
47 +from vllm.model_executor.layers.activation import get_act_fn
48 +from vllm.model_executor.layers.linear import (
49 + ColumnParallelLinear,
50 + QKVParallelLinear,
51 + RowParallelLinear,
52 +)
53 +from vllm.model_executor.layers.quantization import QuantizationConfig
54 +from vllm.model_executor.layers.rotary_embedding.common import (
55 + dispatch_rotary_emb_function,
56 +)
57 +from vllm.model_executor.model_loader.weight_utils import (
58 + default_weight_loader,
59 + maybe_remap_kv_scale_name,
60 +)
61 +from vllm.multimodal import MULTIMODAL_REGISTRY
62 +from vllm.multimodal.inputs import (
63 + MultiModalDataDict,
64 + MultiModalFieldConfig,
65 + MultiModalKwargs,
66 +)
67 +from vllm.multimodal.parse import (
68 + ImageProcessorItems,
69 + ImageSize,
70 + MultiModalDataItems,
71 +)
72 +from vllm.multimodal.processing import (
73 + BaseMultiModalProcessor,
74 + BaseProcessingInfo,
75 + PromptReplacement,
76 + PromptUpdate,
77 +)
78 +from vllm.multimodal.profiling import BaseDummyInputsBuilder
79 +from vllm.sequence import IntermediateTensors
80 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
81 +
82 +from .ernie45 import Ernie4_5ForCausalLM
83 +from .interfaces import MultiModalEmbeddings, SupportsMRoPE, SupportsMultiModal
84 +from .utils import (
85 + AutoWeightsLoader,
86 + PPMissingLayer,
87 + WeightsMapper,
88 + is_pp_missing_parameter,
89 + maybe_prefix,
90 +)
91 +from .vision import get_vit_attn_backend
92 +
93 +
94 +def smart_resize(
95 + height: int,
96 + width: int,
97 + factor: int = 28,
98 + min_pixels: int = 28 * 28 * 130,
99 + max_pixels: int = 28 * 28 * 1280,
100 +):
101 + """Rescales the image so that the following conditions are met:
102 +
103 + 1. Both dimensions (height and width) are divisible by 'factor'.
104 +
105 + 2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
106 +
107 + 3. The aspect ratio of the image is maintained as closely as possible.
108 +
109 + """
110 +
111 + if height < factor:
112 + width = round((width * factor) / height)
113 + height = factor
114 +
115 + if width < factor:
116 + height = round((height * factor) / width)
117 + width = factor
118 +
119 + if max(height, width) / min(height, width) > 200:
120 + raise ValueError(
121 + f"absolute aspect ratio must be smaller than 200, "
122 + f"got {max(height, width) / min(height, width)}"
123 + )
124 + h_bar = round(height / factor) * factor
125 + w_bar = round(width / factor) * factor
126 + if h_bar * w_bar > max_pixels:
127 + beta = math.sqrt((height * width) / max_pixels)
128 + h_bar = math.floor(height / beta / factor) * factor
129 + w_bar = math.floor(width / beta / factor) * factor
130 + elif h_bar * w_bar < min_pixels:
131 + beta = math.sqrt(min_pixels / (height * width))
132 + h_bar = math.ceil(height * beta / factor) * factor
133 + w_bar = math.ceil(width * beta / factor) * factor
134 + return h_bar, w_bar
135 +
136 +
137 +def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:
138 + if not interleaved:
139 + x1, x2 = x.chunk(2, dim=-1)
140 + return torch.cat((-x2, x1), dim=-1)
141 + x1, x2 = x[..., ::2], x[..., 1::2]
142 + return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2)
143 +
144 +
145 +def apply_rotary_emb_torch(
146 + x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
147 +) -> torch.Tensor:
148 + """
149 + x: (batch_size, seqlen, nheads, headdim)
150 + cos, sin: (seqlen, rotary_dim / 2) or (batch_size, seqlen, rotary_dim / 2)
151 + """
152 + ro_dim = cos.shape[-1] * 2
153 + assert ro_dim <= x.shape[-1]
154 + cos = repeat(
155 + cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)"
156 + )
157 + sin = repeat(
158 + sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)"
159 + )
160 + return torch.cat(
161 + [
162 + x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin,
163 + x[..., ro_dim:],
164 + ],
165 + dim=-1,
166 + )
167 +
168 +
169 +def apply_rotary_pos_emb_vision(t: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
170 + rotary_emb_function = dispatch_rotary_emb_function(default=apply_rotary_emb_torch)
171 + t_ = t.float()
172 + cos = freqs.cos()
173 + sin = freqs.sin()
174 + output = rotary_emb_function(t_, cos, sin).type_as(t)
175 + return output
176 +
177 +
178 +class PaddleOCRVLProcessingInfo(BaseProcessingInfo):
179 + def get_hf_config(self):
180 + return self.ctx.get_hf_config()
181 +
182 + def get_hf_processor(self, **kwargs: object):
183 + return self.ctx.get_hf_processor(**kwargs)
184 +
185 + def get_image_processor(self, **kwargs: object):
186 + return self.get_hf_processor(**kwargs).image_processor
187 +
188 + def get_supported_mm_limits(self):
189 + return {"image": None}
190 +
191 + def get_num_image_tokens(
192 + self,
193 + *,
194 + image_width: int,
195 + image_height: int,
196 + image_processor,
197 + ) -> int:
198 + if image_processor is None:
199 + image_processor = self.get_image_processor()
200 +
201 + do_resize = True
202 + hf_config = self.get_hf_config()
203 + vision_config = hf_config.vision_config
204 + patch_size = vision_config.patch_size
205 + merge_size = vision_config.spatial_merge_size
206 +
207 + if do_resize:
208 + resized_height, resized_width = smart_resize(
209 + height=image_height,
210 + width=image_width,
211 + factor=patch_size * merge_size,
212 + min_pixels=image_processor.min_pixels,
213 + max_pixels=image_processor.max_pixels,
214 + )
215 + preprocessed_size = ImageSize(width=resized_width, height=resized_height)
216 + else:
217 + preprocessed_size = ImageSize(width=image_width, height=image_height)
218 +
219 + grid_t = 1
220 + grid_h = preprocessed_size.height // patch_size
221 + grid_w = preprocessed_size.width // patch_size
222 +
223 + num_patches = grid_t * grid_h * grid_w
224 + num_image_tokens = num_patches // (merge_size**2)
225 +
226 + return num_image_tokens
227 +
228 + def get_image_size_with_most_features(self) -> ImageSize:
229 + hf_config = self.get_hf_config()
230 + image_size = hf_config.vision_config.image_size
231 + return ImageSize(height=image_size, width=image_size)
232 +
233 +
234 +class PaddleOCRVLDummyInputsBuilder(BaseDummyInputsBuilder[PaddleOCRVLProcessingInfo]):
235 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
236 + num_images = mm_counts.get("image", 0)
237 +
238 + processor = self.info.get_hf_processor()
239 + image_token = processor.image_token
240 +
241 + return image_token * num_images
242 +
243 + def get_dummy_mm_data(
244 + self,
245 + seq_len: int,
246 + mm_counts: Mapping[str, int],
247 + mm_options: Mapping[str, BaseDummyOptions] | None = None,
248 + ) -> MultiModalDataDict:
249 + num_images = mm_counts.get("image", 0)
250 +
251 + max_image_size = self.info.get_image_size_with_most_features()
252 + image_overrides = mm_options.get("image") if mm_options else None
253 +
254 + return {
255 + "image": self._get_dummy_images(
256 + width=max_image_size.width,
257 + height=max_image_size.height,
258 + num_images=num_images,
259 + overrides=image_overrides,
260 + )
261 + }
262 +
263 +
264 +class PaddleOCRVLMultiModalProcessor(
265 + BaseMultiModalProcessor[PaddleOCRVLProcessingInfo]
266 +):
267 + def _call_hf_processor(
268 + self,
269 + prompt: str,
270 + mm_data: Mapping[str, object],
271 + mm_kwargs: Mapping[str, object],
272 + tok_kwargs: Mapping[str, object],
273 + ) -> BatchFeature:
274 + if mm_data:
275 + processed_outputs = self.info.ctx.call_hf_processor(
276 + self.info.get_hf_processor(**mm_kwargs),
277 + dict(text=prompt, **mm_data),
278 + dict(**mm_kwargs, **tok_kwargs),
279 + )
280 + num_patches_per_image = processed_outputs["image_grid_thw"].prod(-1)
281 + processed_outputs["pixel_values"] = processed_outputs["pixel_values"].split(
282 + num_patches_per_image.tolist()
283 + )
284 + else:
285 + tokenizer = self.info.get_tokenizer()
286 + processed_outputs = tokenizer(
287 + prompt, add_special_tokens=True, return_tensors="pt"
288 + )
289 + return processed_outputs
290 +
291 + def _get_mm_fields_config(
292 + self,
293 + hf_inputs: BatchFeature,
294 + hf_processor_mm_kwargs: Mapping[str, object],
295 + ) -> Mapping[str, MultiModalFieldConfig]:
296 + return dict(
297 + pixel_values=MultiModalFieldConfig.batched("image"),
298 + image_grid_thw=MultiModalFieldConfig.batched("image"),
299 + )
300 +
301 + def _get_prompt_updates(
302 + self,
303 + mm_items: MultiModalDataItems,
304 + hf_processor_mm_kwargs: Mapping[str, object],
305 + out_mm_kwargs: MultiModalKwargs,
306 + ) -> Sequence[PromptUpdate]:
307 + image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs)
308 + hf_config = self.info.get_hf_config()
309 + image_token_id = hf_config.image_token_id
310 +
311 + def get_replacement(item_idx: int, image_processor):
312 + images = mm_items.get_items("image", ImageProcessorItems)
313 +
314 + image_size = images.get_image_size(item_idx)
315 + num_image_tokens = self.info.get_num_image_tokens(
316 + image_width=image_size.width,
317 + image_height=image_size.height,
318 + image_processor=image_processor,
319 + )
320 +
321 + return [image_token_id] * num_image_tokens
322 +
323 + return [
324 + PromptReplacement(
325 + modality="image",
326 + target=[image_token_id],
327 + replacement=partial(get_replacement, image_processor=image_processor),
328 + ),
329 + ]
330 +
331 +
332 +class Projector(nn.Module):
333 + def __init__(
334 + self,
335 + text_config: PretrainedConfig,
336 + vision_config: PretrainedConfig,
337 + prefix: str = "",
338 + ):
339 + super().__init__()
340 + self.text_config = text_config
341 + self.vision_config = vision_config
342 + self.merge_kernel_size = (2, 2)
343 +
344 + self.hidden_size = (
345 + self.vision_config.hidden_size
346 + * self.merge_kernel_size[0]
347 + * self.merge_kernel_size[1]
348 + )
349 +
350 + self.pre_norm = torch.nn.LayerNorm(self.vision_config.hidden_size, eps=1e-05)
351 + self.linear_1 = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
352 + self.act = GELUActivation()
353 + self.linear_2 = nn.Linear(
354 + self.hidden_size, self.text_config.hidden_size, bias=True
355 + )
356 +
357 + def forward(
358 + self,
359 + image_features: torch.Tensor,
360 + image_grid_thw: torch.Tensor,
361 + ) -> torch.Tensor:
362 + m1, m2 = self.merge_kernel_size
363 + if isinstance(image_features, (list, tuple)):
364 + processed_features = list()
365 + for image_feature, image_grid in zip(image_features, image_grid_thw):
366 + image_feature = self.pre_norm(image_feature)
367 + t, h, w = image_grid
368 +
369 + image_feature = rearrange(
370 + image_feature,
371 + "(t h p1 w p2) d -> (t h w) (p1 p2 d)",
372 + t=t,
373 + h=h // m1,
374 + p1=m1,
375 + w=w // m2,
376 + p2=m2,
377 + )
378 + hidden_states = self.linear_1(image_feature)
379 + hidden_states = self.act(hidden_states)
380 + hidden_states = self.linear_2(hidden_states)
381 + processed_features.append(hidden_states)
382 +
383 + return processed_features
384 +
385 + dims = image_features.shape[:-1]
386 + dim = image_features.shape[-1]
387 + image_features = image_features.view(np.prod(dims), dim)
388 + hidden_states = self.pre_norm(image_features).view(-1, self.hidden_size)
389 + hidden_states = self.linear_1(hidden_states)
390 + hidden_states = self.act(hidden_states)
391 + hidden_states = self.linear_2(hidden_states)
392 +
393 + return hidden_states.view(*dims, -1)
394 +
395 +
396 +class PaddleOCRImagePixelInputs(TensorSchema):
397 + type: Literal["pixel_values"]
398 + pixel_values: Annotated[
399 + torch.Tensor,
400 + TensorShape("bn", "p", 3, "patch_size", "patch_size", dynamic_dims={"p"}),
401 + ]
402 + image_grid_thw: Annotated[
403 + torch.Tensor,
404 + TensorShape("bn", 3),
405 + ]
406 +
407 +
408 +class SiglipVisionEmbeddings(nn.Module):
409 + def __init__(self, config: PretrainedConfig):
410 + super().__init__()
411 + self.config = config
412 + self.embed_dim = config.hidden_size
413 + self.image_size = config.image_size
414 + self.patch_size = config.patch_size
415 +
416 + self.patch_embedding = nn.Conv2d(
417 + in_channels=config.num_channels,
418 + out_channels=self.embed_dim,
419 + kernel_size=self.patch_size,
420 + stride=self.patch_size,
421 + padding="valid",
422 + )
423 +
424 + self.num_patches = (self.image_size // self.patch_size) ** 2
425 + self.num_positions = self.num_patches
426 + self.cache_position_embedding = dict()
427 + self.cache_position_count = dict()
428 + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
429 + self.packing_position_embedding = nn.Embedding(32768, self.embed_dim)
430 +
431 + self.register_buffer(
432 + "position_ids",
433 + torch.arange(self.num_positions).expand((1, -1)),
434 + persistent=False,
435 + )
436 +
437 + def interpolate_pos_encoding(
438 + self,
439 + embeddings: torch.Tensor,
440 + height: int,
441 + width: int,
442 + is_after_patchify: bool = False,
443 + ) -> torch.Tensor:
444 + num_positions = self.position_embedding.weight.shape[0]
445 +
446 + patch_pos_embed = self.position_embedding.weight.unsqueeze(0)
447 +
448 + dim = embeddings.shape[-1]
449 +
450 + if is_after_patchify:
451 + new_height = height
452 + new_width = width
453 + else:
454 + new_height = height // self.patch_size
455 + new_width = width // self.patch_size
456 +
457 + sqrt_num_positions = torch_int(num_positions**0.5)
458 + patch_pos_embed = patch_pos_embed.reshape(
459 + 1, sqrt_num_positions, sqrt_num_positions, dim
460 + )
461 + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
462 +
463 + patch_pos_embed = nn.functional.interpolate(
464 + patch_pos_embed,
465 + size=(new_height, new_width),
466 + mode="bilinear",
467 + align_corners=False,
468 + )
469 +
470 + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
471 + return patch_pos_embed
472 +
473 + def fetch_position_embedding_lfu_cache(
474 + self, embeddings: torch.Tensor, h: int, w: int, max_cache: int = 20
475 + ):
476 + grid = (h, w)
477 + if grid in self.cache_position_embedding:
478 + self.cache_position_count[grid] += 1
479 + return self.cache_position_embedding[grid]
480 +
481 + if len(self.cache_position_embedding) >= max_cache:
482 + min_hit_grid = min(
483 + self.cache_position_count,
484 + key=self.cache_position_count.get,
485 + )
486 + self.cache_position_count.pop(min_hit_grid)
487 + self.cache_position_embedding.pop(min_hit_grid)
488 +
489 + position_embedding = self.interpolate_pos_encoding(embeddings, h, w, True)
490 + self.cache_position_count[grid] = 1
491 + self.cache_position_embedding[grid] = position_embedding
492 + return position_embedding
493 +
494 + def forward(
495 + self,
496 + pixel_values: torch.FloatTensor,
497 + position_ids: torch.Tensor | None = None,
498 + image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]]
499 + | None = None,
500 + interpolate_pos_encoding=False,
501 + ) -> torch.Tensor:
502 + if pixel_values.dim() == 4:
503 + pixel_values = pixel_values.unsqueeze(0)
504 + if pixel_values.dim() == 5:
505 + if position_ids is None:
506 + raise ValueError(
507 + "position_ids cannot be None when pixel_values.dim() is 5."
508 + )
509 + (
510 + batch_size,
511 + squence_len,
512 + channel,
513 + height,
514 + width,
515 + ) = pixel_values.shape
516 + target_dtype = self.patch_embedding.weight.dtype
517 + pixel_values = rearrange(pixel_values, "b l c h w -> (b l) c h w")
518 + patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
519 + embeddings = patch_embeds.flatten(-2).squeeze(-1)
520 +
521 + if interpolate_pos_encoding and image_grid_thw is not None:
522 + start = 0
523 + tmp_embeddings = list()
524 + for image_grid in image_grid_thw:
525 + t, h, w = image_grid
526 + end = start + t * h * w
527 + image_embeddings = embeddings[start:end, :]
528 + position_embedding = (
529 + self.interpolate_pos_encoding(image_embeddings, h, w, True)
530 + .squeeze(0)
531 + .repeat(t, 1)
532 + )
533 + image_embeddings = image_embeddings + position_embedding
534 + tmp_embeddings.append(image_embeddings)
535 + start = end
536 + embeddings = torch.concat(tmp_embeddings, dim=0).unsqueeze(0)
537 + else:
538 + embeddings = embeddings + self.packing_position_embedding(position_ids)
539 + return embeddings
540 + else:
541 + raise ValueError(
542 + "Unsupported pixel_values dimension:"
543 + f" {pixel_values.dim()}. Expected 4 or 5."
544 + )
545 +
546 +
547 +def all_gather_interleave(local_tensor: torch.Tensor, hidden_size: int, tp_size: int):
548 + """All-gather the input tensor interleavely across model parallel group."""
549 + import torch.distributed as dist
550 +
551 + gathered_tensors = [torch.zeros_like(local_tensor) for _ in range(tp_size)]
552 + dist.all_gather(
553 + gathered_tensors, local_tensor, group=parallel_state.get_tp_group().device_group
554 + )
555 +
556 + gathered_tensors_split = [
557 + torch.split(tensor, hidden_size // tp_size, -1) for tensor in gathered_tensors
558 + ]
559 + ordered_tensors = [
560 + tensor for pair in zip(*gathered_tensors_split) for tensor in pair
561 + ]
562 + result_tensor = torch.cat(ordered_tensors, dim=-1)
563 + return result_tensor
564 +
565 +
566 +class SiglipAttention(nn.Module):
567 + """SigLIP vision attention adapted from Qwen2.5-VisionAttention."""
568 +
569 + def __init__(
570 + self,
571 + *,
572 + embed_dim: int,
573 + num_heads: int,
574 + projection_size: int,
575 + quant_config: QuantizationConfig | None = None,
576 + prefix: str = "",
577 + attn_backend: _Backend = _Backend.TORCH_SDPA,
578 + attn_backend_override: _Backend | None = None,
579 + use_upstream_fa: bool = False,
580 + ) -> None:
581 + super().__init__()
582 +
583 + self.tp_size = parallel_state.get_tensor_model_parallel_world_size()
584 + self.tp_rank = parallel_state.get_tensor_model_parallel_rank()
585 + self.hidden_size_per_attention_head = dist_utils.divide(
586 + projection_size, num_heads
587 + )
588 + self.num_attention_heads_per_partition = dist_utils.divide(
589 + num_heads, self.tp_size
590 + )
591 +
592 + self.qkv_proj = QKVParallelLinear(
593 + hidden_size=embed_dim,
594 + head_size=self.hidden_size_per_attention_head,
595 + total_num_heads=num_heads,
596 + total_num_kv_heads=num_heads,
597 + bias=True,
598 + quant_config=quant_config,
599 + prefix=f"{prefix}.qkv_proj",
600 + )
601 + self.out_proj = RowParallelLinear(
602 + input_size=projection_size,
603 + output_size=embed_dim,
604 + quant_config=quant_config,
605 + prefix=f"{prefix}.out_proj",
606 + )
607 +
608 + self.attn_backend = attn_backend
609 + self.use_upstream_fa = use_upstream_fa
610 + self.attn_backend, self.flash_attn_varlen_func = (
611 + maybe_get_vit_flash_attn_backend(
612 + self.attn_backend,
613 + self.use_upstream_fa,
614 + attn_backend_override=attn_backend_override,
615 + )
616 + )
617 + self.is_flash_attn_backend = self.attn_backend in {
618 + _Backend.FLASH_ATTN,
619 + _Backend.ROCM_AITER_FA,
620 + }
621 +
622 + def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
623 + seq_len, bs, _ = qkv.shape
624 + if self.tp_size > 1:
625 + qkv = all_gather_interleave(qkv, self.qkv_proj.hidden_size, self.tp_size)
626 +
627 + q, k, v = qkv.chunk(3, dim=2)
628 +
629 + if self.tp_size > 1:
630 + splitter = partial(
631 + dist_utils.split_tensor_along_last_dim, num_partitions=self.tp_size
632 + )
633 + q = splitter(q)[self.tp_rank]
634 + k = splitter(k)[self.tp_rank]
635 + v = splitter(v)[self.tp_rank]
636 +
637 + new_shape = (
638 + seq_len,
639 + bs,
640 + self.num_attention_heads_per_partition,
641 + self.hidden_size_per_attention_head,
642 + )
643 + q, k, v = (x.view(*new_shape) for x in (q, k, v))
644 + return q, k, v
645 +
646 + def forward(
647 + self,
648 + hidden_states: torch.Tensor,
649 + *,
650 + cu_seqlens: torch.Tensor,
651 + rotary_pos_emb: torch.Tensor | None,
652 + max_seqlen: torch.Tensor | None,
653 + seqlens: torch.Tensor | None,
654 + ) -> torch.Tensor:
655 + batch_size, _, _ = hidden_states.shape
656 +
657 + x = rearrange(hidden_states, "b s d -> s b d")
658 + x, _ = self.qkv_proj(x)
659 + q, k, v = self.split_qkv(x)
660 + q, k, v = (rearrange(t, "s b h d -> b s h d") for t in (q, k, v))
661 +
662 + if rotary_pos_emb is not None:
663 + qk_concat = torch.cat([q, k], dim=0)
664 + qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
665 + q, k = torch.chunk(qk_rotated, 2, dim=0)
666 +
667 + if self.is_flash_attn_backend:
668 + if max_seqlen is None:
669 + raise ValueError("Flash attention backend requires max_seqlen.")
670 + context_layer = vit_flash_attn_wrapper(
671 + q,
672 + k,
673 + v,
674 + cu_seqlens,
675 + max_seqlen,
676 + batch_size,
677 + self.attn_backend == _Backend.ROCM_AITER_FA,
678 + self.use_upstream_fa,
679 + )
680 + elif self.attn_backend == _Backend.TORCH_SDPA:
681 + outputs = []
682 + for i in range(1, len(cu_seqlens)):
683 + start_idx = cu_seqlens[i - 1]
684 + end_idx = cu_seqlens[i]
685 + q_i = q[:, start_idx:end_idx]
686 + k_i = k[:, start_idx:end_idx]
687 + v_i = v[:, start_idx:end_idx]
688 + q_i, k_i, v_i = (
689 + rearrange(tensor, "b s h d -> b h s d")
690 + for tensor in (q_i, k_i, v_i)
691 + )
692 + output_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
693 + output_i = rearrange(output_i, "b h s d -> b s h d")
694 + outputs.append(output_i)
695 + context_layer = torch.cat(outputs, dim=1)
696 + context_layer = rearrange(
697 + context_layer, "b s h d -> s b (h d)"
698 + ).contiguous()
699 + elif self.attn_backend == _Backend.XFORMERS:
700 + if seqlens is None:
701 + raise ValueError("xFormers attention backend requires seqlens tensor.")
702 + context_layer = vit_xformers_attn_wrapper(q, k, v, seqlens)
703 + else:
704 + raise RuntimeError(
705 + f"PaddleOCR-VL does not support {self.attn_backend} backend now."
706 + )
707 +
708 + output, _ = self.out_proj(context_layer)
709 + output = rearrange(output, "s b d -> b s d")
710 + return output
711 +
712 +
713 +class SigLIPRotaryEmbedding(nn.Module):
714 + def __init__(self, dim: int, theta: float = 10000.0) -> None:
715 + super().__init__()
716 + self.dim = dim
717 + self.theta = theta
718 + self.rope_init()
719 +
720 + def rope_init(self):
721 + inv_freq = 1.0 / (
722 + self.theta ** (torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim)
723 + )
724 + self.register_buffer("inv_freq", inv_freq, persistent=False)
725 +
726 + def forward(self, seqlen: int) -> torch.Tensor:
727 + seq = torch.arange(
728 + seqlen,
729 + device=self.inv_freq.device,
730 + dtype=self.inv_freq.dtype,
731 + )
732 + freqs = torch.outer(seq, self.inv_freq)
733 + return freqs
734 +
735 +
736 +class SiglipMLP(nn.Module):
737 + def __init__(
738 + self,
739 + config: PretrainedConfig,
740 + quant_config: QuantizationConfig | None = None,
741 + prefix: str = "",
742 + ) -> None:
743 + super().__init__()
744 +
745 + self.config = config
746 + self.activation_fn = get_act_fn(config.hidden_act)
747 + # Special handling for BNB and torchao quantization
748 + if quant_config and quant_config.get_name() in ["bitsandbytes", "torchao"]:
749 + quantizable = True
750 + else:
751 + # For other quantization, we require the hidden size to be a
752 + # multiple of 64
753 + quantizable = (
754 + config.hidden_size % 64 == 0 and config.intermediate_size % 64 == 0
755 + )
756 + self.fc1 = ColumnParallelLinear(
757 + config.hidden_size,
758 + config.intermediate_size,
759 + quant_config=quant_config if quantizable else None,
760 + prefix=f"{prefix}.fc1",
761 + )
762 + self.fc2 = RowParallelLinear(
763 + config.intermediate_size,
764 + config.hidden_size,
765 + quant_config=quant_config if quantizable else None,
766 + prefix=f"{prefix}.fc2",
767 + )
768 +
769 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
770 + hidden_states, _ = self.fc1(hidden_states)
771 + hidden_states = self.activation_fn(hidden_states)
772 + hidden_states, _ = self.fc2(hidden_states)
773 + return hidden_states
774 +
775 +
776 +class SiglipEncoderLayer(nn.Module):
777 + def __init__(
778 + self,
779 + config: PretrainedConfig,
780 + quant_config: QuantizationConfig | None = None,
781 + prefix: str = "",
782 + *,
783 + attn_backend: _Backend = _Backend.TORCH_SDPA,
784 + attn_backend_override: _Backend | None = None,
785 + use_upstream_fa: bool = False,
786 + ):
787 + super().__init__()
788 + self.embed_dim = config.hidden_size
789 + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
790 + self.self_attn = SiglipAttention(
791 + embed_dim=config.hidden_size,
792 + num_heads=config.num_attention_heads,
793 + projection_size=config.hidden_size,
794 + quant_config=quant_config,
795 + prefix=f"{prefix}.self_attn",
796 + attn_backend=attn_backend,
797 + attn_backend_override=attn_backend_override,
798 + use_upstream_fa=use_upstream_fa,
799 + )
800 + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
801 + self.mlp = SiglipMLP(
802 + config,
803 + quant_config=quant_config,
804 + prefix=f"{prefix}.mlp",
805 + )
806 +
807 + def forward(
808 + self,
809 + hidden_states: torch.Tensor,
810 + *,
811 + cu_seqlens: torch.Tensor,
812 + rotary_pos_emb: torch.Tensor | None,
813 + max_seqlen: torch.Tensor | None,
814 + seqlens: torch.Tensor | None,
815 + ) -> torch.Tensor:
816 + residual = hidden_states
817 +
818 + hidden_states = self.layer_norm1(hidden_states)
819 + hidden_states = self.self_attn(
820 + hidden_states=hidden_states,
821 + cu_seqlens=cu_seqlens,
822 + rotary_pos_emb=rotary_pos_emb,
823 + max_seqlen=max_seqlen,
824 + seqlens=seqlens,
825 + )
826 +
827 + hidden_states = residual + hidden_states
828 +
829 + residual = hidden_states
830 + hidden_states = self.layer_norm2(hidden_states)
831 + hidden_states = self.mlp(hidden_states)
832 +
833 + hidden_states = residual + hidden_states
834 +
835 + return hidden_states
836 +
837 +
838 +class SiglipEncoder(nn.Module):
839 + def __init__(
840 + self,
841 + config: PretrainedConfig,
842 + quant_config: QuantizationConfig | None = None,
843 + prefix: str = "",
844 + attn_backend_override: _Backend | None = None,
845 + ):
846 + super().__init__()
847 + self.config = config
848 + embed_dim = config.hidden_size
849 + num_heads = config.num_attention_heads
850 + head_dim = embed_dim // num_heads
851 + self.attn_backend = get_vit_attn_backend(
852 + head_size=head_dim,
853 + dtype=torch.get_default_dtype(),
854 + attn_backend_override=attn_backend_override,
855 + )
856 + self.use_upstream_fa = False
857 + if self.attn_backend not in {
858 + _Backend.FLASH_ATTN,
859 + _Backend.ROCM_AITER_FA,
860 + } and check_upstream_fa_availability(torch.get_default_dtype()):
861 + self.attn_backend = _Backend.FLASH_ATTN
862 + self.use_upstream_fa = True
863 + if self.attn_backend not in {
864 + _Backend.FLASH_ATTN,
865 + _Backend.TORCH_SDPA,
866 + _Backend.XFORMERS,
867 + _Backend.ROCM_AITER_FA,
868 + }:
869 + raise RuntimeError(
870 + f"PaddleOCR-VL does not support {self.attn_backend} backend now."
871 + )
872 + self.layers = nn.ModuleList(
873 + [
874 + SiglipEncoderLayer(
875 + config,
876 + quant_config=quant_config,
877 + prefix=f"{prefix}.layers.{layer_idx}",
878 + attn_backend=self.attn_backend,
879 + attn_backend_override=attn_backend_override,
880 + use_upstream_fa=self.use_upstream_fa,
881 + )
882 + for layer_idx in range(config.num_hidden_layers)
883 + ]
884 + )
885 + self.rotary_pos_emb = SigLIPRotaryEmbedding(head_dim // 2)
886 +
887 + @staticmethod
888 + def flatten_list(image_grid_thw):
889 + tmp_image_grid_thw = list()
890 + for image_grid in image_grid_thw:
891 + if isinstance(image_grid, list):
892 + tmp_image_grid_thw.extend(image_grid)
893 + else:
894 + tmp_image_grid_thw.append(image_grid)
895 + return tmp_image_grid_thw
896 +
897 + def forward(
898 + self,
899 + inputs_embeds,
900 + cu_seqlens: torch.Tensor | None = None,
901 + image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]]
902 + | None = None,
903 + height_position_ids: torch.Tensor | None = None,
904 + width_position_ids: torch.Tensor | None = None,
905 + ) -> torch.Tensor:
906 + device = inputs_embeds.device
907 + hidden_states = inputs_embeds
908 +
909 + flatten_image_grid_thw = self.flatten_list(image_grid_thw)
910 +
911 + if width_position_ids is None or height_position_ids is None:
912 + split_hids = list()
913 + split_wids = list()
914 + for t, h, w in flatten_image_grid_thw:
915 + image_pids = torch.arange(t * h * w, device=device) % (h * w)
916 + sample_hids = image_pids // w
917 + sample_wids = image_pids % w
918 + split_hids.append(sample_hids)
919 + split_wids.append(sample_wids)
920 + width_position_ids = torch.concat(split_wids, dim=0)
921 + height_position_ids = torch.concat(split_hids, dim=0)
922 +
923 + pids = torch.stack(
924 + [height_position_ids, width_position_ids],
925 + dim=-1,
926 + )
927 + max_grid_size = pids.max() + 1
928 + rope_emb_max_grid = self.rotary_pos_emb(max_grid_size)
929 + rotary_pos_emb = rope_emb_max_grid[pids].flatten(1)
930 +
931 + if cu_seqlens is None:
932 + raise ValueError("cu_seqlens cannot be None for SiglipEncoder.")
933 + if not isinstance(cu_seqlens, torch.Tensor):
934 + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
935 + else:
936 + cu_seqlens = cu_seqlens.to(device=device)
937 +
938 + max_seqlen = None
939 + seqlens = None
940 + if self.attn_backend in {_Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA}:
941 + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
942 + elif self.attn_backend == _Backend.XFORMERS:
943 + seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
944 +
945 + hidden_states = inputs_embeds
946 + for encoder_layer in self.layers:
947 + hidden_states = encoder_layer(
948 + hidden_states,
949 + cu_seqlens=cu_seqlens,
950 + rotary_pos_emb=rotary_pos_emb,
951 + max_seqlen=max_seqlen,
952 + seqlens=seqlens,
953 + )
954 + return hidden_states
955 +
956 +
957 +class SiglipVisionTransformer(nn.Module):
958 + def __init__(
959 + self,
960 + config: PretrainedConfig,
961 + quant_config: QuantizationConfig | None = None,
962 + prefix: str = "",
963 + attn_backend_override: _Backend | None = None,
964 + ):
965 + super().__init__()
966 + self.config = config
967 + embed_dim = config.hidden_size
968 +
969 + self.embeddings = SiglipVisionEmbeddings(config)
970 + self.encoder = SiglipEncoder(
971 + config,
972 + quant_config=quant_config,
973 + prefix=f"{prefix}.encoder",
974 + attn_backend_override=attn_backend_override,
975 + )
976 + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
977 +
978 + def forward(
979 + self,
980 + pixel_values: torch.Tensor,
981 + interpolate_pos_encoding: bool | None = False,
982 + position_ids: torch.Tensor | None = None,
983 + height_position_ids: torch.Tensor | None = None,
984 + width_position_ids: torch.Tensor | None = None,
985 + cu_seqlens: torch.Tensor | None = None,
986 + image_grid_thw: torch.Tensor | None = None,
987 + ) -> torch.Tensor:
988 + hidden_states = self.embeddings(
989 + pixel_values,
990 + interpolate_pos_encoding=interpolate_pos_encoding,
991 + position_ids=position_ids,
992 + image_grid_thw=image_grid_thw,
993 + )
994 +
995 + last_hidden_state = self.encoder(
996 + inputs_embeds=hidden_states,
997 + cu_seqlens=cu_seqlens,
998 + image_grid_thw=image_grid_thw,
999 + height_position_ids=height_position_ids,
1000 + width_position_ids=width_position_ids,
1001 + )
1002 +
1003 + last_hidden_state = self.post_layernorm(last_hidden_state)
1004 + return last_hidden_state
1005 +
1006 +
1007 +class SiglipVisionModel(nn.Module):
1008 + def __init__(
1009 + self,
1010 + config,
1011 + quant_config: QuantizationConfig | None = None,
1012 + prefix: str = "",
1013 + attn_backend_override: _Backend | None = None,
1014 + ):
1015 + super().__init__()
1016 +
1017 + self.vision_model = SiglipVisionTransformer(
1018 + config,
1019 + quant_config=quant_config,
1020 + prefix=f"{prefix}.vision_model",
1021 + attn_backend_override=attn_backend_override,
1022 + )
1023 + self.quant_config = quant_config
1024 +
1025 + @property
1026 + def dtype(self) -> torch.dtype:
1027 + return self.vision_model.embeddings.patch_embedding.weight.dtype
1028 +
1029 + @property
1030 + def device(self) -> torch.device:
1031 + return self.vision_model.embeddings.patch_embedding.weight.device
1032 +
1033 + def get_input_embeddings(self) -> nn.Module:
1034 + return self.vision_model.embeddings.patch_embedding
1035 +
1036 + def forward(
1037 + self,
1038 + pixel_values,
1039 + interpolate_pos_encoding: bool = False,
1040 + position_ids: torch.Tensor | None = None,
1041 + image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]]
1042 + | None = None,
1043 + cu_seqlens: torch.Tensor | None = None,
1044 + ) -> BaseModelOutputWithPooling:
1045 + return self.vision_model(
1046 + pixel_values=pixel_values,
1047 + interpolate_pos_encoding=interpolate_pos_encoding,
1048 + position_ids=position_ids,
1049 + image_grid_thw=image_grid_thw,
1050 + cu_seqlens=cu_seqlens,
1051 + )
1052 +
1053 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1054 + stacked_params_mapping = [
1055 + ("qkv_proj", "q_proj", "q"),
1056 + ("qkv_proj", "k_proj", "k"),
1057 + ("qkv_proj", "v_proj", "v"),
1058 + ]
1059 + params_dict = dict(self.named_parameters(remove_duplicate=False))
1060 + loaded_params: set[str] = set()
1061 + for name, loaded_weight in weights:
1062 + if "rotary_emb.inv_freq" in name:
1063 + continue
1064 + if "head.attention" in name or "head.layernorm" in name:
1065 + continue
1066 + if "head.mlp" in name or "head.probe" in name:
1067 + continue
1068 + if self.quant_config is not None and (
1069 + scale_name := self.quant_config.get_cache_scale(name)
1070 + ):
1071 + param = params_dict[scale_name]
1072 + weight_loader = getattr(
1073 + param,
1074 + "weight_loader",
1075 + default_weight_loader,
1076 + )
1077 + loaded_weight = (
1078 + loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
1079 + )
1080 + weight_loader(param, loaded_weight)
1081 + loaded_params.add(scale_name)
1082 + continue
1083 + for (
1084 + param_name,
1085 + weight_name,
1086 + shard_id,
1087 + ) in stacked_params_mapping:
1088 + if weight_name not in name:
1089 + continue
1090 + name = name.replace(weight_name, param_name)
1091 + if name.endswith(".bias") and name not in params_dict:
1092 + continue
1093 + if is_pp_missing_parameter(name, self):
1094 + continue
1095 + param = params_dict[name]
1096 + weight_loader = param.weight_loader
1097 + weight_loader(param, loaded_weight, shard_id)
1098 + break
1099 + else:
1100 + if name.endswith(".bias") and name not in params_dict:
1101 + continue
1102 + name = maybe_remap_kv_scale_name(name, params_dict)
1103 + if name is None:
1104 + continue
1105 + if is_pp_missing_parameter(name, self):
1106 + continue
1107 + param = params_dict[name]
1108 + weight_loader = getattr(
1109 + param,
1110 + "weight_loader",
1111 + default_weight_loader,
1112 + )
1113 + weight_loader(param, loaded_weight)
1114 + loaded_params.add(name)
1115 + return loaded_params
1116 +
1117 +
1118 +@MULTIMODAL_REGISTRY.register_processor(
1119 + PaddleOCRVLMultiModalProcessor,
1120 + info=PaddleOCRVLProcessingInfo,
1121 + dummy_inputs=PaddleOCRVLDummyInputsBuilder,
1122 +)
1123 +class PaddleOCRVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsMRoPE):
1124 + merge_by_field_config = True
1125 +
1126 + hf_to_vllm_mapper = WeightsMapper(
1127 + orig_to_new_prefix={
1128 + "model.": "language_model.model.",
1129 + "lm_head.": "language_model.lm_head.",
1130 + }
1131 + )
1132 +
1133 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1134 + super().__init__()
1135 + config = vllm_config.model_config.hf_config
1136 + quant_config = vllm_config.quant_config
1137 + multimodal_config = vllm_config.model_config.multimodal_config
1138 +
1139 + self.config = config
1140 + self.multimodal_config = multimodal_config
1141 +
1142 + attn_backend_override = (
1143 + multimodal_config.mm_encoder_attn_backend
1144 + if multimodal_config is not None
1145 + else None
1146 + )
1147 +
1148 + self.visual = SiglipVisionModel(
1149 + config=config.vision_config,
1150 + quant_config=quant_config,
1151 + prefix=maybe_prefix(prefix, "visual"),
1152 + attn_backend_override=attn_backend_override,
1153 + )
1154 + self.mlp_AR = Projector(config, config.vision_config)
1155 +
1156 + self.language_model = Ernie4_5ForCausalLM(
1157 + vllm_config=vllm_config,
1158 + prefix=maybe_prefix(prefix, "language_model"),
1159 + )
1160 +
1161 + for layer in self.language_model.model.layers:
1162 + if not isinstance(layer, PPMissingLayer):
1163 + layer.self_attn.rotary_emb.is_neox_style = True
1164 +
1165 + self.make_empty_intermediate_tensors = (
1166 + self.language_model.make_empty_intermediate_tensors
1167 + )
1168 +
1169 + def compute_logits(
1170 + self,
1171 + hidden_states: torch.Tensor,
1172 + ) -> torch.Tensor | None:
1173 + return self.language_model.compute_logits(hidden_states)
1174 +
1175 + def get_mrope_input_positions(
1176 + self,
1177 + input_tokens: list[int],
1178 + hf_config: PretrainedConfig,
1179 + image_grid_thw: list[list[int]] | torch.Tensor,
1180 + video_grid_thw: list[list[int]] | torch.Tensor,
1181 + second_per_grid_ts: list[float],
1182 + context_len: int = 0,
1183 + seq_len: int | None = None,
1184 + audio_feature_lengths: torch.Tensor | None = None,
1185 + use_audio_in_video: bool = False,
1186 + ) -> tuple[torch.Tensor, int]:
1187 + """Get mrope input positions and delta value."""
1188 +
1189 + image_token_id = hf_config.image_token_id
1190 + video_token_id = hf_config.video_token_id
1191 + vision_start_token_id = hf_config.vision_start_token_id
1192 + spatial_merge_size = hf_config.vision_config.spatial_merge_size
1193 + tokens_per_second = getattr(hf_config.vision_config, "tokens_per_second", 1.0)
1194 +
1195 + input_tokens_tensor = torch.tensor(input_tokens)
1196 + vision_start_indices = torch.argwhere(
1197 + input_tokens_tensor == vision_start_token_id
1198 + ).squeeze(1)
1199 + vision_tokens = input_tokens_tensor[vision_start_indices + 1]
1200 + image_nums = (vision_tokens == image_token_id).sum()
1201 + video_nums = (vision_tokens == video_token_id).sum()
1202 + llm_pos_ids_list: list = []
1203 +
1204 + st = 0
1205 + remain_images, remain_videos = image_nums, video_nums
1206 +
1207 + image_index, video_index = 0, 0
1208 + for _ in range(image_nums + video_nums):
1209 + video_second_per_grid_t = 0.0
1210 + if remain_images > 0:
1211 + try:
1212 + ed_image = input_tokens.index(image_token_id, st)
1213 + except ValueError:
1214 + ed_image = len(input_tokens) + 1
1215 + else:
1216 + ed_image = len(input_tokens) + 1
1217 + if remain_videos > 0:
1218 + try:
1219 + ed_video = input_tokens.index(video_token_id, st)
1220 + except ValueError:
1221 + ed_video = len(input_tokens) + 1
1222 + else:
1223 + ed_video = len(input_tokens) + 1
1224 + if ed_image < ed_video:
1225 + t, h, w = (
1226 + image_grid_thw[image_index][0],
1227 + image_grid_thw[image_index][1],
1228 + image_grid_thw[image_index][2],
1229 + )
1230 + image_index += 1
1231 + remain_images -= 1
1232 + ed = ed_image
1233 + else:
1234 + t, h, w = (
1235 + video_grid_thw[video_index][0],
1236 + video_grid_thw[video_index][1],
1237 + video_grid_thw[video_index][2],
1238 + )
1239 + video_second_per_grid_t = 1.0
1240 + if second_per_grid_ts:
1241 + video_second_per_grid_t = second_per_grid_ts[video_index]
1242 + video_index += 1
1243 + remain_videos -= 1
1244 + ed = ed_video
1245 +
1246 + llm_grid_t, llm_grid_h, llm_grid_w = (
1247 + t,
1248 + h // spatial_merge_size,
1249 + w // spatial_merge_size,
1250 + )
1251 + text_len = ed - st
1252 +
1253 + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
1254 + llm_pos_ids_list.append(
1255 + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
1256 + )
1257 +
1258 + t_index = (
1259 + (
1260 + torch.arange(llm_grid_t)
1261 + .view(-1, 1)
1262 + .expand(-1, llm_grid_h * llm_grid_w)
1263 + * video_second_per_grid_t
1264 + * tokens_per_second
1265 + )
1266 + .long()
1267 + .flatten()
1268 + )
1269 +
1270 + h_index = (
1271 + torch.arange(llm_grid_h)
1272 + .view(1, -1, 1)
1273 + .expand(llm_grid_t, -1, llm_grid_w)
1274 + .flatten()
1275 + )
1276 + w_index = (
1277 + torch.arange(llm_grid_w)
1278 + .view(1, 1, -1)
1279 + .expand(llm_grid_t, llm_grid_h, -1)
1280 + .flatten()
1281 + )
1282 + llm_pos_ids_list.append(
1283 + torch.stack([t_index, h_index, w_index]) + text_len + st_idx
1284 + )
1285 + st = ed + llm_grid_t * llm_grid_h * llm_grid_w
1286 +
1287 + if st < len(input_tokens):
1288 + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
1289 + text_len = len(input_tokens) - st
1290 + llm_pos_ids_list.append(
1291 + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
1292 + )
1293 +
1294 + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
1295 + mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
1296 + llm_positions = llm_positions[:, context_len:seq_len]
1297 +
1298 + return llm_positions, mrope_position_delta
1299 +
1300 + def get_language_model(self) -> nn.Module:
1301 + return self.language_model
1302 +
1303 + def _parse_and_validate_image_input(
1304 + self, **kwargs: object
1305 + ) -> PaddleOCRImagePixelInputs | None:
1306 + pixel_values = kwargs.pop("pixel_values", None)
1307 + image_grid_thw = kwargs.pop("image_grid_thw", None)
1308 +
1309 + if pixel_values is None:
1310 + return None
1311 +
1312 + return PaddleOCRImagePixelInputs(
1313 + type="pixel_values",
1314 + pixel_values=pixel_values,
1315 + image_grid_thw=image_grid_thw,
1316 + )
1317 +
1318 + def forward(
1319 + self,
1320 + input_ids: torch.Tensor,
1321 + positions: torch.Tensor,
1322 + intermediate_tensors: IntermediateTensors | None = None,
1323 + inputs_embeds: torch.Tensor | None = None,
1324 + **kwargs,
1325 + ):
1326 + if intermediate_tensors is not None:
1327 + inputs_embeds = None
1328 +
1329 + elif inputs_embeds is None:
1330 + vision_embeddings = self.get_multimodal_embeddings(**kwargs)
1331 + is_multimodal = kwargs.pop("is_multimodal", None)
1332 + handle_oov_mm_token = kwargs.pop("handle_oov_mm_token", False)
1333 + inputs_embeds = self.get_input_embeddings(
1334 + input_ids,
1335 + vision_embeddings,
1336 + is_multimodal=is_multimodal,
1337 + handle_oov_mm_token=handle_oov_mm_token,
1338 + )
1339 + input_ids = None
1340 +
1341 + return self.language_model(
1342 + input_ids, positions, intermediate_tensors, inputs_embeds
1343 + )
1344 +
1345 + @classmethod
1346 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
1347 + if modality.startswith("image"):
1348 + return "<|IMAGE_START|><|IMAGE_PLACEHOLDER|><|IMAGE_END|>"
1349 +
1350 + raise ValueError("Only image modality is supported")
1351 +
1352 + def encode_image(
1353 + self, pixel_values: torch.Tensor, image_grid_thw: torch.Tensor
1354 + ) -> torch.Tensor:
1355 + pixel_values = pixel_values.type(self.visual.dtype)
1356 + siglip_position_ids = list()
1357 + image_grid_hws = list()
1358 + cu_seqlens = [0]
1359 +
1360 + thw_tuple = tuple(image_grid_thw.tolist())
1361 + numel = np.prod(thw_tuple)
1362 + image_grid_hws.append(thw_tuple)
1363 + image_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:])
1364 + siglip_position_ids.append(image_position_ids)
1365 + cu_seqlens.append(cu_seqlens[-1] + numel)
1366 +
1367 + siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to(
1368 + pixel_values.device
1369 + )
1370 + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to(pixel_values.device)
1371 +
1372 + vision_outputs = self.visual(
1373 + pixel_values=pixel_values,
1374 + image_grid_thw=image_grid_hws,
1375 + position_ids=siglip_position_ids,
1376 + interpolate_pos_encoding=True,
1377 + cu_seqlens=cu_seqlens,
1378 + )
1379 + return vision_outputs
1380 +
1381 + def _process_image_input(
1382 + self, image_input: PaddleOCRImagePixelInputs
1383 + ) -> MultiModalEmbeddings:
1384 + pixel_values = image_input.pixel_values
1385 + image_grid_thw = image_input.image_grid_thw
1386 + vision_outputs = tuple(
1387 + self.encode_image(pixel, grid).squeeze(0)
1388 + for pixel, grid in zip(pixel_values, image_grid_thw)
1389 + )
1390 + image_embeds = self.mlp_AR(vision_outputs, image_grid_thw)
1391 + return image_embeds
1392 +
1393 + def get_multimodal_embeddings(self, **kwargs) -> MultiModalEmbeddings:
1394 + image_input = self._parse_and_validate_image_input(**kwargs)
1395 + if image_input is None:
1396 + return ()
1397 +
1398 + multimodal_embeddings: tuple[torch.Tensor, ...] = ()
1399 + image_embeds = self._process_image_input(image_input)
1400 + multimodal_embeddings += tuple(image_embeds)
1401 +
1402 + return multimodal_embeddings
1403 +
1404 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1405 + loader = AutoWeightsLoader(self)
1406 + autoloaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1407 + return autoloaded_weights
vllm/model_executor/models/registry.py
+4
@@ -340,6 +340,10 @@ _MULTIMODAL_MODELS = {
340 "NVLM_D": ("nvlm_d", "NVLM_D_Model"),
341 "Ovis": ("ovis", "Ovis"),
342 "Ovis2_5": ("ovis2_5", "Ovis2_5"),
343 + "PaddleOCRVLForConditionalGeneration": (
344 + "paddleocr_vl",
345 + "PaddleOCRVLForConditionalGeneration",
346 + ),
347 "PaliGemmaForConditionalGeneration": (
348 "paligemma",
349 "PaliGemmaForConditionalGeneration",