interns1.py 29.9 KB
Newer Older
Lyu Han's avatar
Lyu Han committed
1
2
3
4
5
6
7
8
9
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

# --------------------------------------------------------
# InternS1
# Copyright (c) 2025 Shanghai AI Lab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
from collections.abc import Iterable, Mapping, Sequence
10
from typing import Annotated, Literal, TypeAlias
Lyu Han's avatar
Lyu Han committed
11

12
import regex as re
Lyu Han's avatar
Lyu Han committed
13
14
import torch
import torch.nn as nn
15
from transformers import BatchFeature, InternVLProcessor, PretrainedConfig
Lyu Han's avatar
Lyu Han committed
16
17
from transformers.activations import ACT2FN
from transformers.models.got_ocr2.image_processing_got_ocr2_fast import (
18
19
    GotOcr2ImageProcessorFast,
)
20
from transformers.models.internvl.video_processing_internvl import (
21
22
    InternVLVideoProcessor,
)
Lyu Han's avatar
Lyu Han committed
23
24

from vllm.config import VllmConfig
25
from vllm.config.multimodal import BaseDummyOptions
Lyu Han's avatar
Lyu Han committed
26
27
28
29
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.models.interns1_vit import InternS1VisionModel
from vllm.model_executor.models.module_mapping import MultiModelKeys
from vllm.multimodal import MULTIMODAL_REGISTRY
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
from vllm.multimodal.inputs import (
    MultiModalDataDict,
    MultiModalFieldConfig,
    MultiModalKwargsItems,
)
from vllm.multimodal.parse import (
    ImageEmbeddingItems,
    ImageProcessorItems,
    ImageSize,
    MultiModalDataItems,
)
from vllm.multimodal.processing import (
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    PromptReplacement,
    PromptUpdate,
    PromptUpdateDetails,
)
Lyu Han's avatar
Lyu Han committed
48
49
from vllm.multimodal.profiling import BaseDummyInputsBuilder
from vllm.sequence import IntermediateTensors
50
from vllm.transformers_utils.processor import cached_video_processor_from_config
51
from vllm.utils.tensor_schema import TensorSchema, TensorShape
Lyu Han's avatar
Lyu Han committed
52

53
54
55
56
57
58
59
60
61
62
63
64
from .interfaces import (
    MultiModalEmbeddings,
    SupportsLoRA,
    SupportsMultiModal,
    SupportsPP,
)
from .utils import (
    AutoWeightsLoader,
    WeightsMapper,
    init_vllm_registered_model,
    maybe_prefix,
)
Lyu Han's avatar
Lyu Han committed
65
66
67
68
69


class InternS1MultiModalProjector(nn.Module):
    def __init__(self, config):
        super().__init__()
70
71
72
        self.layer_norm = nn.LayerNorm(
            config.vision_config.hidden_size * int(1 / config.downsample_ratio) ** 2
        )
Lyu Han's avatar
Lyu Han committed
73
        self.linear_1 = nn.Linear(
74
75
76
            config.vision_config.hidden_size * int(1 / config.downsample_ratio) ** 2,
            config.text_config.hidden_size,
        )
Lyu Han's avatar
Lyu Han committed
77
        self.act = ACT2FN[config.projector_hidden_act]
78
79
80
        self.linear_2 = nn.Linear(
            config.text_config.hidden_size, config.text_config.hidden_size
        )
