llava_next.py 24.6 KB
Newer Older
1
from functools import cached_property
2
3
from typing import (Iterable, List, Literal, Mapping, Optional, Tuple,
                    TypedDict, Union)
4
5
6

import torch
import torch.nn as nn
7
from PIL import Image
8
from transformers import CLIPVisionConfig, LlavaNextConfig, SiglipVisionConfig
9
10
11
12
13
from transformers.models.llava_next.modeling_llava_next import (
    get_anyres_image_grid_shape, unpad_image)
from typing_extensions import NotRequired

from vllm.attention import AttentionMetadata
14
from vllm.config import CacheConfig, MultiModalConfig
15
from vllm.inputs import INPUT_REGISTRY, DecoderOnlyInputs, InputContext
Cyrus Leung's avatar
Cyrus Leung committed
16
from vllm.model_executor.layers.pooler import Pooler, PoolingType
17
from vllm.model_executor.layers.quantization import QuantizationConfig
18
from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
Cyrus Leung's avatar
Cyrus Leung committed
19
from vllm.model_executor.pooling_metadata import PoolingMetadata
20
from vllm.model_executor.sampling_metadata import SamplingMetadata
21
from vllm.multimodal import MULTIMODAL_REGISTRY
Cyrus Leung's avatar
Cyrus Leung committed
22
from vllm.sequence import IntermediateTensors, PoolerOutput
23
from vllm.utils import is_list_of
24

25
26
from .clip import (CLIPVisionModel, dummy_image_for_clip,
                   dummy_seq_data_for_clip, get_clip_image_feature_size,
27
                   get_clip_patch_grid_length, input_processor_for_clip)
28
from .interfaces import SupportsMultiModal, SupportsPP
29
from .llava import LlavaMultiModalProjector, init_vision_tower_for_llava
30
31
32
from .siglip import (SiglipVisionModel, dummy_image_for_siglip,
                     dummy_seq_data_for_siglip, get_siglip_image_feature_size,
                     get_siglip_patch_grid_length, input_processor_for_siglip)
Cyrus Leung's avatar
Cyrus Leung committed
33
34
from .utils import (AutoWeightsLoader, embed_multimodal, flatten_bn,
                    init_vllm_registered_model)
35

36
37
38
# Result in the max possible feature size (2x2 grid of 336x336px tiles)
MAX_IMAGE_FEATURE_SIZE_HEIGHT = MAX_IMAGE_FEATURE_SIZE_WIDTH = 448

39
40
41

class LlavaNextImagePixelInputs(TypedDict):
    type: Literal["pixel_values"]
42
    data: Union[torch.Tensor, List[torch.Tensor]]
43
    """
44
45
    Shape:
    `(batch_size * num_images, 1 + num_patches, num_channels, height, width)`
46

47
48
    Note that `num_patches` may be different per batch and image,
    in which case the data is passed as a list instead of a batched tensor.
49
    """
50
51

    image_sizes: NotRequired[torch.Tensor]
52
    """
53
    Shape: `(batch_size * num_images, 2)`
54
55
56

    This should be in `(height, width)` format.
    """
57
58


59
60
61
class LlavaNextImageEmbeddingInputs(TypedDict):
    type: Literal["image_embeds"]
    data: torch.Tensor
62
    """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
63
64
65
66
67
68
69

    `hidden_size` must match the hidden size of language model backbone.
    """


LlavaNextImageInputs = Union[LlavaNextImagePixelInputs,
                             LlavaNextImageEmbeddingInputs]
70
71


72
# Based on: https://github.com/huggingface/text-generation-inference/blob/v2.2.0/server/text_generation_server/models/vlm_causal_lm.py#L79
73
def _get_llava_next_num_unpadded_features(
74
75
    original_height: int,
    original_width: int,
76
77
78
79
80
81
82
    npatches: int,
    num_patch_height: int,
    num_patch_width: int,
) -> Tuple[int, int]:
    current_height = npatches * num_patch_height
    current_width = npatches * num_patch_width

83
    original_aspect_ratio = original_width / original_height
84
85
    current_aspect_ratio = current_width / current_height

86
87
88
    if original_aspect_ratio > current_aspect_ratio:
        scale_factor = current_width / original_width
        new_height = int(original_height * scale_factor)
89
        padding = (current_height - new_height) // 2
90
        current_height -= 2 * padding
91
    else:
