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

import torch
7
import torch.nn as nn
8
9
from transformers import (BatchFeature, CLIPVisionConfig, LlavaConfig,
                          PixtralVisionConfig, PretrainedConfig,
10
                          SiglipVisionConfig)
11
from transformers.models.llava import LlavaProcessor
12
from transformers.models.pixtral import PixtralProcessor
13
14

from vllm.attention import AttentionMetadata
15
from vllm.config import VllmConfig
16
from vllm.inputs import InputProcessingContext
17
from vllm.model_executor.layers.activation import get_act_fn
18
19
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               RowParallelLinear)
20
from vllm.model_executor.layers.quantization import QuantizationConfig
Joe Runde's avatar
Joe Runde committed
21
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
22
from vllm.model_executor.sampling_metadata import SamplingMetadata
23
from vllm.multimodal import MULTIMODAL_REGISTRY
24
25
26
from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
                                    MultiModalInputsV2, MultiModalKwargs,
                                    NestedTensors)
27
from vllm.multimodal.parse import (ImageEmbeddingItems, ImageProcessorItems,
28
                                   ImageSize, MultiModalDataItems)
29
from vllm.multimodal.processing import (BaseMultiModalProcessor,
30
31
32
                                        BaseProcessingInfo, ProcessingCache,
                                        PromptReplacement)
from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
33
from vllm.sequence import IntermediateTensors
34

35
from .clip import CLIPVisionModel
36
from .interfaces import SupportsMultiModal, SupportsPP
37
38
39
from .pixtral import (PixtralHFVisionModel,
                      get_pixtral_hf_image_feature_grid_size)
from .siglip import SiglipVisionModel
40
from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,
41
                    maybe_prefix, merge_multimodal_embeddings)
42
from .vision import get_vision_encoder_info
43
44


45
46
class LlavaImagePixelInputs(TypedDict):
    type: Literal["pixel_values"]