Lyu Han's avatar
Lyu Han committed
81
82
83
84
85
86
87
88
89

    def forward(self, image_features):
        hidden_states = self.layer_norm(image_features)
        hidden_states = self.linear_1(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.linear_2(hidden_states)
        return hidden_states


90
class InternS1ImagePixelInputs(TensorSchema):
Lyu Han's avatar
Lyu Han committed
91
    """
92
93
94
95
96
97
    Dimensions:
        - bnp: Batch size * number of images * (1 + num_patches)
        - c: Number of channels (3)
        - h: Height
        - w: Width
        - bn: Batch size * number of images
Lyu Han's avatar
Lyu Han committed
98
    """
99

100
101
102
    type: Literal["pixel_values"] = "pixel_values"
    pixel_values: Annotated[torch.Tensor, TensorShape("bnp", 3, "h", "w")]
    num_patches: Annotated[torch.Tensor, TensorShape("bn")]
Lyu Han's avatar
Lyu Han committed
103
104


105
class InternS1ImageEmbeddingInputs(TensorSchema):
Lyu Han's avatar
Lyu Han committed
106
    """
107
108
109
110
    Dimensions:
        - ni: Number of images
        - tifs: Total image feature size
        - hs: Hidden size (must match language model backbone)
Lyu Han's avatar
Lyu Han committed
111
    """
112

113
    type: Literal["image_embeds"] = "image_embeds"
114
    data: Annotated[torch.Tensor | list[torch.Tensor], TensorShape("ni", "tifs", "hs")]
Lyu Han's avatar
Lyu Han committed
115
116


117
InternS1ImageInputs: TypeAlias = InternS1ImagePixelInputs | InternS1ImageEmbeddingInputs
Lyu Han's avatar
Lyu Han committed
118
119


120
class InternS1VideoPixelInputs(TensorSchema):
Lyu Han's avatar
Lyu Han committed
121
    """
122
123
124
125
126
127
    Dimensions:
        - bnv: Batch size * number of videos * number of frames
        - bn: Batch size * number of images
        - c: Number of channels (3)
        - h: Height
        - w: Width
Lyu Han's avatar
Lyu Han committed
128
    """
129

130
131
132
    type: Literal["pixel_values_videos"] = "pixel_values_videos"
    pixel_values: Annotated[torch.Tensor, TensorShape("bnv", 3, "h", "w")]
    num_patches: Annotated[torch.Tensor, TensorShape("bn")]
Lyu Han's avatar
Lyu Han committed
133
134


135
class InternS1VideoEmbeddingInputs(TensorSchema):
Lyu Han's avatar
Lyu Han committed
136
    """
137
138
139
140
    Dimensions:
        - nv: Number of videos
        - tvfs: Total video feature size
        - hs: Hidden size (must match language model backbone)
Lyu Han's avatar
Lyu Han committed
141
    """
142

143
    type: Literal["video_embeds"] = "video_embeds"
144
    data: Annotated[torch.Tensor | list[torch.Tensor], TensorShape("nv", "tvfs", "hs")]
Lyu Han's avatar
Lyu Han committed
145
146


147
InternS1VideoInputs: TypeAlias = InternS1VideoPixelInputs | InternS1VideoEmbeddingInputs
Lyu Han's avatar
Lyu Han committed
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168


def resolve_interns1_min_max_num(
    min_dynamic_patch: int,
    max_dynamic_patch: int,
    dynamic_image_size: bool,
    use_thumbnail: bool,
) -> tuple[int, int]:
    min_dynamic_patch = min_dynamic_patch if dynamic_image_size else 1
    max_dynamic_patch = max_dynamic_patch if dynamic_image_size else 1

    if use_thumbnail and max_dynamic_patch != 1:
        max_dynamic_patch += 1

    return min_dynamic_patch, max_dynamic_patch


def get_interns1_target_ratios(
    min_num: int,
    max_num: int,
) -> list[tuple[int, int]]:
169
170
171
172
173
174
175
    target_ratios = {
        (i, j)
        for n in range(min_num, max_num + 1)
        for i in range(1, n + 1)
        for j in range(1, n + 1)
        if min_num <= i * j <= max_num
    }
Lyu Han's avatar
Lyu Han committed
176
177
178
179
    return sorted(target_ratios, key=lambda x: x[0] * x[1])


class InternS1ProcessingInfo(BaseProcessingInfo):
180
    """ProcessingInfo for InternS1-style models."""
Lyu Han's avatar
Lyu Han committed
181
182

    def get_hf_processor(self, **kwargs: object) -> InternVLProcessor:
183
184
        hf_processor = self.ctx.get_hf_processor(InternVLProcessor, **kwargs)
        hf_processor.video_processor = cached_video_processor_from_config(
185
186
            self.ctx.model_config, processor_cls=InternVLVideoProcessor, **kwargs
        )
187
        return hf_processor
Lyu Han's avatar
Lyu Han committed
188

189
    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
190
        return {"image": None, "video": None}
Lyu Han's avatar
Lyu Han committed
191
192
193
194
195
196

    def get_num_image_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
197
        processor: GotOcr2ImageProcessorFast | None = None,
Lyu Han's avatar
Lyu Han committed
198
199
200
201
202
    ) -> int:
        if processor is None:
            processor = self.get_hf_processor().image_processor

        if not isinstance(processor, GotOcr2ImageProcessorFast):
