llava.py 22.8 KB
Newer Older
1
from functools import cached_property
2
from typing import (Iterable, List, Literal, Mapping, Optional, Protocol, Set,
3
                    Tuple, TypedDict, Union)
4
5

import torch
6
import torch.nn as nn
7
from PIL import Image
8
from transformers import (CLIPVisionConfig, LlavaConfig, PixtralVisionConfig,
9
                          PretrainedConfig, SiglipVisionConfig)
10
11

from vllm.attention import AttentionMetadata
12
from vllm.config import VllmConfig
13
14
from vllm.inputs import (INPUT_REGISTRY, DecoderOnlyInputs, DummyData,
                         InputContext)
15
from vllm.model_executor.layers.activation import get_act_fn
16
from vllm.model_executor.layers.quantization import QuantizationConfig
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
from vllm.multimodal.inputs import NestedTensors
21
from vllm.sequence import IntermediateTensors
22
from vllm.utils import is_list_of
23

24
25
26
from .clip import (CLIPVisionModel, dummy_image_for_clip,
                   dummy_seq_data_for_clip, get_max_clip_image_tokens,
                   input_processor_for_clip)
27
from .interfaces import SupportsMultiModal, SupportsPP
28
29
30
31
from .pixtral import (PixtralHFVisionModel, dummy_image_for_pixtral_hf,
                      dummy_seq_data_for_pixtral_hf,
                      get_max_pixtral_hf_image_tokens,
                      input_processor_for_pixtral_hf)
32
33
34
from .siglip import (SiglipVisionModel, dummy_image_for_siglip,
                     dummy_seq_data_for_siglip, get_max_siglip_image_tokens,
                     input_processor_for_siglip)
35
from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,
36
                    maybe_prefix, merge_multimodal_embeddings)
37
38


39
40
class LlavaImagePixelInputs(TypedDict):
    type: Literal["pixel_values"]