92
93
        scale_factor = current_height / original_height
        new_width = int(original_width * scale_factor)
94
        padding = (current_width - new_width) // 2
95
        current_width -= 2 * padding
96
97
98
99
100
101

    unpadded_features = current_height * current_width
    newline_features = current_height
    return (unpadded_features, newline_features)


102
# Based on: https://github.com/huggingface/text-generation-inference/blob/v2.2.0/server/text_generation_server/models/vlm_causal_lm.py#L106
103
def get_llava_next_image_feature_size(
104
105
106
107
108
109
110
111
112
113
114
115
    hf_config: LlavaNextConfig,
    *,
    input_height: int,
    input_width: int,
) -> int:
    vision_config = hf_config.vision_config

    if isinstance(vision_config, CLIPVisionConfig):
        num_patches = get_clip_patch_grid_length(
            image_size=vision_config.image_size,
            patch_size=vision_config.patch_size,
        )
116
117
118
119
120
        base_feature_size = get_clip_image_feature_size(vision_config)
    elif isinstance(vision_config, SiglipVisionConfig):
        num_patches = get_siglip_patch_grid_length(
            image_size=vision_config.image_size,
            patch_size=vision_config.patch_size,
121
        )
122
123
124
125
126
127
128
129
130
131
132
133
        base_feature_size = get_siglip_image_feature_size(vision_config)
    else:
        msg = f"Unsupported vision config: {type(vision_config)}"
        raise NotImplementedError(msg)

    strategy = hf_config.vision_feature_select_strategy
    if strategy == "default":
        base_feature_size -= 1
    elif strategy == "full":
        pass
    else:
        raise ValueError(f"Unexpected select feature strategy: {strategy}")
134

135
136
137
138
139
    num_patch_height, num_patch_width = get_anyres_image_grid_shape(
        image_size=(input_height, input_width),
        grid_pinpoints=hf_config.image_grid_pinpoints,
        patch_size=vision_config.image_size,
    )
140

141
142
143
144
145
146
    (
        unpadded_feature_size,
        newline_feature_size,
    ) = _get_llava_next_num_unpadded_features(input_height, input_width,
                                              num_patches, num_patch_height,
                                              num_patch_width)
147

148
    return unpadded_feature_size + newline_feature_size + base_feature_size
149
150


151
152
153
def get_max_llava_next_image_tokens(ctx: InputContext):
    return get_llava_next_image_feature_size(
        ctx.get_hf_config(LlavaNextConfig),
154
155
        input_height=MAX_IMAGE_FEATURE_SIZE_HEIGHT,
        input_width=MAX_IMAGE_FEATURE_SIZE_WIDTH,
156
157
158
    )


159
160
def dummy_data_for_llava_next(ctx: InputContext, seq_len: int,
                              mm_counts: Mapping[str, int]):
161
162
    hf_config = ctx.get_hf_config(LlavaNextConfig)
    vision_config = hf_config.vision_config
163
    num_images = mm_counts["image"]
164

165
    image_feature_size = get_max_llava_next_image_tokens(ctx)
166
167
168
169
170

    if isinstance(vision_config, CLIPVisionConfig):
        seq_data = dummy_seq_data_for_clip(
            vision_config,
            seq_len,
171
            num_images,
172
173
174
175
            image_token_id=hf_config.image_token_index,
            image_feature_size_override=image_feature_size,
        )

176
177
        mm_data = dummy_image_for_clip(
            vision_config,
178
            num_images,
179
180
            image_width_override=MAX_IMAGE_FEATURE_SIZE_WIDTH,
            image_height_override=MAX_IMAGE_FEATURE_SIZE_HEIGHT,
181
        )
182

183
184
185
186
187
        return seq_data, mm_data
    elif isinstance(vision_config, SiglipVisionConfig):
        seq_data = dummy_seq_data_for_siglip(
            vision_config,
            seq_len,
188
            num_images,
189
190
191
192
193
194
            image_token_id=hf_config.image_token_index,
            image_feature_size_override=image_feature_size,
        )

        mm_data = dummy_image_for_siglip(
            vision_config,
195
            num_images,
196
197
198
199
            image_width_override=MAX_IMAGE_FEATURE_SIZE_WIDTH,
            image_height_override=MAX_IMAGE_FEATURE_SIZE_HEIGHT,
        )