203
204
205
            raise ValueError(
                f"GotOcr2ImageProcessorFast is expected but got {type(processor)}"
            )
206
        num_image_patches = processor.get_number_of_image_patches(
207
208
209
            image_height, image_width, images_kwargs=dict()
        )
        num_image_tokens = self.get_hf_processor().image_seq_length * num_image_patches
Lyu Han's avatar
Lyu Han committed
210
211
        return num_image_tokens

212
    def resolve_target_ratios(self, use_thumbnail: bool | None = None):
Lyu Han's avatar
Lyu Han committed
213
214
215
216
217
218
219
220
221
222
223
        image_processor = self.get_hf_processor().image_processor
        min_dynamic_patch = image_processor.min_patches
        max_dynamic_patch = image_processor.max_patches
        # HF format's InternVL processor uses `crop_to_patches` which is
        # equivalent to `use_thumbnail` in original format.
        use_thumbnail = image_processor.crop_to_patches
        dynamic_image_size = True
        min_num, max_num = resolve_interns1_min_max_num(
            min_dynamic_patch,
            max_dynamic_patch,
            dynamic_image_size,
224
225
            use_thumbnail=use_thumbnail,
        )
Lyu Han's avatar
Lyu Han committed
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246

        return get_interns1_target_ratios(min_num, max_num)

    def get_image_size_with_most_features(self) -> ImageSize:
        processor = self.get_hf_processor()

        hf_config = self.ctx.get_hf_config()
        base_height, base_width = hf_config.vision_config.image_size
        target_ratios = self.resolve_target_ratios()

        largest_feature_size, largest_feature_pinpoint = 0, None
        for wr, hr in target_ratios:
            width, height = base_width * wr, base_height * hr

            feat_size = self.get_num_image_tokens(
                image_width=width,
                image_height=height,
                processor=processor.image_processor,
            )
            if feat_size > largest_feature_size:
                largest_feature_size = feat_size
247
                largest_feature_pinpoint = ImageSize(width=width, height=height)
Lyu Han's avatar
Lyu Han committed
248

249
250
251
        assert not (largest_feature_size == 0 or largest_feature_pinpoint is None), (
            "Cannot have a largest feature size of 0!"
        )
Lyu Han's avatar
Lyu Han committed
252
253
254
255
256
257
258
259
260
261
262
263
264

        return largest_feature_pinpoint

    def get_max_image_tokens(self) -> int:
        processor = self.get_hf_processor()
        target_width, target_height = self.get_image_size_with_most_features()

        return self.get_num_image_tokens(
            image_width=target_width,
            image_height=target_height,
            processor=processor.image_processor,
        )

265
266
267
268
269
270
271
272
273
274
275
    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)

        processor = self.get_hf_processor()

        max_image_tokens = self.get_max_image_tokens() * max_images
276
        max_total_frames = (seq_len - max_image_tokens) // processor.image_seq_length
277
278
279
280
        max_frames_per_video = max_total_frames // max(max_videos, 1)

        return max(max_frames_per_video, 1)

Lyu Han's avatar
Lyu Han committed
281

282
class InternS1DummyInputsBuilder(BaseDummyInputsBuilder[InternS1ProcessingInfo]):
283
    """DummyInputsBuilder for InternS1-style models."""
Lyu Han's avatar
Lyu Han committed
284
285
286

    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        num_images = mm_counts.get("image", 0)
287
        num_videos = mm_counts.get("video", 0)
Lyu Han's avatar
Lyu Han committed
288
        image_token = self.info.get_hf_processor().image_token
289
        video_token = self.info.get_hf_processor().video_token
Lyu Han's avatar
Lyu Han committed
290

291
        return image_token * num_images + video_token * num_videos