47
48
49
50
51
52
53
    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.
    """
54
55
56
57
58


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

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


LlavaImageInputs = Union[LlavaImagePixelInputs, LlavaImageEmbeddingInputs]


68
69
class LlavaMultiModalProjector(nn.Module):

70
71
72
73
74
75
    def __init__(self,
                 vision_hidden_size: int,
                 text_hidden_size: int,
                 projector_hidden_act: str,
                 quant_config: Optional[QuantizationConfig] = None,
                 prefix: str = ""):
76
77
        super().__init__()

78
79
80
81
82
        self.linear_1 = ColumnParallelLinear(vision_hidden_size,
                                             text_hidden_size,
                                             bias=True,
                                             quant_config=quant_config,
                                             prefix=f"{prefix}.linear_1")
83
        self.act = get_act_fn(projector_hidden_act)
84
85
86
87
88
        self.linear_2 = RowParallelLinear(text_hidden_size,
                                          text_hidden_size,
                                          bias=True,
                                          quant_config=quant_config,
                                          prefix=f"{prefix}.linear_2")
89

90
    def forward(self, image_features: torch.Tensor) -> torch.Tensor:
91
        hidden_states, _ = self.linear_1(image_features)
92
        hidden_states = self.act(hidden_states)
93
        hidden_states, _ = self.linear_2(hidden_states)
94
95
96
        return hidden_states


97
98
class LlavaLikeConfig(Protocol):
    vision_config: Final[PretrainedConfig]
99
    image_token_index: Final[int]
100
    vision_feature_select_strategy: Final[str]
101
    vision_feature_layer: Final[Union[int, list[int]]]
102

103

104
105
106
107
class LlavaLikeProcessor(Protocol):
    image_token: Final[str]


108
class BaseLlavaProcessingInfo(BaseProcessingInfo):
109

110
    def get_hf_config(self) -> LlavaLikeConfig:
111
        return self.ctx.get_hf_config(LlavaConfig)
112

113
114
    def get_vision_encoder_info(self):
        return get_vision_encoder_info(self.get_hf_config())
115

116
    @abstractmethod
117
    def get_hf_processor(self) -> LlavaLikeProcessor:
118
        raise NotImplementedError
119

120
121
    def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
        return {"image": None}
122

123
124
    def get_mm_max_tokens_per_item(self, seq_len: int) -> Mapping[str, int]:
        return {"image": self.get_max_image_tokens()}
125

126
127
128
129
130
131
132
133
134
135
136
137
138
    def _apply_feature_select_strategy(
        self,
        strategy: str,
        encoder_num_image_tokens: int,
    ) -> int:
        if strategy == "default":
            return encoder_num_image_tokens - 1
        if strategy == "full":
            return encoder_num_image_tokens

        msg = f"Unexpected feature select strategy: {strategy!r}"
        raise NotImplementedError(msg)

139
140
141
142
143
144
145
146
    def get_num_image_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
    ) -> int:
        hf_config = self.get_hf_config()
        vision_encoder_info = self.get_vision_encoder_info()
147

148
149
150
151
152
153
154
        return self._apply_feature_select_strategy(
            hf_config.vision_feature_select_strategy,
            vision_encoder_info.get_num_image_tokens(
                image_width=image_width,
                image_height=image_height,
            ),
        )
155

156
157
    def get_image_size_with_most_features(self) -> ImageSize:
        vision_encoder_info = self.get_vision_encoder_info()
158
159
        width = height = vision_encoder_info.get_image_size()
        return ImageSize(width=width, height=height)
160

161
162
    def get_max_image_tokens(self) -> int:
        target_width, target_height = self.get_image_size_with_most_features()
163

164
        return self.get_num_image_tokens(
165
166
167
168
            image_width=target_width,
            image_height=target_height,
        )

169
170
171
172
173
174

_I = TypeVar("_I", bound=BaseLlavaProcessingInfo)


class LlavaDummyInputsBuilder(BaseDummyInputsBuilder[_I]):

175
    def get_dummy_processor_inputs(
176
        self,
177
        seq_len: int,
178
179
180
181
        mm_counts: Mapping[str, int],
    ) -> ProcessorInputs:
        num_images = mm_counts.get("image", 0)

182
        processor = self.info.get_hf_processor()
183
        image_token = processor.image_token
184
185
        target_width, target_height = \
            self.info.get_image_size_with_most_features()
186
187
188
189
190
191
192
193
194
195
196
197
198
199

        mm_data = {
            "image":
            self._get_dummy_images(width=target_width,
                                   height=target_height,
                                   num_images=num_images)
        }

        return ProcessorInputs(
            prompt_text=image_token * num_images,
            mm_data=mm_data,
        )


200
class LlavaProcessingInfo(BaseLlavaProcessingInfo):
201

202
    def get_hf_processor(self):
203
204
205
        return self.ctx.get_hf_processor(LlavaProcessor)


206
class BaseLlavaMultiModalProcessor(BaseMultiModalProcessor[_I]):
207
208
209
210
211
212
213
214
215

    # 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
216
217
218
219
220
221
222

    def _get_prompt_replacements(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
        out_mm_kwargs: MultiModalKwargs,
    ) -> list[PromptReplacement]:
223
        hf_config = self.info.get_hf_config()
224
225
226
227
228
229
230
231
232
233
        image_token_id = hf_config.image_token_index

        def get_replacement(item_idx: int):
            images = mm_items.get_items(
                "image", (ImageEmbeddingItems, ImageProcessorItems))

            if isinstance(images, ImageEmbeddingItems):
                num_image_tokens = images.get_feature_size(item_idx)
            else:
                image_size = images.get_image_size(item_idx)
234
                num_image_tokens = self.info.get_num_image_tokens(
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
                    image_width=image_size.width,
                    image_height=image_size.height,
                )

            return [image_token_id] * num_image_tokens

        return [
            PromptReplacement(
                modality="image",
                target=[image_token_id],
                replacement=get_replacement,
            ),
        ]


250
251
class LlavaMultiModalProcessor(
        BaseLlavaMultiModalProcessor[LlavaProcessingInfo]):
252

253
254
255
256
257
258
259
260
261
262
263
    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_embeds=MultiModalFieldConfig.batched("image"),
        )


264
class PixtralHFProcessingInfo(BaseLlavaProcessingInfo):
265

266
    def get_hf_processor(self):
267
268
        return self.ctx.get_hf_processor(PixtralProcessor)

269

270
271
class PixtralHFMultiModalProcessor(
        BaseMultiModalProcessor[PixtralHFProcessingInfo]):
272

273
274
275
276
277
278
279
280
281
282
283
    def _call_hf_processor(
        self,
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
    ) -> BatchFeature:
        processed_outputs = super()._call_hf_processor(
            prompt=prompt,
            mm_data=mm_data,
            mm_kwargs=mm_kwargs,
        )
284

285
286
287
288
        pixel_values = processed_outputs.get("pixel_values")
        if pixel_values is not None:
            images = mm_data["images"]
            assert isinstance(images, list)
289

290
291
292
293
294
            # Original output: (1, num_images, C, H, W)
            # New output: (num_images, C, H, W)
            assert (isinstance(pixel_values, list) and len(pixel_values) == 1)
            assert (isinstance(pixel_values[0], list)
                    and len(pixel_values[0]) == len(images))
295

296
            processed_outputs["pixel_values"] = pixel_values[0]
297

298
        return processed_outputs
299

300
301
302
303
304
305
306
307
308
309
    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_embeds=MultiModalFieldConfig.batched("image"),
        )

310
311
312
    def _get_prompt_replacements(
        self,
        mm_items: MultiModalDataItems,
313
314
        hf_processor_mm_kwargs: Mapping[str, object],
        out_mm_kwargs: MultiModalKwargs,
315
    ) -> list[PromptReplacement]:
316
        hf_config = self.info.get_hf_config()
317
318
        image_token_id = hf_config.image_token_index

319
        processor = self.info.get_hf_processor()
320
321
322
        image_token = processor.image_token
        image_break_token = processor.image_break_token
        image_end_token = processor.image_end_token
323

324
325
        vision_config = hf_config.vision_config
        assert isinstance(vision_config, PixtralVisionConfig)
326

327
328
329
        def get_replacement(item_idx: int):
            images = mm_items.get_items("image", ImageProcessorItems)
            image_size = images.get_image_size(item_idx)
330

331
332
333
334
335
            ncols, nrows = get_pixtral_hf_image_feature_grid_size(
                vision_config,
                image_width=image_size.width,
                image_height=image_size.height,
            )
336

337
338
            tokens = ([image_token] * ncols + [image_break_token]) * nrows
            tokens[-1] = image_end_token
339

340
            return "".join(tokens)
341
342
343
344
345

        return [
            PromptReplacement(
                modality="image",
                target=[image_token_id],
346
347
                replacement=get_replacement,
            ),
348
349
        ]

350

351
352
353
354
355
356
357
358
359
360
def _build_llava_or_pixtral_hf_info(
    ctx: InputProcessingContext, ) -> BaseLlavaProcessingInfo:
    hf_config = ctx.get_hf_config(LlavaConfig)

    if isinstance(hf_config.vision_config, PixtralVisionConfig):
        return PixtralHFProcessingInfo(ctx)

    return LlavaProcessingInfo(ctx)


361
def _build_llava_or_pixtral_hf_processor(
362
363
    info: _I,
    dummy_inputs: BaseDummyInputsBuilder[_I],
364
365
366
    *,
    cache: Optional[ProcessingCache] = None,
    enable_sanity_checks: bool = True,
367
) -> BaseMultiModalProcessor:
368
    if isinstance(info, PixtralHFProcessingInfo):
369
        return PixtralHFMultiModalProcessor(
370
371
372
373
374
375
376
377
378
379
            info,
            dummy_inputs,  # type: ignore
            cache=cache,
            enable_sanity_checks=enable_sanity_checks,
        )

    if isinstance(info, LlavaProcessingInfo):
        return LlavaMultiModalProcessor(
            info,
            dummy_inputs,  # type: ignore
380
381
            cache=cache,
            enable_sanity_checks=enable_sanity_checks,
382
        )
383

384
    raise NotImplementedError(type(info))
385
386
387
388
389
390
391
392
393
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


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
419
420
421
422
423
424
425


def init_vision_tower_for_llava(
    hf_config: LlavaLikeConfig,
    quant_config: Optional[QuantizationConfig],
    *,
    require_post_norm: Optional[bool] = None,
426
    prefix: str = "",
427
):
428
429
    vision_config = hf_config.vision_config

430
431
    # Initialize the vision tower only up to the deepest required feature layer
    num_hidden_layers = _get_num_hidden_layers(hf_config)
432
433
434
435

    if isinstance(vision_config, CLIPVisionConfig):
        return CLIPVisionModel(
            vision_config,
436
            quant_config=quant_config,
437
            num_hidden_layers_override=num_hidden_layers,
438
            require_post_norm=require_post_norm,
439
            prefix=prefix,
440
441
442
443
        )
    elif isinstance(vision_config, SiglipVisionConfig):
        return SiglipVisionModel(
            vision_config,
444
            quant_config=quant_config,
445
            num_hidden_layers_override=num_hidden_layers,
446
            require_post_norm=require_post_norm,
447
            prefix=prefix,
448
        )
449
    elif isinstance(vision_config, PixtralVisionConfig):
450
451
        return PixtralHFVisionModel(
            vision_config,
452
            quant_config=quant_config,
453
454
            num_hidden_layers_override=num_hidden_layers,
            require_post_norm=require_post_norm,
455
            prefix=prefix,
456
        )
457
458
459
460
461

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


462
463
464
@MULTIMODAL_REGISTRY.register_processor(_build_llava_or_pixtral_hf_processor,
                                        info=_build_llava_or_pixtral_hf_info,
                                        dummy_inputs=LlavaDummyInputsBuilder)
465
class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
466
467
468
469

    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"]
470
    }
471

472
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
473
        super().__init__()
474

475
476
477
478
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config

479
        self.config = config
480
        self.multimodal_config = multimodal_config
481

482
483
484
485
486
487
488
489
490
        # 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"

491
        # TODO: Optionally initializes this for supporting embeddings.
492
        self.vision_tower = init_vision_tower_for_llava(
493
494
495
            config,
            quant_config,
            require_post_norm=False,
496
            prefix=maybe_prefix(prefix, "vision_tower"))
497
498
499
        self.multi_modal_projector = LlavaMultiModalProjector(
            vision_hidden_size=config.vision_config.hidden_size,
            text_hidden_size=config.text_config.hidden_size,
500
501
502
            projector_hidden_act=config.projector_hidden_act,
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, "multi_modal_projector"))
503

504
        self.language_model = init_vllm_registered_model(
505
            vllm_config=vllm_config,
506
507
508
            hf_config=config.text_config,
            prefix=maybe_prefix(prefix, "language_model"),
        )
509

510
511
512
513
514
515
516
517
        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
518
        return get_sampler()
519

520
521
522
523
524
525
526
    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))
527
            raise ValueError(
528
529
                f"The expected shape of pixel values is {expected_expr}. "
                f"You supplied {tuple(data.shape)}.")
530
531
532
533

        return data

    def _parse_and_validate_image_input(
534
535
            self, **kwargs: object) -> Optional[LlavaImageInputs]:
        pixel_values = kwargs.pop("pixel_values", None)
536
        image_embeds = kwargs.pop("image_embeds", None)
537

538
        if pixel_values is None and image_embeds is None:
539
            return None
540

541
        if pixel_values is not None:
542
            if not isinstance(pixel_values, (torch.Tensor, list)):
543
544
                raise ValueError("Incorrect type of pixel values. "
                                 f"Got type: {type(pixel_values)}")
545

546
547
548
549
550
551
            if self.config.vision_config.model_type == "pixtral":
                return LlavaImagePixelInputs(
                    type="pixel_values",
                    data=flatten_bn(pixel_values),
                )

552
553
            return LlavaImagePixelInputs(
                type="pixel_values",
554
555
                data=self._validate_pixel_values(
                    flatten_bn(pixel_values, concat=True)),
556
557
558
            )

        if image_embeds is not None:
559
            if not isinstance(image_embeds, (torch.Tensor, list)):
560
561
                raise ValueError("Incorrect type of image embeddings. "
                                 f"Got type: {type(image_embeds)}")
562

563
564
            return LlavaImageEmbeddingInputs(
                type="image_embeds",
565
                data=flatten_bn(image_embeds, concat=True),
566
567
568
            )

        raise AssertionError("This line should be unreachable.")
569
570
571
572
573
574
575
576
577
578
579

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

580
581
    def _image_pixels_to_features(
        self,
582
583
        vision_tower: Union[CLIPVisionModel, SiglipVisionModel,
                            PixtralHFVisionModel],
584
585
        pixel_values: torch.Tensor,
    ) -> torch.Tensor:
586

587
588
        # NOTE: we skip the step to select the vision feature layer since
        # this is already done inside the vision tower
589
        image_features = vision_tower(pixel_values)
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605

        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:
606
607
608
609

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

610
611
        assert self.vision_tower is not None
        image_features = self._process_image_pixels(image_input)
612
613
        return self.multi_modal_projector(image_features)

614
    def get_multimodal_embeddings(self, **kwargs) -> Optional[NestedTensors]:
615
616
617
618
619
620
621
622
623
        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,
624
        multimodal_embeddings: Optional[NestedTensors] = None,
625
626
    ) -> torch.Tensor:
        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
627
        if multimodal_embeddings is not None:
628
            inputs_embeds = merge_multimodal_embeddings(
629
                input_ids, inputs_embeds, multimodal_embeddings,
630
631
632
                self.config.image_token_index)
        return inputs_embeds

633
634
635
636
637
638
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
639
        intermediate_tensors: Optional[IntermediateTensors] = None,
640
        inputs_embeds: Optional[torch.Tensor] = None,
641
        **kwargs: object,
642
    ) -> Union[torch.Tensor, IntermediateTensors]:
Cyrus Leung's avatar
Cyrus Leung committed
643
        """Run forward pass for LLaVA-1.5.
