@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 675aa2ec64

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 14, 05:30 UTC

[Model] Upstream Deepseek-OCR model (#27247)

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

Isotr0py committed Oct 22, 2025 at 22:59 UTC 675aa2ec64b2d8ab45948f45cef80f74ebfadbbb

10 files changed +1821 -40

docs/models/supported_models.md

@@ -639,6 +639,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

639 | `ChameleonForConditionalGeneration` | Chameleon | T + I | `facebook/chameleon-7b`, etc. | | ✅︎ |

640 | `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I+ | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ |

641 | `DeepseekVLV2ForCausalLM`^ | DeepSeek-VL2 | T + I+ | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |

642 +| `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I+ | `deepseek-ai/DeepSeek-OCR`, etc. | | ✅︎ |

643 | `Ernie4_5_VLMoeForConditionalGeneration` | Ernie4.5-VL | T + I+/ V+ | `baidu/ERNIE-4.5-VL-28B-A3B-PT`, `baidu/ERNIE-4.5-VL-424B-A47B-PT` | | ✅︎ |

644 | `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ |

examples/offline_inference/vision_language.py

+69 -20

@@ -30,6 +30,7 @@ class ModelRequestData(NamedTuple):

30 prompts: list[str]

31 stop_token_ids: list[int] | None = None

32 lora_requests: list[LoRARequest] | None = None

33 + sampling_params: list[SamplingParams] | None = None

36 # NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on

@@ -153,23 +154,6 @@ def run_chameleon(questions: list[str], modality: str) -> ModelRequestData:

154 )

155

156

156 -# Dots-OCR

157 -def run_dots_ocr(questions: list[str], modality: str) -> ModelRequestData:

158 - assert modality == "image"

159 -

160 - prompts = [f"<|img|><|imgpad|><|endofimg|>{question}" for question in questions]

161 - engine_args = EngineArgs(

162 - model="rednote-hilab/dots.ocr",

163 - limit_mm_per_prompt={modality: 1},

164 - trust_remote_code=True,

165 - )

166 -

167 - return ModelRequestData(

168 - engine_args=engine_args,

169 - prompts=prompts,

170 - )

171 -

172 -

157 def run_command_a_vision(questions: list[str], modality: str) -> ModelRequestData:

158 assert modality == "image"

159

@@ -217,6 +201,66 @@ def run_deepseek_vl2(questions: list[str], modality: str) -> ModelRequestData:

201 )

202

203

204 +def run_deepseek_ocr(questions: list[str], modality: str) -> ModelRequestData:

205 + from vllm.model_executor.models.deepseek_ocr import NGramPerReqLogitsProcessor

206 +

207 + assert modality == "image"

208 +

209 + model_name = "deepseek-ai/DeepSeek-OCR"

210 +

211 + engine_args = EngineArgs(

212 + model=model_name,

213 + limit_mm_per_prompt={modality: 1},

214 + logits_processors=[NGramPerReqLogitsProcessor],

215 + )

216 +

217 + # deepseek-ocr use plain prompt template

218 + prompts = [f"<image>\n{question}" for question in questions]

219 +

220 + # The following sampling params config is taken from

221 + # the official Deepseek-OCR inference example.

222 + # (IMPORTANT) Use the custom logits processor and avoid skipping

223 + # special tokens for this model for the optimal OCR performance.

224 + sampling_params = [

225 + SamplingParams(

226 + temperature=0.0,

227 + max_tokens=8192,

228 + # ngram logit processor args

229 + extra_args=dict(

230 + ngram_size=30,

231 + window_size=90,

232 + # whitelist: <td>, </td>

233 + whitelist_token_ids={128821, 128822},

234 + ),

235 + skip_special_tokens=False,

236 + )

237 + for _ in questions

238 + ]

239 +

240 + return ModelRequestData(

241 + engine_args=engine_args,

242 + prompts=prompts,

243 + sampling_params=sampling_params,

244 + )

245 +

246 +

247 +# Dots-OCR

248 +def run_dots_ocr(questions: list[str], modality: str) -> ModelRequestData:

249 + assert modality == "image"

250 +

251 + prompts = [f"<|img|><|imgpad|><|endofimg|>{question}" for question in questions]

252 + engine_args = EngineArgs(

253 + model="rednote-hilab/dots.ocr",

254 + limit_mm_per_prompt={modality: 1},

255 + trust_remote_code=True,

256 + )

257 +

258 + return ModelRequestData(

259 + engine_args=engine_args,

260 + prompts=prompts,

261 + )

262 +

263 +

264 # Ernie4.5-VL

265 def run_ernie45_vl(questions: list[str], modality: str) -> ModelRequestData:

266 model_name = "baidu/ERNIE-4.5-VL-28B-A3B-PT"

@@ -1738,9 +1782,10 @@ model_example_map = {

1782 "bee": run_bee,

1783 "blip-2": run_blip2,

1784 "chameleon": run_chameleon,

1741 - "dots_ocr": run_dots_ocr,

1785 "command_a_vision": run_command_a_vision,

1786 "deepseek_vl_v2": run_deepseek_vl2,

1787 + "deepseek_ocr": run_deepseek_ocr,

1788 + "dots_ocr": run_dots_ocr,

1789 "ernie45_vl": run_ernie45_vl,

1790 "fuyu": run_fuyu,

1791 "gemma3": run_gemma3,

@@ -2003,8 +2048,12 @@ def main(args):

2048

2049 # We set temperature to 0.2 so that outputs can be different

2050 # even when all prompts are identical when running batch inference.

2006 - sampling_params = SamplingParams(

2007 - temperature=0.2, max_tokens=64, stop_token_ids=req_data.stop_token_ids

2051 + sampling_params = (

2052 + SamplingParams(

2053 + temperature=0.2, max_tokens=64, stop_token_ids=req_data.stop_token_ids

2054 + )

2055 + if req_data.sampling_params is None

2056 + else req_data.sampling_params

2057 )

2058

2059 assert args.num_prompts > 0

tests/models/registry.py

@@ -585,6 +585,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {

585 transformers_version_reason="HF model is not compatible.",

586 hf_overrides={"architectures": ["DeepseekVLV2ForCausalLM"]},

587 ),

588 + "DeepseekOCRForCausalLM": _HfExamplesInfo(

589 + "deepseek-ai/DeepSeek-OCR",

590 + ),

591 "DotsOCRForCausalLM": _HfExamplesInfo(

592 "rednote-hilab/dots.ocr", trust_remote_code=True

593 ),

vllm/model_executor/models/deepencoder.py new

+673

@@ -0,0 +1,673 @@

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/blob/main/DeepSeek-OCR-master/DeepSeek-OCR-vllm/deepencoder/sam_vary_sdpa.py

6 +

7 +# Copyright (c) Meta Platforms, Inc. and affiliates.

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 +import math

13 +from collections.abc import Iterable

14 +from functools import partial

15 +

16 +import torch

17 +import torch.nn as nn

18 +import torch.nn.functional as F

19 +from transformers import CLIPVisionConfig

20 +

21 +from vllm.attention.layer import MultiHeadAttention

22 +from vllm.model_executor.layers.quantization import QuantizationConfig

23 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

24 +

25 +from .clip import CLIPEncoder, CLIPVisionEmbeddings

26 +

27 +

28 +class MLPBlock(nn.Module):

29 + def __init__(

30 + self,

31 + embedding_dim: int,

32 + mlp_dim: int,

33 + act: type[nn.Module] = nn.GELU,

34 + ) -> None:

35 + super().__init__()

36 + self.lin1 = nn.Linear(embedding_dim, mlp_dim)

37 + self.lin2 = nn.Linear(mlp_dim, embedding_dim)

38 + self.act = act()

39 +

40 + def forward(self, x: torch.Tensor) -> torch.Tensor:

41 + return self.lin2(self.act(self.lin1(x)))

42 +

43 +

44 +# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa

45 +# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa

46 +class LayerNorm2d(nn.Module):

47 + def __init__(self, num_channels: int, eps: float = 1e-6) -> None:

48 + super().__init__()

49 + self.weight = nn.Parameter(torch.ones(num_channels))

50 + self.bias = nn.Parameter(torch.zeros(num_channels))

51 + self.eps = eps

52 +

53 + def forward(self, x: torch.Tensor) -> torch.Tensor:

54 + u = x.mean(1, keepdim=True)

55 + s = (x - u).pow(2).mean(1, keepdim=True)

56 + x = (x - u) / torch.sqrt(s + self.eps)

57 + x = self.weight[:, None, None] * x + self.bias[:, None, None]

58 + return x

59 +

60 +

61 +# This class and its supporting functions below lightly adapted from the ViTDet backbone available at: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py # noqa

62 +class ImageEncoderViT(nn.Module):

63 + def __init__(

64 + self,

65 + img_size: int = 1024,

66 + patch_size: int = 16,

67 + in_chans: int = 3,

68 + embed_dim: int = 768,

69 + depth: int = 12,

70 + num_heads: int = 12,

71 + mlp_ratio: float = 4.0,

72 + out_chans: int = 256,

73 + qkv_bias: bool = True,

74 + norm_layer: type[nn.Module] = nn.LayerNorm,

75 + act_layer: type[nn.Module] = nn.GELU,

76 + use_abs_pos: bool = True,

77 + use_rel_pos: bool = False,

78 + rel_pos_zero_init: bool = True,

79 + window_size: int = 0,

80 + global_attn_indexes: tuple[int, ...] = (),

81 + ) -> None:

82 + """

83 + Args:

84 + img_size (int): Input image size.

85 + patch_size (int): Patch size.

86 + in_chans (int): Number of input image channels.

87 + embed_dim (int): Patch embedding dimension.

88 + depth (int): Depth of ViT.

89 + num_heads (int): Number of attention heads in each ViT block.

90 + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.

91 + qkv_bias (bool): If True, add a learnable bias to query, key, value.

92 + norm_layer (nn.Module): Normalization layer.

93 + act_layer (nn.Module): Activation layer.

94 + use_abs_pos (bool): If True, use absolute positional embeddings.

95 + use_rel_pos (bool): If True, add relative positional embeddings to the attention map.

96 + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.

97 + window_size (int): Window size for window attention blocks.

98 + global_attn_indexes (list): Indexes for blocks using global attention.

99 + """ # noqa: E501

100 + super().__init__()

101 + self.img_size = img_size

102 +

103 + self.patch_embed = PatchEmbed(

104 + kernel_size=(patch_size, patch_size),

105 + stride=(patch_size, patch_size),

106 + in_chans=in_chans,

107 + embed_dim=embed_dim,

108 + )

109 +

110 + self.pos_embed: nn.Parameter | None = None

111 + if use_abs_pos:

112 + # Initialize absolute positional embedding with pretrain image size.

113 + self.pos_embed = nn.Parameter(

114 + torch.zeros(

115 + 1, img_size // patch_size, img_size // patch_size, embed_dim

116 + )

117 + )

118 +

119 + self.blocks = nn.ModuleList()

120 + for i in range(depth):

121 + block = Block(

122 + dim=embed_dim,

123 + num_heads=num_heads,

124 + mlp_ratio=mlp_ratio,

125 + qkv_bias=qkv_bias,

126 + norm_layer=norm_layer,

127 + act_layer=act_layer,

128 + use_rel_pos=use_rel_pos,

129 + rel_pos_zero_init=rel_pos_zero_init,

130 + window_size=window_size if i not in global_attn_indexes else 0,

131 + input_size=(img_size // patch_size, img_size // patch_size),

132 + )

133 + self.blocks.append(block)

134 +

135 + self.neck = nn.Sequential(

136 + nn.Conv2d(

137 + embed_dim,

138 + out_chans,

139 + kernel_size=1,

140 + bias=False,

141 + ),

142 + LayerNorm2d(out_chans),

143 + nn.Conv2d(

144 + out_chans,

145 + out_chans,

146 + kernel_size=3,

147 + padding=1,

148 + bias=False,

149 + ),

150 + LayerNorm2d(out_chans),

151 + )

152 +

153 + self.net_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1, bias=False)

154 + self.net_3 = nn.Conv2d(

155 + 512, 1024, kernel_size=3, stride=2, padding=1, bias=False

156 + )

157 +

158 + def get_abs_pos(self, abs_pos: torch.Tensor, tgt_size: int):

159 + dtype = abs_pos.dtype

160 +

161 + src_size = abs_pos.size(1)

162 +

163 + if src_size != tgt_size:

164 + old_pos_embed = abs_pos.permute(0, 3, 1, 2)

165 + old_pos_embed = old_pos_embed.to(torch.float32)

166 + new_pos_embed = F.interpolate(

167 + old_pos_embed,

168 + size=(tgt_size, tgt_size),

169 + mode="bicubic",

170 + antialias=True,

171 + align_corners=False,

172 + ).to(dtype)

173 + new_pos_embed = new_pos_embed.permute(0, 2, 3, 1)

174 + return new_pos_embed

175 + else:

176 + return abs_pos

177 +

178 + def forward(self, x: torch.Tensor) -> torch.Tensor:

179 + x = self.patch_embed(x)

180 + if self.pos_embed is not None:

181 + x = x + self.get_abs_pos(self.pos_embed, x.size(1))

182 +

183 + for blk in self.blocks:

184 + x = blk(x)

185 +

186 + neck_output = self.neck(x.permute(0, 3, 1, 2))

187 + conv2_output = self.net_2(neck_output)

188 + conv3_output = self.net_3(conv2_output)

189 +

190 + return conv3_output

191 +

192 +

193 +class Block(nn.Module):

194 + """Transformer blocks with support of window attention and residual propagation

195 + blocks"""

196 +

197 + def __init__(

198 + self,

199 + dim: int,

200 + num_heads: int,

201 + mlp_ratio: float = 4.0,

202 + qkv_bias: bool = True,

203 + norm_layer: type[nn.Module] = nn.LayerNorm,

204 + act_layer: type[nn.Module] = nn.GELU,

205 + use_rel_pos: bool = False,

206 + rel_pos_zero_init: bool = True,

207 + window_size: int = 0,

208 + input_size: tuple[int, int] | None = None,

209 + ) -> None:

210 + """

211 + Args:

212 + dim (int): Number of input channels.

213 + num_heads (int): Number of attention heads in each ViT block.

214 + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.

215 + qkv_bias (bool): If True, add a learnable bias to query, key, value.

216 + norm_layer (nn.Module): Normalization layer.

217 + act_layer (nn.Module): Activation layer.

218 + use_rel_pos (bool): If True, add relative positional embeddings to the attention map.

219 + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.

220 + window_size (int): Window size for window attention blocks. If it equals 0, then

221 + use global attention.

222 + input_size (tuple(int, int) or None): Input resolution for calculating the relative

223 + positional parameter size.

224 + """ # noqa: E501

225 + super().__init__()

226 + self.norm1 = norm_layer(dim)

227 + self.attn = RelPosAttention(

228 + dim,

229 + num_heads=num_heads,

230 + qkv_bias=qkv_bias,

231 + use_rel_pos=use_rel_pos,

232 + rel_pos_zero_init=rel_pos_zero_init,

233 + input_size=input_size if window_size == 0 else (window_size, window_size),

234 + )

235 +

236 + self.norm2 = norm_layer(dim)

237 + self.mlp = MLPBlock(

238 + embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer

239 + )

240 +

241 + self.window_size = window_size

242 +

243 + def forward(self, x: torch.Tensor) -> torch.Tensor:

244 + shortcut = x

245 + x = self.norm1(x)

246 + # Window partition

247 + if self.window_size > 0:

248 + H, W = x.shape[1], x.shape[2]

249 + x, pad_hw = window_partition(x, self.window_size)

250 +

251 + x = self.attn(x)

252 + # Reverse window partition

253 + if self.window_size > 0:

254 + x = window_unpartition(x, self.window_size, pad_hw, (H, W))

255 +

256 + x = shortcut + x

257 + x = x + self.mlp(self.norm2(x))

258 +

259 + return x

260 +

261 +

262 +class RelPosAttention(nn.Module):

263 + """Multi-head Attention block with relative position embeddings."""

264 +

265 + def __init__(

266 + self,

267 + dim: int,

268 + num_heads: int = 8,

269 + qkv_bias: bool = True,

270 + use_rel_pos: bool = False,

271 + rel_pos_zero_init: bool = True,

272 + input_size: tuple[int, int] | None = None,

273 + ) -> None:

274 + """

275 + Args:

276 + dim (int): Number of input channels.

277 + num_heads (int): Number of attention heads.

278 + qkv_bias (bool): If True, add a learnable bias to query, key, value.

279 + rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.

280 + input_size (tuple(int, int) or None): Input resolution for calculating the relative

281 + positional parameter size.

282 + """ # noqa: E501

283 + super().__init__()

284 + self.num_heads = num_heads

285 + head_dim = dim // num_heads

286 + self.scale = head_dim**-0.5

287 +

288 + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)

289 + self.proj = nn.Linear(dim, dim)

290 +

291 + self.use_rel_pos = use_rel_pos

292 + if self.use_rel_pos:

293 + assert input_size is not None, (

294 + "Input size must be provided if using relative positional encoding."

295 + )

296 + # initialize relative positional embeddings

297 + self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))

298 + self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))

299 +

300 + def forward(self, x: torch.Tensor) -> torch.Tensor:

301 + B, H, W, _ = x.shape

302 + # qkv with shape (3, B, nHead, H * W, C)

303 + qkv = (

304 + self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)

305 + )

306 + # q, k, v with shape (B * nHead, H * W, C)

307 + q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)

308 +

309 + rel_h, rel_w = None, None

310 + if self.use_rel_pos:

311 + rel_h, rel_w = add_decomposed_rel_pos(

312 + q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W)

313 + )

314 +

315 + q = q.view(B, self.num_heads, H * W, -1)

316 + k = k.view(B, self.num_heads, H * W, -1)

317 + v = v.view(B, self.num_heads, H * W, -1)

318 +

319 + if self.use_rel_pos:

320 + rel_h = rel_h.view(

321 + B, self.num_heads, rel_h.size(1), rel_h.size(2), rel_h.size(3)

322 + )

323 + rel_w = rel_w.view(

324 + B, self.num_heads, rel_w.size(1), rel_w.size(2), rel_w.size(3)

325 + )

326 + attn_bias = (rel_h + rel_w).view(

327 + B, self.num_heads, rel_h.size(2), rel_h.size(3) * rel_w.size(4)

328 + )

329 + x = torch.nn.functional.scaled_dot_product_attention(

330 + q, k, v, attn_mask=attn_bias

331 + )

332 + else:

333 + x = torch.nn.functional.scaled_dot_product_attention(q, k, v)

334 +

335 + x = (

336 + x.view(B, self.num_heads, H, W, -1)

337 + .permute(0, 2, 3, 1, 4)

338 + .reshape(B, H, W, -1)

339 + )

340 +

341 + x = self.proj(x)

342 +

343 + return x

344 +

345 +

346 +def window_partition(

347 + x: torch.Tensor, window_size: int

348 +) -> tuple[torch.Tensor, tuple[int, int]]:

349 + """

350 + Partition into non-overlapping windows with padding if needed.

351 + Args:

352 + x (tensor): input tokens with [B, H, W, C].

353 + window_size (int): window size.

354 +

355 + Returns:

356 + windows: windows after partition with [B * num_windows, window_size, window_size, C].

357 + (Hp, Wp): padded height and width before partition

358 + """ # noqa: E501

359 + B, H, W, C = x.shape

360 +

361 + pad_h = (window_size - H % window_size) % window_size

362 + pad_w = (window_size - W % window_size) % window_size

363 + if pad_h > 0 or pad_w > 0:

364 + x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))

365 + Hp, Wp = H + pad_h, W + pad_w

366 +

367 + x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)

368 + windows = (

369 + x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)

370 + )

371 + return windows, (Hp, Wp)

372 +

373 +

374 +def window_unpartition(

375 + windows: torch.Tensor,

376 + window_size: int,

377 + pad_hw: tuple[int, int],

378 + hw: tuple[int, int],

379 +) -> torch.Tensor:

380 + """

381 + Window unpartition into original sequences and removing padding.

382 + Args:

383 + windows (tensor): input tokens with [B * num_windows, window_size, window_size, C].

384 + window_size (int): window size.

385 + pad_hw (Tuple): padded height and width (Hp, Wp).

386 + hw (Tuple): original height and width (H, W) before padding.

387 +

388 + Returns:

389 + x: unpartitioned sequences with [B, H, W, C].

390 + """ # noqa: E501

391 + Hp, Wp = pad_hw

392 + H, W = hw

393 + B = windows.shape[0] // (Hp * Wp // window_size // window_size)

394 + x = windows.view(

395 + B, Hp // window_size, Wp // window_size, window_size, window_size, -1

396 + )

397 + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1)

398 +

399 + if Hp > H or Wp > W:

400 + x = x[:, :H, :W, :].contiguous()

401 + return x

402 +

403 +

404 +def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:

405 + """

406 + Get relative positional embeddings according to the relative positions of

407 + query and key sizes.

408 + Args:

409 + q_size (int): size of query q.

410 + k_size (int): size of key k.

411 + rel_pos (Tensor): relative position embeddings (L, C).

412 +

413 + Returns:

414 + Extracted positional embeddings according to relative positions.

415 + """

416 + max_rel_dist = int(2 * max(q_size, k_size) - 1)

417 + # Interpolate rel pos if needed.

418 + if rel_pos.shape[0] != max_rel_dist:

419 + # Interpolate rel pos.

420 + dtype = rel_pos.dtype

421 + rel_pos = rel_pos.to(torch.float32)

422 + rel_pos_resized = F.interpolate(

423 + rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),

424 + size=max_rel_dist,

425 + mode="linear",

426 + ).to(dtype)

427 + rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)

428 + else:

429 + rel_pos_resized = rel_pos

430 +

431 + # Scale the coords with short length if shapes for q and k are different.

432 + q_coords = torch.arange(q_size, device=rel_pos.device)[:, None] * max(

433 + k_size / q_size, 1.0

434 + )

435 + k_coords = torch.arange(k_size, device=rel_pos.device)[None, :] * max(

436 + q_size / k_size, 1.0

437 + )

438 + relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)

439 +

440 + return rel_pos_resized[relative_coords.long()]

441 +

442 +

443 +def add_decomposed_rel_pos(

444 + q: torch.Tensor,

445 + rel_pos_h: torch.Tensor,

446 + rel_pos_w: torch.Tensor,

447 + q_size: tuple[int, int],

448 + k_size: tuple[int, int],

449 +) -> torch.Tensor:

450 + """

451 + Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.

452 + https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py

453 + Args:

454 + q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).

455 + rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis.

456 + rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis.

457 + q_size (Tuple): spatial sequence size of query q with (q_h, q_w).

458 + k_size (Tuple): spatial sequence size of key k with (k_h, k_w).

459 +

460 + Returns:

461 + attn (Tensor): attention map with added relative positional embeddings.

462 + """ # noqa: E501

463 + q_h, q_w = q_size

464 + k_h, k_w = k_size

465 + Rh = get_rel_pos(q_h, k_h, rel_pos_h)

466 + Rw = get_rel_pos(q_w, k_w, rel_pos_w)

467 +

468 + B, _, dim = q.shape

469 + r_q = q.reshape(B, q_h, q_w, dim)

470 + rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh)

471 + rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw)

472 + rel_h = rel_h.unsqueeze(-1)

473 + rel_w = rel_w.unsqueeze(-2)

474 + rel_h = rel_h.reshape(B, q_h * q_w, k_h, 1)

475 + rel_w = rel_w.reshape(B, q_h * q_w, 1, k_w)

476 +

477 + return rel_h, rel_w

478 +

479 +

480 +class PatchEmbed(nn.Module):

481 + """

482 + Image to Patch Embedding.

483 + """

484 +

485 + def __init__(

486 + self,

487 + kernel_size: tuple[int, int] = (16, 16),

488 + stride: tuple[int, int] = (16, 16),

489 + padding: tuple[int, int] = (0, 0),

490 + in_chans: int = 3,

491 + embed_dim: int = 768,

492 + ) -> None:

493 + """

494 + Args:

495 + kernel_size (Tuple): kernel size of the projection layer.

496 + stride (Tuple): stride of the projection layer.

497 + padding (Tuple): padding size of the projection layer.

498 + in_chans (int): Number of input image channels.

499 + embed_dim (int): Patch embedding dimension.

500 + """

501 + super().__init__()

502 +

503 + self.proj = nn.Conv2d(

504 + in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding

505 + )

506 +

507 + def forward(self, x: torch.Tensor) -> torch.Tensor:

508 + x = self.proj(x)

509 + # B C H W -> B H W C

510 + x = x.permute(0, 2, 3, 1)

511 + return x

512 +

513 +

514 +# TODO(Isotr0py): use vision_config to build sam model

515 +def build_sam_vit_b():

516 + return _build_sam(

517 + encoder_embed_dim=768,

518 + encoder_depth=12,

519 + encoder_num_heads=12,

520 + encoder_global_attn_indexes=[2, 5, 8, 11],

521 + )

522 +

523 +

524 +def _build_sam(

525 + encoder_embed_dim,

526 + encoder_depth,

527 + encoder_num_heads,

528 + encoder_global_attn_indexes,

529 +):

530 + prompt_embed_dim = 256

531 + image_size = 1024

532 + vit_patch_size = 16

533 + image_encoder = ImageEncoderViT(

534 + depth=encoder_depth,

535 + embed_dim=encoder_embed_dim,

536 + img_size=image_size,

537 + mlp_ratio=4,

538 + norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),

539 + num_heads=encoder_num_heads,

540 + patch_size=vit_patch_size,

541 + qkv_bias=True,

542 + use_rel_pos=True,

543 + global_attn_indexes=encoder_global_attn_indexes,

544 + window_size=14,

545 + out_chans=prompt_embed_dim,

546 + )

547 + return image_encoder

548 +

549 +

550 +class DeepCLIPVisionEmbeddings(CLIPVisionEmbeddings):

551 + def get_abs_pos(self, abs_pos: torch.Tensor, tgt_size: int):

552 + # abs_pos: L, C

553 + # tgt_size: M

554 + # return: M, C

555 + dim = abs_pos.size(-1)

556 + abs_pos_new = abs_pos.squeeze(0)

557 + cls_token, old_pos_embed = abs_pos_new[:1], abs_pos_new[1:]

558 +

559 + src_size = int(math.sqrt(abs_pos_new.shape[0] - 1))

560 + tgt_size = int(math.sqrt(tgt_size))

561 + dtype = abs_pos.dtype

562 +

563 + if src_size != tgt_size:

564 + old_pos_embed = (

565 + old_pos_embed.view(1, src_size, src_size, dim)

566 + .permute(0, 3, 1, 2)

567 + .contiguous()

568 + )

569 + old_pos_embed = old_pos_embed.to(torch.float32)

570 + new_pos_embed = F.interpolate(

571 + old_pos_embed,

572 + size=(tgt_size, tgt_size),

573 + mode="bicubic",

574 + antialias=True,

575 + align_corners=False,

576 + ).to(dtype)

577 + new_pos_embed = new_pos_embed.permute(0, 2, 3, 1)

578 + new_pos_embed = new_pos_embed.view(tgt_size * tgt_size, dim)

579 + vision_pos_embed = torch.cat([cls_token, new_pos_embed], dim=0)

580 + vision_pos_embed = vision_pos_embed.view(1, tgt_size * tgt_size + 1, dim)

581 + return vision_pos_embed

582 + else:

583 + return abs_pos

584 +

585 + def forward(

586 + self, pixel_values: torch.Tensor, patch_embeds: torch.Tensor | None = None

587 + ) -> torch.Tensor:

588 + batch_size = pixel_values.shape[0]

589 + if patch_embeds is not None:

590 + patch_embeds = patch_embeds

591 + else:

592 + patch_embeds = self.patch_embedding(pixel_values)

593 + patch_embeds = patch_embeds.flatten(2).transpose(1, 2)

594 +

595 + class_embeds = self.class_embedding.expand(batch_size, 1, -1)

596 + embeddings = torch.cat([class_embeds, patch_embeds], dim=1)

597 + embeddings = embeddings + self.get_abs_pos(

598 + self.position_embedding(self.position_ids), embeddings.size(1)

599 + )

600 + return embeddings

601 +

602 +

603 +class DeepCLIPVisionTransformer(nn.Module):

604 + def __init__(

605 + self,

606 + config: CLIPVisionConfig,

607 + quant_config: QuantizationConfig | None = None,

608 + *,

609 + num_hidden_layers_override: int | None = None,

610 + prefix: str = "",

611 + ) -> None:

612 + super().__init__()

613 +

614 + self.config = config

615 + embed_dim = config.hidden_size

616 +

617 + self.embeddings = DeepCLIPVisionEmbeddings(config)

618 +

619 + # NOTE: This typo of "layrnorm" is not fixed on purpose to match

620 + # the original transformers code and name of the model weights.

621 + self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

622 +

623 + self.transformer = CLIPEncoder(

624 + config=config,

625 + quant_config=quant_config,

626 + num_hidden_layers_override=num_hidden_layers_override,

627 + prefix=f"{prefix}.encoder",

628 + attn_cls=MultiHeadAttention,

629 + )

630 +

631 + num_hidden_layers = config.num_hidden_layers

632 + if len(self.transformer.layers) > config.num_hidden_layers:

633 + raise ValueError(

634 + f"The original encoder only has {num_hidden_layers} "

635 + f"layers, but you requested {len(self.transformer.layers)} layers."

636 + )

637 +

638 + @property

639 + def dtype(self):

640 + return next(self.parameters()).dtype

641 +

642 + @property

643 + def device(self):

644 + return next(self.parameters()).device

645 +

646 + def forward(

647 + self,

648 + pixel_values: torch.Tensor,

649 + patch_embeds: torch.Tensor | None = None,

650 + *,

651 + select_layers: list[int] | None = None,

652 + ) -> torch.Tensor:

653 + hidden_states = self.embeddings(pixel_values, patch_embeds)

654 + hidden_states = self.pre_layrnorm(hidden_states)

655 +

656 + # Produces either the last layer output or all of the hidden states,

657 + # depending on if we have select_layers or not

658 + encoder_outputs = self.transformer(

659 + inputs_embeds=hidden_states,

660 + return_all_hidden_states=select_layers is not None,

661 + )

662 + return encoder_outputs

663 +

664 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

665 + params_dict = dict(self.named_parameters())

666 + loaded_params: set[str] = set()

667 +

668 + for name, loaded_weight in weights:

669 + param = params_dict[name]

670 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

671 + weight_loader(param, loaded_weight)

672 + loaded_params.add(name)

673 + return loaded_params

vllm/model_executor/models/deepseek_ocr.py new

+594

@@ -0,0 +1,594 @@

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 +

8 +import torch

9 +import torch.nn as nn

10 +from transformers import BatchFeature, CLIPVisionConfig

11 +

12 +from vllm.config import VllmConfig

13 +from vllm.config.multimodal import BaseDummyOptions

14 +from vllm.model_executor.models.interfaces import (

15 + MultiModalEmbeddings,

16 + SupportsMultiModal,

17 + SupportsPP,

18 +)

19 +from vllm.model_executor.models.utils import (

20 + AutoWeightsLoader,

21 + WeightsMapper,

22 + init_vllm_registered_model,

23 + maybe_prefix,

24 +)

25 +from vllm.multimodal import MULTIMODAL_REGISTRY

26 +from vllm.multimodal.inputs import (

27 + MultiModalDataDict,

28 + MultiModalFieldConfig,

29 + MultiModalKwargs,

30 + NestedTensors,

31 +)

32 +from vllm.multimodal.parse import (

33 + ImageEmbeddingItems,

34 + ImageProcessorItems,

35 + ImageSize,

36 + MultiModalDataItems,

37 +)

38 +from vllm.multimodal.processing import (

39 + BaseMultiModalProcessor,

40 + BaseProcessingInfo,

41 + PromptReplacement,

42 + PromptUpdate,

43 +)

44 +from vllm.multimodal.profiling import BaseDummyInputsBuilder

45 +from vllm.sampling_params import SamplingParams

46 +from vllm.sequence import IntermediateTensors

47 +from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config

48 +from vllm.transformers_utils.processors.deepseek_ocr import (

49 + BASE_SIZE,

50 + CROP_MODE,

51 + IMAGE_SIZE,

52 + DeepseekOCRProcessor,

53 + count_tiles,

54 +)

55 +from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config

56 +from vllm.v1.sample.logits_processor import (

57 + AdapterLogitsProcessor,

58 + RequestLogitsProcessor,

59 +)

60 +

61 +from .deepencoder import DeepCLIPVisionTransformer, build_sam_vit_b

62 +from .deepseek_vl2 import MlpProjector

63 +

64 +# The image token id may be various

65 +_IMAGE_TOKEN = "<image>"

66 +

67 +

68 +class NoRepeatNGramLogitsProcessor:

69 + def __init__(

70 + self,

71 + ngram_size: int,

72 + window_size: int,

73 + whitelist_token_ids: set[int] | None = None,

74 + ):

75 + self.ngram_size = ngram_size

76 + self.window_size = window_size

77 + self.whitelist_token_ids = whitelist_token_ids or set()

78 +

79 + def __call__(

80 + self,

81 + output_ids: list[int],

82 + logits: torch.Tensor,

83 + ) -> torch.Tensor:

84 + if len(output_ids) < self.ngram_size:

85 + return logits

86 +

87 + current_prefix = tuple(output_ids[-(self.ngram_size - 1) :])

88 +

89 + search_start = max(0, len(output_ids) - self.window_size)

90 + search_end = len(output_ids) - self.ngram_size + 1

91 +

92 + banned_tokens = set()

93 + for i in range(search_start, search_end):

94 + ngram = tuple(output_ids[i : i + self.ngram_size])

95 + if ngram[:-1] == current_prefix:

96 + banned_tokens.add(ngram[-1])

97 +

98 + banned_tokens = banned_tokens - self.whitelist_token_ids

99 +

100 + if banned_tokens:

101 + logits[list(banned_tokens)] = -float("inf")

102 +

103 + return logits

104 +

105 +

106 +class NGramPerReqLogitsProcessor(AdapterLogitsProcessor):

107 + """Example of overriding the wrapper class `__init__()` in order to utilize

108 + info about the device type"""

109 +

110 + def __init__(

111 + self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool

112 + ):

113 + super().__init__(vllm_config, device, is_pin_memory)

114 +

115 + def is_argmax_invariant(self) -> bool:

116 + return True

117 +

118 + def new_req_logits_processor(

119 + self,

120 + params: SamplingParams,

121 + ) -> RequestLogitsProcessor | None:

122 + ngram_size = params.extra_args and params.extra_args.get("ngram_size")

123 + window_size = params.extra_args and params.extra_args.get("window_size", 100)

124 + whitelist_token_ids = params.extra_args and params.extra_args.get(

125 + "whitelist_token_ids", None

126 + )

127 + if ngram_size is None:

128 + return None

129 + if not isinstance(ngram_size, int) or ngram_size <= 0:

130 + raise ValueError(

131 + f"`ngram_size` has to be a strictly positive integer, got {ngram_size}."

132 + )

133 + if not isinstance(window_size, int) or window_size <= 0:

134 + raise ValueError(

135 + "`window_size` has to be a strictly positive integer, "

136 + f"got {window_size}."

137 + )

138 + if whitelist_token_ids is not None and not isinstance(

139 + whitelist_token_ids, Iterable

140 + ):

141 + raise ValueError(

142 + "`whitelist_token_ids` has to be a set of integers, "

143 + f"got {whitelist_token_ids}."

144 + )

145 + else:

146 + whitelist_token_ids = (

147 + set(whitelist_token_ids) if whitelist_token_ids else None

148 + )

149 + return NoRepeatNGramLogitsProcessor(

150 + ngram_size=ngram_size,

151 + window_size=window_size,

152 + whitelist_token_ids=whitelist_token_ids,

153 + )

154 +

155 +

156 +class DeepseekOCRProcessingInfo(BaseProcessingInfo):

157 + def get_hf_config(self):

158 + return self.ctx.get_hf_config(DeepseekVLV2Config)

159 +

160 + def get_hf_processor(self, **kwargs: object):

161 + return self.ctx.get_hf_processor(DeepseekOCRProcessor, **kwargs)

162 +

163 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

164 + return {"image": None}

165 +

166 + def get_num_image_tokens(

167 + self, *, image_width: int, image_height: int, cropping: bool = True

168 + ) -> int:

169 + image_size = IMAGE_SIZE

170 + base_size = BASE_SIZE

171 + patch_size = 16

172 + downsample_ratio = 4

173 +

174 + if CROP_MODE:

175 + if image_width <= 640 and image_height <= 640:

176 + crop_ratio = [1, 1]

177 + else:

178 + # find the closest aspect ratio to the target

179 + crop_ratio = count_tiles(

180 + image_width, image_height, image_size=IMAGE_SIZE

181 + )

182 +

183 + num_width_tiles, num_height_tiles = crop_ratio

184 + else:

185 + num_width_tiles = num_height_tiles = 1

186 +

187 + h = w = math.ceil((base_size // patch_size) / downsample_ratio)

188 +

189 + h2 = w2 = math.ceil((image_size // patch_size) / downsample_ratio)

190 +

191 + global_views_tokens = h * (w + 1)

192 + if num_width_tiles > 1 or num_height_tiles > 1:

193 + local_views_tokens = (num_height_tiles * h2) * (num_width_tiles * w2 + 1)

194 + else:

195 + local_views_tokens = 0

196 +

197 + return global_views_tokens + local_views_tokens + 1

198 +

199 + def get_image_size_with_most_features(self) -> ImageSize:

200 + if IMAGE_SIZE == 1024 and BASE_SIZE == 1280:

201 + return ImageSize(width=1024 * 2, height=1024 * 2)

202 + return ImageSize(width=640 * 2, height=640 * 2)

203 +

204 +

205 +class DeepseekOCRDummyInputsBuilder(BaseDummyInputsBuilder[DeepseekOCRProcessingInfo]):

206 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

207 + num_images = mm_counts.get("image", 0)

208 +

209 + processor = self.info.get_hf_processor()

210 + image_token = processor.image_token

211 +

212 + return image_token * num_images

213 +

214 + def get_dummy_mm_data(

215 + self,

216 + seq_len: int,

217 + mm_counts: Mapping[str, int],

218 + mm_options: Mapping[str, BaseDummyOptions] | None = None,

219 + ) -> MultiModalDataDict:

220 + num_images = mm_counts.get("image", 0)

221 +

222 + max_image_size = self.info.get_image_size_with_most_features()

223 +

224 + return {

225 + "image": self._get_dummy_images(

226 + width=max_image_size.width,

227 + height=max_image_size.height,

228 + num_images=num_images,

229 + )

230 + }

231 +

232 +

233 +class DeepseekOCRMultiModalProcessor(

234 + BaseMultiModalProcessor[DeepseekOCRProcessingInfo]

235 +):

236 + def _call_hf_processor(

237 + self,

238 + prompt: str,

239 + mm_data: Mapping[str, object],

240 + mm_kwargs: Mapping[str, object],

241 + tok_kwargs: Mapping[str, object],

242 + ) -> BatchFeature:

243 + if mm_data:

244 + processed_outputs = self.info.ctx.call_hf_processor(

245 + self.info.get_hf_processor(**mm_kwargs),

246 + dict(prompt=prompt, **mm_data),

247 + mm_kwargs,

248 + )

249 +

250 + else:

251 + tokenizer = self.info.get_tokenizer()

252 + processed_outputs = tokenizer(

253 + prompt, add_special_tokens=True, return_tensors="pt"

254 + )

255 +

256 + return processed_outputs

257 +

258 + def _get_mm_fields_config(

259 + self,

260 + hf_inputs: BatchFeature,

261 + hf_processor_mm_kwargs: Mapping[str, object],

262 + ) -> Mapping[str, MultiModalFieldConfig]:

263 + return dict(

264 + pixel_values=MultiModalFieldConfig.batched("image"),

265 + images_spatial_crop=MultiModalFieldConfig.batched("image"),

266 + images_crop=MultiModalFieldConfig.batched("image"),

267 + )

268 +

269 + def _get_prompt_updates(

270 + self,

271 + mm_items: MultiModalDataItems,

272 + hf_processor_mm_kwargs: Mapping[str, object],

273 + out_mm_kwargs: MultiModalKwargs,

274 + ) -> Sequence[PromptUpdate]:

275 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

276 +

277 + image_token_id = hf_processor.image_token_id

278 + assert isinstance(image_token_id, int)

279 +

280 + def get_replacement_deepseek_vl2(item_idx: int):

281 + images = mm_items.get_items(

282 + "image", (ImageEmbeddingItems, ImageProcessorItems)

283 + )

284 +

285 + if isinstance(images, ImageEmbeddingItems):

286 + num_image_tokens = images.get_feature_size(item_idx)

287 + else:

288 + size = images.get_image_size(item_idx)

289 +

290 + num_image_tokens = self.info.get_num_image_tokens(

291 + image_width=size.width,

292 + image_height=size.height,

293 + cropping=CROP_MODE,

294 + )

295 + return [image_token_id] * num_image_tokens

296 +

297 + return [

298 + PromptReplacement(

299 + modality="image",

300 + target=[image_token_id],

301 + replacement=get_replacement_deepseek_vl2,

302 + )

303 + ]

304 +

305 + # TODO(Isotr0py): Check if we still need this workaround for

306 + # deepseek-ocr processor.

307 + # def _cached_apply_hf_processor(

308 + # self,

309 + # prompt: str | list[int],

310 + # mm_data_items: MultiModalDataItems,

311 + # hf_processor_mm_kwargs: Mapping[str, object],

312 + # tokenization_kwargs: Mapping[str, object],

313 + # mm_uuids: MultiModalUUIDDict | None = None,

314 + # ) -> tuple[list[int], MultiModalKwargs, bool]:

315 + # # The processor logic is different for len(images) <= 2 vs > 2

316 + # # Since the processing cache assumes that the processor output is

317 + # # invariant of how many images are passed per prompt, we only

318 + # # perform caching for the most common case

319 + # if mm_data_items.get_count("image", strict=False) > 2:

320 + # # This code path corresponds to the cache being disabled

321 + # return self._apply_hf_processor_main(

322 + # prompt=prompt,

323 + # mm_items=mm_data_items,

324 + # hf_processor_mm_kwargs=hf_processor_mm_kwargs,

325 + # enable_hf_prompt_update=True,

326 + # )

327 +

328 + # return super()._cached_apply_hf_processor(

329 + # prompt=prompt,

330 + # mm_data_items=mm_data_items,

331 + # hf_processor_mm_kwargs=hf_processor_mm_kwargs,

332 + # )

333 +

334 +

335 +@MULTIMODAL_REGISTRY.register_processor(

336 + DeepseekOCRMultiModalProcessor,

337 + info=DeepseekOCRProcessingInfo,

338 + dummy_inputs=DeepseekOCRDummyInputsBuilder,

339 +)

340 +class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

341 + hf_to_vllm_mapper = WeightsMapper(

342 + orig_to_new_prefix={

343 + # map prefix for language backbone

344 + "model.embed_tokens.": "language_model.model.embed_tokens.",

345 + "model.layers.": "language_model.model.layers.",

346 + "model.norm.": "language_model.model.norm.",

347 + "lm_head.": "language_model.lm_head.",

348 + # remove "model." prefix for other components

349 + "model.": "",

350 + }

351 + )

352 +

353 + @classmethod

354 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

355 + if modality.startswith("image"):

356 + return "<image>"

357 +

358 + raise ValueError("Only image modality is supported")

359 +

360 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

361 + super().__init__()

362 +

363 + config: DeepseekVLV2Config = vllm_config.model_config.hf_config

364 + quant_config = vllm_config.quant_config

365 + multimodal_config = vllm_config.model_config.multimodal_config

366 +

367 + self.config = config

368 + self.multimodal_config = multimodal_config

369 +

370 + self.vision_config = config.vision_config

371 + self.projector_config = config.projector_config

372 + self.text_config = config.text_config

373 +

374 + model_config = vllm_config.model_config

375 + tokenizer = cached_tokenizer_from_config(model_config)

376 + self.image_token_id = tokenizer.vocab[_IMAGE_TOKEN]

377 +

378 + self.sam_model = build_sam_vit_b()

379 + clip_vision_config = CLIPVisionConfig(

380 + hidden_size=1024,

381 + intermediate_size=4096,

382 + num_attention_heads=16,

383 + num_hidden_layers=24,

384 + image_size=224,

385 + patch_size=14,

386 + projection_dim=512,

387 + layer_norm_eps=1e-5,

388 + )

389 + self.vision_model = DeepCLIPVisionTransformer(

390 + config=clip_vision_config,

391 + quant_config=quant_config,

392 + )

393 +

394 + self.projector = MlpProjector(self.projector_config)

395 + self.tile_tag = config.tile_tag

396 + self.global_view_pos = config.global_view_pos

397 +

398 + # special token for image token sequence format

399 + n_embed = self.projector_config.n_embed

400 + embed_std = 1 / torch.sqrt(torch.tensor(n_embed, dtype=torch.float32))

401 + if self.tile_tag == "2D":

402 + # <|view_separator|>, <|\n|>

403 + self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std)

404 + # This is a typo in original implementation

405 + self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std)

406 + else:

407 + raise ValueError(

408 + f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"

409 + )

410 +

411 + if self.text_config.topk_method == "noaux_tc":

412 + architectures = ["DeepseekV3ForCausalLM"]

413 + elif not self.text_config.use_mla:

414 + architectures = ["DeepseekForCausalLM"]

415 + else:

416 + architectures = ["DeepseekV2ForCausalLM"]

417 +

418 + self.language_model = init_vllm_registered_model(

419 + vllm_config=vllm_config,

420 + hf_config=self.text_config,

421 + prefix=maybe_prefix(prefix, "language_model"),

422 + architectures=architectures,

423 + )

424 +

425 + self.make_empty_intermediate_tensors = (

426 + self.language_model.make_empty_intermediate_tensors

427 + )

428 +

429 + def _parse_and_validate_image_input(self, **kwargs: object):

430 + pixel_values = kwargs.pop("pixel_values", None)

431 + images_spatial_crop = kwargs.pop("images_spatial_crop", None)

432 + images_crop = kwargs.pop("images_crop", None)

433 +

434 + if pixel_values is None or torch.sum(pixel_values).item() == 0:

435 + return None

436 +

437 + if pixel_values is not None:

438 + if not isinstance(pixel_values, (torch.Tensor, list)):

439 + raise ValueError(

440 + f"Incorrect type of pixel values. Got type: {type(pixel_values)}"

441 + )

442 +

443 + if not isinstance(images_spatial_crop, (torch.Tensor, list)):

444 + raise ValueError(

445 + "Incorrect type of image sizes. "

446 + f"Got type: {type(images_spatial_crop)}"

447 + )

448 +

449 + if not isinstance(images_crop, (torch.Tensor, list)):

450 + raise ValueError(

451 + f"Incorrect type of image crop. Got type: {type(images_crop)}"

452 + )

453 +

454 + return [pixel_values, images_crop, images_spatial_crop]

455 +

456 + raise AssertionError("This line should be unreachable.")

457 +

458 + def _encode_global_features(self, image_tensor: torch.Tensor) -> torch.Tensor:

459 + global_features_1 = self.sam_model(image_tensor)

460 + global_features_2 = self.vision_model(image_tensor, global_features_1)

461 + features = torch.cat(

462 + (

463 + global_features_2[:, 1:],

464 + global_features_1.flatten(2).permute(0, 2, 1),

465 + ),

466 + dim=-1,

467 + )

468 + features = self.projector(features)

469 +

470 + _, hw, dim = features.shape

471 + side = int(hw**0.5)

472 +

473 + features = features.view(side, side, dim)

474 + newline = self.image_newline[None, None, :].expand(side, 1, dim)

475 + features = torch.cat([features, newline], dim=1)

476 + return features.view(-1, dim)

477 +

478 + def _encode_local_features(

479 + self, patches: torch.Tensor, crop_shape: torch.Tensor

480 + ) -> torch.Tensor | None:

481 + if torch.sum(patches).item() == 0:

482 + return None

483 +

484 + local_features_1 = self.sam_model(patches)

485 + local_features_2 = self.vision_model(patches, local_features_1)

486 + features = torch.cat(

487 + (

488 + local_features_2[:, 1:],

489 + local_features_1.flatten(2).permute(0, 2, 1),

490 + ),

491 + dim=-1,

492 + )

493 + features = self.projector(features)

494 +

495 + _, hw, dim = features.shape

496 + patch_side = int(hw**0.5)

497 +

498 + width_tiles = int(crop_shape[0].item())

499 + height_tiles = int(crop_shape[1].item())

500 +

501 + features = (

502 + features.view(height_tiles, width_tiles, patch_side, patch_side, dim)

503 + .permute(0, 2, 1, 3, 4)

504 + .reshape(height_tiles * patch_side, width_tiles * patch_side, dim)

505 + )

506 + newline = self.image_newline[None, None, :].expand(

507 + height_tiles * patch_side, 1, dim

508 + )

509 + features = torch.cat([features, newline], dim=1)

510 +

511 + return features.view(-1, dim)

512 +

513 + def _pixel_values_to_embedding(

514 + self,

515 + pixel_values: torch.Tensor,

516 + images_crop: torch.Tensor,

517 + images_spatial_crop: torch.Tensor,

518 + ) -> NestedTensors:

519 + images_in_this_batch = []

520 +

521 + for jdx in range(images_spatial_crop.size(0)):

522 + patches = images_crop[jdx][0].to(torch.bfloat16)

523 + image_ori = pixel_values[jdx]

524 + crop_shape = images_spatial_crop[jdx][0]

525 +

526 + global_features = self._encode_global_features(image_ori)

527 + local_features = self._encode_local_features(patches, crop_shape)

528 +

529 + if local_features is not None:

530 + combined = torch.cat(

531 + [local_features, global_features, self.view_seperator[None, :]],

532 + dim=0,

533 + )

534 + else:

535 + combined = torch.cat(

536 + [global_features, self.view_seperator[None, :]], dim=0

537 + )

538 +

539 + images_in_this_batch.append(combined)

540 +

541 + return images_in_this_batch

542 +

543 + def _process_image_input(self, image_input) -> torch.Tensor:

544 + pixel_values = image_input[0].to(torch.bfloat16)

545 + images_crop = image_input[1]

546 + images_spatial_crop = image_input[2].to(dtype=torch.long)

547 +

548 + vision_features = self._pixel_values_to_embedding(

549 + pixel_values=pixel_values,

550 + images_crop=images_crop,

551 + images_spatial_crop=images_spatial_crop,

552 + )

553 +

554 + return vision_features

555 +

556 + def get_language_model(self) -> torch.nn.Module:

557 + return self.language_model

558 +

559 + def get_multimodal_embeddings(

560 + self, **kwargs: object

561 + ) -> MultiModalEmbeddings | None:

562 + image_input = self._parse_and_validate_image_input(**kwargs)

563 + if image_input is None:

564 + return None

565 + vision_embeddings = self._process_image_input(image_input)

566 + return vision_embeddings

567 +

568 + def forward(

569 + self,

570 + input_ids: torch.Tensor,

571 + positions: torch.Tensor,

572 + intermediate_tensors: IntermediateTensors | None = None,

573 + inputs_embeds: torch.Tensor | None = None,

574 + **kwargs: object,

575 + ):

576 + if intermediate_tensors is not None:

577 + inputs_embeds = None

578 +

579 + hidden_states = self.language_model(

580 + input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds

581 + )

582 +

583 + return hidden_states

584 +

585 + def compute_logits(

586 + self,

587 + hidden_states: torch.Tensor,

588 + ) -> torch.Tensor | None:

589 + return self.language_model.compute_logits(hidden_states)

590 +

591 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

592 + loader = AutoWeightsLoader(self)

593 + autoloaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

594 + return autoloaded_weights

vllm/model_executor/models/deepseek_vl2.py

+23 -20

@@ -101,9 +101,10 @@ class MlpProjector(nn.Module):

101 super().__init__()

102

103 self.cfg = cfg

104 + self.projector_type = cfg.projector_type

105 assert not cfg.token_pooling, "Token pooling is not supported currently."

106

106 - if cfg.projector_type == "downsample_mlp_gelu":

107 + if self.projector_type == "downsample_mlp_gelu":

108 mlp_depth = cfg.depth

109 mlp_ratio = cfg.mlp_ratio

110 modules = [

@@ -120,7 +121,8 @@ class MlpProjector(nn.Module):

121 modules.append(nn.GELU())

122 modules.append(nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed))

123 modules = nn.Sequential(*modules)

123 -

124 + elif self.projector_type == "linear":

125 + modules = nn.Linear(cfg.input_dim, cfg.n_embed)

126 else:

127 raise NotImplementedError(

128 f"Unsupported projector type: {cfg.projector_type}"

@@ -130,24 +132,25 @@ class MlpProjector(nn.Module):

132

133 def forward(self, x):

134 bs, hw, input_dim = x.shape

133 - h = w = int((hw) ** 0.5)

134 - """compute padding"""

135 - if h % self.cfg.downsample_ratio:

136 - pad = self.cfg.downsample_ratio - h % self.cfg.downsample_ratio

137 - else:

138 - pad = 0

139 - x = x.reshape(bs, h, w, input_dim)

140 - if pad > 0:

141 - x = F.pad(x, (0, 0, 0, pad, 0, pad), "constant", 0)

142 - """4 to 1 concat"""

143 - x = x.permute(0, 3, 1, 2) # B, C, H, W

144 - x = F.unfold(

145 - x,

146 - kernel_size=self.cfg.downsample_ratio,

147 - stride=self.cfg.downsample_ratio,

148 - padding=0,

149 - ) # B, C*4, HW // 4

150 - x = x.permute(0, 2, 1)

135 + if self.projector_type == "downsample_mlp_gelu":

136 + h = w = int((hw) ** 0.5)

137 + """compute padding"""

138 + if h % self.cfg.downsample_ratio:

139 + pad = self.cfg.downsample_ratio - h % self.cfg.downsample_ratio

140 + else:

141 + pad = 0

142 + x = x.reshape(bs, h, w, input_dim)

143 + if pad > 0:

144 + x = F.pad(x, (0, 0, 0, pad, 0, pad), "constant", 0)

145 + """4 to 1 concat"""

146 + x = x.permute(0, 3, 1, 2) # B, C, H, W

147 + x = F.unfold(

148 + x,

149 + kernel_size=self.cfg.downsample_ratio,

150 + stride=self.cfg.downsample_ratio,

151 + padding=0,

152 + ) # B, C*4, HW // 4

153 + x = x.permute(0, 2, 1)

154

155 return self.layers(x)

156

vllm/model_executor/models/registry.py

@@ -258,6 +258,7 @@ _MULTIMODAL_MODELS = {

258 "Cohere2VisionForConditionalGeneration",

259 ),

260 "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),

261 + "DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"),

262 "DotsOCRForCausalLM": ("dots_ocr", "DotsOCRForCausalLM"),

263 "Ernie4_5_VLMoeForConditionalGeneration": (

264 "ernie45_vl",

vllm/transformers_utils/chat_templates/registry.py

@@ -34,6 +34,7 @@ _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK: dict[str, ChatTemplatePath] = {

34 "clip": CHAT_TEMPLATES_DIR / "template_basic.jinja",

35 "chameleon": CHAT_TEMPLATES_DIR / "template_basic.jinja",

36 "deepseek_vl_v2": CHAT_TEMPLATES_DIR / "template_deepseek_vl2.jinja",

37 + "deepseek_ocr": CHAT_TEMPLATES_DIR / "template_deepseek_ocr.jinja",

38 "fuyu": CHAT_TEMPLATES_DIR / "template_fuyu.jinja",

39 "minicpmv": _get_minicpmv_chat_template_fallback,

40 "paligemma": CHAT_TEMPLATES_DIR / "template_basic.jinja",

vllm/transformers_utils/chat_templates/template_deepseek_ocr.jinja new

+14

@@ -0,0 +1,14 @@

1 +{%- if messages[0]['role'] == 'system' -%}

2 + {%- set system_message = messages[0]['content'] -%}

3 + {%- set messages = messages[1:] -%}

4 +{%- else -%}

5 + {% set system_message = '' -%}

6 +{%- endif -%}

7 +

8 +{{ bos_token + system_message }}

9 +{%- for message in messages -%}

10 + {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}

11 + {{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}

12 + {%- endif -%}

13 + {{ message['content'] }}

14 +{%- endfor -%}

vllm/transformers_utils/processors/deepseek_ocr.py new

+442

@@ -0,0 +1,442 @@

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 +import torchvision.transforms as T

8 +from PIL import Image, ImageOps

9 +from transformers import AutoProcessor, BatchFeature, LlamaTokenizerFast

10 +from transformers.processing_utils import ProcessorMixin

11 +

12 +# TODO(Isotr0py): change modes for variants

13 +# see: https://github.com/deepseek-ai/DeepSeek-OCR/blob/8cf003d38821fa1b19c73da3bd1b0dc262ea8136/DeepSeek-OCR-master/DeepSeek-OCR-vllm/config.py#L1-L6

14 +# Tiny: base_size = 512, image_size = 512, crop_mode = False

15 +# Small: base_size = 640, image_size = 640, crop_mode = False

16 +# Base: base_size = 1024, image_size = 1024, crop_mode = False

17 +# Large: base_size = 1280, image_size = 1280, crop_mode = False

18 +# Gundam: base_size = 1024, image_size = 640, crop_mode = True

19 +BASE_SIZE = 1024

20 +IMAGE_SIZE = 640

21 +CROP_MODE = True

22 +

23 +# TODO(Isotr0py): Expose as mm_kwargs

24 +MIN_CROPS = 2

25 +MAX_CROPS = 6 # max:9; If your GPU memory is small, it is recommended to set it to 6.

26 +

27 +

28 +def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):

29 + best_ratio_diff = float("inf")

30 + best_ratio = (1, 1)

31 + area = width * height

32 + for ratio in target_ratios:

33 + target_aspect_ratio = ratio[0] / ratio[1]

34 + ratio_diff = abs(aspect_ratio - target_aspect_ratio)

35 + if ratio_diff < best_ratio_diff:

36 + best_ratio_diff = ratio_diff

37 + best_ratio = ratio

38 + elif ratio_diff == best_ratio_diff:

39 + if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:

40 + best_ratio = ratio

41 + return best_ratio

42 +

43 +

44 +def calculate_aspect_ratios(

45 + min_num: int = MIN_CROPS, max_num: int = MAX_CROPS

46 +) -> list[tuple[int, int]]:

47 + target_ratios: set[tuple[int, int]] = set(

48 + (i, j)

49 + for n in range(min_num, max_num + 1)

50 + for i in range(1, n + 1)

51 + for j in range(1, n + 1)

52 + if i * j <= max_num and i * j >= min_num

53 + )

54 + sorted_target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

55 + return sorted_target_ratios

56 +

57 +

58 +def count_tiles(

59 + orig_width,

60 + orig_height,

61 + min_num=MIN_CROPS,

62 + max_num=MAX_CROPS,

63 + image_size=640,

64 + use_thumbnail=False,

65 +):

66 + aspect_ratio = orig_width / orig_height

67 +

68 + # calculate the existing image aspect ratio

69 + target_ratios = calculate_aspect_ratios(min_num, max_num)

70 +

71 + # find the closest aspect ratio to the target

72 + target_aspect_ratio = find_closest_aspect_ratio(

73 + aspect_ratio, target_ratios, orig_width, orig_height, image_size

74 + )

75 +

76 + return target_aspect_ratio

77 +

78 +

79 +def dynamic_preprocess(

80 + image, min_num=MIN_CROPS, max_num=MAX_CROPS, image_size=640, use_thumbnail=False

81 +):

82 + orig_width, orig_height = image.size

83 + aspect_ratio = orig_width / orig_height

84 +

85 + # calculate the existing image aspect ratio

86 + target_ratios = calculate_aspect_ratios(min_num, max_num)

87 +

88 + # find the closest aspect ratio to the target

89 + target_aspect_ratio = find_closest_aspect_ratio(

90 + aspect_ratio, target_ratios, orig_width, orig_height, image_size

91 + )

92 +

93 + # calculate the target width and height

94 + target_width = image_size * target_aspect_ratio[0]

95 + target_height = image_size * target_aspect_ratio[1]

96 + blocks = target_aspect_ratio[0] * target_aspect_ratio[1]

97 +

98 + # resize the image

99 + resized_img = image.resize((target_width, target_height))

100 + processed_images = []

101 + for i in range(blocks):

102 + box = (

103 + (i % (target_width // image_size)) * image_size,

104 + (i // (target_width // image_size)) * image_size,

105 + ((i % (target_width // image_size)) + 1) * image_size,

106 + ((i // (target_width // image_size)) + 1) * image_size,

107 + )

108 + # split the image

109 + split_img = resized_img.crop(box)

110 + processed_images.append(split_img)

111 + assert len(processed_images) == blocks

112 + if use_thumbnail and len(processed_images) != 1:

113 + thumbnail_img = image.resize((image_size, image_size))

114 + processed_images.append(thumbnail_img)

115 + return processed_images, target_aspect_ratio

116 +

117 +

118 +class ImageTransform:

119 + def __init__(

120 + self,

121 + mean: tuple[float, float, float] = (0.5, 0.5, 0.5),

122 + std: tuple[float, float, float] = (0.5, 0.5, 0.5),

123 + normalize: bool = True,

124 + ):

125 + self.mean = mean

126 + self.std = std

127 + self.normalize = normalize

128 +

129 + transform_pipelines = [T.ToTensor()]

130 +

131 + if normalize:

132 + transform_pipelines.append(T.Normalize(mean, std))

133 +

134 + self.transform = T.Compose(transform_pipelines)

135 +

136 + def __call__(self, pil_img: Image.Image):

137 + x = self.transform(pil_img)

138 + return x

139 +

140 +

141 +class DeepseekOCRProcessor(ProcessorMixin):

142 + tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast")

143 + attributes = ["tokenizer"]

144 +

145 + def __init__(

146 + self,

147 + tokenizer: LlamaTokenizerFast,

148 + patch_size: int = 16,

149 + downsample_ratio: int = 4,

150 + image_mean: tuple[float, float, float] = (0.5, 0.5, 0.5),

151 + image_std: tuple[float, float, float] = (0.5, 0.5, 0.5),

152 + normalize: bool = True,

153 + image_token: str = "<image>",

154 + pad_token: str = "<｜▁pad▁｜>",

155 + add_special_token: bool = False,

156 + sft_format: str = "deepseek",

157 + mask_prompt: bool = True,

158 + ignore_id: int = -100,

159 + **kwargs,

160 + ):

161 + self.image_size = IMAGE_SIZE

162 + self.base_size = BASE_SIZE

163 + self.patch_size = 16

164 + self.image_mean = image_mean

165 + self.image_std = image_std

166 + self.normalize = normalize

167 + self.downsample_ratio = 4

168 +

169 + self.image_transform = ImageTransform(

170 + mean=image_mean, std=image_std, normalize=normalize

171 + )

172 +

173 + self.tokenizer = tokenizer

174 + self.tokenizer.padding_side = "left" # must set this，padding side with make a difference in batch inference # noqa: E501

175 +

176 + # add the pad_token as special token to use 'tokenizer.pad_token'

177 + # and 'tokenizer.pad_token_id'

178 + if self.tokenizer.pad_token is None:

179 + self.tokenizer.add_special_tokens({"pad_token": pad_token})

180 +

181 + # add image token

182 + self.image_token_id = self.tokenizer.vocab.get(image_token)

183 + self.image_token = image_token

184 + self.pad_token = pad_token

185 + self.add_special_token = add_special_token

186 + self.sft_format = sft_format

187 + self.mask_prompt = mask_prompt

188 + self.ignore_id = ignore_id

189 +

190 + super().__init__(

191 + tokenizer,

192 + **kwargs,

193 + )

194 +

195 + @property

196 + def bos_id(self):

197 + return self.tokenizer.bos_token_id

198 +

199 + @property

200 + def eos_id(self):

201 + return self.tokenizer.eos_token_id

202 +

203 + @property

204 + def pad_id(self):

205 + return self.tokenizer.pad_token_id

206 +

207 + def encode(self, text: str, bos: bool = True, eos: bool = False):

208 + t = self.tokenizer.encode(text, add_special_tokens=False)

209 + if bos:

210 + t = [self.bos_id] + t

211 + if eos:

212 + t = t + [self.eos_id]

213 + return t

214 +

215 + def decode(self, t: list[int], **kwargs) -> str:

216 + return self.tokenizer.decode(t, **kwargs)

217 +

218 + def process_one(

219 + self,

220 + prompt: str,

221 + images: list[Image.Image],

222 + crop_mode: bool = CROP_MODE,

223 + ):

224 + """

225 +

226 + Args:

227 + prompt (str): the formatted prompt;

228 + images (List[ImageType]): the list of images;

229 + crop_mode (bool): if True, then crop the image;

230 +

231 + Returns:

232 + outputs (BaseProcessorOutput): the output of the processor,

233 + - input_ids (torch.LongTensor): [N + image tokens]

234 + - target_ids (torch.LongTensor): [N + image tokens]

235 + - pixel_values (torch.FloatTensor): [n_patches, 3, H, W]

236 + - image_id (int): the id of the image token

237 + - num_image_tokens (List[int]): the number of image tokens

238 + """

239 +

240 + assert prompt is not None and images is not None, (

241 + "prompt and images must be used at the same time."

242 + )

243 +

244 + sft_format = prompt

245 +

246 + (

247 + input_ids,

248 + pixel_values,

249 + images_crop,

250 + images_seq_mask,

251 + images_spatial_crop,

252 + num_image_tokens,

253 + _,

254 + ) = self.tokenize_with_images(

255 + conversation=sft_format,

256 + images=images,

257 + bos=True,

258 + eos=True,

259 + cropping=crop_mode,

260 + )

261 +

262 + prepare = BatchFeature(

263 + data=dict(

264 + input_ids=input_ids,

265 + pixel_values=pixel_values,

266 + images_crop=images_crop,

267 + images_seq_mask=images_seq_mask,

268 + images_spatial_crop=images_spatial_crop,

269 + num_image_tokens=num_image_tokens,

270 + ),

271 + tensor_type="pt",

272 + )

273 + return prepare

274 +

275 + def __call__(

276 + self,

277 + *,

278 + prompt: str,

279 + images: list[Image.Image],

280 + crop_mode: bool = CROP_MODE,

281 + **kwargs,

282 + ):

283 + prepare = self.process_one(

284 + prompt=prompt,

285 + images=images,

286 + crop_mode=crop_mode,

287 + )

288 +

289 + return prepare

290 +

291 + def tokenize_with_images(

292 + self,

293 + conversation: str,

294 + images: list[Image.Image],

295 + bos: bool = True,

296 + eos: bool = True,

297 + cropping: bool = True,

298 + ):

299 + """Tokenize text with <image> tags."""

300 +

301 + assert conversation.count(self.image_token) == len(images)

302 + text_splits = conversation.split(self.image_token)

303 + images_list, images_crop_list, images_seq_mask, images_spatial_crop = (

304 + [],

305 + [],

306 + [],

307 + [],

308 + )

309 + image_shapes = []

310 + num_image_tokens = []

311 + tokenized_str = []

312 + for text_sep, image in zip(text_splits, images):

313 + tokenized_sep = self.encode(text_sep, bos=False, eos=False)

314 + tokenized_str += tokenized_sep

315 + images_seq_mask += [False] * len(tokenized_sep)

316 +

317 + image_shapes.append(image.size)

318 +

319 + images_crop_raw = []

320 + if image.size[0] <= 640 and image.size[1] <= 640:

321 + crop_ratio = [1, 1]

322 + elif cropping:

323 + images_crop_raw, crop_ratio = dynamic_preprocess(

324 + image, image_size=IMAGE_SIZE

325 + )

326 + else:

327 + crop_ratio = [1, 1]

328 +

329 + if self.image_size <= 640 and not cropping:

330 + image = image.resize((self.image_size, self.image_size))

331 +

332 + global_view = ImageOps.pad(

333 + image,

334 + (self.base_size, self.base_size),

335 + color=tuple(int(x * 255) for x in self.image_transform.mean),

336 + )

337 + images_list.append(self.image_transform(global_view))

338 +

339 + num_width_tiles, num_height_tiles = crop_ratio

340 + images_spatial_crop.append([num_width_tiles, num_height_tiles])

341 +

342 + if num_width_tiles > 1 or num_height_tiles > 1:

343 + for cropped_image in images_crop_raw:

344 + images_crop_list.append(self.image_transform(cropped_image))

345 +

346 + num_queries = math.ceil(

347 + (self.image_size // self.patch_size) / self.downsample_ratio

348 + )

349 + num_queries_base = math.ceil(

350 + (self.base_size // self.patch_size) / self.downsample_ratio

351 + )

352 +

353 + tokenized_image = (

354 + [self.image_token_id] * num_queries_base + [self.image_token_id]

355 + ) * num_queries_base

356 + tokenized_image += [self.image_token_id]

357 + if num_width_tiles > 1 or num_height_tiles > 1:

358 + local_row = [self.image_token_id] * (num_queries * num_width_tiles + 1)

359 + tokenized_image += local_row * (num_queries * num_height_tiles)

360 + tokenized_str += tokenized_image

361 + images_seq_mask += [True] * len(tokenized_image)

362 + num_image_tokens.append(len(tokenized_image))

363 +

364 + """process the last text split"""

365 + tokenized_sep = self.encode(text_splits[-1], bos=False, eos=False)

366 + tokenized_str += tokenized_sep

367 + images_seq_mask += [False] * len(tokenized_sep)

368 +

369 + """add the bos and eos tokens"""

370 + if bos:

371 + tokenized_str = [self.bos_id] + tokenized_str

372 + images_seq_mask = [False] + images_seq_mask

373 + if eos:

374 + tokenized_str = tokenized_str + [self.eos_id]

375 + images_seq_mask = images_seq_mask + [False]

376 +

377 + assert len(tokenized_str) == len(images_seq_mask), (

378 + f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} "

379 + f"is not equal to images_seq_mask's length {len(images_seq_mask)}."

380 + )

381 +

382 + masked_tokenized_str = []

383 + for token_index in tokenized_str:

384 + if token_index != self.image_token_id:

385 + masked_tokenized_str.append(token_index)

386 + else:

387 + masked_tokenized_str.append(self.ignore_id)

388 +

389 + assert (

390 + len(tokenized_str) == len(images_seq_mask) == len(masked_tokenized_str)

391 + ), (

392 + f"tokenized_str's length {len(tokenized_str)}, "

393 + f"input_ids' length {len(masked_tokenized_str)}, "

394 + f"images_seq_mask's length {len(images_seq_mask)}, are not equal."

395 + )

396 +

397 + input_ids = torch.LongTensor(tokenized_str)

398 + target_ids = torch.LongTensor(masked_tokenized_str)

399 + images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool)

400 +

401 + # set input_ids < 0 | input_ids == self.image_token_id as ignore_id

402 + target_ids[(input_ids < 0) | (input_ids == self.image_token_id)] = (

403 + self.ignore_id

404 + )

405 + input_ids[input_ids < 0] = self.pad_id

406 +

407 + # Remove the ending eos token

408 + assert input_ids[-1] == self.eos_id

409 + input_ids = input_ids[:-1]

410 + target_ids = target_ids[:-1]

411 + images_seq_mask = images_seq_mask[:-1]

412 +

413 + if len(images_list) == 0:

414 + pixel_values = torch.zeros((1, 3, self.base_size, self.base_size))

415 + images_spatial_crop = torch.zeros((1, 1), dtype=torch.long)

416 + images_crop = torch.zeros(

417 + (1, 3, self.image_size, self.image_size)

418 + ).unsqueeze(0)

419 + else:

420 + pixel_values = torch.stack(images_list, dim=0)

421 + images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long)

422 + if images_crop_list:

423 + images_crop = torch.stack(images_crop_list, dim=0).unsqueeze(0)

424 + else:

425 + images_crop = torch.zeros(

426 + (1, 3, self.image_size, self.image_size)

427 + ).unsqueeze(0)

428 +

429 + input_ids = input_ids.unsqueeze(0)

430 +

431 + return (

432 + input_ids,

433 + pixel_values,

434 + images_crop,

435 + images_seq_mask,

436 + images_spatial_crop,

437 + num_image_tokens,

438 + image_shapes,

439 + )

440 +

441 +

442 +AutoProcessor.register("DeepseekOCRProcessor", DeepseekOCRProcessor)