Lyu Han's avatar
Lyu Han committed
292
293
294
295
296

    def get_dummy_mm_data(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
297
        mm_options: Mapping[str, BaseDummyOptions] | None = None,
Lyu Han's avatar
Lyu Han committed
298
    ) -> MultiModalDataDict:
299
300
301
302
        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
        )
Lyu Han's avatar
Lyu Han committed
303
        num_images = mm_counts.get("image", 0)
304
305
306
307
        num_videos = mm_counts.get("video", 0)

        config = self.info.get_hf_config()
        image_size_h, image_size_w = config.vision_config.image_size
Lyu Han's avatar
Lyu Han committed
308

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

Lyu Han's avatar
Lyu Han committed
312
        return {
313
314
315
316
317
318
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(
                width=image_size_w,
                height=image_size_h,
                num_frames=target_num_frames,
                num_videos=num_videos,
                overrides=video_overrides,
            ),
Lyu Han's avatar
Lyu Han committed
326
327
328
        }


329
330
class InternS1MultiModalProcessor(BaseMultiModalProcessor[InternS1ProcessingInfo]):
    """Basic image-only MultiModalProcessor for InternS1-style models."""
Lyu Han's avatar
Lyu Han committed
331
332
333
334
335
336
337

    def _call_hf_processor(
        self,
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
        tok_kwargs: Mapping[str, object],
338
    ) -> BatchFeature:
339
340
341
342
343
        mm_data = dict(mm_data)
        videos = mm_data.pop("videos", [])
        images = mm_data.pop("images", [])
        assert isinstance(videos, list)
        assert isinstance(images, list)
Lyu Han's avatar
Lyu Han committed
344
345

        hf_processor = self.info.get_hf_processor(**mm_kwargs)
346
        tokenizer = hf_processor.tokenizer
347
348
349
        video_token_id = tokenizer.encode(
            hf_processor.video_token, add_special_tokens=False
        )
350
351
352
        assert len(video_token_id) == 1
        video_token_id = video_token_id[0]

353
354
        prompt = re.sub(hf_processor.image_token, "<image_placeholder>", prompt)
        prompt = re.sub(hf_processor.video_token, "<video_placeholder>", prompt)
355
356
357
358
359
360
361
362
363
364
365

        image_outputs = {}
        if images:
            image_pixel_values = []
            for image in images:
                processed_outputs = super()._call_hf_processor(
                    prompt=hf_processor.image_token,
                    mm_data={"images": image},
                    mm_kwargs=mm_kwargs,
                    tok_kwargs=tok_kwargs,
                )
366
                image_pixel_values.append(processed_outputs.pop("pixel_values"))
367
368
369

                input_ids = processed_outputs.pop("input_ids")
                image_placeholder = tokenizer.batch_decode(input_ids)[0]
370
                prompt = prompt.replace("<image_placeholder>", image_placeholder, 1)
371
372

            num_patches = [len(item) for item in image_pixel_values]
373
            image_outputs = {
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
                "pixel_values": torch.concat(image_pixel_values),
                "image_num_patches": torch.tensor(num_patches),
                "image_token_id": torch.tensor(hf_processor.image_token_id),
            }

        video_outputs = {}
        if videos:
            video_pixel_values = []
            for video in videos:
                processed_outputs = super()._call_hf_processor(
                    prompt=hf_processor.video_token,
                    mm_data={"videos": video},
                    mm_kwargs=mm_kwargs,
                    tok_kwargs=tok_kwargs,
                )
389
                video_pixel_values.append(processed_outputs.pop("pixel_values"))
390
391

                input_ids = processed_outputs.pop("input_ids")
392
                input_ids[input_ids == hf_processor.image_token_id] = video_token_id
393
394

                video_placeholder = tokenizer.batch_decode(input_ids)[0]
395
                prompt = prompt.replace("<video_placeholder>", video_placeholder, 1)
396
397

            num_frames = [len(item) for item in video_pixel_values]
398
            video_outputs = {
399
400
401
402
403
                "pixel_values_videos": torch.concat(video_pixel_values),
                "video_num_patches": torch.tensor(num_frames),
                "video_token_id": torch.tensor(video_token_id),
            }

404
405
        prompt = re.sub("<image_placeholder>", hf_processor.image_token, prompt)
        prompt = re.sub("<video_placeholder>", hf_processor.video_token, prompt)