644
645
646

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

648
        Concretely, consider a text prompt:
649
650
        `"USER: <image>\\nWhat's the content of the image?\\nASSISTANT:"`.

651
        Tokenizer outputs:
652
653
654
655
        `[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
656
        before they are inputted to the model, so the input processor prepends
657
658
659
660
661
662
663
664
665
        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.
666
667
668
669
670
671
672

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

675
676
        See also:
            :class:`LlavaImageInputs`
677
        """
678
679
        if intermediate_tensors is not None:
            inputs_embeds = None
680
681
682

        # NOTE: In v1, inputs_embeds is always generated at model runner, this
        # condition is for v0 compatibility.
683
        elif inputs_embeds is None:
684
            vision_embeddings = self.get_multimodal_embeddings(**kwargs)
685
686
687
            inputs_embeds = self.get_input_embeddings(input_ids,
                                                      vision_embeddings)
            input_ids = None
688

689
690
691
692
        hidden_states = self.language_model.model(input_ids,
                                                  positions,
                                                  kv_caches,
                                                  attn_metadata,
693
                                                  intermediate_tensors,
694
                                                  inputs_embeds=inputs_embeds)
695
696
697

        return hidden_states

698
699
700
701
702
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
703
704
        return self.language_model.compute_logits(hidden_states,
                                                  sampling_metadata)
705
706
707
708
709
710

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

713
714
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
715
        loader = AutoWeightsLoader(self)
716
        return loader.load_weights(weights)
717
718


719
class MantisMultiModalProcessor(LlavaMultiModalProcessor):
720

721
722
    def apply(
        self,
723
        prompt: Union[str, list[int]],
724
725
726
        mm_data: MultiModalDataDict,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> MultiModalInputsV2:
727
        hf_config = self.info.get_hf_config()
728
        image_token_id = hf_config.image_token_index
729
730

        # Assume that it doesn't depend on the image size
731
        num_image_tokens = self.info.get_num_image_tokens(
732
733
734
            image_width=-1,
            image_height=-1,
        )
735

736
        result = super().apply(prompt, mm_data, hf_processor_mm_kwargs)
737

738
739
        mm_items = self._to_mm_items(mm_data)
        mm_item_counts = mm_items.get_all_counts()
740
741
742
743
744
745
746
        mm_kwargs = result["mm_kwargs"]

        # We reimplement the functionality of MLlavaProcessor from
        # https://github.com/TIGER-AI-Lab/Mantis.git
        def get_replacement_mantis(item_idx: int):
            return "".join([
                f"(image {item_idx+1}: <Image>",  # 7 tokens
747
                "<image>" * num_image_tokens,
748
749
750
                "</Image>)",  # 3 tokens
            ])

751
        mantis_mm_repls = self._bind_and_group_repls([
752
753
            PromptReplacement(
                modality="image",
754
                target=[image_token_id] * num_image_tokens,
755
756
757
758
                replacement=get_replacement_mantis,
            )
        ])

759
        prompt_ids, prompt, _ = self._apply_prompt_replacements(
760
            result["prompt_token_ids"],
761
            mantis_mm_repls,
762
763
764
765
766
767
768
769
            mm_item_counts,
        )

        unbound_orig_repls = self._get_prompt_replacements(
            mm_items,
            hf_processor_mm_kwargs,
            mm_kwargs,
        )
770
771
772
773
774
775
776
        orig_repls = self._bind_and_group_repls(unbound_orig_repls)

        mm_placeholders = self._find_mm_placeholders(
            orig_repls,
            prompt_ids,
            mm_item_counts,
        )
777

778
        self._validate_mm_placeholders(mm_placeholders, mm_item_counts)
779

780
781
782
        mm_placeholder_ranges = {
            modality: [item.to_range() for item in placeholders]
            for modality, placeholders in mm_placeholders.items()
783
784
785
786
        }

        return MultiModalInputsV2(
            type="multimodal",
787
            prompt=prompt,
788
789
            prompt_token_ids=prompt_ids,
            mm_kwargs=mm_kwargs,
790
            mm_placeholders=mm_placeholder_ranges,
791
        )
792
793
794
795


# To use this model, please use
# `--hf_overrides '{"architectures": ["MantisForConditionalGeneration"]}'`
796
797
798
@MULTIMODAL_REGISTRY.register_processor(MantisMultiModalProcessor,
                                        info=LlavaProcessingInfo,
                                        dummy_inputs=LlavaDummyInputsBuilder)
799
800
class MantisForConditionalGeneration(LlavaForConditionalGeneration):
    pass