200
201
202
203
204
205
        return seq_data, mm_data

    msg = f"Unsupported vision config: {type(vision_config)}"
    raise NotImplementedError(msg)


206
207
208
def input_processor_for_llava_next(ctx: InputContext,
                                   inputs: DecoderOnlyInputs):
    multi_modal_data = inputs.get("multi_modal_data")
209
    if multi_modal_data is None or "image" not in multi_modal_data:
210
        return inputs
211

212
213
214
    model_config = ctx.model_config
    hf_config = ctx.get_hf_config(LlavaNextConfig)
    vision_config = hf_config.vision_config
215

216
217
218
219
220
221
222
223
224
    image_data = multi_modal_data["image"]
    if isinstance(image_data, Image.Image):
        width, height = image_data.size

        image_feature_size = get_llava_next_image_feature_size(
            hf_config,
            input_height=height,
            input_width=width,
        )
225
226
227
228
229
230
231
    elif is_list_of(image_data, Image.Image):
        image_feature_size = [
            get_llava_next_image_feature_size(hf_config,
                                              input_height=img.height,
                                              input_width=img.width)
            for img in image_data
        ]
232
    elif isinstance(image_data, torch.Tensor):
233
234
235
        num_images, image_feature_size, hidden_size = image_data.shape
    elif is_list_of(image_data, torch.Tensor):
        image_feature_size = [item.shape[1] for item in image_data]
236
237
    else:
        raise TypeError(f"Invalid image type: {type(image_data)}")
238

239
    vision_config = hf_config.vision_config
240

241
242
243
244
    if isinstance(vision_config, CLIPVisionConfig):
        return input_processor_for_clip(
            model_config,
            vision_config,
245
            inputs,
246
247
248
            image_token_id=hf_config.image_token_index,
            image_feature_size_override=image_feature_size,
        )
249
250
251
252
    elif isinstance(vision_config, SiglipVisionConfig):
        return input_processor_for_siglip(
            model_config,
            vision_config,
253
            inputs,
254
255
256
257
258
259
260
261
            image_token_id=hf_config.image_token_index,
            image_feature_size_override=image_feature_size,
        )

    msg = f"Unsupported vision config: {type(vision_config)}"
    raise NotImplementedError(msg)


262
@MULTIMODAL_REGISTRY.register_image_input_mapper()
263
@MULTIMODAL_REGISTRY.register_max_image_tokens(get_max_llava_next_image_tokens)
264
@INPUT_REGISTRY.register_dummy_data(dummy_data_for_llava_next)
265
@INPUT_REGISTRY.register_input_processor(input_processor_for_llava_next)
266
267
class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal,
                                        SupportsPP):
268

269
270
    def __init__(self,
                 config: LlavaNextConfig,
271
                 multimodal_config: MultiModalConfig,
272
273
                 cache_config: Optional[CacheConfig] = None,
                 quant_config: Optional[QuantizationConfig] = None) -> None:
274
        super().__init__()
275
276

        self.config = config
277
        self.multimodal_config = multimodal_config
278

279
        # TODO: Optionally initializes this for supporting embeddings.
280
281
        self.vision_tower = init_vision_tower_for_llava(
            config, quant_config, require_post_norm=False)
282
283
        self.image_newline = nn.Parameter(
            torch.empty(config.text_config.hidden_size))
284
285
286
287
288
        self.multi_modal_projector = LlavaMultiModalProjector(
            vision_hidden_size=config.vision_config.hidden_size,
            text_hidden_size=config.text_config.hidden_size,
            projector_hidden_act=config.projector_hidden_act)

289
290
        self.language_model = init_vllm_registered_model(
            config.text_config, cache_config, quant_config)
291

Cyrus Leung's avatar
Cyrus Leung committed
292
293
294
295
        # The same model class supports both language generation and embedding
        # because the architecture name is the same
        self._pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)

296
297
298
299
300
301
302
303
304
        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors)

    @cached_property
    def sampler(self):
        if hasattr(self.language_model, "sampler"):
            return self.language_model.sampler

        return Sampler()
305
306

    def _validate_image_sizes(self, data: torch.Tensor) -> torch.Tensor:
307
308
309
310
311
312
313
314
315
316
317
318
319
        expected_dims = (2, )

        def _validate_shape(d: torch.Tensor):
            actual_dims = tuple(d.shape)

            if actual_dims != expected_dims:
                expected_expr = str(expected_dims)
                raise ValueError(
                    f"The expected shape of image sizes per image per batch "
                    f"is {expected_expr}. You supplied {tuple(d.shape)}.")

        for d in data:
            _validate_shape(d)