406
407
        text_outputs = tokenizer(prompt, **tok_kwargs, return_tensors="pt")

408
        return BatchFeature({**text_outputs, **image_outputs, **video_outputs})
Lyu Han's avatar
Lyu Han committed
409
410
411

    def _get_mm_fields_config(
        self,
412
        hf_inputs: BatchFeature,
Lyu Han's avatar
Lyu Han committed
413
414
415
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
        image_num_patches = hf_inputs.get("image_num_patches", torch.empty(0))
416
        video_num_patches = hf_inputs.get("video_num_patches", torch.empty(0))
Lyu Han's avatar
Lyu Han committed
417
        num_images = len(image_num_patches)
418
        num_videos = len(video_num_patches)
Lyu Han's avatar
Lyu Han committed
419
420
421

        return dict(
            pixel_values=MultiModalFieldConfig.flat_from_sizes(
422
423
                "image", image_num_patches
            ),
Lyu Han's avatar
Lyu Han committed
424
425
426
            image_num_patches=MultiModalFieldConfig.batched("image"),
            image_embeds=MultiModalFieldConfig.batched("image"),
            image_token_id=MultiModalFieldConfig.shared("image", num_images),
427
            pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
428
429
                "video", video_num_patches
            ),
430
431
            video_num_patches=MultiModalFieldConfig.batched("video"),
            video_token_id=MultiModalFieldConfig.shared("video", num_videos),
Lyu Han's avatar
Lyu Han committed
432
433
434
435
436
437
        )

    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
438
        out_mm_kwargs: MultiModalKwargsItems,
Lyu Han's avatar
Lyu Han committed
439
440
441
442
443
    ) -> Sequence[PromptUpdate]:
        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
        img_context_token = hf_processor.image_token
        start_image_token = hf_processor.start_image_token
        end_image_token = hf_processor.end_image_token
444
445
        video_token = hf_processor.video_token

446
447
448
        out_mm_data = out_mm_kwargs.get_data()
        if "video_num_patches" in out_mm_data:
            video_num_patches = out_mm_data["video_num_patches"]
449
450
451
452
453
            assert isinstance(video_num_patches, torch.Tensor)
            video_num_patches = video_num_patches.tolist()
        else:
            video_num_patches = []

454
455
        if "image_num_patches" in out_mm_data:
            image_num_patches = out_mm_data["image_num_patches"]
456
457
458
459
460
461
            assert isinstance(image_num_patches, torch.Tensor)
            image_num_patches = image_num_patches.tolist()
        else:
            image_num_patches = []

        def get_replacement_interns1_image(item_idx: int):
Lyu Han's avatar
Lyu Han committed
462
            images = mm_items.get_items(
463
464
                "image", (ImageEmbeddingItems, ImageProcessorItems)
            )
Lyu Han's avatar
Lyu Han committed
465
466
467
468

            if isinstance(images, ImageEmbeddingItems):
                feature_size = images.get_feature_size(item_idx)
            else:
469
470
                num_patches = image_num_patches[item_idx]
                feature_size = num_patches * hf_processor.image_seq_length
Lyu Han's avatar
Lyu Han committed
471
472
473

            repl_features = img_context_token * feature_size
            repl_full = start_image_token + repl_features + end_image_token
474
            return PromptUpdateDetails.select_text(repl_full, img_context_token)
Lyu Han's avatar
Lyu Han committed
475

476
477
478
        def get_replacement_interns1_video(item_idx: int):
            num_patches = video_num_patches[item_idx]
            repl_features = video_token * hf_processor.image_seq_length
479
            repl_features_with_sep = start_image_token + repl_features + end_image_token
480
            # num_patches is equal to num_frames
481
482
483
            repl_full = "\n".join(
                [f"Frame{i + 1}: {repl_features_with_sep}" for i in range(num_patches)]
            )
484
485
486

            return PromptUpdateDetails.select_text(repl_full, video_token)

Lyu Han's avatar
Lyu Han committed
487
488
489
490
        return [
            PromptReplacement(
                modality="image",
                target=img_context_token,
491
492
493
494
495
496
497
                replacement=get_replacement_interns1_image,
            ),
            PromptReplacement(
                modality="video",
                target=video_token,
                replacement=get_replacement_interns1_video,
            ),
Lyu Han's avatar
Lyu Han committed
498
499
500
501
502
503
        ]


