ovis2_5.py 24.1 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
"""PyTorch Ovis model."""

5
6
from collections.abc import Iterable, Mapping
from functools import partial
7
from typing import Annotated, Literal
8
9
10
11
12

import torch
import torch.nn as nn
from transformers import BaseImageProcessor, BatchFeature, PretrainedConfig

13
from vllm.config import VllmConfig
14
from vllm.config.multimodal import BaseDummyOptions
15
from vllm.model_executor.layers.linear import ReplicatedLinear
16
from vllm.model_executor.layers.quantization import QuantizationConfig
17
from vllm.model_executor.models.ovis import VisualEmbedding
18
from vllm.model_executor.models.siglip2navit import Siglip2NavitModel
19
20
21
22
23
24
from vllm.model_executor.models.utils import (
    AutoWeightsLoader,
    flatten_bn,
    init_vllm_registered_model,
    maybe_prefix,
)
25
from vllm.multimodal import MULTIMODAL_REGISTRY
26
27
28
29
30
from vllm.multimodal.inputs import (
    MultiModalDataDict,
    MultiModalFieldConfig,
    MultiModalKwargsItems,
)
31
from vllm.multimodal.parse import ImageSize, MultiModalDataItems
32
from vllm.multimodal.processing import (
33
    BaseDummyInputsBuilder,
34
35
36
37
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    PromptReplacement,
)
38
39
from vllm.sequence import IntermediateTensors
from vllm.transformers_utils.processors.ovis2_5 import Ovis2_5Processor
40
from vllm.utils.tensor_schema import TensorSchema, TensorShape
41

42
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

IMAGE_TOKEN = "<image>"
VIDEO_TOKEN = "<video>"
INDICATOR_IDS = [-301, -302, -303, -304]

IMAGE_PAD_TOKEN_MAP = {
    "gemma2": "<unused0>",
    "llama": "<|reserved_special_token_0|>",
    "qwen2": "<|image_pad|>",
    "qwen3": "<|image_pad|>",
}
IMAGE_PAD_TOKEN_ID_MAP = {
    "gemma2": 7,
    "llama": 128002,
    "qwen2": 151655,
    "qwen3": 151655,
}


62
class Ovis2_5ImagePatchInputs(TensorSchema):
63
    """
64
65
66
67
68
    Dimensions:
        - bnp: Batch size * number of images * number of patches
        - patch_size: patch_size_x * patch_size_y * num_channels
        - patch_indicators: Batch size * (number of patches + 1)
        - bn: Batch size * number of images
69
70
    """

71
72
73
74
75
76
77
    type: Literal["image_patches"]
    flat_data: Annotated[torch.Tensor, TensorShape("bnp", "patch_size")]
    indicator_tokens: Annotated[torch.Tensor, TensorShape("patch_indicators")]
    patches_per_item: Annotated[list[int], TensorShape("bn")]
    grids: Annotated[torch.Tensor, TensorShape("bn", 3)]
    # This is used to restore the first two dimensions of `flat_data`.

78

79
class Ovis2_5VideoPatchInputs(TensorSchema):
80
    """
81
82
83
84
85
    Dimensions:
        - bnp: Batch size * number of videos * number of patches
        - patch_size: patch_size_x * patch_size_y * num_channels
        - patch_indicators: Batch size * (number of patches + 1)
        - bn: Batch size * number of videos
86
87
    """

88
89
90
91
92
93
    type: Literal["video_patches"]
    flat_data: Annotated[torch.Tensor, TensorShape("bnp", "patch_size")]
    indicator_tokens: Annotated[torch.Tensor, TensorShape("patch_indicators")]
    patches_per_item: Annotated[list[int], TensorShape("bn")]
    grids: Annotated[torch.Tensor, TensorShape("bn", 3)]
    # This is used to restore the first two dimensions of `flat_data`.
94
95
96
97
98
99
100
101
102
103
104