41
42
43
44
45
46
47
    data: Union[torch.Tensor, List[torch.Tensor]]
    """
    Shape: `(batch_size * num_images, num_channels, height, width)`

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


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

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


LlavaImageInputs = Union[LlavaImagePixelInputs, LlavaImageEmbeddingInputs]


62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
# TODO(xwjiang): Run benchmark and decide if TP.
class LlavaMultiModalProjector(nn.Module):

    def __init__(self, vision_hidden_size: int, text_hidden_size: int,
                 projector_hidden_act: str):
        super().__init__()

        self.linear_1 = nn.Linear(vision_hidden_size,
                                  text_hidden_size,
                                  bias=True)
        self.act = get_act_fn(projector_hidden_act)
        self.linear_2 = nn.Linear(text_hidden_size,
                                  text_hidden_size,
                                  bias=True)

77
    def forward(self, image_features: torch.Tensor) -> torch.Tensor:
78
79
80
81
82
83
        hidden_states = self.linear_1(image_features)
        hidden_states = self.act(hidden_states)
        hidden_states = self.linear_2(hidden_states)
        return hidden_states


84
85
86
87
88
def get_max_llava_image_tokens(ctx: InputContext):
    hf_config = ctx.get_hf_config(LlavaConfig)
    vision_config = hf_config.vision_config

    if isinstance(vision_config, CLIPVisionConfig):
89
90
91
        num_image_tokens = get_max_clip_image_tokens(vision_config)
    elif isinstance(vision_config, SiglipVisionConfig):
        num_image_tokens = get_max_siglip_image_tokens(vision_config)
92
93
    elif isinstance(vision_config, PixtralVisionConfig):
        num_image_tokens = get_max_pixtral_hf_image_tokens(vision_config)
94
95
96
97
98
99
100
101
102
103
104
    else:
        msg = f"Unsupported vision config: {type(vision_config)}"
        raise NotImplementedError(msg)

    strategy = hf_config.vision_feature_select_strategy
    if strategy == "default":
        return num_image_tokens - 1
    elif strategy == "full":
        return num_image_tokens
    else:
        raise ValueError(f"Unexpected select feature strategy: {strategy}")
105
106


107
108
def dummy_data_for_llava(ctx: InputContext, seq_len: int,
                         mm_counts: Mapping[str, int]):
109
110
    hf_config = ctx.get_hf_config(LlavaConfig)
    vision_config = hf_config.vision_config
111
    num_images = mm_counts["image"]
112

113
114
    image_feature_size = get_max_llava_image_tokens(ctx)

115
    if isinstance(vision_config, CLIPVisionConfig):
116
        seq_data, ranges = dummy_seq_data_for_clip(
117
118
            vision_config,
            seq_len,
119
            num_images,
120
            image_token_id=hf_config.image_token_index,
121
            image_feature_size_override=image_feature_size,
122
123
        )

124
        mm_data = dummy_image_for_clip(vision_config, num_images)
125
        return DummyData(seq_data, mm_data, ranges)
126
    elif isinstance(vision_config, SiglipVisionConfig):
127
        seq_data, ranges = dummy_seq_data_for_siglip(
128
129
            vision_config,
            seq_len,
130
            num_images,
131
132
133
134
            image_token_id=hf_config.image_token_index,
            image_feature_size_override=image_feature_size,
        )

135
        mm_data = dummy_image_for_siglip(vision_config, num_images)
136
        return DummyData(seq_data, mm_data, ranges)
137
    elif isinstance(vision_config, PixtralVisionConfig):
138
        seq_data, ranges = dummy_seq_data_for_pixtral_hf(
139
140
141
142
143
144
145
146
            vision_config,
            seq_len,
            num_images,
            image_token_id=hf_config.image_token_index,
            image_feature_size_override=image_feature_size,
        )

        mm_data = dummy_image_for_pixtral_hf(vision_config, num_images)
147
        return DummyData(seq_data, mm_data, ranges)
148
149
150
151
152

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


153
154
def input_processor_for_llava(ctx: InputContext, inputs: DecoderOnlyInputs):
    multi_modal_data = inputs.get("multi_modal_data")
155
    if multi_modal_data is None or "image" not in multi_modal_data:
156
        return inputs
157
158
159
160
161

    model_config = ctx.model_config
    hf_config = ctx.get_hf_config(LlavaConfig)
    vision_config = hf_config.vision_config

162
163
164
165
166
167
168
169
170
171
172
173
    image_data = multi_modal_data["image"]
    if isinstance(image_data, Image.Image):
        image_feature_size = get_max_llava_image_tokens(ctx)
    elif is_list_of(image_data, Image.Image):
        image_feature_size = [get_max_llava_image_tokens(ctx)
                              ] * len(image_data)
    elif isinstance(image_data, torch.Tensor):
        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]
    else:
        raise TypeError(f"Invalid image type: {type(image_data)}")
174

175
176
177
178
    if isinstance(vision_config, CLIPVisionConfig):
        return input_processor_for_clip(
            model_config,
            vision_config,
179
            inputs,
180
            image_token_id=hf_config.image_token_index,
181
182
183
184
185
186
            image_feature_size_override=image_feature_size,
        )
    elif isinstance(vision_config, SiglipVisionConfig):
        return input_processor_for_siglip(
            model_config,
            vision_config,
187
            inputs,
188
189
190
            image_token_id=hf_config.image_token_index,
            image_feature_size_override=image_feature_size,
        )
191
192
193
194
195
196
197
198
199
    elif isinstance(vision_config, PixtralVisionConfig):
        # We ignore image_feature_size_override since we have non-uniform
        # image sizes for Pixtral
        return input_processor_for_pixtral_hf(
            model_config,
            vision_config,
            inputs,
            image_token_id=hf_config.image_token_index,
        )
200
201
202
203
204

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


205
206
class LlavaLikeConfig(Protocol):
    vision_config: PretrainedConfig
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
235
236
237
238
239
240
241
    vision_feature_layer: Union[int, List[int]]


def _get_num_hidden_layers(hf_config: LlavaLikeConfig) -> int:
    """Determine the number of hidden layers to initialize up to in the
    visual encoder.
    
    Args:
        hf_config: Model config with vision feature layer(s).
    """
    feature_layers = hf_config.vision_feature_layer
    num_hidden_layers = hf_config.vision_config.num_hidden_layers
    # If we have one feature layer, initialize up to that layer
    if isinstance(feature_layers, int):
        return _get_layer_index(feature_layers, num_hidden_layers)
    # If we have multiple feature layers, initialize up to the deepest one
    elif isinstance(feature_layers, (list, tuple)):
        return max(
            _get_layer_index(idx, num_hidden_layers) for idx in feature_layers)
    raise TypeError(f"vision_layer_feature type: {type(feature_layers)}"
                    " is not supported")


def _get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int:
    """Given an signed vision feature layer, get the number of hidden layers
    needed to leverage it.

    Args:
        feature_layer_index: Index of a required layer in the visual encoder.
        num_hidden_layers: The total number of hidden layers in the visual
            encoder.
    """
    if feature_layer_index < 0:
        return num_hidden_layers + feature_layer_index + 1
    return feature_layer_index + 1
242
243
244
245
246
247
248


def init_vision_tower_for_llava(
    hf_config: LlavaLikeConfig,
    quant_config: Optional[QuantizationConfig],
    *,
    require_post_norm: Optional[bool] = None,
249
    prefix: str = "",
250
):
251
252
    vision_config = hf_config.vision_config

253
254
    # Initialize the vision tower only up to the deepest required feature layer
    num_hidden_layers = _get_num_hidden_layers(hf_config)
255
256
257
258

    if isinstance(vision_config, CLIPVisionConfig):
        return CLIPVisionModel(
            vision_config,
259
            quant_config=quant_config,
260
            num_hidden_layers_override=num_hidden_layers,
261
            require_post_norm=require_post_norm,
262
            prefix=prefix,
263
264
265
266
        )
    elif isinstance(vision_config, SiglipVisionConfig):
        return SiglipVisionModel(
            vision_config,
267
            quant_config=quant_config,
268
            num_hidden_layers_override=num_hidden_layers,
269
            require_post_norm=require_post_norm,
270
            prefix=prefix,
271
        )
272
    elif isinstance(vision_config, PixtralVisionConfig):
273
274
        return PixtralHFVisionModel(
            vision_config,
275
            quant_config=quant_config,
276
277
            num_hidden_layers_override=num_hidden_layers,
            require_post_norm=require_post_norm,
278
            prefix=prefix,
279
        )
280
281
282
283
284

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


285
@MULTIMODAL_REGISTRY.register_image_input_mapper()
286
@MULTIMODAL_REGISTRY.register_max_image_tokens(get_max_llava_image_tokens)
287
@INPUT_REGISTRY.register_dummy_data(dummy_data_for_llava)
288
@INPUT_REGISTRY.register_input_processor(input_processor_for_llava)
289
class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
290

291
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
292
        super().__init__()
293

294
295
296
297
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config

298
        self.config = config
299
        self.multimodal_config = multimodal_config
300

301
302
303
304
305
306
307
308
309
        # NOTE: These are special cases for Pixtral-12B in the HF-format
        # https://huggingface.co/mistral-community/pixtral-12b/blob/main/config.json  # noqa
        if (config.text_config.architectures is None
                and config.text_config.model_type == "mistral"):
            config.text_config.architectures = ["MistralForCausalLM"]
        if (config.projector_hidden_act is None
                and config.vision_config.hidden_act == "gelu"):
            config.projector_hidden_act = "gelu"

310
        # TODO: Optionally initializes this for supporting embeddings.
311
        self.vision_tower = init_vision_tower_for_llava(
312
313
314
            config,
            quant_config,
            require_post_norm=False,
315
            prefix=maybe_prefix(prefix, "vision_tower"))
316
317
318
319
320
        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)

321
        self.language_model = init_vllm_registered_model(
322
            config.text_config,
323
            vllm_config=vllm_config,
324
            prefix=maybe_prefix(prefix, "language_model"))
325

326
327
328
329
330
331
332
333
        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
334
        return get_sampler()
335

336
337
338
339
340
341
342
    def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
        h = w = self.config.vision_config.image_size
        expected_dims = (3, h, w)
        actual_dims = tuple(data.shape[1:])

        if actual_dims != expected_dims:
            expected_expr = ("batch_size", *map(str, expected_dims))
343
            raise ValueError(
344
345
                f"The expected shape of pixel values is {expected_expr}. "
                f"You supplied {tuple(data.shape)}.")
346
347
348

        return data

349
350
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
    def _validate_image_sizes(self, images: List[torch.Tensor],
                              sizes: List[torch.Tensor]) -> List[torch.Tensor]:
        if not isinstance(sizes, list):
            sizes = [sizes]

        total_images = sum(size.numel() // 2 for size in sizes)
        if total_images != len(images):
            raise ValueError("Mismatch in number of images. "
                             f"Expected {total_images}, got {len(images)}")
        img_idx = 0
        for size in sizes:
            # Flatten the size tensor to a list of (height, width) pairs
            size = size.view(-1, 2).tolist()
            for expected_h, expected_w in size:
                if img_idx >= len(images):
                    raise ValueError("Ran out of images before sizes. "
                                     f"{img_idx} >= {len(images)}")
                img = images[img_idx]
                if img.shape[-2:] != (expected_h, expected_w):
                    raise ValueError(
                        "Image size mismatch. Expected "
                        f"{(expected_h, expected_w)}, got {img.shape[-2:]}")
                if img.shape[-3] != 3:
                    raise ValueError("Image channel mismatch. Expected 3, "
                                     f"got {img.shape[-3]}")
                img_idx += 1
        return images

377
    def _parse_and_validate_image_input(
378
379
            self, **kwargs: object) -> Optional[LlavaImageInputs]:
        pixel_values = kwargs.pop("pixel_values", None)
380
        image_sizes = kwargs.pop("image_sizes", None)
381
        image_embeds = kwargs.pop("image_embeds", None)
382

383
        if pixel_values is None and image_embeds is None:
384
            return None
385

386
        if pixel_values is not None:
387
            if not isinstance(pixel_values, (torch.Tensor, list)):
388
389
                raise ValueError("Incorrect type of pixel values. "
                                 f"Got type: {type(pixel_values)}")
390

391
392
393
394
            # Case for models like PixtralHF that have dynamic image sizes
            # so we need to produce a list of tensors
            if image_sizes is not None:
                images = pixel_values
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413

                def flatten_to_3d_tensors(item):
                    if isinstance(item, torch.Tensor):
                        if item.dim() >= 3:
                            return [t for t in item.view(-1, *item.shape[-3:])]
                        else:
                            raise ValueError(
                                f"Unexpected tensor dimension: {item.dim()}")
                    elif isinstance(item, list):
                        return [
                            t for subitem in item
                            for t in flatten_to_3d_tensors(subitem)
                        ]
                    else:
                        raise ValueError(f"Unexpected type: {type(item)}")

                # Restructure the batched images into a list of lists of images
                images = flatten_to_3d_tensors(pixel_values)

414
415
                return LlavaImagePixelInputs(
                    type="pixel_values",
416
                    data=self._validate_image_sizes(images, image_sizes),
417
418
                )

419
420
            return LlavaImagePixelInputs(
                type="pixel_values",
421
422
                data=self._validate_pixel_values(
                    flatten_bn(pixel_values, concat=True)),
423
424
425
            )

        if image_embeds is not None:
426
            if not isinstance(image_embeds, (torch.Tensor, list)):
427
428
                raise ValueError("Incorrect type of image embeddings. "
                                 f"Got type: {type(image_embeds)}")
429

430
431
            return LlavaImageEmbeddingInputs(
                type="image_embeds",
432
                data=flatten_bn(image_embeds, concat=True),
433
434
435
            )

        raise AssertionError("This line should be unreachable.")
436
437
438
439
440
441
442
443
444
445
446

    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}")

447
448
    def _image_pixels_to_features(
        self,
449
450
        vision_tower: Union[CLIPVisionModel, SiglipVisionModel,
                            PixtralHFVisionModel],
451
452
        pixel_values: torch.Tensor,
    ) -> torch.Tensor:
453

454
455
        # NOTE: we skip the step to select the vision feature layer since
        # this is already done inside the vision tower
456
        image_features = vision_tower(pixel_values)
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472

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

    def _process_image_pixels(self,
                              inputs: LlavaImagePixelInputs) -> torch.Tensor:
        assert self.vision_tower is not None

        pixel_values = inputs["data"]

        return self._image_pixels_to_features(self.vision_tower, pixel_values)

    def _process_image_input(self,
                             image_input: LlavaImageInputs) -> torch.Tensor:
473
474
475
476

        if image_input["type"] == "image_embeds":
            return image_input["data"]

477
478
        assert self.vision_tower is not None
        image_features = self._process_image_pixels(image_input)
479
480
        return self.multi_modal_projector(image_features)

481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
    def process_mm_inputs(self, **kwargs):
        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,
        vision_embeddings: Optional[NestedTensors] = None,
    ) -> torch.Tensor:
        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
        if vision_embeddings is not None:
            inputs_embeds = merge_multimodal_embeddings(
                input_ids, inputs_embeds, vision_embeddings,
                self.config.image_token_index)
        return inputs_embeds

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

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

515
        Concretely, consider a text prompt:
516
517
        `"USER: <image>\\nWhat's the content of the image?\\nASSISTANT:"`.

518
        Tokenizer outputs:
519
520
521
522
        `[1, 3148, 1001, 29901, 29871, 32000, 29871, 13, 5618, 29915, 29879,
        278, 2793, 310, 278, 1967, 29973, 13, 22933, 9047, 13566, 29901]`.

        To reserve space in KV cache, we have to insert placeholder tokens
523
        before they are inputted to the model, so the input processor prepends
524
525
526
527
528
529
530
531
532
        additional image tokens (denoted as `32000`), resulting in:
        `[1, 3148, 1001, 29901, 29871, 32000, ..., 32000, 29871, 13, 5618,
        29915, 29879, 278, 2793, 310, 278, 1967, 29973, 13, 22933, 9047, 13566,
        29901]`.

        We insert 575 tokens so that including the original image token in the
        input, there are a total of 576 (24 * 24) image tokens, which
        corresponds to the number of image tokens inputted to the language
        model, i.e. the number of image tokens outputted by the visual encoder.
533
534
535
536
537
538
539

        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
540
            pixel_values: The pixels in each input image.
541

542
543
        See also:
            :class:`LlavaImageInputs`
544
        """
545
546
        if intermediate_tensors is not None:
            inputs_embeds = None
547
548
549
550
551
552
553
        elif inputs_embeds is None:
            vision_embeddings = self.process_mm_inputs(**kwargs)
            # always pass the input via `inputs_embeds`
            # to make sure the computation graph is consistent
            inputs_embeds = self.get_input_embeddings(input_ids,
                                                      vision_embeddings)
            input_ids = None
554

555
556
557
558
        hidden_states = self.language_model.model(input_ids,
                                                  positions,
                                                  kv_caches,
                                                  attn_metadata,
559
                                                  intermediate_tensors,
560
                                                  inputs_embeds=inputs_embeds)
561
562
563

        return hidden_states

564
565
566
567
568
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
569
570
        return self.language_model.compute_logits(hidden_states,
                                                  sampling_metadata)
571
572
573
574
575
576

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

579
580
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
581
        loader = AutoWeightsLoader(self)
582
        return loader.load_weights(weights)