@MULTIMODAL_REGISTRY.register_processor(
    InternS1MultiModalProcessor,
    info=InternS1ProcessingInfo,
504
505
506
507
508
    dummy_inputs=InternS1DummyInputsBuilder,
)
class InternS1ForConditionalGeneration(
    nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA
):
509
    merge_by_field_config = True
Lyu Han's avatar
Lyu Han committed
510
511
512
513
514
515
516
517

    # To ensure correct weight loading and mapping.
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            "lm_head.": "language_model.lm_head.",
            "model.language_model.": "language_model.model.",
            "model.vision_tower.": "vision_tower.",
            "model.multi_modal_projector.": "multi_modal_projector.",
518
519
        }
    )
Lyu Han's avatar
Lyu Han committed
520
521

    @classmethod
522
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
co63oc's avatar
co63oc committed
523
        # transformers InternVLProcessor uses <IMG_CONTEXT> as the separator
Lyu Han's avatar
Lyu Han committed
524
525
        # refer to https://github.com/huggingface/transformers/blob/f90de364c2484c7c325bbe05befdcf487bd75b63/src/transformers/models/internvl/processing_internvl.py#L116
        if modality.startswith("image"):
526
            return "<IMG_CONTEXT>"
Lyu Han's avatar
Lyu Han committed
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
        if modality.startswith("video"):
            return "<video>"

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

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config

        self.config = config
        self.multimodal_config = multimodal_config

        image_size = config.vision_config.image_size[0]
        patch_size = config.vision_config.patch_size[0]
        self.patch_size = patch_size
        self.num_image_token = int(
545
546
            (image_size // patch_size) ** 2 * (config.downsample_ratio**2)
        )
Lyu Han's avatar
Lyu Han committed
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
        self.downsample_ratio = config.downsample_ratio

        self.llm_arch_name = config.text_config.architectures[0]
        self.vision_tower = self._init_vision_model(
            config,
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, "vision_tower"),
        )

        self.language_model = init_vllm_registered_model(
            vllm_config=vllm_config,
            hf_config=config.text_config,
            prefix=maybe_prefix(prefix, "language_model"),
        )

        self.multi_modal_projector = self._init_mlp1(config)

        self.img_context_token_id = None
        self.video_context_token_id = None

        self.visual_token_mask = None
        self.make_empty_intermediate_tensors = (
569
570
            self.language_model.make_empty_intermediate_tensors
        )
Lyu Han's avatar
Lyu Han committed
571
572
573
574

    def _init_vision_model(
        self,
        config: PretrainedConfig,
575
        quant_config: QuantizationConfig | None,
Lyu Han's avatar
Lyu Han committed
576
577
578
579
580
581
582
583
584
585
586
        *,
        prefix: str,
    ):
        num_hidden_layers = config.vision_config.num_hidden_layers
        return InternS1VisionModel(
            config.vision_config,
            quant_config=quant_config,
            num_hidden_layers_override=num_hidden_layers,
            prefix=prefix,
        )

587
    def _init_mlp1(self, config: PretrainedConfig) -> nn.Module:
Lyu Han's avatar
Lyu Han committed
588
589
590
591
592
593
594
595
        return InternS1MultiModalProjector(config)

    def pixel_shuffle(self, x, scale_factor=0.5):
        n, w, h, c = x.size()
        # N, W, H, C --> N, W, H * scale, C // scale
        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
        x = x.permute(0, 2, 1, 3).contiguous()
596
597
598
599
600
601
        x = x.view(
            n,
            int(h * scale_factor),
            int(w * scale_factor),
            int(c / (scale_factor * scale_factor)),
        )
Lyu Han's avatar
Lyu Han committed
602
603
604
605
606
607
608
        x = x.permute(0, 2, 1, 3).contiguous()
        return x

    def extract_feature(self, pixel_values: torch.Tensor) -> torch.Tensor:
        vit_embeds = self.vision_tower(pixel_values=pixel_values)
        vit_embeds = vit_embeds[:, 1:, :]