class VisualTokenizer(torch.nn.Module):
    """
    VIT
    """

    def __init__(
        self,
        config: PretrainedConfig,
        visual_vocab_size: int,
105
        quant_config: QuantizationConfig | None = None,
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
        prefix: str = "",
    ):
        super().__init__()
        self.config = config
        self.vit = self._init_backbone(
            config=config,
            quant_config=quant_config,
            prefix=f"{prefix}.vit",
        )
        # reserved tokens for INDICATOR_IDS
        head_dim = visual_vocab_size - len(INDICATOR_IDS)
        self.head = torch.nn.Sequential(
            ReplicatedLinear(
                self.config.hidden_size * self.config.hidden_stride**2,
                head_dim,
                bias=False,
                return_bias=False,
123
124
125
            ),
            torch.nn.LayerNorm(head_dim),
        )
126
127
128
129

    def _init_backbone(
        self,
        config: PretrainedConfig,
130
        quant_config: QuantizationConfig | None = None,
131
132
133
134
        prefix: str = "",
    ):
        model_type = config.model_type
        if model_type == "siglip2_navit":
135
136
137
138
139
140
            return Siglip2NavitModel(
                config=config,
                quant_config=quant_config,
                prefix=prefix,
            )
        raise ValueError(f"Unsupported visual tokenizer model_type: {model_type}")
141
142

    @property
143
    def dtype(self) -> torch.dtype:
144
145
146
        return next(self.head.parameters()).dtype

    @property
147
    def device(self) -> torch.device:
148
149
        return next(self.head.parameters()).device

150
    def tokenize(self, logits: torch.Tensor) -> torch.Tensor:
151
        tokens = torch.softmax(logits, dim=-1, dtype=torch.float32).to(logits.dtype)
152
153
        return tokens

154
155
156
    def encode(
        self, pixel_values: torch.Tensor, grid_thws: torch.Tensor
    ) -> torch.Tensor:
157
        features = self.vit(pixel_values, grid_thws)
158
159
        # refer to qwen2.5-vl patchmerger
        seq_len, _ = features.shape
