@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 54072f315f
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 15, 12:15 UTC

[MODEL ADDITION] Ovis2 Model Addition (#15826)

Signed-off-by: Marco <121761685+mlinmg@users.noreply.github.com> Signed-off-by: Isotr0py <2037008807@qq.com> Signed-off-by: isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Marco committed Apr 30, 2025 at 09:33 UTC 54072f315f8f39f906e59b2153e05e9cf7cd524b
17 files changed +1349 -7
docs/source/models/supported_models.md
+7
@@ -1014,6 +1014,13 @@ See [this page](#generative-models) for more information on how to use generativ
1014 *
1015 * ✅︎
1016 * ✅︎
1017 +- * `Ovis2ForConditionalGeneration`<sup>^</sup>
1018 + * Ovis2
1019 + * T + I<sup>+</sup>
1020 + * `AIDC-AI/Ovis2-1B`, `AIDC-AI/Ovis2-2B`, etc.
1021 + *
1022 + *
1023 + * ✅︎
1024 - * `PaliGemmaForConditionalGeneration`
1025 * PaliGemma, PaliGemma 2
1026 * T + I<sup>E</sup>
examples/offline_inference/vision_language.py
+31
@@ -725,6 +725,36 @@ def run_nvlm_d(questions: list[str], modality: str) -> ModelRequestData:
725 )
726
727
728 +# Ovis2
729 +def run_ovis2(questions: list[str], modality: str) -> ModelRequestData:
730 + assert modality == "image"
731 +
732 + model_name = "AIDC-AI/Ovis2-1B"
733 + tokenizer = "Isotr0py/Ovis2-tokenizer"
734 +
735 + engine_args = EngineArgs(
736 + model=model_name,
737 + tokenizer=tokenizer,
738 + max_model_len=4096,
739 + max_num_seqs=2,
740 + trust_remote_code=True,
741 + dtype="half",
742 + hf_overrides={"architectures": ["Ovis2ForConditionalGeneration"]},
743 + limit_mm_per_prompt={"image": 1},
744 + )
745 +
746 + placeholder = "<image>\n"
747 + prompts = [("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
748 + f"<|im_start|>user\n{placeholder}"
749 + f"{question}<|im_end|>\n"
750 + "<|im_start|>assistant\n") for question in questions]
751 +
752 + return ModelRequestData(
753 + engine_args=engine_args,
754 + prompts=prompts,
755 + )
756 +
757 +
758 # PaliGemma
759 def run_paligemma(questions: list[str], modality: str) -> ModelRequestData:
760 assert modality == "image"
@@ -1041,6 +1071,7 @@ model_example_map = {
1071 "llama4": run_llama4,
1072 "molmo": run_molmo,
1073 "NVLM_D": run_nvlm_d,
1074 + "ovis2": run_ovis2,
1075 "paligemma": run_paligemma,
1076 "paligemma2": run_paligemma2,
1077 "phi3_v": run_phi3v,
examples/offline_inference/vision_language_multi_image.py
+31
@@ -436,6 +436,36 @@ def load_nvlm_d(question: str, image_urls: list[str]) -> ModelRequestData:
436 )
437
438
439 +# Ovis2
440 +def load_ovis2(question: str, image_urls: list[str]) -> ModelRequestData:
441 + model_name = "AIDC-AI/Ovis2-1B"
442 + tokenizer = "Isotr0py/Ovis2-tokenizer"
443 +
444 + engine_args = EngineArgs(
445 + model=model_name,
446 + tokenizer=tokenizer,
447 + max_model_len=8192,
448 + max_num_seqs=2,
449 + trust_remote_code=True,
450 + dtype="half",
451 + limit_mm_per_prompt={"image": len(image_urls)},
452 + hf_overrides={"architectures": ["Ovis2ForConditionalGeneration"]},
453 + )
454 +
455 + placeholder = '\n'.join(
456 + [f'Image {i+1}: <image>' for i in range(len(image_urls))]) + '\n'
457 + prompt = ("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
458 + f"<|im_start|>user\n{placeholder}"
459 + f"{question}<|im_end|>\n"
460 + "<|im_start|>assistant\n")
461 +
462 + return ModelRequestData(
463 + engine_args=engine_args,
464 + prompt=prompt,
465 + image_data=[fetch_image(url) for url in image_urls],
466 + )
467 +
468 +
469 def load_pixtral_hf(question: str, image_urls: list[str]) -> ModelRequestData:
470 model_name = "mistral-community/pixtral-12b"
471
@@ -685,6 +715,7 @@ model_example_map = {
715 "mistral3": load_mistral3,
716 "mllama": load_mllama,
717 "NVLM_D": load_nvlm_d,
718 + "ovis2": load_ovis2,
719 "phi3_v": load_phi3v,
720 "phi4_mm": load_phi4mm,
721 "pixtral_hf": load_pixtral_hf,
tests/models/decoder_only/vision_language/test_models.py
+12
@@ -467,6 +467,18 @@ VLM_TEST_SETTINGS = {
467 max_num_seqs=2,
468 patch_hf_runner=model_utils.molmo_patch_hf_runner,
469 ),
470 + "ovis2": VLMTestInfo(
471 + models=["AIDC-AI/Ovis2-1B"],
472 + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
473 + prompt_formatter=lambda img_prompt: f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{img_prompt}<|im_end|>\n<|im_start|>assistant\n", # noqa: E501
474 + img_idx_to_prompt=lambda idx: "<image>\n", # noqa: E501
475 + max_model_len=4096,
476 + max_num_seqs=2,
477 + dtype="half",
478 + # use sdpa mode for hf runner since ovis2 didn't work with flash_attn
479 + hf_model_kwargs={"llm_attn_implementation": "sdpa"},
480 + patch_hf_runner=model_utils.ovis2_patch_hf_runner,
481 + ),
482 "phi3v": VLMTestInfo(
483 models=["microsoft/Phi-3.5-vision-instruct"],
484 test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
tests/models/decoder_only/vision_language/vlm_utils/core.py
+1 -1
@@ -67,7 +67,7 @@ def run_test(
67 "disable_mm_preprocessor_cache": True,
68 }
69 if model_info.tokenizer:
70 - vllm_runner_kwargs_["tokenizer"] = model_info.tokenizer
70 + vllm_runner_kwargs_["tokenizer_name"] = model_info.tokenizer
71 if model_info.tokenizer_mode:
72 vllm_runner_kwargs_["tokenizer_mode"] = model_info.tokenizer_mode
73 if model_info.hf_overrides:
tests/models/decoder_only/vision_language/vlm_utils/model_utils.py
+30
@@ -676,3 +676,33 @@ def molmo_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
676 hf_model.model.generate = types.MethodType(_generate, hf_model.model)
677
678 return hf_model
679 +
680 +
681 +def ovis2_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
682 + """Patches and returns an instance of the HfRunner to use for Ovis2."""
683 + hf_model.model.visual_tokenizer.to(hf_model.dtype)
684 + hf_model.model.vte.to(hf_model.dtype)
685 + hf_model.model.llm.to(hf_model.dtype)
686 +
687 + hf_model.model.get_output_embeddings = lambda: \
688 + hf_model.model.llm.get_output_embeddings()
689 +
690 + def processor(*args, text="", images=None, **kwargs):
691 + text_tokenizer = hf_model.model.get_text_tokenizer()
692 + images = [images] if isinstance(images, Image) else images
693 +
694 + text = text.split("<|im_start|>user\n")[1].split("<|im_end|>\n")[0]
695 +
696 + prompt, input_ids, pixel_values = hf_model.model.preprocess_inputs(
697 + text_or_conversations=text, images=images)
698 + attention_mask = torch.ne(input_ids, text_tokenizer.pad_token_id)
699 +
700 + inputs = {
701 + "inputs": input_ids.unsqueeze(0),
702 + "pixel_values": pixel_values.unsqueeze(0),
703 + "attention_mask": attention_mask.unsqueeze(0),
704 + }
705 + return BatchFeature(data=inputs, tensor_type="pt")
706 +
707 + hf_model.processor = processor
708 + return hf_model
tests/models/multimodal/processing/test_common.py
+1
@@ -274,6 +274,7 @@ def _test_processing_correctness_mistral(
274 "allenai/Molmo-7B-D-0924",
275 "allenai/Molmo-7B-O-0924",
276 "nvidia/NVLM-D-72B",
277 + "AIDC-AI/Ovis2-1B",
278 "google/paligemma-3b-mix-224",
279 "google/paligemma2-3b-ft-docci-448",
280 "microsoft/Phi-4-multimodal-instruct",
tests/models/registry.py
+4
@@ -348,6 +348,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
348 max_transformers_version="4.48",
349 transformers_version_reason="Use of deprecated imports which have been removed.", # noqa: E501
350 extras={"phi3.5": "microsoft/Phi-3.5-vision-instruct"}), # noqa: E501
351 + "Ovis2ForConditionalGeneration": _HfExamplesInfo("AIDC-AI/Ovis2-1B",
352 + tokenizer="Isotr0py/Ovis2-tokenizer",
353 + trust_remote_code=True,
354 + hf_overrides={"architectures": ["Ovis2ForConditionalGeneration"]}), # noqa: E501
355 "Phi4MMForCausalLM": _HfExamplesInfo("microsoft/Phi-4-multimodal-instruct",
356 trust_remote_code=True),
357 "PixtralForConditionalGeneration": _HfExamplesInfo("mistralai/Pixtral-12B-2409", # noqa: E501
vllm/entrypoints/chat_utils.py
+3 -2
@@ -496,9 +496,10 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
496 if model_type.startswith("llava"):
497 return self._cached_token_str(self._tokenizer,
498 hf_config.image_token_index)
499 +
500 if model_type in ("aya_vision", "chameleon", "deepseek_vl_v2",
500 - "internvl_chat", "skywork_chat", "NVLM_D",
501 - "h2ovl_chat", "idefics3", "smolvlm"):
501 + "internvl_chat", "ovis2", "skywork_chat",
502 + "NVLM_D", "h2ovl_chat", "idefics3", "smolvlm"):
503 return "<image>"
504 if model_type in ("mllama", "llama4"):
505 return "<|image|>"
vllm/model_executor/models/aimv2.py new
+322
@@ -0,0 +1,322 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# A modified implementation of the AIMv2 Transformer
4 +# inserted here also the image tokenizer used by Ovis2
5 +from typing import Optional
6 +
7 +import torch
8 +from torch import nn, softmax
9 +from torch.nn import functional as F
10 +from torch.nn.functional import gumbel_softmax, pad
11 +
12 +from vllm.model_executor.layers.layernorm import RMSNorm
13 +from vllm.model_executor.layers.linear import ReplicatedLinear
14 +from vllm.model_executor.layers.quantization.base_config import (
15 + QuantizationConfig)
16 +from vllm.transformers_utils.configs.ovis2 import (AIMv2Config,
17 + Aimv2VisualTokenizerConfig)
18 +
19 +IMAGE_INDICATOR_IDS = [-301, -302, -303, -304,
20 + -305] # kept for vocab prefixed tokens
21 +
22 +
23 +def st_argmax(y_soft: torch.Tensor, dim: int): # straight-through softmax
24 + index = y_soft.max(dim, keepdim=True)[1]
25 + y_hard = torch.zeros_like(
26 + y_soft, memory_format=torch.legacy_contiguous_format).scatter_(
27 + dim, index, 1.0)
28 + ret = y_hard - y_soft.detach() + y_soft
29 + return ret
30 +
31 +
32 +class Aimv2VisualTokenizer(torch.nn.Module):
33 +
34 + def __init__(self,
35 + config: Aimv2VisualTokenizerConfig,
36 + quant_config: Optional[QuantizationConfig] = None,
37 + prefix: str = "",
38 + **kwargs):
39 + super().__init__()
40 + self.config = config
41 + self.backbone = AIMv2Model(
42 + config=config.backbone_config, # noqa
43 + quant_config=quant_config,
44 + prefix=f"{prefix}.visual_tokenizer")
45 + # reserved tokens for IMAGE_INDICATORS
46 + head_dim = config.vocab_size - len(IMAGE_INDICATOR_IDS)
47 + self.head = torch.nn.Sequential(
48 + ReplicatedLinear(
49 + config.backbone_config.hidden_size * config.hidden_stride *
50 + config.hidden_stride,
51 + head_dim,
52 + bias=False,
53 + ), torch.nn.LayerNorm(head_dim))
54 +
55 + @property
56 + def dtype(self):
57 + return self.backbone.dtype
58 +
59 + @property
60 + def device(self):
61 + return self.backbone.device
62 +
63 + def tokenize(self, logits):
64 + if self.config.tokenize_function == 'softmax':
65 + tokens = softmax(logits, dim=-1)
66 + elif self.config.tokenize_function == 'gumbel_argmax':
67 + tokens = gumbel_softmax(logits, tau=self.config.tau, hard=True)
68 + elif self.config.tokenize_function == 'st_argmax':
69 + tokens = st_argmax(logits, dim=-1)
70 + else:
71 + raise ValueError(
72 + 'Invalid `max_type`, expected softmax or gumbel_argmax '
73 + f'or st_argmax, but got {self.config.tokenize_function}')
74 + return tokens
75 +
76 + def encode(self, pixel_values):
77 + features = self.backbone(pixel_values)
78 + if self.config.drop_cls_token:
79 + features = features[:, 1:, :]
80 +
81 + # merge number of `hidden_stride * hidden_stride` hidden states together
82 + # to reduce token sequence length
83 + # e.g., for hidden_stride=2, this leads to a token length reduction:
84 + # 1024 -> 256 for aimv2
85 + if self.config.hidden_stride > 1:
86 + # this `d` maybe different from the above `d``
87 + n, L, d = features.shape
88 + sqrt_l = int(L**0.5)
89 + assert sqrt_l**2 == L, (
90 + "The token sequence length should be a perfect square.")
91 + features = features.reshape(n, sqrt_l, sqrt_l, d)
92 + pl = (self.config.hidden_stride -
93 + (sqrt_l %
94 + self.config.hidden_stride)) % self.config.hidden_stride
95 + features = pad(features, (0, 0, 0, pl, 0, pl), "constant", 0)
96 + sqrt_l += pl
97 + features = features.reshape(n, sqrt_l // self.config.hidden_stride,
98 + self.config.hidden_stride,
99 + sqrt_l // self.config.hidden_stride,
100 + self.config.hidden_stride, d)
101 + # [n, sqrt_l/hs, sqrt_l/hs, hs, hs, d]
102 + features = features.permute(0, 1, 3, 2, 4, 5)
103 + # [n, sqrt_l/hs, sqrt_l/hs, hs*hs*d]
104 + features = features.flatten(3)
105 + # [n, sqrt_l/hs*sqrt_l/hs, hs*hs*d]
106 + features = features.reshape(
107 + n, -1,
108 + self.config.hidden_stride * self.config.hidden_stride * d)
109 +
110 + return features
111 +
112 + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
113 + """[BatchSize, ImageShape] -> [BatchSize, Token, VocabSize]"""
114 + features = self.encode(pixel_values)
115 + logits, _ = self.head[0](
116 + features) # we spllit the sequncial here for not throwing an error
117 + logits = self.head[1](logits)
118 + tokens = self.tokenize(logits)
119 + # tokens' shape is [BatchSize, #Token, VocabSize-5], so padding with
120 + # [BatchSize, #Token, 5], after which, tokens' shape should become
121 + # [BatchSize, #Token, VocabSize]
122 + batch_size, token_len, _ = tokens.shape
123 + padding_tensor = torch.zeros(size=(batch_size, token_len,
124 + len(IMAGE_INDICATOR_IDS)),
125 + dtype=tokens.dtype,
126 + device=tokens.device,
127 + layout=tokens.layout,
128 + requires_grad=False)
129 + tokens = torch.cat((tokens, padding_tensor), dim=2)
130 + return tokens
131 +
132 +
133 +class AIMv2SwiGLUFFN(nn.Module):
134 +
135 + def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
136 + prefix: str):
137 + super().__init__()
138 + hidden_features = config.intermediate_size
139 + in_features = config.hidden_size
140 + bias = config.use_bias
141 +
142 + # TODO(Isotr0py): investigate if we can add TP to visual tokenizer
143 + self.fc1 = ReplicatedLinear(in_features,
144 + hidden_features,
145 + bias=bias,
146 + quant_config=quant_config,
147 + prefix=f"{prefix}.fc1")
148 + self.fc2 = ReplicatedLinear(hidden_features,
149 + in_features,
150 + bias=bias,
151 + quant_config=quant_config,
152 + prefix=f"{prefix}.fc2")
153 + self.fc3 = ReplicatedLinear(in_features,
154 + hidden_features,
155 + bias=bias,
156 + quant_config=quant_config,
157 + prefix=f"{prefix}.fc3")
158 +
159 + def forward(self, x: torch.Tensor) -> torch.Tensor:
160 + x_parallel, _ = self.fc1(x)
161 + gate, _ = self.fc3(x)
162 + x_parallel = F.silu(x_parallel) * gate
163 + out, _ = self.fc2(x_parallel)
164 + return out
165 +
166 +
167 +class AIMv2PatchEmbed(nn.Module):
168 +
169 + def __init__(self, config: AIMv2Config):
170 + super().__init__()
171 + self.proj = nn.Conv2d(
172 + config.num_channels,
173 + config.hidden_size,
174 + kernel_size=(config.patch_size, config.patch_size),
175 + stride=(config.patch_size, config.patch_size),
176 + )
177 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
178 +
179 + def forward(self, x: torch.Tensor) -> torch.Tensor:
180 + x = self.proj(x).flatten(2).transpose(1, 2)
181 + x = self.norm.forward_native(x)
182 + return x
183 +
184 +
185 +class AIMv2ViTPreprocessor(nn.Module):
186 +
187 + def __init__(self, config: AIMv2Config):
188 + super().__init__()
189 + num_patches = (config.image_size // config.patch_size)**2
190 +
191 + self.patchifier = AIMv2PatchEmbed(config)
192 + self.pos_embed = nn.Parameter(
193 + torch.zeros((1, num_patches, config.hidden_size)))
194 +
195 + def forward(self, x: torch.Tensor) -> torch.Tensor:
196 + tokens = self.patchifier(x)
197 + _, N, _ = tokens.shape
198 + pos_embed = self.pos_embed.to(tokens.device)
199 + tokens = tokens + pos_embed[:, :N]
200 + return tokens
201 +
202 +
203 +class AIMv2Attention(nn.Module):
204 +
205 + def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
206 + prefix: str):
207 + super().__init__()
208 + dim = config.hidden_size
209 +
210 + # TODO(Isotr0py): investigate if we can add TP to visual tokenizer
211 + self.num_heads = config.num_attention_heads
212 + self.qkv = ReplicatedLinear(dim, dim * 3, bias=config.qkv_bias)
213 + # self.qkv = QKVParallelLinear(
214 + # hidden_size=dim,
215 + # head_size=dim // config.num_attention_heads,
216 + # total_num_heads=config.num_attention_heads,
217 + # bias=config.qkv_bias,
218 + # quant_config=quant_config,
219 + # prefix=f"{prefix}.qkv")
220 + self.proj = ReplicatedLinear(dim, dim, bias=config.use_bias)
221 + # self.proj = RowParallelLinear(input_size=dim,
222 + # output_size=dim,
223 + # bias = config.use_bias,
224 + # quant_config=quant_config,
225 + # prefix=f"{prefix}.proj")
226 +
227 + def forward( # todo might implement multiple attn implementations
228 + self,
229 + x: torch.Tensor,
230 + mask: Optional[torch.Tensor] = None) -> torch.Tensor:
231 + B, N, C = x.shape
232 + qkv, _ = self.qkv(x)
233 +
234 + qkv = qkv.reshape(B, N, 3, self.num_heads,
235 + C // self.num_heads).permute(2, 0, 3, 1, 4)
236 +
237 + q, k, v = qkv.unbind(0)
238 +
239 + x = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
240 + x = x.transpose(1, 2).contiguous().reshape(B, N, C)
241 + x, _ = self.proj(x)
242 + return x
243 +
244 +
245 +class AIMv2Block(nn.Module):
246 +
247 + def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
248 + prefix: str):
249 + super().__init__()
250 + self.attn = AIMv2Attention(config,
251 + quant_config=quant_config,
252 + prefix=f"{prefix}.attn")
253 + self.norm_1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
254 + self.mlp = AIMv2SwiGLUFFN(config,
255 + quant_config=quant_config,
256 + prefix=f"{prefix}.mlp")
257 + self.norm_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
258 +
259 + def forward(self,
260 + x: torch.Tensor,
261 + mask: Optional[torch.Tensor] = None) -> torch.Tensor:
262 + x = x + self.attn(self.norm_1.forward_native(x), mask)
263 + x = x + self.mlp(self.norm_2.forward_native(x))
264 + return x
265 +
266 +
267 +class AIMv2Transformer(nn.Module):
268 +
269 + def __init__(self, config: AIMv2Config, quant_config: QuantizationConfig,
270 + prefix: str):
271 + super().__init__()
272 +
273 + self.blocks = nn.ModuleList([
274 + AIMv2Block(config, quant_config, prefix=f"{prefix}.blocks.{i}")
275 + for i in range(config.num_hidden_layers)
276 + ])
277 + self.post_trunk_norm = RMSNorm(config.hidden_size,
278 + eps=config.rms_norm_eps)
279 +
280 + def forward(
281 + self,
282 + tokens: torch.Tensor,
283 + mask: Optional[torch.Tensor] = None,
284 + ) -> torch.Tensor:
285 + # they take the -1 as the ref embeddings, like a clip skip
286 + for block in self.blocks:
287 + tokens = block(tokens, mask)
288 + # NO NORM IN THE OG IMPLEMENTATION
289 + # tokens = self.post_trunk_norm(tokens)
290 + return tokens
291 +
292 +
293 +class AIMv2Model(torch.nn.Module):
294 +
295 + def __init__(self,
296 + config: AIMv2Config,
297 + quant_config: QuantizationConfig,
298 + prefix: str = ""):
299 + super().__init__()
300 + self.preprocessor = AIMv2ViTPreprocessor(config)
301 + self.trunk = AIMv2Transformer(config,
302 + quant_config=quant_config,
303 + prefix=f"{prefix}.trunk")
304 +
305 + @property
306 + def dtype(self):
307 + return self.trunk.blocks[0].attn.qkv.weight.dtype
308 +
309 + @property
310 + def device(self):
311 + return self.trunk.blocks[0].attn.qkv.device
312 +
313 + def forward(
314 + self,
315 + pixel_values: torch.Tensor,
316 + mask: Optional[torch.Tensor] = None,
317 + ) -> torch.Tensor:
318 +
319 + x = self.preprocessor(pixel_values)
320 + x = self.trunk(x, mask)
321 +
322 + return x
vllm/model_executor/models/ovis2.py new
+331
@@ -0,0 +1,331 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# adapted from https://github.com/huggingface/transformers/blob/v4.39.3/src/transformers/models/ovis/modeling_ovis.py
4 +# Copyright 2023 The vLLM team.
5 +# Copyright 2023 HuggingFace Inc. team. All rights reserved.
6 +#
7 +# Licensed under the Apache License, Version 2.0 (the "License");
8 +# you may not use this file except in compliance with the License.
9 +# You may obtain a copy of the License at
10 +#
11 +# http://www.apache.org/licenses/LICENSE-2.0
12 +#
13 +# Unless required by applicable law or agreed to in writing, software
14 +# distributed under the License is distributed on an "AS IS" BASIS,
15 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 +# See the License for the specific language governing permissions and
17 +# limitations under the License.
18 +""" PyTorch Ovis2 model."""
19 +from typing import (Iterable, List, Literal, Mapping, Optional, Set, Tuple,
20 + TypedDict, Union)
21 +
22 +import torch
23 +import torch.nn as nn
24 +from torch import Tensor
25 +from transformers import BatchFeature
26 +
27 +from vllm.config import VllmConfig
28 +from vllm.model_executor.models.aimv2 import Aimv2VisualTokenizer
29 +from vllm.model_executor.models.utils import (AutoWeightsLoader, flatten_bn,
30 + init_vllm_registered_model,
31 + maybe_prefix)
32 +from vllm.model_executor.sampling_metadata import SamplingMetadata
33 +from vllm.multimodal import MULTIMODAL_REGISTRY
34 +from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
35 + MultiModalKwargs)
36 +from vllm.multimodal.parse import ImageSize, MultiModalDataItems
37 +from vllm.multimodal.processing import (BaseMultiModalProcessor,
38 + BaseProcessingInfo, PromptReplacement)
39 +from vllm.multimodal.profiling import BaseDummyInputsBuilder
40 +from vllm.sequence import IntermediateTensors
41 +from vllm.transformers_utils.configs.ovis2 import OvisConfig
42 +from vllm.transformers_utils.processors.ovis2 import OvisProcessor
43 +
44 +from .interfaces import MultiModalEmbeddings, SupportsMultiModal
45 +from .utils import merge_multimodal_embeddings
46 +
47 +# Cannot find the following number from hf config.
48 +IMAGE_TOKEN = "<image>"
49 +IMAGE_ATOM_TOKEN_ID = 151666
50 +IMAGE_PAD_TOKEN_ID = 151672
51 +NUMBER_OF_TOKEN_TO_RESERVE_FOR_SEGMENT = 256
52 +
53 +
54 +class Ovis2ImagePatchInputs(TypedDict):
55 + type: Literal["image_patches"]
56 + flat_data: torch.Tensor
57 + """
58 + Shape:
59 + `(batch_size * num_patches, patch_size_x * patch_size_y * num_channels)`
60 + """
61 +
62 + patches_per_image: List[int]
63 + """
64 + List of number of total patches for each image in the batch.
65 + This is used to restore the first two dimensions of `flat_data`.
66 + """
67 +
68 +
69 +class VisualEmbedding(torch.nn.Embedding):
70 +
71 + def __init__(self, *args, **kwargs):
72 + super().__init__(*args, **kwargs)
73 +
74 + def forward(self, visual_tokens: Tensor) -> Tensor:
75 + if visual_tokens.dtype in [
76 + torch.int8, torch.int16, torch.int32, torch.int64, torch.long
77 + ]:
78 + return super().forward(visual_tokens)
79 + return torch.matmul(visual_tokens, self.weight)
80 +
81 + @property
82 + def device(self):
83 + return self.weight.device
84 +
85 + @property
86 + def dtype(self):
87 + return self.weight.dtype
88 +
89 +
90 +class Ovis2ProcessingInfo(BaseProcessingInfo):
91 +
92 + def get_hf_config(self):
93 + return self.ctx.get_hf_config(OvisConfig)
94 +
95 + def get_hf_processor(self, **kwargs):
96 + return self.ctx.get_hf_processor(OvisProcessor)
97 +
98 + def get_image_processor(self) -> OvisProcessor:
99 + return self.get_hf_processor().image_processor # type: ignore
100 +
101 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
102 + return { # 32k is model token limit at the moment
103 + "image":
104 + self.get_hf_config().multimodal_max_length //
105 + ((9 + 1) * NUMBER_OF_TOKEN_TO_RESERVE_FOR_SEGMENT)
106 + }
107 +
108 + def get_image_size_with_most_features(self) -> ImageSize:
109 + image_processor = self.get_image_processor()
110 + return ImageSize(width=image_processor.size['shortest_edge'] * 9 * 2,
111 + height=image_processor.size['shortest_edge'] * 9 * 2)
112 +
113 +
114 +class Ovis2DummyInputsBuilder(BaseDummyInputsBuilder[Ovis2ProcessingInfo]):
115 +
116 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
117 + num_images = mm_counts.get("image", 0)
118 + return IMAGE_TOKEN * num_images
119 +
120 + def get_dummy_mm_data(
121 + self,
122 + seq_len: int,
123 + mm_counts: Mapping[str, int],
124 + ) -> MultiModalDataDict:
125 + num_images = mm_counts.get("image", 0)
126 +
127 + target_width, target_height = \
128 + self.info.get_image_size_with_most_features()
129 +
130 + mm_data = {
131 + "image":
132 + self._get_dummy_images(width=target_width,
133 + height=target_height,
134 + num_images=num_images),
135 + }
136 + return mm_data
137 +
138 +
139 +class Ovis2MultiModalProcessor(BaseMultiModalProcessor[Ovis2ProcessingInfo]):
140 +
141 + def _call_hf_processor(
142 + self,
143 + prompt: str,
144 + mm_data: Mapping[str, object],
145 + mm_kwargs: Mapping[str, object],
146 + ) -> BatchFeature:
147 + if not mm_data:
148 + # # Avoid warning from HF logger for text-only input
149 + prompt_ids = self.info.get_tokenizer().encode(prompt)
150 + # prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids) nope
151 + return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
152 +
153 + processed_outputs = super()._call_hf_processor(
154 + prompt=prompt,
155 + mm_data=mm_data,
156 + mm_kwargs=mm_kwargs,
157 + )
158 +
159 + return processed_outputs
160 +
161 + def _apply_hf_processor_tokens_only(
162 + self,
163 + prompt_tokens: list[int],
164 + ) -> list[int]:
165 +
166 + return prompt_tokens
167 +
168 + def _get_mm_fields_config(
169 + self,
170 + hf_inputs: BatchFeature,
171 + hf_processor_mm_kwargs: Mapping[str, object],
172 + ) -> Mapping[str, MultiModalFieldConfig]:
173 + return dict(pixel_values=MultiModalFieldConfig.batched("image"),
174 + grids=MultiModalFieldConfig.batched("image"))
175 +
176 + def _get_prompt_updates(
177 + self,
178 + mm_items: MultiModalDataItems,
179 + hf_processor_mm_kwargs: Mapping[str, object],
180 + out_mm_kwargs: MultiModalKwargs,
181 + ) -> list[PromptReplacement]:
182 +
183 + def get_replacement_ovis(item_idx):
184 + grid = out_mm_kwargs["grids"][item_idx]
185 +
186 + hf_processor = self.info.get_hf_processor()
187 + return hf_processor.construct_image_placeholders(grid)
188 +
189 + return [
190 + PromptReplacement(
191 + modality="image",
192 + target=IMAGE_TOKEN,
193 + replacement=get_replacement_ovis,
194 + ),
195 + ]
196 +
197 +
198 +@MULTIMODAL_REGISTRY.register_processor(Ovis2MultiModalProcessor,
199 + info=Ovis2ProcessingInfo,
200 + dummy_inputs=Ovis2DummyInputsBuilder)
201 +class Ovis2ForConditionalGeneration(nn.Module, SupportsMultiModal):
202 +
203 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
204 + super().__init__()
205 + config = vllm_config.model_config.hf_config
206 + quant_config = vllm_config.quant_config
207 +
208 + self.config: OvisConfig = config
209 + self.llm = init_vllm_registered_model(
210 + vllm_config=vllm_config.with_hf_config(config.get_text_config()),
211 + prefix=maybe_prefix(prefix, "llm"),
212 + )
213 +
214 + self.visual_tokenizer = Aimv2VisualTokenizer(
215 + config=config.visual_tokenizer_config,
216 + quant_config=quant_config,
217 + prefix=f"{prefix}.visual_tokenizer",
218 + image_processor_name_or_path=config.visual_tokenizer_config.
219 + backbone_config.name_or_path,
220 + )
221 +
222 + self.vte = VisualEmbedding(
223 + self.config.visual_tokenizer_config.vocab_size,
224 + self.config.hidden_size)
225 +
226 + # TODO(Isotr0py): PP support
227 + # self.make_empty_intermediate_tensors = (
228 + # self.language_model.make_empty_intermediate_tensors)
229 +
230 + def _parse_and_validate_image_input(
231 + self, **kwargs: object) -> Optional[Ovis2ImagePatchInputs]:
232 + pixel_values = kwargs.pop("pixel_values", None)
233 + if pixel_values is None:
234 + return None
235 +
236 + if pixel_values is not None:
237 + if not isinstance(pixel_values, (torch.Tensor, list)):
238 + raise ValueError("Incorrect type of pixel values. "
239 + f"Got type: {type(pixel_values)}")
240 +
241 + return Ovis2ImagePatchInputs(
242 + type="image_patches",
243 + flat_data=flatten_bn(flatten_bn(pixel_values), concat=True),
244 + patches_per_image=[
245 + x.shape[0] for x in flatten_bn(pixel_values)
246 + ],
247 + )
248 +
249 + raise AssertionError("This line should be unreachable.")
250 +
251 + def _process_image_input(
252 + self, image_input: Ovis2ImagePatchInputs) -> MultiModalEmbeddings:
253 + image_patches_flat = image_input["flat_data"]
254 + patches_per_image = image_input["patches_per_image"]
255 +
256 + target_dtype = self.visual_tokenizer.dtype
257 + visual_tokens = self.visual_tokenizer(
258 + image_patches_flat.to(target_dtype))
259 + visual_embeds = self.vte(visual_tokens) # 1:1 numeric eq.
260 +
261 + return tuple(
262 + x.flatten(0, 1)
263 + for x in visual_embeds.split(patches_per_image, dim=0))
264 +
265 + def get_multimodal_embeddings(
266 + self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
267 + image_input = self._parse_and_validate_image_input(**kwargs)
268 + if image_input is None:
269 + return None
270 +
271 + image_features = self._process_image_input(image_input)
272 +
273 + return image_features
274 +
275 + def get_input_embeddings(
276 + self,
277 + input_ids: torch.Tensor,
278 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
279 + ) -> torch.Tensor:
280 + inputs_embeds = self.llm.get_input_embeddings(input_ids)
281 + if multimodal_embeddings is not None:
282 + inputs_embeds = merge_multimodal_embeddings(
283 + input_ids, inputs_embeds, multimodal_embeddings,
284 + [IMAGE_ATOM_TOKEN_ID, IMAGE_PAD_TOKEN_ID])
285 + return inputs_embeds
286 +
287 + def forward(
288 + self,
289 + input_ids: torch.Tensor,
290 + positions: torch.Tensor,
291 + intermediate_tensors: Optional[IntermediateTensors] = None,
292 + inputs_embeds: Optional[torch.Tensor] = None,
293 + **kwargs: object,
294 + ) -> Union[torch.Tensor, IntermediateTensors]:
295 + if intermediate_tensors is not None:
296 + inputs_embeds = None
297 +
298 + # NOTE: In v1, inputs_embeds is always generated at model runner, this
299 + # condition is for v0 compatibility.
300 + elif inputs_embeds is None:
301 + vision_embeddings = self.get_multimodal_embeddings(**kwargs)
302 + inputs_embeds = self.get_input_embeddings(input_ids,
303 + vision_embeddings)
304 + input_ids = None
305 +
306 + # up until here we have a inputs_embeds 100% numerical identity
307 + # between the OG HF Transformers implementation and ours
308 + hidden_states = self.llm(
309 + input_ids=input_ids,
310 + positions=positions,
311 + intermediate_tensors=intermediate_tensors,
312 + inputs_embeds=inputs_embeds,
313 + )
314 + return hidden_states
315 +
316 + def compute_logits(
317 + self,
318 + hidden_states: torch.Tensor,
319 + sampling_metadata: SamplingMetadata,
320 + ) -> Optional[torch.Tensor]:
321 + logits = self.llm.logits_processor(self.llm.lm_head, hidden_states,
322 + sampling_metadata)
323 + return logits
324 +
325 + def load_weights(self, weights: Iterable[Tuple[str,
326 + torch.Tensor]]) -> Set[str]:
327 + loader = AutoWeightsLoader(self)
328 + return loader.load_weights(weights)
329 +
330 + def get_language_model(self) -> torch.nn.Module:
331 + return self.llm
vllm/model_executor/models/registry.py
+1
@@ -195,6 +195,7 @@ _MULTIMODAL_MODELS = {
195 "Mistral3ForConditionalGeneration": ("mistral3", "Mistral3ForConditionalGeneration"), # noqa: E501
196 "MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),
197 "NVLM_D": ("nvlm_d", "NVLM_D_Model"),
198 + "Ovis2ForConditionalGeneration": ("ovis2", "Ovis2ForConditionalGeneration"),
199 "PaliGemmaForConditionalGeneration": ("paligemma", "PaliGemmaForConditionalGeneration"), # noqa: E501
200 "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
201 "PixtralForConditionalGeneration": ("pixtral", "PixtralForConditionalGeneration"), # noqa: E501
vllm/transformers_utils/config.py
+4 -3
@@ -38,9 +38,9 @@ from vllm.transformers_utils.configs import (ChatGLMConfig, Cohere2Config,
38 MiniMaxVL01Config, MllamaConfig,
39 MLPSpeculatorConfig, MPTConfig,
40 NemotronConfig, NVLM_D_Config,
41 - RWConfig, SkyworkR1VChatConfig,
42 - SolarConfig, Telechat2Config,
43 - UltravoxConfig)
41 + OvisConfig, RWConfig,
42 + SkyworkR1VChatConfig, SolarConfig,
43 + Telechat2Config, UltravoxConfig)
44 # yapf: enable
45 from vllm.transformers_utils.utils import check_gguf_file
46 from vllm.utils import resolve_obj_by_qualname
@@ -79,6 +79,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
79 "minimax_vl_01": MiniMaxVL01Config,
80 "nemotron": NemotronConfig,
81 "NVLM_D": NVLM_D_Config,
82 + "ovis": OvisConfig,
83 "solar": SolarConfig,
84 "skywork_chat": SkyworkR1VChatConfig,
85 "telechat": Telechat2Config,
vllm/transformers_utils/configs/__init__.py
+2
@@ -23,6 +23,7 @@ from vllm.transformers_utils.configs.moonvit import MoonViTConfig
23 from vllm.transformers_utils.configs.mpt import MPTConfig
24 from vllm.transformers_utils.configs.nemotron import NemotronConfig
25 from vllm.transformers_utils.configs.nvlm_d import NVLM_D_Config
26 +from vllm.transformers_utils.configs.ovis2 import OvisConfig
27 from vllm.transformers_utils.configs.skyworkr1v import SkyworkR1VChatConfig
28 from vllm.transformers_utils.configs.solar import SolarConfig
29 from vllm.transformers_utils.configs.telechat2 import Telechat2Config
@@ -49,6 +50,7 @@ __all__ = [
50 "KimiVLConfig",
51 "NemotronConfig",
52 "NVLM_D_Config",
53 + "OvisConfig",
54 "SkyworkR1VChatConfig",
55 "SolarConfig",
56 "Telechat2Config",
vllm/transformers_utils/configs/ovis2.py new
+170
@@ -0,0 +1,170 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# yapf: disable
4 +# ruff: noqa: E501
5 +# copied from https://huggingface.co/AIDC-AI/Ovis2-1B/blob/main/configuration_aimv2.py
6 +# and https://huggingface.co/AIDC-AI/Ovis2-1B/blob/main/configuration_ovis.py
7 +from typing import Any, Optional, Union
8 +
9 +from transformers import AutoConfig, PretrainedConfig
10 +
11 +
12 +class AIMv2Config(PretrainedConfig):
13 + """This is the configuration class to store the configuration of an [`AIMv2Model`].
14 +
15 + Instantiating a configuration with the defaults will yield a similar configuration
16 + to that of the [apple/aimv2-large-patch14-224](https://huggingface.co/apple/aimv2-large-patch14-224).
17 +
18 + Args:
19 + hidden_size: Dimension of the hidden representations.
20 + intermediate_size: Dimension of the SwiGLU representations.
21 + num_hidden_layers: Number of hidden layers in the Transformer.
22 + num_attention_heads: Number of attention heads for each attention layer
23 + in the Transformer.
24 + num_channels: Number of input channels.
25 + image_size: Image size.
26 + patch_size: Patch size.
27 + rms_norm_eps: Epsilon value used for the RMS normalization layer.
28 + attention_dropout: Dropout ratio for attention probabilities.
29 + projection_dropout: Dropout ratio for the projection layer after the attention.
30 + qkv_bias: Whether to add a bias to the queries, keys and values.
31 + use_bias: Whether to add a bias in the feed-forward and projection layers.
32 + kwargs: Keyword arguments for the [`PretrainedConfig`].
33 + """
34 +
35 + model_type: str = "aimv2"
36 +
37 + def __init__(
38 + self,
39 + hidden_size: int = 1024,
40 + intermediate_size: int = 2816,
41 + num_hidden_layers: int = 24,
42 + num_attention_heads: int = 8,
43 + num_channels: int = 3,
44 + image_size: int = 224,
45 + patch_size: int = 14,
46 + rms_norm_eps: float = 1e-5,
47 + attention_dropout: float = 0.0,
48 + projection_dropout: float = 0.0,
49 + qkv_bias: bool = False,
50 + use_bias: bool = False,
51 + **kwargs: Any,
52 + ):
53 + super().__init__(**kwargs)
54 + self.hidden_size = hidden_size
55 + self.intermediate_size = intermediate_size
56 + self.num_hidden_layers = num_hidden_layers
57 + self.num_attention_heads = num_attention_heads
58 + self.num_channels = num_channels
59 + self.patch_size = patch_size
60 + self.image_size = image_size
61 + self.attention_dropout = attention_dropout
62 + self.rms_norm_eps = rms_norm_eps
63 +
64 + self.projection_dropout = projection_dropout
65 + self.qkv_bias = qkv_bias
66 + self.use_bias = use_bias
67 +
68 +
69 +IGNORE_ID = -100
70 +IMAGE_TOKEN_ID = -200
71 +IMAGE_TOKEN = "<image>"
72 +IMAGE_ATOM_ID = -300
73 +IMAGE_INDICATOR_IDS = [-301, -302, -303, -304, -305]
74 +
75 +AutoConfig.register("aimv2", AIMv2Config)
76 +
77 +
78 +# ----------------------------------------------------------------------
79 +# Visual Tokenizer Configuration
80 +# ----------------------------------------------------------------------
81 +class BaseVisualTokenizerConfig(PretrainedConfig):
82 +
83 + def __init__(self,
84 + vocab_size=16384,
85 + tokenize_function="softmax",
86 + tau=1.0,
87 + depths=None,
88 + drop_cls_token=False,
89 + backbone_config: Optional[Union[PretrainedConfig,
90 + dict]] = None,
91 + hidden_stride: int = 1,
92 + **kwargs):
93 + super().__init__(**kwargs)
94 + self.vocab_size = vocab_size
95 + self.tokenize_function = tokenize_function
96 + self.tau = tau
97 + if isinstance(depths, str):
98 + depths = [int(x) for x in depths.split('|')]
99 + self.depths = depths
100 + self.backbone_kwargs = dict[str, Any]()
101 + self.drop_cls_token = drop_cls_token
102 + if backbone_config is not None:
103 + assert isinstance(backbone_config, (PretrainedConfig, dict)), \
104 + f"expect `backbone_config` to be instance of PretrainedConfig or dict, but got {type(backbone_config)} type"
105 + if not isinstance(backbone_config, PretrainedConfig):
106 + model_type = backbone_config['model_type']
107 + backbone_config.pop('model_type')
108 + backbone_config = AutoConfig.for_model(model_type,
109 + **backbone_config)
110 + self.backbone_config = backbone_config
111 + self.hidden_stride = hidden_stride
112 +
113 +
114 +class Aimv2VisualTokenizerConfig(BaseVisualTokenizerConfig):
115 + model_type = "aimv2_visual_tokenizer"
116 +
117 + def __init__(self, **kwargs):
118 + super().__init__(**kwargs)
119 + if self.drop_cls_token:
120 + self.drop_cls_token = False
121 + if self.depths:
122 + assert len(self.depths) == 1
123 + self.backbone_kwargs['num_hidden_layers'] = self.depths[0]
124 +
125 +
126 +AutoConfig.register("aimv2_visual_tokenizer", Aimv2VisualTokenizerConfig)
127 +
128 +
129 +# ----------------------------------------------------------------------
130 +# Ovis Configuration
131 +# ----------------------------------------------------------------------
132 +class OvisConfig(PretrainedConfig):
133 + model_type = "ovis"
134 +
135 + def __init__(self,
136 + llm_config: Optional[Union[PretrainedConfig, dict]] = None,
137 + visual_tokenizer_config: Optional[Union[PretrainedConfig,
138 + dict]] = None,
139 + multimodal_max_length=8192,
140 + hidden_size=None,
141 + conversation_formatter_class=None,
142 + llm_attn_implementation=None,
143 + disable_tie_weight=False,
144 + **kwargs):
145 + super().__init__(**kwargs)
146 + if llm_config is not None:
147 + assert isinstance(llm_config, (PretrainedConfig, dict)), \
148 + f"expect `llm_config` to be instance of PretrainedConfig or dict, but got {type(llm_config)} type"
149 + if not isinstance(llm_config, PretrainedConfig):
150 + model_type = llm_config['model_type']
151 + llm_config.pop('model_type')
152 + llm_config = AutoConfig.for_model(model_type, **llm_config)
153 +
154 + # map llm_config to text_config
155 + self.text_config = llm_config
156 + if visual_tokenizer_config is not None:
157 + assert isinstance(visual_tokenizer_config, (PretrainedConfig, dict)), \
158 + f"expect `visual_tokenizer_config` to be instance of PretrainedConfig or dict, but got {type(visual_tokenizer_config)} type"
159 + if not isinstance(visual_tokenizer_config, PretrainedConfig):
160 + model_type = visual_tokenizer_config['model_type']
161 + visual_tokenizer_config.pop('model_type')
162 + visual_tokenizer_config = AutoConfig.for_model(
163 + model_type, **visual_tokenizer_config)
164 +
165 + self.visual_tokenizer_config = visual_tokenizer_config
166 + self.multimodal_max_length = multimodal_max_length
167 + self.hidden_size = hidden_size
168 + self.conversation_formatter_class = conversation_formatter_class
169 + self.llm_attn_implementation = llm_attn_implementation
170 + self.disable_tie_weight = disable_tie_weight
vllm/transformers_utils/processors/__init__.py
+2 -1
@@ -2,5 +2,6 @@
2
3 from vllm.transformers_utils.processors.deepseek_vl2 import (
4 DeepseekVLV2Processor)
5 +from vllm.transformers_utils.processors.ovis2 import OvisProcessor
6
6 -__all__ = ["DeepseekVLV2Processor"]
7 +__all__ = ["DeepseekVLV2Processor", "OvisProcessor"]
vllm/transformers_utils/processors/ovis2.py new
+397
@@ -0,0 +1,397 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +# yapf: disable
4 +# ruff: noqa: E501
5 +# coding=utf-8
6 +# adapted from https://github.com/AIDC-AI/Ovis/blob/35ab51a1a1e3542fa6db260a1084cefbc8f164bb/ovis/vllm/processing_ovis.py
7 +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.
8 +#
9 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10 +# and OPT implementations in this library. It has been modified from its
11 +# original forms to accommodate minor architectural differences compared
12 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13 +#
14 +# Licensed under the Apache License, Version 2.0 (the "License");
15 +# you may not use this file except in compliance with the License.
16 +# You may obtain a copy of the License at
17 +#
18 +# http://www.apache.org/licenses/LICENSE-2.0
19 +#
20 +# Unless required by applicable law or agreed to in writing, software
21 +# distributed under the License is distributed on an "AS IS" BASIS,
22 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23 +# See the License for the specific language governing permissions and
24 +# limitations under the License.
25 +from typing import List, Union
26 +
27 +import PIL
28 +import torch
29 +from transformers import AutoProcessor, BatchFeature
30 +from transformers.image_utils import ImageInput
31 +from transformers.processing_utils import (ProcessingKwargs, ProcessorMixin,
32 + Unpack)
33 +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
34 +
35 +__all__ = [ 'OvisProcessor']
36 +IGNORE_ID = -100
37 +
38 +class OvisProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call-arg]
39 + _defaults = {
40 + "text_kwargs": {
41 + "padding": False,
42 + },
43 + "images_kwargs": {
44 + 'max_partition':9,
45 + 'covering_threshold':0.9,
46 + 'convert_to_rgb':True,
47 + 'return_tensors':'pt'},
48 + }
49 +
50 +
51 +
52 +class OvisProcessor(ProcessorMixin):
53 + r"""
54 + Constructs a Ovis processor which wraps a Ovis image processor and a Qwen2 tokenizer into a single processor.
55 + [`OvisProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`] and [`Qwen2TokenizerFast`]. See the
56 + [`~OvisProcessor.__call__`] and [`~OvisProcessor.decode`] for more information.
57 + Args:
58 + image_processor ([`Qwen2VLImageProcessor`], *optional*):
59 + The image processor is a required input.
60 + tokenizer ([`Qwen2TokenizerFast`], *optional*):
61 + The tokenizer is a required input.
62 + chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
63 + in a chat into a tokenizable string.
64 + """
65 +
66 + attributes = ["image_processor", "tokenizer"]
67 + valid_kwargs = ["chat_template"]
68 +
69 + image_processor_class = "AutoImageProcessor"
70 + tokenizer_class = "Qwen2Tokenizer"
71 +
72 + def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):
73 + self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
74 + self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token
75 + super().__init__(image_processor, tokenizer, chat_template=chat_template)
76 +
77 + self.extra_special_tokens = {
78 + "image_token": "<image>",
79 + "image_atom": "<image_atom>",
80 + "image_start": "<img>",
81 + "image_prefix": "<pre>",
82 + "image_col_sep": "<col>",
83 + "image_row_sep": "<row>",
84 + "image_end": "</img>",
85 + 'image_pad': '<image_pad>',
86 + }
87 +
88 + def __call__(
89 + self,
90 + images: ImageInput = None,
91 + text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
92 + **kwargs: Unpack[OvisProcessorKwargs],
93 + ) -> BatchFeature:
94 + """
95 + Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
96 + and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
97 + the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwrags` arguments to
98 + Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`.
99 + Args:
100 + images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
101 + The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
102 + tensor. Both channels-first and channels-last formats are supported.
103 + text (`str`, `List[str]`, `List[List[str]]`):
104 + The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
105 + (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
106 + `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
107 + videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`):
108 + The image or batch of videos to be prepared. Each video can be a 4D NumPy array or PyTorch
109 + tensor, or a nested list of 3D frames. Both channels-first and channels-last formats are supported.
110 + return_tensors (`str` or [`~utils.TensorType`], *optional*):
111 + If set, will return tensors of a particular framework. Acceptable values are:
112 + - `'tf'`: Return TensorFlow `tf.constant` objects.
113 + - `'pt'`: Return PyTorch `torch.Tensor` objects.
114 + - `'np'`: Return NumPy `np.ndarray` objects.
115 + - `'jax'`: Return JAX `jnp.ndarray` objects.
116 + Returns:
117 + [`BatchFeature`]: A [`BatchFeature`] with the following fields:
118 + - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
119 + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
120 + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
121 + `None`).
122 + - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
123 + - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.
124 + - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`.
125 + - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`.
126 + - **second_per_grid_ts** -- List of video seconds per time grid. Returned when `videos` is not `None`.
127 + """
128 + output_kwargs = self._merge_kwargs(
129 + OvisProcessorKwargs,
130 + tokenizer_init_kwargs=self.tokenizer.init_kwargs,
131 + **kwargs,
132 + )
133 +
134 + # Process all images first
135 + image_features = {}
136 + if images is not None:
137 + processed_images = []
138 + image_placeholders_list = []
139 + grids = []
140 +
141 + # Process each image
142 + for image in images if isinstance(images, list) else [images]:
143 + pixel_values, image_placeholders, grid = self.preprocess_image(
144 + image=image, **output_kwargs["images_kwargs"]
145 + )
146 + processed_images.append(pixel_values)
147 + image_placeholders_list.append(image_placeholders)
148 + grids.append(grid)
149 +
150 + # assign all processed images
151 + if processed_images:
152 + image_features["image_placeholders"] = image_placeholders_list
153 +
154 + # Process text input
155 + if text is not None:
156 +
157 + if not isinstance(text, list):
158 + text = [text]
159 +
160 + tokenized_batched_text = self.tokenizer.batch_encode_plus(
161 + text,
162 + **output_kwargs["text_kwargs"]
163 + )
164 + image_token_id = self.get_token_value("image_token")
165 + replaced_ids_list = []
166 + replaced_attn_mask_list = []
167 + idx = 0
168 + for ids_tensor, attn_mask in zip(tokenized_batched_text['input_ids'],
169 + tokenized_batched_text['attention_mask']):
170 + if image_token_id in ids_tensor and "image_placeholders" in image_features:
171 + if idx < len(image_features["image_placeholders"]):
172 + # Converts in list for ease of use
173 + ids_list = ids_tensor.tolist()
174 + attn_list = attn_mask.tolist()
175 +
176 + new_ids = []
177 + new_attn = []
178 +
179 + # replace placeholders
180 + for i, token_id in enumerate(ids_list):
181 + if token_id == image_token_id:
182 + placeholder_ids = image_features["image_placeholders"][idx]
183 + new_ids.extend(placeholder_ids)
184 + new_attn.extend([1] * len(placeholder_ids))
185 + idx += 1
186 + else:
187 + new_ids.append(token_id)
188 + new_attn.append(attn_list[i])
189 +
190 + # Converts back to tensors
191 + ids_tensor = torch.tensor(new_ids, dtype=torch.long)
192 + attn_mask = torch.tensor(new_attn, dtype=torch.long)
193 + else:
194 + raise RuntimeError(
195 + 'Mismatch between the images you provided and the number of placeholder present in the text')
196 +
197 + replaced_ids_list.append(ids_tensor)
198 + replaced_attn_mask_list.append(attn_mask)
199 +
200 + if replaced_ids_list:
201 + replaced_and_tokenized_ids = torch.stack(replaced_ids_list)
202 + replaced_and_tokenized_attn_mask = torch.stack(replaced_attn_mask_list)
203 + else:
204 + replaced_and_tokenized_ids = torch.tensor([], dtype=torch.long)
205 + replaced_and_tokenized_attn_mask = torch.tensor([], dtype=torch.long)
206 +
207 + # Create the output with text features
208 + output = BatchFeature(
209 + data={
210 + "input_ids": replaced_and_tokenized_ids,
211 + "attention_mask": replaced_and_tokenized_attn_mask,
212 + }
213 + )
214 +
215 + # Add image features if present
216 + if image_features:
217 + output["pixel_values"] = processed_images
218 + output['grids'] = grids
219 +
220 + return output
221 +
222 +
223 + # If only images were provided
224 + return BatchFeature(data=image_features)
225 +
226 +
227 + def get_image_size(self):
228 + height = self.image_processor.crop_size["height"]
229 + width = self.image_processor.crop_size["width"]
230 + return height, width
231 +
232 + def get_token_value(self, tok):
233 + return self.tokenizer.get_vocab()[self.extra_special_tokens[tok]]
234 +
235 + def construct_image_placeholders(self, grid):
236 +
237 + image_placeholders = [self.get_token_value('image_start'),
238 + self.get_token_value('image_atom'),
239 + self.get_token_value('image_prefix')]
240 + if grid[0] * grid[1] > 1:
241 + for r in range(grid[0]):
242 + for c in range(grid[1]):
243 + image_placeholders.append(self.get_token_value('image_atom') )
244 + if c < grid[1] - 1:
245 + image_placeholders.append(self.get_token_value('image_col_sep'))
246 + if r < grid[0] - 1:
247 + image_placeholders.append(self.get_token_value('image_row_sep'))
248 + image_placeholders.append(self.get_token_value('image_end'))
249 + # return image_placeholders
250 +
251 + image_atom_token_id = self.get_token_value('image_atom')
252 + # Extract the padding token ID from tokenizer
253 + image_padding_token_id = self.get_token_value('image_pad')
254 +
255 + # Create a new list with padding tokens inserted
256 + padded_placeholder_tokens = []
257 + for token in image_placeholders:
258 + padded_placeholder_tokens.append(token)
259 + if token == image_atom_token_id:
260 + # Add 255 padding tokens after each image atom token
261 + padded_placeholder_tokens.extend([image_padding_token_id] * 255)
262 + return padded_placeholder_tokens
263 +
264 + def preprocess_image(self, image: PIL.Image.Image, max_partition, covering_threshold, convert_to_rgb, return_tensors):
265 + def _preprocess(img: PIL.Image.Image, side):
266 + # first resize and preprocess
267 + w, h = img.size
268 + if w == h:
269 + new_width = new_height = side
270 + elif w > h:
271 + new_width = side
272 + new_height = int(h / w * new_width)
273 + else:
274 + new_height = side
275 + new_width = int(w / h * new_height)
276 + new_size = dict(height=new_height, width=new_width)
277 + pixel_values = self.image_processor.preprocess(img, size=new_size, return_tensors=return_tensors)['pixel_values']
278 +
279 + # then pad to square
280 + square_values = torch.zeros([1, 3, side, side], dtype=pixel_values.dtype, device=pixel_values.device)
281 + new_height, new_width = pixel_values.shape[2:]
282 + if new_height == new_width:
283 + square_values[:, :, :, :] = pixel_values
284 + elif new_height > new_width:
285 + from_index = (side - new_width) // 2
286 + square_values[:, :, :, from_index:from_index + new_width] = pixel_values
287 + else:
288 + from_index = (side - new_height) // 2
289 + square_values[:, :, from_index:from_index + new_height, :] = pixel_values
290 +
291 + return square_values
292 +
293 + def _partition(img, grid) -> list[tuple[int, int, int, int]]:
294 + w, h = img.size
295 + row_height = h // grid[0]
296 + col_width = w // grid[1]
297 +
298 + partition = []
299 + for row in range(grid[0]):
300 + for col in range(grid[1]):
301 + left = col * col_width
302 + upper = row * row_height
303 + right = w if col == grid[1] - 1 else (col + 1) * col_width
304 + lower = h if row == grid[0] - 1 else (row + 1) * row_height
305 + partition.append((left, upper, right, lower))
306 +
307 + return partition
308 +
309 + def _covering_area(left, upper, right, lower, side):
310 + w = right - left
311 + h = lower - upper
312 + w, h = max(w, h), min(w, h)
313 + if w > side:
314 + h = h / w * side
315 + w = side
316 + return w * h
317 +
318 + def _get_best_grid(img, side):
319 + img_area = img.size[0] * img.size[1]
320 +
321 + candidate_grids = []
322 + for i in range(1, max_partition + 1):
323 + for j in range(1, max_partition + 1):
324 + if i * j <= max_partition:
325 + candidate_grids.append((i, j))
326 +
327 + all_grids = []
328 + good_grids = []
329 + for grid in candidate_grids:
330 + partition = _partition(img, grid)
331 + covering_ratio = sum([_covering_area(*p, side) for p in partition]) / img_area
332 + assert covering_ratio <= 1.0
333 + all_grids.append((grid, covering_ratio))
334 + if covering_ratio > covering_threshold:
335 + good_grids.append((grid, covering_ratio))
336 +
337 + if len(good_grids) > 0:
338 + # pick the good partition with minimum #sub_images and break the tie using covering_ratio
339 + return sorted(good_grids, key=lambda x: (x[0][0] * x[0][1], -x[1]))[0][0]
340 + else:
341 + # pick the partition with maximum covering_ratio and break the tie using #sub_images
342 + return sorted(all_grids, key=lambda x: (-x[1], x[0][0] * x[0][1]))[0][0]
343 +
344 + if convert_to_rgb and image.mode != 'RGB':
345 + image = image.convert('RGB')
346 +
347 +
348 + sides = self.get_image_size()
349 + if sides[0] != sides[1]:
350 + raise ValueError('get_image_size() returns non-square size')
351 + side = sides[0]
352 + grid = _get_best_grid(image, side)
353 + partition = _partition(image, grid)
354 + crops = [image.crop(p) for p in partition]
355 + if len(crops) > 1:
356 + crops.insert(0, image)
357 + pixel_values = torch.cat([_preprocess(crop, side) for crop in crops], dim=0)
358 + image_placeholders = self.construct_image_placeholders(grid)
359 + return pixel_values, image_placeholders, grid
360 +
361 + def batch_decode(self, *args, **kwargs):
362 + """
363 + This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
364 + refer to the docstring of this method for more information.
365 + """
366 + return self.tokenizer.batch_decode(*args, **kwargs)
367 +
368 + def decode(self, *args, **kwargs):
369 + """
370 + This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
371 + the docstring of this method for more information.
372 + """
373 + return self.tokenizer.decode(*args, **kwargs)
374 +
375 + def post_process_image_text_to_text(self, generated_outputs):
376 + """
377 + Post-process the output of the model to decode the text.
378 + Args:
379 + generated_outputs (`torch.Tensor` or `np.ndarray`):
380 + The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
381 + or `(sequence_length,)`.
382 + Returns:
383 + `List[str]`: The decoded text.
384 + """
385 + return self.tokenizer.batch_decode(
386 + generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False
387 + )
388 +
389 + @property
390 + def model_input_names(self):
391 + tokenizer_input_names = self.tokenizer.model_input_names
392 + image_processor_input_names = self.image_processor.model_input_names
393 + names_from_processor = list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
394 + return names_from_processor + ["second_per_grid_ts"]
395 +
396 +
397 +AutoProcessor.register("OvisProcessor", OvisProcessor)
\ No newline at end of file