llava_next.py 23.3 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
from abc import abstractmethod
4
from functools import cached_property
5
from typing import (Final, Iterable, List, Literal, Mapping, Optional,
6
                    Protocol, Set, Tuple, TypedDict, TypeVar, Union)
7
8
9

import torch
import torch.nn as nn
10
from transformers import BatchFeature, LlavaNextConfig, LlavaNextProcessor
11
12
13
14
15
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
16
from vllm.config import VllmConfig
Joe Runde's avatar
Joe Runde committed
17
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
18
from vllm.model_executor.sampling_metadata import SamplingMetadata
19
from vllm.multimodal import MULTIMODAL_REGISTRY
20
21
from vllm.multimodal.inputs import MultiModalFieldConfig, NestedTensors
from vllm.multimodal.parse import ImageSize
22
from vllm.sequence import IntermediateTensors
23

24
from .clip import CLIPVisionModel
25
from .interfaces import SupportsMultiModal, SupportsPP
26
27
from .llava import (BaseLlavaMultiModalProcessor, BaseLlavaProcessingInfo,
                    LlavaDummyInputsBuilder, LlavaLikeConfig,
28
                    LlavaMultiModalProjector, init_vision_tower_for_llava)
29
from .siglip import SiglipVisionModel
Cyrus Leung's avatar
Cyrus Leung committed
30
from .utils import (AutoWeightsLoader, embed_multimodal, flatten_bn,
31
                    init_vllm_registered_model, maybe_prefix)
32
33
34
35


class LlavaNextImagePixelInputs(TypedDict):
    type: Literal["pixel_values"]
36
    data: Union[torch.Tensor, List[torch.Tensor]]
37
    """
38
39
    Shape:
    `(batch_size * num_images, 1 + num_patches, num_channels, height, width)`
40

41
42
    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.
43
    """
44
45

    image_sizes: NotRequired[torch.Tensor]
46
    """
47
    Shape: `(batch_size * num_images, 2)`
48
49
50

    This should be in `(height, width)` format.
    """
51
52


53
54
55
class LlavaNextImageEmbeddingInputs(TypedDict):
    type: Literal["image_embeds"]
    data: torch.Tensor