320
321
322

        return data

323
324
325
326
    def _validate_pixel_values(
        self, data: Union[torch.Tensor, List[torch.Tensor]]
    ) -> Union[torch.Tensor, List[torch.Tensor]]:

327
328
329
330
331
332
333
334
        h = w = self.config.vision_config.image_size
        expected_dims = (3, h, w)

        def _validate_shape(d: torch.Tensor):
            actual_dims = tuple(d.shape[1:])

            if actual_dims != expected_dims:
                expected_expr = ("num_patches", *map(str, expected_dims))
335
                raise ValueError(
336
                    "The expected shape of pixel values per image per batch "
337
                    f"is {expected_expr}. You supplied {tuple(d.shape)}.")
338

339
340
        for d in data:
            _validate_shape(d)
341
342
343

        return data

344
    def _parse_and_validate_image_input(
345
            self, **kwargs: object) -> Optional[LlavaNextImageInputs]:
346
347
        pixel_values = kwargs.pop("pixel_values", None)
        image_sizes = kwargs.pop("image_sizes", None)
348
        image_embeds = kwargs.pop("image_embeds", None)
349

350
        if pixel_values is None and image_embeds is None:
351
            return None
352

353
354
355
356
        if pixel_values is not None:
            if not isinstance(pixel_values, (torch.Tensor, list)):
                raise ValueError("Incorrect type of pixel values. "
                                 f"Got type: {type(pixel_values)}")
357

358
            if not isinstance(image_sizes, (torch.Tensor, list)):
359
360
                raise ValueError("Incorrect type of image sizes. "
                                 f"Got type: {type(image_sizes)}")
361

362
363
            return LlavaNextImagePixelInputs(
                type="pixel_values",
364
365
366
                data=self._validate_pixel_values(flatten_bn(pixel_values)),
                image_sizes=self._validate_image_sizes(
                    flatten_bn(image_sizes, concat=True)),
367
368
369
370
371
372
373
374
375
            )

        if image_embeds is not None:
            if not isinstance(image_embeds, torch.Tensor):
                raise ValueError("Incorrect type of image embeds. "
                                 f"Got type: {type(image_embeds)}")

            return LlavaNextImageEmbeddingInputs(
                type="image_embeds",
376
                data=flatten_bn(image_embeds),
377
378
379
            )

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

Cyrus Leung's avatar
Cyrus Leung committed
381
382
383
384
385
386
387
388
389
390
    def _select_image_features(self, image_features: torch.Tensor, *,
                               strategy: str) -> torch.Tensor:
        # Copied from https://github.com/huggingface/transformers/blob/39c3c0a72af6fbda5614dde02ff236069bb79827/src/transformers/models/llava/modeling_llava.py#L421  # noqa
        if strategy == "default":
            return image_features[:, 1:]
        elif strategy == "full":
            return image_features

        raise ValueError(f"Unexpected select feature strategy: {strategy}")

391
392
393
394
395
    def _image_pixels_to_features(
        self,
        vision_tower: Union[CLIPVisionModel, SiglipVisionModel],
        pixel_values: torch.Tensor,
    ) -> torch.Tensor:
Cyrus Leung's avatar
Cyrus Leung committed
396

397
398
        # NOTE: we skip the step to select the vision feature layer since
        # this is already done inside the vision tower
399
        image_features = vision_tower(pixel_values)
Cyrus Leung's avatar
Cyrus Leung committed
400
401
402
403
404
405

        return self._select_image_features(
            image_features,
            strategy=self.config.vision_feature_select_strategy,
        )

406
    # Based on: https://github.com/haotian-liu/LLaVA/blob/main/llava/model/llava_arch.py
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
    def _merge_image_patch_embeddings(self, image_size: torch.Tensor,
                                      patch_embeddings: torch.Tensor, *,
                                      strategy: str) -> torch.Tensor:
        if strategy == "flat":
            return patch_embeddings.flatten(0, 1)

        if strategy.startswith("spatial"):
            height = width = self.config.vision_config.image_size \
                // self.config.vision_config.patch_size

            base_patch_embeds = patch_embeddings[0]
            if height * width != base_patch_embeds.shape[0]:
                raise ValueError(
                    "The number of patches is not consistent with the "
                    "image size.")

            if patch_embeddings.shape[0] > 1:
                other_patch_embeds = patch_embeddings[1:]