160
        features = features.reshape(seq_len // (self.config.hidden_stride**2), -1)
161
162
163

        return features

164
165
166
    def forward(
        self, pixel_values: torch.Tensor, grid_thws: torch.Tensor
    ) -> torch.Tensor:
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
        features = self.encode(pixel_values, grid_thws)
        logits = self.head(features)
        tokens = self.tokenize(logits)
        # tokens' shape is [#Token, VocabSize-4],
        # so padding with [#Token, 4], after which,
        # tokens' shape should become [#Token, VocabSize];
        tokens = torch.nn.functional.pad(
            tokens,
            (0, len(INDICATOR_IDS)),
            mode="constant",
            value=0,
        )
        return tokens


class Ovis2_5ProcessingInfo(BaseProcessingInfo):
    def get_hf_config(self):
        return self.ctx.get_hf_config()

    def get_hf_processor(self, **kwargs):
        vit_config = self.get_hf_config().vit_config
        return self.ctx.get_hf_processor(
            Ovis2_5Processor,
            image_pad_token=self.get_image_pad_token(),
            patch_size=vit_config.patch_size,
            hidden_stride=vit_config.hidden_stride,
            temporal_patch_size=vit_config.temporal_patch_size,
        )

    def get_image_pad_token(self) -> str:
        hf_text_config = self.get_hf_config().get_text_config()
        text_model_type = hf_text_config.model_type
        return IMAGE_PAD_TOKEN_MAP.get(text_model_type)

    def get_image_processor(self) -> BaseImageProcessor:
        return self.get_hf_processor().image_processor  # type: ignore

204
    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
        return {"image": None, "video": 1}

    def get_image_size_with_most_features(self) -> ImageSize:
        # NOTE(myselvess): max_pixels 1792 * 1792 hardcoded in original code
        # TODO(myselvess): Be adjusted based on the max_pixels
        return ImageSize(width=1792, height=1792)

    def get_num_image_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
        num_frames: int = 1,
    ) -> tuple[ImageSize, int]:
        hf_config = self.get_hf_config()
        vit_config = hf_config.vit_config
        patch_size = vit_config.patch_size
        temporal_patch_size = vit_config.temporal_patch_size
        # NOTE: Frames are padded to be divisible by `temporal_patch_size`
        # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294
        padded_num_frames = num_frames + (-num_frames % temporal_patch_size)
        grid_t = max(padded_num_frames // temporal_patch_size, 1)
        grid_h = image_height // patch_size
        grid_w = image_width // patch_size
        num_patches = grid_t * grid_h * grid_w
        num_vision_tokens = num_patches
        return num_vision_tokens

    def get_max_image_tokens(self) -> int:
        target_width, target_height = self.get_image_size_with_most_features()
235
236
237
        return self.get_num_image_tokens(
            image_width=target_width, image_height=target_height
        )
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262

    def _get_max_video_frames(self, max_tokens: int) -> int:
        target_width, target_height = self.get_image_size_with_most_features()
        num_frames = 0
        while True:
            next_num_frames = num_frames + 1
            next_max_tokens = self.get_num_video_tokens(
                image_width=target_width,
                image_height=target_height,
                num_frames=next_num_frames,
                image_processor=None,
            )
            if next_max_tokens > max_tokens:
                break
            num_frames = next_num_frames
        return num_frames

    def get_num_frames_with_most_features(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> int:
        max_images = mm_counts.get("image", 0)
        max_videos = mm_counts.get("video", 0)
        max_image_tokens = self.get_max_image_tokens() * max_images
263
        max_total_frames = self._get_max_video_frames(seq_len - max_image_tokens)
264
265
266
267
268
269
270
271
272
        max_frames_per_video = max_total_frames // max(max_videos, 1)
        return max(max_frames_per_video, 1)

    def get_num_video_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
        num_frames: int,
273
        image_processor: BaseImageProcessor | None,
274
    ) -> int:
275
276
277
        num_video_tokens = self.get_num_image_tokens(
            image_width=image_width, image_height=image_height, num_frames=num_frames
        )
278
279
280
281
282
283
284
285
286
287
288
        return num_video_tokens

    def get_max_video_tokens(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> int:
        target_width, target_height = self.get_image_size_with_most_features()
        return self.get_num_video_tokens(
            image_width=target_width,
            image_height=target_height,
289
            num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts),
290
291
292
293
294
295
296
297
298
299
300
301
302
303
            image_processor=None,
        )


class Ovis2_5DummyInputsBuilder(BaseDummyInputsBuilder[Ovis2_5ProcessingInfo]):
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        num_images = mm_counts.get("image", 0)
        num_videos = mm_counts.get("video", 0)
        return IMAGE_TOKEN * num_images + VIDEO_TOKEN * num_videos

    def get_dummy_mm_data(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
304
        mm_options: Mapping[str, BaseDummyOptions] | None = None,
305
        mm_processor_kwargs: Mapping[str, object] | None = None,
306
307
308
309
    ) -> MultiModalDataDict:
        num_images = mm_counts.get("image", 0)
        num_videos = mm_counts.get("video", 0)

310
311
312
313
        target_width, target_height = self.info.get_image_size_with_most_features()
        target_num_frames = self.info.get_num_frames_with_most_features(
            seq_len, mm_counts
        )
314
315
316
317

        image_overrides = mm_options.get("image") if mm_options else None
        video_overrides = mm_options.get("video") if mm_options else None

318
        mm_data = {
319
320
321
322
323
324
325
            "image": self._get_dummy_images(
                width=target_width,
                height=target_height,
                num_images=num_images,
                overrides=image_overrides,
            ),
            "video": self._get_dummy_videos(
326
327
328
329
                width=target_width,
                height=target_height,
                num_frames=target_num_frames,
                num_videos=num_videos,
330
                overrides=video_overrides,
331
            ),
332
333
334
335
        }
        return mm_data


336
class Ovis2_5MultiModalProcessor(BaseMultiModalProcessor[Ovis2_5ProcessingInfo]):
337
338
339
340
341
    def visual_indicators_to_visual_tokens(
        self,
        visual_indicators: list[int],
    ) -> list[int]:
        """
342
        Filter image indicators placeholders and convert them to corresponding
343
344
345
346
347
348
        tokens in visual tokenizer.
        """
        hf_config = self.info.get_hf_config()
        vte_vocab_size = hf_config.visual_vocab_size
        return [
            vte_vocab_size - len(INDICATOR_IDS) + abs(x + 300) - 1
349
350
            for x in visual_indicators
            if x < -300
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
        ]

    def _call_hf_processor(
        self,
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
        tok_kwargs: Mapping[str, object],
    ) -> BatchFeature:
        if not mm_data:
            # Avoid warning from HF logger for text-only input
            tokenizer = self.info.get_tokenizer()
            prompt_ids = tokenizer.encode(prompt, add_special_tokens=False)
            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")

        processed_outputs = super()._call_hf_processor(
            prompt=prompt,
            mm_data=mm_data,
            mm_kwargs=mm_kwargs,
            tok_kwargs=tok_kwargs,
        )
        hf_processor = self.info.get_hf_processor()

        if "videos" in mm_data:
            visual_indicators = [
                hf_processor.construct_visual_indicators((1, 1, 1), True)
                for grid in processed_outputs["video_grids"]
            ]
            indicator_tokens = [
                self.visual_indicators_to_visual_tokens(indicator)
                for indicator in visual_indicators
            ]
383
            processed_outputs["video_indicator_tokens"] = torch.tensor(indicator_tokens)
384
385
386
387
388
389
390
391
392
393
        if "images" in mm_data:
            visual_indicators = [
                hf_processor.construct_visual_indicators((1, 1, 1), False)
                for grid in processed_outputs["grids"]
            ]
            indicator_tokens = [
                self.visual_indicators_to_visual_tokens(indicator)
                for indicator in visual_indicators
            ]

394
            processed_outputs["indicator_tokens"] = torch.tensor(indicator_tokens)
395
396
397
398
399
400
401
402
403
404
405
406
407
        return processed_outputs

    def _apply_hf_processor_tokens_only(
        self,
        prompt_tokens: list[int],
    ) -> list[int]:
        return prompt_tokens

    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
408
409
410
411
412
413
414
415
        return dict(
            pixel_values=MultiModalFieldConfig.batched("image"),
            grids=MultiModalFieldConfig.batched("image"),
            indicator_tokens=MultiModalFieldConfig.batched("image"),
            video_pixel_values=MultiModalFieldConfig.batched("video"),
            video_indicator_tokens=MultiModalFieldConfig.batched("video"),
            video_grids=MultiModalFieldConfig.batched("video"),
        )
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430

    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
        out_mm_kwargs: MultiModalKwargsItems,
    ) -> list[PromptReplacement]:
        def get_replacement_ovis(item_idx, modality: str):
            if modality == "image":
                out_item = out_mm_kwargs["image"][item_idx]
                grid = out_item["grids"].data
            elif modality == "video":
                out_item = out_mm_kwargs["video"][item_idx]
                grid = out_item["video_grids"].data
            hf_processor = self.info.get_hf_processor()