56
    """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
57
58
59
60
61
62
63

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


LlavaNextImageInputs = Union[LlavaNextImagePixelInputs,
                             LlavaNextImageEmbeddingInputs]
64
65


66
67
class LlavaNextLikeConfig(LlavaLikeConfig, Protocol):
    image_grid_pinpoints: Final[list[list[int]]]
68

69

70
class LlavaNextProcessingInfo(BaseLlavaProcessingInfo):
71

72
    def get_hf_config(self) -> LlavaNextLikeConfig:
73
        return self.ctx.get_hf_config(LlavaNextConfig)
74

75
    def get_hf_processor(self):
76
        return self.ctx.get_hf_processor(LlavaNextProcessor)
77

78
    # Based on: https://github.com/huggingface/text-generation-inference/blob/v3.0.1/server/text_generation_server/models/vlm_causal_lm.py#L113
79
    def get_num_image_tokens(
80
81
82
83
84
        self,
        *,
        image_width: int,
        image_height: int,
    ) -> int:
85
86
        hf_config = self.get_hf_config()
        vision_encoder_info = self.get_vision_encoder_info()
87
88
89

        base_feature_size = self._apply_feature_select_strategy(
            hf_config.vision_feature_select_strategy,
90
            vision_encoder_info.get_num_image_tokens(
91
92
93
                image_width=image_width,
                image_height=image_height,
            ),
94
        )
95
96
97
98

        num_patch_height, num_patch_width = get_anyres_image_grid_shape(
            image_size=(image_height, image_width),
            grid_pinpoints=hf_config.image_grid_pinpoints,
99
            patch_size=vision_encoder_info.get_image_size(),
100
101
        )

102
103
104
105
106
107
        (
            unpadded_feature_size,
            newline_feature_size,
        ) = self._get_num_unpadded_features(
            original_height=image_height,
            original_width=image_width,
108
            npatches=vision_encoder_info.get_patch_grid_length(),
109
110
111
            num_patch_height=num_patch_height,
            num_patch_width=num_patch_width,
        )
112

113
        return unpadded_feature_size + newline_feature_size + base_feature_size
114

115
    # Based on: https://github.com/huggingface/text-generation-inference/blob/v3.0.1/server/text_generation_server/models/vlm_causal_lm.py#L86
116
117
118
119
120
121
122
123
124
    def _get_num_unpadded_features(
        self,
        *,
        original_height: int,
        original_width: int,
        npatches: int,
        num_patch_height: int,
        num_patch_width: int,
    ) -> tuple[int, int]:
125
126
        current_height = npatches * num_patch_height
        current_width = npatches * num_patch_width
127

128
129
        aspect_ratio = original_width / original_height
        current_aspect_ratio = current_width / current_height
130

131
132
133
134
        if aspect_ratio > current_aspect_ratio:
            new_height = (original_height * current_width) // original_width
            padding = (current_height - new_height) // 2
            current_height = current_height - (2 * padding)
135
        else:
136
137
138
            new_width = (original_width * current_height) // original_height
            padding = (current_width - new_width) // 2
            current_width = current_width - (2 * padding)
139

140
141
        unpadded_features = current_height * current_width
        newline_features = current_height
142

143
144
        return (unpadded_features, newline_features)

145
146
    def get_image_size_with_most_features(self) -> ImageSize:
        hf_config = self.get_hf_config()
147
148
149

        largest_feature_size, largest_feature_pinpoint = 0, None
        for (height, width) in hf_config.image_grid_pinpoints:
150
151
            feat_size = self.get_num_image_tokens(image_width=width,
                                                  image_height=height)
152
153
154
155
156
157
158
159
            if feat_size > largest_feature_size:
                largest_feature_size = feat_size
                largest_feature_pinpoint = ImageSize(width=width,
                                                     height=height)

        if largest_feature_size == 0 or largest_feature_pinpoint is None:
            raise ValueError("Cannot have a largest feature size of 0!")

160
161
162
        return largest_feature_pinpoint


163
164
165
166
167
168
169
170
171
172
173
174
175
176
_I = TypeVar("_I", bound=LlavaNextProcessingInfo)


class BaseLlavaNextMultiModalProcessor(BaseLlavaMultiModalProcessor[_I]):

    # Copied from BaseMultiModalProcessor
    @abstractmethod
    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
        raise NotImplementedError

177

178
179
class LlavaNextMultiModalProcessor(
        BaseLlavaNextMultiModalProcessor[LlavaNextProcessingInfo]):
180
181
182
183
184
185
186
187
188
189
190

    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
        return dict(
            pixel_values=MultiModalFieldConfig.batched("image"),
            image_sizes=MultiModalFieldConfig.batched("image"),
            image_embeds=MultiModalFieldConfig.batched("image"),
        )
191
192


193
194
195
@MULTIMODAL_REGISTRY.register_processor(LlavaNextMultiModalProcessor,
                                        info=LlavaNextProcessingInfo,
                                        dummy_inputs=LlavaDummyInputsBuilder)
196
197
class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal,
                                        SupportsPP):
198

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

205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
        vision_feature_layer = config.vision_feature_layer
        # Determine the layer up to which we will initialize the vision tower
        if isinstance(vision_feature_layer, int):
            vision_hidden_size = config.vision_config.hidden_size
            self.feature_sample_layers = None
        # Used for multimodal granite models to control encoder outputs
        elif isinstance(vision_feature_layer, (list, tuple)):
            vision_hidden_size = config.vision_config.hidden_size * len(
                vision_feature_layer)
            self.feature_sample_layers = vision_feature_layer
        else:
            raise TypeError(
                f"vision_layer_feature type: {type(vision_feature_layer)}"
                " is not supported")

220
        self.config = config
221
        self.multimodal_config = multimodal_config
222

223
        # TODO: Optionally initializes this for supporting embeddings.
224
        self.vision_tower = init_vision_tower_for_llava(
225
226
227
            config,
            quant_config,
            require_post_norm=False,
228
            prefix=maybe_prefix(prefix, "vision_tower"))
229
230
        self.image_newline = nn.Parameter(
            torch.empty(config.text_config.hidden_size))
231
        self.multi_modal_projector = LlavaMultiModalProjector(
232
            vision_hidden_size=vision_hidden_size,
233
234
235
            text_hidden_size=config.text_config.hidden_size,
            projector_hidden_act=config.projector_hidden_act)

236
        self.language_model = init_vllm_registered_model(
237
            vllm_config=vllm_config,
238
239
240
241
            hf_config=config.text_config,
            prefix=maybe_prefix(prefix, "language_model"),
        )

242
243
244
245
246
247
248
249
        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

Joe Runde's avatar
Joe Runde committed
250
        return get_sampler()
251
252

    def _validate_image_sizes(self, data: torch.Tensor) -> torch.Tensor:
253
254
255
256
257
258
259
260
261
262
263
264
265
        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)
266
267
268

        return data

269
270
271
272
    def _validate_pixel_values(
        self, data: Union[torch.Tensor, List[torch.Tensor]]
    ) -> Union[torch.Tensor, List[torch.Tensor]]:

273
274
275
276
277
278
279
280
        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))
281
                raise ValueError(
282
                    "The expected shape of pixel values per image per batch "
283
                    f"is {expected_expr}. You supplied {tuple(d.shape)}.")
284

285
286
        for d in data:
            _validate_shape(d)
287
288
289

        return data

290
    def _parse_and_validate_image_input(
291
            self, **kwargs: object) -> Optional[LlavaNextImageInputs]:
292
293
        pixel_values = kwargs.pop("pixel_values", None)
        image_sizes = kwargs.pop("image_sizes", None)
294
        image_embeds = kwargs.pop("image_embeds", None)
295

296
        if pixel_values is None and image_embeds is None:
297
            return None
298

299
300
301
302
        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)}")
303

304
            if not isinstance(image_sizes, (torch.Tensor, list)):
305
306
                raise ValueError("Incorrect type of image sizes. "
                                 f"Got type: {type(image_sizes)}")
307

308
309
            return LlavaNextImagePixelInputs(
                type="pixel_values",
310
311
312
                data=self._validate_pixel_values(flatten_bn(pixel_values)),
                image_sizes=self._validate_image_sizes(
                    flatten_bn(image_sizes, concat=True)),
313
314
315
316
317
318
319
320
321
            )

        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",
322
                data=flatten_bn(image_embeds),
323
324
325
            )

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

Cyrus Leung's avatar
Cyrus Leung committed
327
328
329
330
331
332
333
334
335
336
    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}")

337
338
339
340
341
    def _image_pixels_to_features(
        self,
        vision_tower: Union[CLIPVisionModel, SiglipVisionModel],
        pixel_values: torch.Tensor,
    ) -> torch.Tensor:
Cyrus Leung's avatar
Cyrus Leung committed
342

343
344
        # NOTE: we skip the step to select the vision feature layer since
        # this is already done inside the vision tower
345
346
        image_features = vision_tower(
            pixel_values, feature_sample_layers=self.feature_sample_layers)
Cyrus Leung's avatar
Cyrus Leung committed
347
348
349
350
351
352

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

353
    # Based on: https://github.com/haotian-liu/LLaVA/blob/main/llava/model/llava_arch.py
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
    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:]

373
374
375
                # Move to CPU to avoid floating-point errors
                orig_height, orig_width = image_size.tolist()

376
                # image_aspect_ratio == "anyres"
377
378
                num_patch_height, num_patch_width = get_anyres_image_grid_shape(
                    (orig_height, orig_width),
379
380
381
                    self.config.image_grid_pinpoints,
                    self.config.vision_config.image_size,
                )
382
383
384
385
                num_patches = num_patch_height * num_patch_width

                # Image patches might be padded for batch processing
                other_patch_embeds = other_patch_embeds[:num_patches] \
386
                    .view(num_patch_height, num_patch_width, height, width, -1)
387
388
389
390
391
392

                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,
393
                                                     (orig_height, orig_width))
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
                    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(
424
425
        self,
        inputs: LlavaNextImagePixelInputs,
426
    ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:
427
428
429
430
        assert self.vision_tower is not None

        pixel_values = inputs["data"]

431
432
433
434
435
436
437
        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)
438

439
440
441
442
443
            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)
444
445
446
        stacked_image_features = self._image_pixels_to_features(
            self.vision_tower, stacked_pixel_values)

447
448
        return torch.split(self.multi_modal_projector(stacked_image_features),
                           num_patches_per_batch)
449
450

    def _process_image_input(
451
452
453
        self,
        image_input: LlavaNextImageInputs,
    ) -> Union[torch.Tensor, List[torch.Tensor]]:
454
455
456
        if image_input["type"] == "image_embeds":
            return [image_input["data"]]

457
        patch_embeddings = self._process_image_pixels(image_input)
458
459
460

        image_sizes = image_input.get("image_sizes")
        if image_sizes is None:
461
            batch_size = len(image_input["data"])
462
            vision_config = self.config.vision_config
463
464
            default_height = default_width = vision_config.image_size
            image_sizes = torch.as_tensor([[default_height, default_width]
465
466
                                           for _ in range(batch_size)])

467
        return [
468
            self._merge_image_patch_embeddings(image_sizes[i],
469
                                               patch_features_batch,
470
                                               strategy="spatial_unpad")
471
            for i, patch_features_batch in enumerate(patch_embeddings)
472
473
        ]

474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
    def get_multimodal_embeddings(self, **kwargs) -> Optional[NestedTensors]:
        image_input = self._parse_and_validate_image_input(**kwargs)
        if image_input is None:
            return None
        vision_embeddings = self._process_image_input(image_input)
        return vision_embeddings

    def get_input_embeddings(
        self,
        input_ids: torch.Tensor,
        multimodal_embeddings: Optional[NestedTensors] = None,
    ) -> torch.Tensor:

        if multimodal_embeddings is None:
            return self.language_model.get_input_embeddings(input_ids)

        inputs_embeds = embed_multimodal(
            input_ids,
            self.config.image_token_index,
            self.language_model.model.get_input_embeddings,
            multimodal_embeddings,
        )
        return inputs_embeds

498
499
500
501
502
503
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
504
        intermediate_tensors: Optional[IntermediateTensors] = None,
505
        inputs_embeds: Optional[torch.Tensor] = None,
506
        **kwargs: object,
507
    ) -> Union[torch.Tensor, IntermediateTensors]:
Cyrus Leung's avatar
Cyrus Leung committed
508
        """Run forward pass for LlaVA-NeXT.