609
        h = w = int(vit_embeds.shape[1] ** 0.5)
Lyu Han's avatar
Lyu Han committed
610
        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
611
612
        vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
Lyu Han's avatar
Lyu Han committed
613
614
615
616
617

        vit_embeds = self.multi_modal_projector(vit_embeds)
        return vit_embeds

    def _parse_and_validate_image_input(
618
        self, **kwargs: object
619
    ) -> InternS1ImageInputs | None:
Lyu Han's avatar
Lyu Han committed
620
621
622
623
624
625
626
627
628
629
        pixel_values = kwargs.pop("pixel_values", None)
        image_num_patches = kwargs.pop("image_num_patches", None)
        image_embeds = kwargs.pop("image_embeds", None)

        if pixel_values is None and image_embeds is None:
            return None

        if image_embeds is not None:
            return InternS1ImageEmbeddingInputs(
                type="image_embeds",
630
                data=image_embeds,
Lyu Han's avatar
Lyu Han committed
631
632
633
            )

        image_token_id = kwargs["image_token_id"]
634
635
636
637
638
        if isinstance(image_token_id, torch.Tensor):
            image_token_id = image_token_id.flatten().unique().item()

        assert isinstance(image_token_id, int)
        self.img_context_token_id = image_token_id
Lyu Han's avatar
Lyu Han committed
639
640

        if pixel_values is not None:
641
            h, w = self.config.vision_config.image_size
Lyu Han's avatar
Lyu Han committed
642
643
            return InternS1ImagePixelInputs(
                type="pixel_values",
644
                pixel_values=pixel_values,
Lyu Han's avatar
Lyu Han committed
645
                num_patches=image_num_patches,
646
647
648
649
                resolve_bindings={
                    "h": h,
                    "w": w,
                },
Lyu Han's avatar
Lyu Han committed
650
651
652
653
654
            )

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

    def _parse_and_validate_video_input(
655
        self, **kwargs: object
656
    ) -> InternS1VideoInputs | None:
657
        pixel_values_flat_video = kwargs.pop("pixel_values_videos", None)
Lyu Han's avatar
Lyu Han committed
658
659
660
661
662
663
664
        video_num_patches = kwargs.pop("video_num_patches", None)
        video_embeds = kwargs.pop("video_embeds", None)

        if pixel_values_flat_video is None and video_embeds is None:
            return None

        if video_embeds is not None:
665
            return InternS1VideoEmbeddingInputs(
Lyu Han's avatar
Lyu Han committed
666
                type="video_embeds",
667
                data=video_embeds,
Lyu Han's avatar
Lyu Han committed
668
669
670
            )

        video_token_id = kwargs["video_token_id"]
671
672
673
674
675
        if isinstance(video_token_id, torch.Tensor):
            video_token_id = video_token_id.flatten().unique().item()

        assert isinstance(video_token_id, int)
        self.video_context_token_id = video_token_id
Lyu Han's avatar
Lyu Han committed
676
677

        if pixel_values_flat_video is not None:
678
            h, w = self.config.vision_config.image_size
Lyu Han's avatar
Lyu Han committed
679
680
681
            return InternS1VideoPixelInputs(
                type="pixel_values_videos",
                num_patches=video_num_patches,
682
683
684
685
686
                pixel_values=pixel_values_flat_video,
                resolve_bindings={
                    "h": h,
                    "w": w,
                },
Lyu Han's avatar
Lyu Han committed
687
688
689
690
            )

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

691
    def _process_vision_input(
Lyu Han's avatar
Lyu Han committed
692
        self,
693
        image_input: InternS1ImageInputs | InternS1VideoInputs,
Lyu Han's avatar
Lyu Han committed
694
    ) -> tuple[torch.Tensor, ...]:
695
696
697
698
        if (
            image_input["type"] == "image_embeds"
            or image_input["type"] == "video_embeds"
        ):
Lyu Han's avatar
Lyu Han committed
699
700
701
702
703
704
705
706
707
708
            return image_input["data"]

        assert self.vision_tower is not None

        image_embeds = self.extract_feature(image_input["pixel_values"])

        num_patches = image_input["num_patches"]

        # Only one image in the current batch
        if len(num_patches) == 1:
