paligemma.py 14.3 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterable, Mapping, Sequence
4
from typing import Literal, Optional, TypedDict, Union
Roger Wang's avatar
Roger Wang committed
5
6
7

import torch
from torch import nn
8
from transformers import BatchFeature, PaliGemmaConfig
Roger Wang's avatar
Roger Wang committed
9

10
from vllm.config import VllmConfig
Roger Wang's avatar
Roger Wang committed
11
12
13
from vllm.logger import init_logger
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.multimodal import MULTIMODAL_REGISTRY
14
from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
15
                                    MultiModalInputs, MultiModalKwargs)
16
17
from vllm.multimodal.parse import (ImageEmbeddingItems, ImageProcessorItems,
                                   MultiModalDataItems)
18
19
from vllm.multimodal.processing import (BaseMultiModalProcessor,
                                        BaseProcessingInfo, PromptIndexTargets,
20
                                        PromptInsertion, PromptUpdate,
21
                                        PromptUpdateDetails)
22
from vllm.multimodal.profiling import BaseDummyInputsBuilder
23
from vllm.sequence import IntermediateTensors
Roger Wang's avatar
Roger Wang committed
24

25
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
26
from .siglip import SiglipVisionModel
27
from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,
28
                    maybe_prefix, merge_multimodal_embeddings)
29
from .vision import get_vision_encoder_info
Roger Wang's avatar
Roger Wang committed
30
31
32
33

logger = init_logger(__name__)


34
35
36
class PaliGemmaImagePixelInputs(TypedDict):
    type: Literal["pixel_values"]
    data: torch.Tensor
37
    """Shape: `(batch_size * num_images, num_channels, height, width)`"""
38
39
40
41
42


class PaliGemmaImageEmbeddingInputs(TypedDict):
    type: Literal["image_embeds"]
    data: torch.Tensor