431
432
433
            return hf_processor.construct_visual_placeholders(
                grid[0],
            )
434
435
436
437
438
439

        return [
            PromptReplacement(
                modality=modality,
                target=IMAGE_TOKEN if modality == "image" else VIDEO_TOKEN,
                replacement=partial(get_replacement_ovis, modality=modality),
440
441
            )
            for modality in ("image", "video")
442
443
444
        ]


445
446
447
448
449
@MULTIMODAL_REGISTRY.register_processor(
    Ovis2_5MultiModalProcessor,
    info=Ovis2_5ProcessingInfo,
    dummy_inputs=Ovis2_5DummyInputsBuilder,
)
450
class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):
451
452
453
454
455
456
457
458
459
    @classmethod
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
        if modality.startswith("image"):
            return IMAGE_TOKEN
        if modality.startswith("video"):
            return VIDEO_TOKEN

        raise ValueError("Only image or video modality is supported")

460
461
462
463
464
465
466
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config

        self.config: PretrainedConfig = config

467
468
469
470
471
        with self._mark_language_model(vllm_config):
            self.llm = init_vllm_registered_model(
                vllm_config=vllm_config.with_hf_config(config.text_config),
                prefix=maybe_prefix(prefix, "llm"),
            )
472

473
474
475
476
477
478
479
480
        with self._mark_tower_model(vllm_config, {"image", "video"}):
            self.visual_tokenizer = VisualTokenizer(
                config=config.vit_config,
                visual_vocab_size=config.visual_vocab_size,
                quant_config=quant_config,
                prefix=f"{prefix}.visual_tokenizer",
            )
            self.vte = VisualEmbedding(config.visual_vocab_size, config.hidden_size)
