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, 07:15 UTC
Add tarsier model support (#18985)
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
汪志鹏 committed
Jun 3, 2025 at 13:13 UTC
1282bd812ea4e1511378bad5b918d609280d2b89
7 files changed
+689
docs/models/supported_models.md
+1
@@ -550,6 +550,7 @@ Specified using `--task generate`.
550
| `Qwen2_5OmniThinkerForConditionalGeneration` | Qwen2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `Qwen/Qwen2.5-Omni-7B` | | ✅︎ | ✅︎\* |
551
| `SkyworkR1VChatModel` | Skywork-R1V-38B | T + I | `Skywork/Skywork-R1V-38B` | | ✅︎ | ✅︎ |
552
| `SmolVLMForConditionalGeneration` | SmolVLM2 | T + I | `SmolVLM2-2.2B-Instruct` | ✅︎ | | ✅︎ |
553
+| `TarsierForConditionalGeneration` | Tarsier | T + I<sup>E+</sup> | `omni-search/Tarsier-7b`,`omni-search/Tarsier-34b` | | ✅︎ | ✅︎ |
554
555
<sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.
556
• For example, to use DeepSeek-VL2 series models:
examples/offline_inference/vision_language.py
+20
@@ -333,6 +333,25 @@ def run_smolvlm(questions: list[str], modality: str) -> ModelRequestData:
333
)
334
335
336
+# omni-research/Tarsier-7b
337
+def run_tarsier(questions: list[str], modality: str) -> ModelRequestData:
338
+ assert modality == "image"
339
+ model_name = "omni-research/Tarsier-7b"
340
+
341
+ engine_args = EngineArgs(
342
+ model=model_name,
343
+ trust_remote_code=True,
344
+ max_model_len=4096,
345
+ limit_mm_per_prompt={modality: 1},
346
+ )
347
+ prompts = [(f"USER: <image>\n{question} ASSISTANT:") for question in questions]
348
+
349
+ return ModelRequestData(
350
+ engine_args=engine_args,
351
+ prompts=prompts,
352
+ )
353
+
354
+
355
# InternVL
356
def run_internvl(questions: list[str], modality: str) -> ModelRequestData:
357
model_name = "OpenGVLab/InternVL3-2B"
@@ -1091,6 +1110,7 @@ model_example_map = {
1110
"qwen2_5_omni": run_qwen2_5_omni,
1111
"skywork_chat": run_skyworkr1v,
1112
"smolvlm": run_smolvlm,
1113
+ "tarsier": run_tarsier,
1114
}
1115
1116
examples/offline_inference/vision_language_multi_image.py
+21
@@ -691,6 +691,26 @@ def load_qwen2_5_vl(question: str, image_urls: list[str]) -> ModelRequestData:
691
)
692
693
694
+def load_tarsier(question: str, image_urls: list[str]) -> ModelRequestData:
695
+ model_name = "omni-research/Tarsier-7b"
696
+
697
+ engine_args = EngineArgs(
698
+ model=model_name,
699
+ trust_remote_code=True,
700
+ max_model_len=4096,
701
+ limit_mm_per_prompt={"image": len(image_urls)},
702
+ )
703
+
704
+ prompt = f"USER: {'<image>' * len(image_urls)}\n{question}\n ASSISTANT:"
705
+ image_data = [fetch_image(url) for url in image_urls]
706
+
707
+ return ModelRequestData(
708
+ engine_args=engine_args,
709
+ prompt=prompt,
710
+ image_data=image_data,
711
+ )
712
+
713
+
714
model_example_map = {
715
"aria": load_aria,
716
"aya_vision": load_aya_vision,
@@ -712,6 +732,7 @@ model_example_map = {
732
"qwen2_vl": load_qwen2_vl,
733
"qwen2_5_vl": load_qwen2_5_vl,
734
"smolvlm": load_smolvlm,
735
+ "tarsier": load_tarsier,
736
}
737
738
tests/models/multimodal/processing/test_common.py
+1
@@ -282,6 +282,7 @@ def _test_processing_correctness_one(
282
"Skywork/Skywork-R1V-38B",
283
"fixie-ai/ultravox-v0_5-llama-3_2-1b",
284
"openai/whisper-large-v3",
285
+ "omni-research/Tarsier-7b",
286
])
287
@pytest.mark.parametrize("hit_rate", [0.3, 0.5, 1.0])
288
@pytest.mark.parametrize("num_batches", [32])
tests/models/registry.py
+2
@@ -406,6 +406,8 @@ _MULTIMODAL_EXAMPLE_MODELS = {
406
"SmolVLMForConditionalGeneration": _HfExamplesInfo("HuggingFaceTB/SmolVLM2-2.2B-Instruct"), # noqa: E501
407
"UltravoxModel": _HfExamplesInfo("fixie-ai/ultravox-v0_5-llama-3_2-1b", # noqa: E501
408
trust_remote_code=True),
409
+ "TarsierForConditionalGeneration": _HfExamplesInfo("omni-research/Tarsier-7b", # noqa: E501
410
+ hf_overrides={"architectures": ["TarsierForConditionalGeneration"]}), # noqa: E501
411
# [Encoder-decoder]
412
# Florence-2 uses BartFastTokenizer which can't be loaded from AutoTokenizer
413
# Therefore, we borrow the BartTokenizer from the original Bart model
vllm/model_executor/models/registry.py
+1
@@ -211,6 +211,7 @@ _MULTIMODAL_MODELS = {
211
"Qwen2_5OmniForConditionalGeneration": ("qwen2_5_omni_thinker", "Qwen2_5OmniThinkerForConditionalGeneration"), # noqa: E501
212
"UltravoxModel": ("ultravox", "UltravoxModel"),
213
"Phi4MMForCausalLM": ("phi4mm", "Phi4MMForCausalLM"),
214
+ "TarsierForConditionalGeneration": ("tarsier", "TarsierForConditionalGeneration"), # noqa: E501
215
# [Encoder-decoder]
216
"Florence2ForConditionalGeneration": ("florence2", "Florence2ForConditionalGeneration"), # noqa: E501
217
"MllamaForConditionalGeneration": ("mllama", "MllamaForConditionalGeneration"), # noqa: E501
vllm/model_executor/models/tarsier.py
new
+643
@@ -0,0 +1,643 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import math
4
+from collections.abc import Iterable, Mapping, Sequence
5
+from typing import (Final, Literal, Optional, Protocol, TypedDict, TypeVar,
6
+ Union, cast)
7
+
8
+import torch
9
+import torch.nn as nn
10
+from transformers import BatchFeature, CLIPVisionConfig
11
+from transformers import LlavaConfig as HfLlavaConfig
12
+from transformers import PretrainedConfig, SiglipVisionConfig
13
+from transformers.image_utils import ImageInput, get_image_size, to_numpy_array
14
+from transformers.models.llava import LlavaProcessor
15
+from transformers.processing_utils import (ProcessingKwargs, Unpack,
16
+ _validate_images_text_input_order)
17
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
18
+
19
+from vllm.config import VllmConfig
20
+from vllm.inputs import InputProcessingContext
21
+from vllm.jsontree import json_map_leaves
22
+from vllm.model_executor.layers.activation import get_act_fn
23
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
24
+ RowParallelLinear)
25
+from vllm.model_executor.layers.quantization import QuantizationConfig
26
+from vllm.model_executor.models.llava import LlavaDummyInputsBuilder
27
+from vllm.model_executor.sampling_metadata import SamplingMetadata
28
+from vllm.multimodal import MULTIMODAL_REGISTRY
29
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargs
30
+from vllm.multimodal.parse import (ImageEmbeddingItems, ImageProcessorItems,
31
+ ImageSize, MultiModalDataItems)
32
+from vllm.multimodal.processing import (BaseMultiModalProcessor,
33
+ BaseProcessingInfo, ProcessingCache,
34
+ PromptReplacement, PromptUpdate)
35
+from vllm.multimodal.profiling import BaseDummyInputsBuilder
36
+from vllm.sequence import IntermediateTensors
37
+
38
+from .clip import CLIPVisionModel
39
+from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
40
+from .siglip import SiglipVisionModel
41
+from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,
42
+ maybe_prefix, merge_multimodal_embeddings)
43
+from .vision import VisionEncoderInfo, get_vision_encoder_info
44
+
45
+
46
+class TarsierImagePixelInputs(TypedDict):
47
+ type: Literal["pixel_values"]
48
+ pixel_values: torch.Tensor
49
+
50
+
51
+class TarsierImageEmbeddingInputs(TypedDict):
52
+ type: Literal["image_embeds"]
53
+ data: torch.Tensor
54
+
55
+
56
+TarsierImageInputs = Union[TarsierImagePixelInputs,
57
+ TarsierImageEmbeddingInputs]
58
+
59
+
60
+class TarsierHfConfig(Protocol): # Based on the Tarsier's LlavaConfig
61
+ vision_config: Final[PretrainedConfig]
62
+ text_config: Final[PretrainedConfig] # Added from Tarsier's LlavaConfig
63
+ image_token_index: Final[int]
64
+ vision_feature_select_strategy: Final[str]
65
+ vision_feature_layer: Final[Union[int, list[int]]]
66
+ projector_hidden_act: Final[str]
67
+ image_newline_idx: Final[int]
68
+ image_new_idx: Final[int]
69
+ multimodal_projector_bias: bool = True
70
+
71
+
72
+class TarsierProcessorKwargs(ProcessingKwargs, total=False):
73
+ _defaults = {
74
+ "text_kwargs": {
75
+ "padding": False,
76
+ },
77
+ "images_kwargs": {},
78
+ }
79
+
80
+
81
+class TarsierProcessor(LlavaProcessor):
82
+
83
+ def __call__(
84
+ self,
85
+ images: ImageInput = None,
86
+ text: Union[TextInput, PreTokenizedInput, list[TextInput],
87
+ list[PreTokenizedInput]] = None,
88
+ audio=None,
89
+ videos=None,
90
+ **kwargs: Unpack[TarsierProcessorKwargs],
91
+ ) -> BatchFeature:
92
+ if images is None and text is None:
93
+ raise ValueError(
94
+ "You have to specify at least one of `images` or `text`.")
95
+
96
+ # check if images and text inputs are reversed for BC
97
+ images, text = _validate_images_text_input_order(images, text)
98
+
99
+ output_kwargs = self._merge_kwargs(
100
+ TarsierProcessorKwargs,
101
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
102
+ **kwargs,
103
+ )
104
+ if images is not None:
105
+ image_inputs = self.image_processor(
106
+ images, **output_kwargs["images_kwargs"])
107
+ else:
108
+ image_inputs = {}
109
+
110
+ if isinstance(text, str):
111
+ text = [text]
112
+ elif not isinstance(text, list) and not isinstance(text[0], str):
113
+ raise ValueError("Invalid input text. Please provide a string,"
114
+ " or a list of strings")
115
+
116
+ # try to expand inputs in processing if we have the necessary parts
117
+ prompt_strings = text
118
+ if image_inputs.get("pixel_values") is not None:
119
+ # Replace the image token with the expanded image token sequence
120
+ pixel_values = image_inputs["pixel_values"]
121
+ height, width = get_image_size(to_numpy_array(pixel_values[0]))
122
+ num_image_tokens = (height // self.patch_size) * (
123
+ width // self.patch_size +
124
+ 1) + self.num_additional_image_tokens + 1
125
+ if self.vision_feature_select_strategy == "default":
126
+ num_image_tokens -= 1
127
+
128
+ prompt_strings = []
129
+ for sample in text:
130
+ sample = sample.replace(self.image_token,
131
+ self.image_token * num_image_tokens)
132
+ prompt_strings.append(sample)
133
+
134
+ return_tensors = output_kwargs["text_kwargs"].pop(
135
+ "return_tensors", None)
136
+ text_inputs = self.tokenizer(prompt_strings,
137
+ **output_kwargs["text_kwargs"])
138
+ return BatchFeature(data={
139
+ **text_inputs,
140
+ **image_inputs
141
+ },
142
+ tensor_type=return_tensors)
143
+
144
+
145
+class TarsierMultiModalProjector(nn.Module):
146
+
147
+ def __init__(self,
148
+ vision_hidden_size: int,
149
+ text_hidden_size: int,
150
+ projector_hidden_act: str,
151
+ multimodal_projector_bias: bool,
152
+ quant_config: Optional[QuantizationConfig] = None,
153
+ prefix: str = ""):
154
+ super().__init__()
155
+
156
+ self.linear_1 = ColumnParallelLinear(vision_hidden_size,
157
+ text_hidden_size,
158
+ bias=multimodal_projector_bias,
159
+ quant_config=quant_config,
160
+ prefix=f"{prefix}.linear_1")
161
+ self.act = get_act_fn(projector_hidden_act)
162
+ self.linear_2 = RowParallelLinear(text_hidden_size,
163
+ text_hidden_size,
164
+ bias=multimodal_projector_bias,
165
+ quant_config=quant_config,
166
+ prefix=f"{prefix}.linear_2")
167
+
168
+ def forward(self, image_features: torch.Tensor) -> torch.Tensor:
169
+ hidden_states, _ = self.linear_1(image_features)
170
+ hidden_states = self.act(hidden_states)
171
+ hidden_states, _ = self.linear_2(hidden_states)
172
+ return hidden_states
173
+
174
+
175
+class TarsierProcessingInfo(BaseProcessingInfo):
176
+
177
+ def get_hf_config(self) -> TarsierHfConfig:
178
+ return self.ctx.get_hf_config(HfLlavaConfig)
179
+
180
+ def get_vision_encoder_info(self) -> VisionEncoderInfo:
181
+ return get_vision_encoder_info(self.get_hf_config())
182
+
183
+ def get_hf_processor(self, **kwargs: object) -> TarsierProcessor:
184
+ hf_processor = self.ctx.get_hf_processor(TarsierProcessor, **kwargs)
185
+ # Patch for patch_size if needed (copied from vLLM LLaVA)
186
+ if hasattr(hf_processor,
187
+ 'patch_size') and hf_processor.patch_size is None:
188
+ patch_size = self.get_vision_encoder_info().get_patch_size()
189
+ hf_processor.patch_size = patch_size
190
+ return hf_processor
191
+
192
+ def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
193
+ return {"image": None}
194
+
195
+ def _apply_feature_select_strategy(
196
+ self,
197
+ strategy: str,
198
+ encoder_num_image_tokens: int,
199
+ ) -> int:
200
+ if strategy == "default":
201
+ return encoder_num_image_tokens - 1
202
+ if strategy == "full":
203
+ return encoder_num_image_tokens
204
+ msg = f"Unexpected feature select strategy: {strategy!r}"
205
+ raise NotImplementedError(msg)
206
+
207
+ def get_num_image_tokens(
208
+ self,
209
+ *,
210
+ image_width: int,
211
+ image_height: int,
212
+ ) -> int:
213
+ hf_config = self.get_hf_config()
214
+ vision_encoder_info = self.get_vision_encoder_info()
215
+ num_projected_patches = self._apply_feature_select_strategy(
216
+ hf_config.vision_feature_select_strategy,
217
+ vision_encoder_info.get_num_image_tokens(
218
+ image_width=image_width,
219
+ image_height=image_height,
220
+ ),
221
+ )
222
+ if num_projected_patches <= 0:
223
+ default_size = self.get_image_size_with_most_features()
224
+ num_projected_patches_default = self._apply_feature_select_strategy(
225
+ hf_config.vision_feature_select_strategy,
226
+ vision_encoder_info.get_num_image_tokens(
227
+ image_width=default_size.width,
228
+ image_height=default_size.height,
229
+ ),
230
+ )
231
+ if num_projected_patches_default <= 0:
232
+ raise ValueError(
233
+ "Could not determine a valid number of image patches.")
234
+ num_projected_patches = num_projected_patches_default
235
+ num_height_patches = int(math.sqrt(num_projected_patches))
236
+ total_image_tokens_for_llm = num_projected_patches \
237
+ + num_height_patches + 1
238
+ return total_image_tokens_for_llm
239
+
240
+ def get_image_size_with_most_features(self) -> ImageSize:
241
+ vision_encoder_info = self.get_vision_encoder_info()
242
+ width = height = vision_encoder_info.get_image_size()
243
+ return ImageSize(width=width, height=height)
244
+
245
+ def get_max_image_tokens(self) -> int:
246
+ target_width, target_height = self.get_image_size_with_most_features()
247
+ return self.get_num_image_tokens(
248
+ image_width=target_width,
249
+ image_height=target_height,
250
+ )
251
+
252
+ def get_image_newline_idx(self) -> int:
253
+ return self.get_hf_config().image_newline_idx
254
+
255
+ def get_image_new_idx(self) -> int:
256
+ return self.get_hf_config().image_new_idx
257
+
258
+
259
+_I_Tarsier = TypeVar("_I_Tarsier", bound=TarsierProcessingInfo)
260
+
261
+
262
+class TarsierDummyInputsBuilder(LlavaDummyInputsBuilder[_I_Tarsier]):
263
+
264
+ pass
265
+
266
+
267
+class TarsierMultiModalProcessor(BaseMultiModalProcessor[_I_Tarsier]):
268
+
269
+ def _get_mm_fields_config(
270
+ self,
271
+ hf_inputs: BatchFeature,
272
+ hf_processor_mm_kwargs: Mapping[str, object],
273
+ ) -> Mapping[str, MultiModalFieldConfig]:
274
+ return dict(
275
+ pixel_values=MultiModalFieldConfig.batched("image"),
276
+ image_embeds=MultiModalFieldConfig.batched("image"),
277
+ )
278
+
279
+ def _get_prompt_updates(
280
+ self,
281
+ mm_items: MultiModalDataItems,
282
+ hf_processor_mm_kwargs: Mapping[str, object],
283
+ out_mm_kwargs: MultiModalKwargs,
284
+ ) -> Sequence[PromptUpdate]:
285
+ hf_config = self.info.get_hf_config()
286
+ image_token_id = hf_config.image_token_index # The <IMAGE> token ID
287
+
288
+ def get_replacement(item_idx: int):
289
+ images = mm_items.get_items(
290
+ "image", (ImageEmbeddingItems, ImageProcessorItems))
291
+
292
+ if isinstance(images, ImageEmbeddingItems):
293
+ num_projected_patches = images.get_feature_size(item_idx)
294
+ # This assumes num_projected_patches is a perfect square
295
+ num_height_patches = int(math.sqrt(num_projected_patches))
296
+ num_final_image_tokens = num_projected_patches \
297
+ + num_height_patches + 1
298
+ else:
299
+ image_size = images.get_image_size(item_idx)
300
+ num_final_image_tokens = self.info.get_num_image_tokens(
301
+ image_width=image_size.width,
302
+ image_height=image_size.height,
303
+ )
304
+
305
+ return [image_token_id] * num_final_image_tokens
306
+
307
+ return [
308
+ PromptReplacement(
309
+ modality="image",
310
+ target=[image_token_id], # Replace each single <IMAGE> token
311
+ replacement=get_replacement,
312
+ ),
313
+ ]
314
+
315
+
316
+def _build_tarsier_hf_info(
317
+ ctx: InputProcessingContext) -> TarsierProcessingInfo:
318
+ return TarsierProcessingInfo(ctx)
319
+
320
+
321
+def _build_tarsier_hf_processor(
322
+ info: _I_Tarsier,
323
+ dummy_inputs: BaseDummyInputsBuilder[_I_Tarsier],
324
+ *,
325
+ cache: Optional[ProcessingCache] = None,
326
+) -> BaseMultiModalProcessor:
327
+ if isinstance(info, TarsierProcessingInfo):
328
+ return TarsierMultiModalProcessor(
329
+ info,
330
+ dummy_inputs,
331
+ cache=cache,
332
+ )
333
+ raise NotImplementedError(type(info))
334
+
335
+
336
+def init_vision_tower_for_tarsier(
337
+ hf_config: TarsierHfConfig, # Use the Tarsier specific config protocol
338
+ quant_config: Optional[QuantizationConfig],
339
+ *,
340
+ require_post_norm: Optional[bool] = None,
341
+ prefix: str = "",
342
+) -> Union[CLIPVisionModel, SiglipVisionModel]:
343
+ vision_config = hf_config.vision_config
344
+
345
+ feature_layers = hf_config.vision_feature_layer
346
+ base_num_hidden_layers = vision_config.num_hidden_layers
347
+
348
+ def _get_layer_index(feature_layer_index: int,
349
+ num_hidden_layers_total: int) -> int:
350
+ if feature_layer_index < 0:
351
+ return num_hidden_layers_total + feature_layer_index + 1
352
+ return feature_layer_index
353
+
354
+ if isinstance(feature_layers, int):
355
+ num_hidden_layers_to_init = _get_layer_index(feature_layers,
356
+ base_num_hidden_layers)
357
+ elif isinstance(feature_layers, (list, tuple)):
358
+ num_hidden_layers_to_init = max(
359
+ _get_layer_index(idx, base_num_hidden_layers)
360
+ for idx in feature_layers)
361
+ else:
362
+ raise TypeError(f"vision_layer_feature type: {type(feature_layers)}"
363
+ " is not supported")
364
+
365
+ if isinstance(vision_config, CLIPVisionConfig):
366
+ return CLIPVisionModel(
367
+ vision_config,
368
+ quant_config=quant_config,
369
+ num_hidden_layers_override=num_hidden_layers_to_init,
370
+ require_post_norm=require_post_norm,
371
+ prefix=prefix,
372
+ )
373
+ elif isinstance(vision_config, SiglipVisionConfig):
374
+ return SiglipVisionModel(
375
+ vision_config,
376
+ quant_config=quant_config,
377
+ num_hidden_layers_override=num_hidden_layers_to_init,
378
+ require_post_norm=require_post_norm,
379
+ prefix=prefix,
380
+ )
381
+
382
+ msg = f"Unsupported vision config for Tarsier: {type(vision_config)}"
383
+ raise NotImplementedError(msg)
384
+
385
+
386
+@MULTIMODAL_REGISTRY.register_processor(_build_tarsier_hf_processor,
387
+ info=_build_tarsier_hf_info,
388
+ dummy_inputs=TarsierDummyInputsBuilder)
389
+class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal,
390
+ SupportsPP):
391
+ packed_modules_mapping = {
392
+ "qkv_proj": ["q_proj", "k_proj", "v_proj"],
393
+ "gate_up_proj": ["gate_proj", "up_proj"]
394
+ }
395
+
396
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
397
+ super().__init__()
398
+ config: TarsierHfConfig = vllm_config.model_config.hf_config
399
+ quant_config = vllm_config.quant_config
400
+ self.config = config # Storing the Tarsier-specific HF config
401
+ self.vision_tower = init_vision_tower_for_tarsier(
402
+ config,
403
+ quant_config,
404
+ require_post_norm=False,
405
+ prefix=maybe_prefix(prefix, "vision_tower"))
406
+ projector_bias = getattr(config, "multimodal_projector_bias", True)
407
+
408
+ self.multi_modal_projector = TarsierMultiModalProjector(
409
+ vision_hidden_size=config.vision_config.hidden_size,
410
+ text_hidden_size=config.text_config.hidden_size,
411
+ projector_hidden_act=config.projector_hidden_act,
412
+ multimodal_projector_bias=projector_bias,
413
+ quant_config=quant_config,
414
+ prefix=maybe_prefix(prefix, "multi_modal_projector"))
415
+ self.language_model = init_vllm_registered_model(
416
+ vllm_config=vllm_config,
417
+ hf_config=config.
418
+ text_config, # Use text_config from Tarsier's main config
419
+ prefix=maybe_prefix(prefix, "language_model"),
420
+ )
421
+ self.register_buffer('image_newline_idx_tensor',
422
+ torch.tensor([config.image_newline_idx],
423
+ dtype=torch.long),
424
+ persistent=False)
425
+ self.register_buffer('image_new_idx_tensor',
426
+ torch.tensor([config.image_new_idx],
427
+ dtype=torch.long),
428
+ persistent=False)
429
+
430
+ self.make_empty_intermediate_tensors = (
431
+ self.language_model.make_empty_intermediate_tensors)
432
+
433
+ def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
434
+ h = w = self.config.vision_config.image_size
435
+ expected_dims = (3, h, w) # Assuming 3 channels
436
+ actual_dims = tuple(data.shape[1:])
437
+
438
+ if actual_dims != expected_dims:
439
+ expected_expr = ("batch_size", *map(str, expected_dims))
440
+ raise ValueError(
441
+ f"The expected shape of pixel values is {expected_expr}. "
442
+ f"You supplied {tuple(data.shape)}.")
443
+ return data
444
+
445
+ def _parse_and_validate_image_input(
446
+ self, **kwargs: object) -> Optional[TarsierImageInputs]:
447
+ pixel_values = kwargs.pop("pixel_values", None)
448
+ image_embeds = kwargs.pop("image_embeds", None)
449
+
450
+ if pixel_values is None and image_embeds is None:
451
+ return None
452
+
453
+ if pixel_values is not None:
454
+ if not isinstance(pixel_values, (torch.Tensor, list)):
455
+ raise ValueError("Incorrect type of pixel values. "
456
+ f"Got type: {type(pixel_values)}")
457
+
458
+ return TarsierImagePixelInputs(
459
+ type="pixel_values",
460
+ pixel_values=self._validate_pixel_values(
461
+ flatten_bn(pixel_values, concat=True)),
462
+ )
463
+
464
+ if image_embeds is not None:
465
+ if not isinstance(image_embeds, (torch.Tensor, list)):
466
+ raise ValueError("Incorrect type of image embeddings. "
467
+ f"Got type: {type(image_embeds)}")
468
+ return TarsierImageEmbeddingInputs(
469
+ type="image_embeds",
470
+ data=flatten_bn(image_embeds, concat=True),
471
+ )
472
+
473
+ raise AssertionError("This line should be unreachable.")
474
+
475
+ def _select_image_features(self, image_features: torch.Tensor, *,
476
+ strategy: str) -> torch.Tensor:
477
+ if strategy == "default":
478
+ return image_features[:, 1:]
479
+ elif strategy == "full":
480
+ return image_features
481
+ raise ValueError(f"Unexpected select feature strategy: {strategy}")
482
+
483
+ def _image_pixels_to_features(
484
+ self,
485
+ vision_tower: Union[CLIPVisionModel, SiglipVisionModel],
486
+ pixel_values: Union[torch.Tensor, list[torch.Tensor]],
487
+ ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:
488
+ # From vLLM LLaVA, vision tower output handling
489
+ image_hidden_states = vision_tower(pixel_values)
490
+ if not isinstance(image_hidden_states, torch.Tensor):
491
+ raise TypeError(
492
+ f"image_hidden_states type: {type(image_hidden_states)}"
493
+ " is not supported")
494
+
495
+ def select_features_fn(leaf: torch.Tensor):
496
+ return self._select_image_features(
497
+ leaf,
498
+ strategy=self.config.vision_feature_select_strategy,
499
+ )
500
+
501
+ selected_features = cast(
502
+ Union[torch.Tensor, tuple[torch.Tensor, ...]],
503
+ json_map_leaves(select_features_fn, image_hidden_states),
504
+ )
505
+ return selected_features
506
+
507
+ def _add_tarsier_split_tokens(
508
+ self, projected_image_features: torch.Tensor) -> torch.Tensor:
509
+ """
510
+ Implements Tarsier's `add_split_tokens` logic.
511
+ """
512
+ num_images, num_projected_patches, embed_dim = \
513
+ projected_image_features.shape
514
+ num_height_patches = int(math.sqrt(num_projected_patches))
515
+ num_width_patches = num_projected_patches // num_height_patches
516
+ device = projected_image_features.device
517
+ embedding_layer = self.language_model.model.embed_tokens
518
+ image_newline_emb = embedding_layer(
519
+ self.image_newline_idx_tensor.to(device)).squeeze(0)
520
+ image_new_emb = embedding_layer(
521
+ self.image_new_idx_tensor.to(device)).squeeze(0)
522
+ try:
523
+ current_image_features_grid = projected_image_features.view(
524
+ num_images, num_height_patches, num_width_patches, embed_dim)
525
+ except RuntimeError as e:
526
+ raise RuntimeError(
527
+ "Cannot reshape projected_image_features"
528
+ f" with shape {projected_image_features.shape} "
529
+ f"to ({num_images}, {num_height_patches},"
530
+ f" {num_width_patches}, {embed_dim}). "
531
+ "Ensure num_projected_patches is compatible"
532
+ " with a grid structure. "
533
+ f"num_projected_patches={num_projected_patches}, "
534
+ f"derived num_height_patches={num_height_patches}. ") from e
535
+
536
+ image_newline_expanded = image_newline_emb.expand(
537
+ (num_images, num_height_patches, 1, embed_dim))
538
+ features_with_newlines = torch.cat(
539
+ [current_image_features_grid, image_newline_expanded],
540
+ dim=2 # Concatenate along width dim
541
+ )
542
+ new_num_patches_after_newline = num_projected_patches \
543
+ + num_height_patches
544
+ features_with_newlines_flat = features_with_newlines.view(
545
+ num_images, new_num_patches_after_newline, embed_dim)
546
+ image_new_expanded = image_new_emb.expand((num_images, 1, embed_dim))
547
+ final_image_features = torch.cat(
548
+ [features_with_newlines_flat, image_new_expanded],
549
+ dim=1 # Concatenate along patch sequence dim
550
+ )
551
+ return final_image_features
552
+
553
+ def _process_image_pixels(
554
+ self,
555
+ inputs: TarsierImagePixelInputs,
556
+ ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:
557
+ assert self.vision_tower is not None
558
+ pixel_values = inputs["pixel_values"]
559
+ image_features_selected = self._image_pixels_to_features(
560
+ self.vision_tower, pixel_values) # type: ignore
561
+ if isinstance(image_features_selected, torch.Tensor):
562
+ projected_features = self.multi_modal_projector(
563
+ image_features_selected)
564
+ final_features = self._add_tarsier_split_tokens(projected_features)
565
+ return final_features
566
+ else:
567
+ raise TypeError(
568
+ f"_image_pixels_to_features type:"
569
+ f" {type(image_features_selected)} is not supported")
570
+
571
+ def _process_image_input(
572
+ self,
573
+ image_input: TarsierImageInputs,
574
+ ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:
575
+ if image_input["type"] == "image_embeds":
576
+ projected_features = image_input["data"]
577
+ if isinstance(projected_features, torch.Tensor):
578
+ return self._add_tarsier_split_tokens(projected_features)
579
+ else:
580
+ raise ValueError("Incorrect type of image_embeds. "
581
+ f"Got type: {type(projected_features)}. ")
582
+ assert self.vision_tower is not None
583
+ return self._process_image_pixels(image_input)
584
+
585
+ def get_language_model(self) -> torch.nn.Module:
586
+ return self.language_model
587
+
588
+ def get_multimodal_embeddings(
589
+ self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
590
+ image_input = self._parse_and_validate_image_input(**kwargs)
591
+ if image_input is None:
592
+ return None
593
+ return self._process_image_input(image_input)
594
+
595
+ def get_input_embeddings(
596
+ self,
597
+ input_ids: torch.Tensor,
598
+ multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
599
+ ) -> torch.Tensor:
600
+ inputs_embeds = self.language_model.get_input_embeddings(input_ids)
601
+ if multimodal_embeddings is not None:
602
+ inputs_embeds = merge_multimodal_embeddings(
603
+ input_ids,
604
+ inputs_embeds,
605
+ multimodal_embeddings,
606
+ self.config.image_token_index,
607
+ )
608
+ return inputs_embeds
609
+
610
+ def forward(
611
+ self,
612
+ input_ids: torch.Tensor,
613
+ positions: torch.Tensor,
614
+ intermediate_tensors: Optional[IntermediateTensors] = None,
615
+ inputs_embeds: Optional[torch.Tensor] = None,
616
+ **kwargs: object,
617
+ ) -> Union[torch.Tensor, IntermediateTensors]:
618
+ if intermediate_tensors is not None:
619
+ inputs_embeds = None
620
+ elif inputs_embeds is None:
621
+ vision_embeddings = self.get_multimodal_embeddings(**kwargs)
622
+ inputs_embeds = self.get_input_embeddings(input_ids,
623
+ vision_embeddings)
624
+ input_ids = None
625
+ hidden_states = self.language_model.model(
626
+ input_ids=input_ids,
627
+ positions=positions,
628
+ intermediate_tensors=intermediate_tensors,
629
+ inputs_embeds=inputs_embeds)
630
+ return hidden_states
631
+
632
+ def compute_logits(
633
+ self,
634
+ hidden_states: torch.Tensor,
635
+ sampling_metadata: SamplingMetadata,
636
+ ) -> Optional[torch.Tensor]:
637
+ return self.language_model.compute_logits(hidden_states,
638
+ sampling_metadata)
639
+
640
+ def load_weights(self, weights: Iterable[tuple[str,
641
+ torch.Tensor]]) -> set[str]:
642
+ loader = AutoWeightsLoader(self)
643
+ return loader.load_weights(weights)