paligemma.py 14 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
3
from collections.abc import Iterable, Mapping, Sequence
from typing import Literal, Optional, Set, Tuple, TypedDict, Union
Roger Wang's avatar
Roger Wang committed
4
5
6

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

9
from vllm.config import VllmConfig
Roger Wang's avatar
Roger Wang committed
10
from vllm.logger import init_logger
11
from vllm.model_executor.layers.sampler import SamplerOutput
Roger Wang's avatar
Roger Wang committed
12
13
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
18
from vllm.multimodal.parse import MultiModalDataItems
from vllm.multimodal.processing import (BaseMultiModalProcessor,
                                        BaseProcessingInfo, PromptIndexTargets,
19
                                        PromptInsertion, PromptUpdate,
20
21
                                        PromptUpdateDetails)
from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
22
from vllm.sequence import IntermediateTensors
Roger Wang's avatar
Roger Wang committed
23

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

logger = init_logger(__name__)


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


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

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


PaliGemmaImageInputs = Union[PaliGemmaImagePixelInputs,
                             PaliGemmaImageEmbeddingInputs]


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

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

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

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


64
65
66
67
68
class PaliGemmaProcessingInfo(BaseProcessingInfo):

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

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

72
73
74
75
76
77
78
79
80
81
82
    def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
        return {"image": 1}

    def get_mm_max_tokens_per_item(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> Mapping[str, int]:
        return {"image": self.get_num_image_tokens()}

    def get_num_image_tokens(self) -> int:
83
84
        vision_encoder_info = self.get_vision_encoder_info()
        return vision_encoder_info.get_max_image_tokens()
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
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


class PaliGemmaDummyInputsBuilder(
        BaseDummyInputsBuilder[PaliGemmaProcessingInfo]):

    def get_dummy_processor_inputs(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> ProcessorInputs:
        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)

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

        return ProcessorInputs(
            prompt_text="",
            mm_data=mm_data,
        )


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,
146
    ) -> Sequence[PromptUpdate]:
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
        hf_config = self.info.get_hf_config()
        image_token_id = hf_config.image_token_index

        tokenizer = self.info.get_tokenizer()
        num_image_tokens = self.info.get_num_image_tokens()
        image_tokens = [image_token_id] * num_image_tokens

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

        # 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 []),
165
166
167
                insertion=PromptUpdateDetails.select_token_id(
                    image_tokens + [bos_token_id],
                    embed_token_id=image_token_id,
168
169
170
171
172
173
174
175
176
                ),
            )
        ]

    def apply(
        self,
        prompt: Union[str, list[int]],
        mm_data: MultiModalDataDict,
        hf_processor_mm_kwargs: Mapping[str, object],
177
        return_mm_hashes: bool = False,
178
    ) -> MultiModalInputs:
179
180
        mm_inputs = super().apply(prompt, mm_data, hf_processor_mm_kwargs,
                                  return_mm_hashes)
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
        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)
200
class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal,
201
                                        SupportsPP):
202
203
204
205
206
207
208
209
210
211
212
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }
213

214
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Roger Wang's avatar
Roger Wang committed
215
        super().__init__()
216
217
218
        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
219
220
221
        self.config = config
        self.multimodal_config = multimodal_config

222
        self.vision_tower = SiglipVisionModel(config.vision_config,
223
                                              quant_config,
224
225
                                              prefix=maybe_prefix(
                                                  prefix, "vision_tower"))
Roger Wang's avatar
Roger Wang committed
226
227
228
229
230
        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
231
232
233
234
235

        if config.text_config.model_type == "gemma":
            config.text_config.architectures = ["GemmaForCausalLM"]
        else:
            config.text_config.architectures = ["Gemma2ForCausalLM"]
236
237
        self.language_model = init_vllm_registered_model(
            vllm_config=vllm_config,
238
239
240
            hf_config=config.text_config,
            prefix=maybe_prefix(prefix, "language_model"),
        )
Roger Wang's avatar
Roger Wang committed
241
        logit_scale = getattr(config, "logit_scale", 1.0)
242
243
244
245
246
247
248
249
        self.language_model.logits_processor.scale *= logit_scale

        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors)

    @property
    def sampler(self):
        return self.language_model.sampler
Roger Wang's avatar
Roger Wang committed
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266

    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)
267
        image_embeds = kwargs.pop("image_embeds", None)
Roger Wang's avatar
Roger Wang committed
268

269
        if pixel_values is None and image_embeds is None:
Roger Wang's avatar
Roger Wang committed
270
271
            return None

272
        if pixel_values is not None:
273
            if not isinstance(pixel_values, (torch.Tensor, list)):
274
275
                raise ValueError("Incorrect type of pixel values. "
                                 f"Got type: {type(pixel_values)}")
276

277
            pixel_values = flatten_bn(pixel_values, concat=True)
278

279
280
281
282
283
284
285
286
287
            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)}")
288

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

291
292
293
294
295
296
            return PaliGemmaImageEmbeddingInputs(
                type="image_embeds",
                data=image_embeds,
            )

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

298
299
300
301
302
    def _image_pixels_to_features(
        self,
        vision_tower: SiglipVisionModel,
        pixel_values: torch.Tensor,
    ) -> torch.Tensor:
Roger Wang's avatar
Roger Wang committed
303

304
        target_dtype = vision_tower.get_input_embeddings().weight.dtype
305
        image_features = vision_tower(pixel_values.to(dtype=target_dtype))
Roger Wang's avatar
Roger Wang committed
306

307
        return image_features
Roger Wang's avatar
Roger Wang committed
308

309
    def _process_image_input(
310
        self,
311
        image_input: PaliGemmaImageInputs,
312
    ) -> torch.Tensor:
Roger Wang's avatar
Roger Wang committed
313

314
315
        if image_input["type"] == "image_embeds":
            return image_input["data"]
Roger Wang's avatar
Roger Wang committed
316

317
318
319
        assert self.vision_tower is not None
        pixel_values = image_input["data"]
        image_features = self._image_pixels_to_features(
320
321
322
            self.vision_tower,
            pixel_values,
        )
Roger Wang's avatar
Roger Wang committed
323
324
325

        return self.multi_modal_projector(image_features)

326
    def get_multimodal_embeddings(
327
            self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
328
329
330
331
332
333
334
335
336
337
338
        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,
339
        multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
340
341
342
343
344
345
346
347
    ) -> 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

348
349
350
351
    def forward(self,
                input_ids: torch.Tensor,
                positions: torch.Tensor,
                intermediate_tensors: Optional[IntermediateTensors] = None,
352
                inputs_embeds: Optional[torch.Tensor] = None,
353
354
355
                **kwargs: object) -> Union[SamplerOutput, IntermediateTensors]:
        if intermediate_tensors is not None:
            inputs_embeds = None
356
357
358
359
360
361
362
363

        # 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
364

365
366
        hidden_states = self.language_model.model(input_ids,
                                                  positions,
367
                                                  intermediate_tensors,
368
                                                  inputs_embeds=inputs_embeds)
Roger Wang's avatar
Roger Wang committed
369
370
371

        return hidden_states

372
373
374
375
376
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
377
378
        return self.language_model.compute_logits(hidden_states,
                                                  sampling_metadata)
Roger Wang's avatar
Roger Wang committed
379
380
381
382
383
384

    def sample(
        self,
        logits: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
385
        return self.language_model.sample(logits, sampling_metadata)
Roger Wang's avatar
Roger Wang committed
386

387
388
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
389
        loader = AutoWeightsLoader(self)
390
        return loader.load_weights(weights)