509
510
511

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

513
        Concretely, consider a text prompt:
514
515
516
517
518
        `"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:"`.

519
        Tokenizer outputs:
520
521
522
523
524
525
526
        `[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
527
        before they are inputted to the model, so the input processor prepends
528
529
530
531
532
533
534
535
536
537
538
        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`.
539
540
541
542
543
544
545

        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
546
            pixel_values: The pixels in each grid patch for each input image.
547
            image_sizes: The original `(height, width)` for each input image.
548

Cyrus Leung's avatar
Cyrus Leung committed
549
        See also:
550
            :class:`LlavaNextImageInputs`
551
        """
552
553
        if intermediate_tensors is not None:
            inputs_embeds = None
554

555
556
557
558
559
560
561
        # NOTE: In v1, inputs_embeds is always generated at model runner, this
        # condition is for v0 compatibility.
        elif inputs_embeds is None:
            vision_embeddings = self.get_multimodal_embeddings(**kwargs)
            inputs_embeds = self.get_input_embeddings(input_ids,
                                                      vision_embeddings)
            input_ids = None
562

563
564
565
566
        hidden_states = self.language_model.model(input_ids,
                                                  positions,
                                                  kv_caches,
                                                  attn_metadata,
567
                                                  intermediate_tensors,
568
                                                  inputs_embeds=inputs_embeds)
569
570
        return hidden_states

571
572
573
574
575
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
576
577
        return self.language_model.compute_logits(hidden_states,
                                                  sampling_metadata)
578
579
580
581
582
583

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

586
587
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
588
        loader = AutoWeightsLoader(self)
589
        return loader.load_weights(weights)