43
    """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
44
45
46
47
48
49
50
51
52

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


PaliGemmaImageInputs = Union[PaliGemmaImagePixelInputs,
                             PaliGemmaImageEmbeddingInputs]


Roger Wang's avatar
Roger Wang committed
53
54
55
56
57
class PaliGemmaMultiModalProjector(nn.Module):

    def __init__(self, vision_hidden_size: int, projection_dim: int):
        super().__init__()

58
        self.linear = nn.Linear(vision_hidden_size, projection_dim, bias=True)
Roger Wang's avatar
Roger Wang committed
59
60

    def forward(self, image_features: torch.Tensor) -> torch.Tensor:
61
        hidden_states = self.linear(image_features)
Roger Wang's avatar
Roger Wang committed
62
63
64
        return hidden_states


65
66
67
68
69
class PaliGemmaProcessingInfo(BaseProcessingInfo):

    def get_hf_config(self):
        return self.ctx.get_hf_config(PaliGemmaConfig)

70
71
72
    def get_vision_encoder_info(self):
        return get_vision_encoder_info(self.get_hf_config())

73
74
75
    def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
        return {"image": 1}

76
    def get_num_image_tokens(
77
        self,
78
79
80
81
        *,
        image_width: int,
        image_height: int,
    ) -> int:
82
        vision_encoder_info = self.get_vision_encoder_info()
83
84
85
86
87

        return vision_encoder_info.get_num_image_tokens(
            image_width=image_width,
            image_height=image_height,
        )
88
89
90
91
92


class PaliGemmaDummyInputsBuilder(
        BaseDummyInputsBuilder[PaliGemmaProcessingInfo]):

93
94
95
96
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        return ""

    def get_dummy_mm_data(
97
98
99
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
100
    ) -> MultiModalDataDict:
101
102
103
104
105
106
        hf_config = self.info.get_hf_config()
        vision_config = hf_config.vision_config
        max_image_size = vision_config.image_size

        num_images = mm_counts.get("image", 0)

107
        return {
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
            "image":
            self._get_dummy_images(width=max_image_size,
                                   height=max_image_size,
                                   num_images=num_images)
        }


class PaliGemmaMultiModalProcessor(
        BaseMultiModalProcessor[PaliGemmaProcessingInfo]):

    def _call_hf_processor(
        self,
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
    ) -> BatchFeature:
        tokenizer = self.info.get_tokenizer()
        if not mm_data:
            prompt_ids = tokenizer.encode(prompt)
            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")

        return super()._call_hf_processor(
            prompt=prompt,
            mm_data=mm_data,
            mm_kwargs=mm_kwargs,
        )

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

    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
        out_mm_kwargs: MultiModalKwargs,
147
    ) -> Sequence[PromptUpdate]:
148
149
150
151
152
153
154
155
        hf_config = self.info.get_hf_config()
        image_token_id = hf_config.image_token_index

        tokenizer = self.info.get_tokenizer()

        bos_token_id = tokenizer.bos_token_id
        assert isinstance(bos_token_id, int)

156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
        def get_insertion(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)
                num_image_tokens = self.info.get_num_image_tokens(
                    image_width=image_size.width,
                    image_height=image_size.height,
                )

            image_tokens = [image_token_id] * num_image_tokens

            return PromptUpdateDetails.select_token_id(
                image_tokens + [bos_token_id],
                embed_token_id=image_token_id,
            )

176
177
178
179
180
181
182
183
        # Paligemma 1 and 2 have different tokenizer.add_bos_token
        # Insert <image>*n + <bos> after <bos> for Paligemma 1
        # Insert <image>*n + <bos> for Paligemma 2
        return [
            PromptInsertion(
                modality="image",
                target=PromptIndexTargets.prefix(
                    [bos_token_id] if tokenizer.add_bos_token else []),
184
                insertion=get_insertion,
185
186
187
188
189
190
191
192
            )
        ]

    def apply(
        self,
        prompt: Union[str, list[int]],
        mm_data: MultiModalDataDict,
        hf_processor_mm_kwargs: Mapping[str, object],
193
        return_mm_hashes: bool = False,
194
    ) -> MultiModalInputs:
195
196
        mm_inputs = super().apply(prompt, mm_data, hf_processor_mm_kwargs,
                                  return_mm_hashes)
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
        prompt_token_ids = mm_inputs["prompt_token_ids"]

        tokenizer = self.info.get_tokenizer()
        newline_prompt = "\n"
        newline_token_id = tokenizer.encode(newline_prompt)[-1]  # 108
        # Force to add newline at the end of prompt for paligemma's format
        # This step can NOT be replacemented by current PromptUpdate methods
        if len(prompt_token_ids) and prompt_token_ids[-1] != newline_token_id:
            prompt_token_ids.append(newline_token_id)
            mm_inputs["prompt_token_ids"] = prompt_token_ids
            mm_inputs["prompt"] += newline_prompt

        return mm_inputs


@MULTIMODAL_REGISTRY.register_processor(
    PaliGemmaMultiModalProcessor,
    info=PaliGemmaProcessingInfo,
    dummy_inputs=PaliGemmaDummyInputsBuilder)
216
class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal,
217
                                        SupportsPP):
218
219
220
221
222
223
224
225
226
227
228
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }
229

230
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Roger Wang's avatar
Roger Wang committed
231
        super().__init__()
232
233
234
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config
Roger Wang's avatar
Roger Wang committed
235
236
237
        self.config = config
        self.multimodal_config = multimodal_config

238
        self.vision_tower = SiglipVisionModel(config.vision_config,
239
                                              quant_config,
240
241
                                              prefix=maybe_prefix(
                                                  prefix, "vision_tower"))
Roger Wang's avatar
Roger Wang committed
242
243
244
245
246
        self.multi_modal_projector = PaliGemmaMultiModalProjector(
            vision_hidden_size=config.vision_config.hidden_size,
            projection_dim=config.vision_config.projection_dim)

        self.quant_config = quant_config
Jani Monoses's avatar
Jani Monoses committed
247
248
249
250
251

        if config.text_config.model_type == "gemma":
            config.text_config.architectures = ["GemmaForCausalLM"]
        else:
            config.text_config.architectures = ["Gemma2ForCausalLM"]
252
253
        self.language_model = init_vllm_registered_model(
            vllm_config=vllm_config,
254
255
256
            hf_config=config.text_config,
            prefix=maybe_prefix(prefix, "language_model"),
        )
Roger Wang's avatar
Roger Wang committed
257
        logit_scale = getattr(config, "logit_scale", 1.0)
258
259
260
261
262
        self.language_model.logits_processor.scale *= logit_scale

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors)

Roger Wang's avatar
Roger Wang committed
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
    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))
            raise ValueError(
                f"The expected shape of pixel values is {expected_expr}. "
                f"You supplied {tuple(data.shape)}.")

        return data

    def _parse_and_validate_image_input(
            self, **kwargs: object) -> Optional[PaliGemmaImageInputs]:
        pixel_values = kwargs.pop("pixel_values", None)
279
        image_embeds = kwargs.pop("image_embeds", None)
Roger Wang's avatar
Roger Wang committed
280

281
        if pixel_values is None and image_embeds is None:
Roger Wang's avatar
Roger Wang committed
282
283
            return None

284
        if pixel_values is not None:
285
            if not isinstance(pixel_values, (torch.Tensor, list)):
286
287
                raise ValueError("Incorrect type of pixel values. "
                                 f"Got type: {type(pixel_values)}")
288

289
            pixel_values = flatten_bn(pixel_values, concat=True)
290

291
292
293
294
295
296
297
298
299
            return PaliGemmaImagePixelInputs(
                type="pixel_values",
                data=self._validate_pixel_values(pixel_values),
            )

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

301
            image_embeds = flatten_bn(image_embeds, concat=True)
302

303
304
305
306
307
308
            return PaliGemmaImageEmbeddingInputs(
                type="image_embeds",
                data=image_embeds,
            )

        raise AssertionError("This line should be unreachable.")
Roger Wang's avatar
Roger Wang committed
309

310
311
312
313
314
    def _image_pixels_to_features(
        self,
        vision_tower: SiglipVisionModel,
        pixel_values: torch.Tensor,
    ) -> torch.Tensor:
Roger Wang's avatar
Roger Wang committed
315

316
        target_dtype = vision_tower.get_input_embeddings().weight.dtype
317
        image_features = vision_tower(pixel_values.to(dtype=target_dtype))
Roger Wang's avatar
Roger Wang committed
318

319
        return image_features
Roger Wang's avatar
Roger Wang committed
320

321
    def _process_image_input(
322
        self,
323
        image_input: PaliGemmaImageInputs,
324
    ) -> torch.Tensor:
Roger Wang's avatar
Roger Wang committed
325

326
327
        if image_input["type"] == "image_embeds":
            return image_input["data"]
Roger Wang's avatar
Roger Wang committed
328

329
330
331
        assert self.vision_tower is not None
        pixel_values = image_input["data"]
        image_features = self._image_pixels_to_features(
332
333
334
            self.vision_tower,
            pixel_values,
        )
Roger Wang's avatar
Roger Wang committed
335
336
337

        return self.multi_modal_projector(image_features)

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

341
    def get_multimodal_embeddings(
342
            self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
343
344
345
346
347
348
349
350
351
352
353
        image_input = self._parse_and_validate_image_input(**kwargs)
        if image_input is None:
            return None
        vision_embeddings = self._process_image_input(image_input)
        # https://github.com/huggingface/transformers/blob/main/src/transformers/models/paligemma/modeling_paligemma.py#L294 # noqa
        vision_embeddings = vision_embeddings * (self.config.hidden_size**-0.5)
        return vision_embeddings

    def get_input_embeddings(
        self,
        input_ids: torch.Tensor,
354
        multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
355
356
357
358
359
360
361
362
    ) -> torch.Tensor:
        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
        if multimodal_embeddings is not None:
            inputs_embeds = merge_multimodal_embeddings(
                input_ids, inputs_embeds, multimodal_embeddings,
                self.config.image_token_index)
        return inputs_embeds

363
364
365
366
    def forward(self,
                input_ids: torch.Tensor,
                positions: torch.Tensor,
                intermediate_tensors: Optional[IntermediateTensors] = None,
367
                inputs_embeds: Optional[torch.Tensor] = None,
368
                **kwargs: object) -> IntermediateTensors:
369
370
        if intermediate_tensors is not None:
            inputs_embeds = None
371
372
373
374
375
376
377
378

        # 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
Roger Wang's avatar
Roger Wang committed
379

380
381
        hidden_states = self.language_model.model(input_ids,
                                                  positions,
382
                                                  intermediate_tensors,
383
                                                  inputs_embeds=inputs_embeds)
Roger Wang's avatar
Roger Wang committed
384
385
386

        return hidden_states

387
388
389
390
391
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
392
393
        return self.language_model.compute_logits(hidden_states,
                                                  sampling_metadata)
Roger Wang's avatar
Roger Wang committed
394

395
396
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
397
        loader = AutoWeightsLoader(self)
398
        return loader.load_weights(weights)