481
482
483
484

        text_model_type = self.config.get_text_config().model_type
        self.image_pad_token_id = IMAGE_PAD_TOKEN_ID_MAP[text_model_type]

485
        self.make_empty_intermediate_tensors = (
486
487
            self.get_language_model().make_empty_intermediate_tensors
        )
488

489
    def _parse_and_validate_image_input(
490
        self, **kwargs: object
491
    ) -> Ovis2_5ImagePatchInputs | None:
492
493
494
        pixel_values = kwargs.pop("pixel_values", None)
        indicator_tokens = kwargs.pop("indicator_tokens", None)
        grids = kwargs.pop("grids", None)
495
496
497
498
499
        if pixel_values is None and indicator_tokens is None:
            return None

        if pixel_values is not None and indicator_tokens is not None:
            if not isinstance(pixel_values, (torch.Tensor, list)):
500
501
502
                raise ValueError(
                    f"Incorrect type of pixel values. Got type: {type(pixel_values)}"
                )
503
504

            if not isinstance(indicator_tokens, (torch.Tensor, list)):
505
506
507
508
                raise ValueError(
                    "Incorrect type of indicator_tokens. "
                    f"Got type: {type(indicator_tokens)}"
                )
509

510
            return Ovis2_5ImagePatchInputs(
511
                type="image_patches",
512
513
                flat_data=flatten_bn(pixel_values, concat=True),
                patches_per_item=[
514
                    x.shape[0] // (self.config.vit_config.hidden_stride**2)
515
                    for x in pixel_values
516
                ],
517
518
                indicator_tokens=flatten_bn(indicator_tokens, concat=True),
                grids=flatten_bn(grids, concat=True),
519
520
521
522
            )

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

523
    def _parse_and_validate_video_input(
524
        self, **kwargs: object
525
    ) -> Ovis2_5VideoPatchInputs | None:
526
527
528
529
530
531
532
533
        pixel_values = kwargs.pop("video_pixel_values", None)
        indicator_tokens = kwargs.pop("video_indicator_tokens", None)
        grids = kwargs.pop("video_grids", None)
        if pixel_values is None and indicator_tokens is None:
            return None

        if pixel_values is not None and indicator_tokens is not None:
            if not isinstance(pixel_values, (torch.Tensor, list)):
534
535
536
                raise ValueError(
                    f"Incorrect type of pixel values. Got type: {type(pixel_values)}"
                )
537
538

            if not isinstance(indicator_tokens, (torch.Tensor, list)):
539
540
541
542
                raise ValueError(
                    "Incorrect type of indicator_tokens. "
                    f"Got type: {type(indicator_tokens)}"
                )
543

544
            return Ovis2_5VideoPatchInputs(
545
                type="video_patches",
546
547
                flat_data=flatten_bn(pixel_values, concat=True),
                patches_per_item=[
548
                    x.shape[0] // (self.config.vit_config.hidden_stride**2)
549
                    for x in pixel_values
550
                ],
551
552
                indicator_tokens=flatten_bn(indicator_tokens, concat=True),
                grids=flatten_bn(grids, concat=True),
553
554
555
556
            )

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