426
427
428
                # Move to CPU to avoid floating-point errors
                orig_height, orig_width = image_size.tolist()

429
                # image_aspect_ratio == "anyres"
430
431
                num_patch_height, num_patch_width = get_anyres_image_grid_shape(
                    (orig_height, orig_width),
432
433
434
                    self.config.image_grid_pinpoints,
                    self.config.vision_config.image_size,
                )
435
436
437
438
                num_patches = num_patch_height * num_patch_width

                # Image patches might be padded for batch processing
                other_patch_embeds = other_patch_embeds[:num_patches] \
439
                    .view(num_patch_height, num_patch_width, height, width, -1)
440
441
442
443
444
445

                if "unpad" in strategy:
                    other_patch_embeds = other_patch_embeds \
                        .permute(4, 0, 2, 1, 3).contiguous() \
                        .flatten(1, 2).flatten(2, 3)
                    other_patch_embeds = unpad_image(other_patch_embeds,
446
                                                     (orig_height, orig_width))
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
                    other_patch_embeds = torch.cat((
                        other_patch_embeds,
                        self.image_newline[:, None, None] \
                            .expand(*other_patch_embeds.shape[:-1], 1) \
                            .to(other_patch_embeds.device),
                    ), dim=-1)
                    other_patch_embeds = other_patch_embeds \
                        .flatten(1, 2).transpose(0, 1)
                else:
                    other_patch_embeds = other_patch_embeds \
                        .permute(0, 2, 1, 3, 4).contiguous() \
                        .flatten(0, 3)

                merged_patch_embeddings = torch.cat(
                    (base_patch_embeds, other_patch_embeds), dim=0)
            else:
                if "unpad" in strategy:
                    merged_patch_embeddings = torch.cat(
                        (base_patch_embeds,
                         self.image_newline[None] \
                            .to(base_patch_embeds.device)
                    ), dim=0)
                else:
                    merged_patch_embeddings = base_patch_embeds

            return merged_patch_embeddings

        raise ValueError(f"Unexpected patch merge strategy: {strategy}")

    def _process_image_pixels(
477
478
        self,
        inputs: LlavaNextImagePixelInputs,
479
    ) -> Union[torch.Tensor, List[torch.Tensor]]:
480
481
482
483
        assert self.vision_tower is not None

        pixel_values = inputs["data"]

484
485
486
487
488
489
490
        if isinstance(pixel_values, torch.Tensor):
            b, num_patches, c, h, w = pixel_values.shape
            stacked_pixel_values = pixel_values.view(b * num_patches, c, h, w)
            stacked_image_features = self._image_pixels_to_features(
                self.vision_tower, stacked_pixel_values)
            stacked_patch_embeddings = self.multi_modal_projector(
                stacked_image_features)
491

492
493
494
495
496
            return stacked_patch_embeddings.view(
                b, num_patches, *stacked_patch_embeddings.shape[1:])

        num_patches_per_batch = [v.shape[0] for v in pixel_values]
        stacked_pixel_values = torch.cat(pixel_values)
497
498
499
        stacked_image_features = self._image_pixels_to_features(
            self.vision_tower, stacked_pixel_values)

500
501
502
503
        return [
            self.multi_modal_projector(image_features) for image_features in
            torch.split(stacked_image_features, num_patches_per_batch)
        ]
504
505

    def _process_image_input(
506
507
508
        self,
        image_input: LlavaNextImageInputs,
    ) -> Union[torch.Tensor, List[torch.Tensor]]:
509
510
511
        if image_input["type"] == "image_embeds":
            return [image_input["data"]]

512
        patch_embeddings = self._process_image_pixels(image_input)
513
514
515

        image_sizes = image_input.get("image_sizes")
        if image_sizes is None:
516
            batch_size = len(image_input["data"])
517
            vision_config = self.config.vision_config
518
519
            default_height = default_width = vision_config.image_size
            image_sizes = torch.as_tensor([[default_height, default_width]
520
521
                                           for _ in range(batch_size)])

522
        return [
523
            self._merge_image_patch_embeddings(image_sizes[i],
524
                                               patch_features_batch,
525
                                               strategy="spatial_unpad")
526
            for i, patch_features_batch in enumerate(patch_embeddings)
527
528
529
530
531
532
533
534
        ]

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
535
        intermediate_tensors: Optional[IntermediateTensors] = None,
536
        **kwargs: object,
537
    ) -> Union[torch.Tensor, IntermediateTensors]:
Cyrus Leung's avatar
Cyrus Leung committed
538
        """Run forward pass for LlaVA-NeXT.
539
540
541

        One key thing to understand is the `input_ids` already accounts for the
        positions of the to-be-inserted image embeddings.
542

543
        Concretely, consider a text prompt:
544
545
546
547
548
        `"A chat between a curious human and an artificial intelligence
        assistant. The assistant gives helpful, detailed, and polite answers to
        the human's questions.
        USER: <image>\\nWhat is shown in this image? ASSISTANT:"`.

549
        Tokenizer outputs:
550
551
552
553
554
555
556
        `[1, 319, 13563, 1546, 263, 12758, 5199, 322, 385, 23116, 21082, 20255,
        29889, 450, 20255, 4076, 8444, 29892, 13173, 29892, 322, 1248, 568,
        6089, 304, 278, 5199, 29915, 29879, 5155, 29889, 3148, 1001, 29901,
        29871, 32000, 13, 5618, 338, 4318, 297, 445, 1967, 29973, 319, 1799,
        9047, 13566, 29901]`.

        To reserve space in KV cache, we have to insert placeholder tokens
557
        before they are inputted to the model, so the input processor prepends
558
559
560
561
562
563
564
565
566
567
568
        additional image tokens (denoted as `32000`), resulting in:
        `[1, 319, 13563, 1546, 263, 12758, 5199, 322, 385, 23116, 21082, 20255,
        29889, 450, 20255, 4076, 8444, 29892, 13173, 29892, 322, 1248, 568,
        6089, 304, 278, 5199, 29915, 29879, 5155, 29889, 3148, 1001, 29901,
        29871, 32000, ..., 32000, 13, 5618, 338, 4318, 297, 445, 1967, 29973,
        319, 1799, 9047, 13566, 29901]`.

        Unlike in LLaVA-1.5, the number of image tokens inputted to the language
        model depends on the original size of the input image. Including the
        original image token in the input, the required number of image tokens
        is given by :func:`get_llava_next_image_feature_size`.
569
570
571
572
573
574
575

        This way, the `positions` and `attn_metadata` are consistent
        with the `input_ids`.

        Args:
            input_ids: Flattened (concatenated) input_ids corresponding to a
                batch.
Cyrus Leung's avatar
Cyrus Leung committed
576
            pixel_values: The pixels in each grid patch for each input image.
577
            image_sizes: The original `(height, width)` for each input image.
578

Cyrus Leung's avatar
Cyrus Leung committed
579
        See also:
580
            :class:`LlavaNextImageInputs`
581
        """
582
583
584
585
586
        if intermediate_tensors is not None:
            input_ids = None
            inputs_embeds = None
        else:
            image_input = self._parse_and_validate_image_input(**kwargs)
587

588
            if image_input is not None:
Cyrus Leung's avatar
Cyrus Leung committed
589
590
591
592
593
594
                inputs_embeds = embed_multimodal(
                    input_ids,
                    self.config.image_token_index,
                    self.language_model.model.get_input_embeddings,
                    lambda _: self._process_image_input(image_input),
                )
595
596
597
                input_ids = None
            else:
                inputs_embeds = None
598

599
600
601
602
        hidden_states = self.language_model.model(input_ids,
                                                  positions,
                                                  kv_caches,
                                                  attn_metadata,
603
                                                  intermediate_tensors,
604
                                                  inputs_embeds=inputs_embeds)
605
606
607

        return hidden_states

608
609
610
611
612
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
613
614
        return self.language_model.compute_logits(hidden_states,
                                                  sampling_metadata)
615
616
617
618
619
620

    def sample(
        self,
        logits: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
621
        return self.language_model.sample(logits, sampling_metadata)
622

Cyrus Leung's avatar
Cyrus Leung committed
623
624
625
626
627
628
629
    def pooler(
        self,
        hidden_states: torch.Tensor,
        pooling_metadata: PoolingMetadata,
    ) -> Optional[PoolerOutput]:
        return self._pooler(hidden_states, pooling_metadata)

630
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
631
632
        loader = AutoWeightsLoader(self)
        loader.load_weights(weights)