709
            return (image_embeds.view(-1, self.config.text_config.hidden_size),)
Lyu Han's avatar
Lyu Han committed
710
711
712
713

        # NOTE: Image embeddings are split into separate tensors for each image
        # by the size of each embedding.
        feature_size = image_embeds.shape[1]
714
        image_embeds = image_embeds.view(-1, self.config.text_config.hidden_size)
Lyu Han's avatar
Lyu Han committed
715
716
717
718
719
720
721
722
723
724
725
        image_feature_sizes = [
            num_patches * feature_size for num_patches in num_patches
        ]
        return image_embeds.split(image_feature_sizes)

    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:
726
727
728
729
730
731
732
            if (
                input_key in ("pixel_values", "image_embeds")
                and "images" not in modalities
            ):
                modalities["images"] = self._parse_and_validate_image_input(**kwargs)
            if input_key in ("pixel_values_videos",) and "videos" not in modalities:
                modalities["videos"] = self._parse_and_validate_video_input(**kwargs)
Lyu Han's avatar
Lyu Han committed
733
734
735
736
737
738
739
740
741

        return modalities

    def _set_visual_token_mask(self, input_ids: torch.Tensor) -> None:
        self.visual_token_mask = None

    def get_language_model(self) -> torch.nn.Module:
        return self.language_model

742
    def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
Lyu Han's avatar
Lyu Han committed
743
744
745
746
747
        modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
        if not modalities:
            return []

        # The result multimodal_embeddings is tuple of tensors, with each
748
        # tensor corresponding to a multimodal data item (image or video).
Lyu Han's avatar
Lyu Han committed
749
750
751
752
753
754
755
        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"]
756
                vision_embeddings = self._process_vision_input(image_input)
Lyu Han's avatar
Lyu Han committed
757
758
759
                multimodal_embeddings += vision_embeddings
            if modality == "videos":
                video_input = modalities["videos"]
760
                video_embeddings = self._process_vision_input(video_input)
Lyu Han's avatar
Lyu Han committed
761
762
763
764
765
766
767
                multimodal_embeddings += video_embeddings

        return multimodal_embeddings

    def get_input_embeddings(
        self,
        input_ids: torch.Tensor,
768
        multimodal_embeddings: MultiModalEmbeddings | None = None,
769
        *,
770
        is_multimodal: torch.Tensor | None = None,
771
        handle_oov_mm_token: bool = False,
Lyu Han's avatar
Lyu Han committed
772
    ) -> torch.Tensor:
773
        if multimodal_embeddings is not None and len(multimodal_embeddings) > 0:
Lyu Han's avatar
Lyu Han committed
774
            self._set_visual_token_mask(input_ids)
775
776
777
778
779
780
781
782
783
784
785

        # This is to satisfy the type checker for each overload
        if multimodal_embeddings is None or is_multimodal is None:
            return super().get_input_embeddings(input_ids)

        return super().get_input_embeddings(
            input_ids,
            multimodal_embeddings=multimodal_embeddings,
            is_multimodal=is_multimodal,
            handle_oov_mm_token=handle_oov_mm_token,
        )
Lyu Han's avatar
Lyu Han committed
786
787
788
789
790

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
791
792
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
Lyu Han's avatar
Lyu Han committed
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
        **kwargs: object,
    ) -> IntermediateTensors:
        if intermediate_tensors is not None:
            input_ids = None
            inputs_embeds = None

        forward_kwargs = {
            "input_ids": input_ids,
            "positions": positions,
            "intermediate_tensors": intermediate_tensors,
            "inputs_embeds": inputs_embeds,
        }

        hidden_states = self.language_model.model(**forward_kwargs)
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
812
    ) -> torch.Tensor | None:
813
        return self.language_model.compute_logits(hidden_states)
Lyu Han's avatar
Lyu Han committed
814

815
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
Lyu Han's avatar
Lyu Han committed
816
817
818
819
820
821
822
823
824
825
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

    def get_mm_mapping(self) -> MultiModelKeys:
        """
        Get the module prefix in multimodal models
        """
        return MultiModelKeys.from_string_field(
            language_model="language_model",
            connector="multi_modal_projector",
826
827
            tower_model="vision_tower",
        )