557
    def _process_visual_input(
558
        self, visual_input: Ovis2_5ImagePatchInputs | Ovis2_5VideoPatchInputs
559
    ) -> MultiModalEmbeddings:
560
561
562
563
        image_patches_flat = visual_input["flat_data"]
        patches_per_image = visual_input["patches_per_item"]
        indicator_tokens = visual_input["indicator_tokens"]
        grid_thws = visual_input["grids"]
564
565

        indicator_per_image = list(
566
567
            map(lambda x: 2 if x > 1 else x + 2, patches_per_image)
        )
568
569
570

        target_dtype = self.visual_tokenizer.dtype
        visual_tokens = self.visual_tokenizer(
571
572
            image_patches_flat.to(target_dtype), grid_thws
        )
573
574
575
576
577

        visual_embeds = self.vte(visual_tokens)  # 1:1 numeric eq.
        indicator_embeds = self.vte(indicator_tokens)

        visual_embeds_per_image = visual_embeds.split(patches_per_image, dim=0)
578
        indicator_embeds_per_image = indicator_embeds.split(indicator_per_image)
579
580

        vision_embeddings = []
581
582
583
        for indicator, visual in zip(
            indicator_embeds_per_image, visual_embeds_per_image
        ):
584
585
586
587
            vision_embeddings_per_image = []
            visual = visual.unsqueeze(0)
            for i in range(visual.shape[0]):
                vision_embeddings_per_image.append(
588
589
590
591
                    torch.cat([indicator[i : i + 1], visual[i]], dim=0)
                )
            vision_embeddings_per_image.append(indicator[i + 1 :])
            vision_embeddings.append(torch.cat(vision_embeddings_per_image, dim=0))
592
593
        return tuple(vision_embeddings)

594
595
596
597
598
599
    def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
        modalities = {}

        # Preserve the order of modalities if there are multiple of them
        # from the order of kwargs.
        for input_key in kwargs:
600
601
602
603
604
605
606
607
608
609
610
            if (
                input_key in ("pixel_values", "indicator_tokens", "grids")
                and "images" not in modalities
            ):
                modalities["images"] = self._parse_and_validate_image_input(**kwargs)
            if (
                input_key
                in ("video_pixel_values", "video_indicator_tokens", "video_grids")
                and "videos" not in modalities
            ):
                modalities["videos"] = self._parse_and_validate_video_input(**kwargs)
611
612
613

        return modalities

614
    def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
615
616
617
618
619
620
621
622
623
624
        modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
        if not modalities:
            return []

        multimodal_embeddings: tuple[torch.Tensor, ...] = ()
        # NOTE: It is important to iterate over the keys in this dictionary
        # to preserve the order of the modalities.
        for modality in modalities:
            if modality == "images":
                image_input = modalities["images"]
625
626
                image_embeddings = self._process_visual_input(image_input)
                multimodal_embeddings += tuple(image_embeddings)
627
628
            if modality == "videos":
                video_input = modalities["videos"]
629
                video_embeddings = self._process_visual_input(video_input)
630
                multimodal_embeddings += tuple(video_embeddings)
631
632

        return multimodal_embeddings
633
634
635

    def forward(
        self,
636
        input_ids: torch.Tensor | None,
637
        positions: torch.Tensor,
638
639
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
640
        **kwargs: object,
641
    ) -> torch.Tensor | IntermediateTensors:
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
        if intermediate_tensors is not None:
            inputs_embeds = None

        # up until here we have a inputs_embeds 100% numerical identity
        # between the OG HF Transformers implementation and ours
        hidden_states = self.llm(
            input_ids=input_ids,
            positions=positions,
            intermediate_tensors=intermediate_tensors,
            inputs_embeds=inputs_embeds,
        )
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
658
    ) -> torch.Tensor | None:
659
        return self.llm.compute_logits(hidden_states)
660

661
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
662
663
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights)