paligemma.py 11.6 KB
Newer Older
1
import itertools
2
3
from typing import (Iterable, List, Literal, Mapping, Optional, Tuple,
                    TypedDict, Union)
Roger Wang's avatar
Roger Wang committed
4
5
6

import torch
from torch import nn
7
from transformers import PaliGemmaConfig
Roger Wang's avatar
Roger Wang committed
8
9
10
11
12
13

from vllm.attention import AttentionMetadata
from vllm.config import CacheConfig, MultiModalConfig
from vllm.inputs import INPUT_REGISTRY, InputContext, LLMInputs
from vllm.logger import init_logger
from vllm.model_executor.layers.logits_processor import LogitsProcessor
14
from vllm.model_executor.layers.quantization import QuantizationConfig
15
from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
Roger Wang's avatar
Roger Wang committed
16
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
17
from vllm.model_executor.models.gemma import GemmaForCausalLM
Roger Wang's avatar
Roger Wang committed
18
19
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.multimodal import MULTIMODAL_REGISTRY
20
from vllm.multimodal.utils import cached_get_tokenizer
21
from vllm.sequence import IntermediateTensors
Roger Wang's avatar
Roger Wang committed
22

23
from .interfaces import SupportsMultiModal
24
25
from .siglip import (SiglipVisionModel, dummy_image_for_siglip,
                     dummy_seq_data_for_siglip, get_max_siglip_image_tokens)
26
from .utils import filter_weights, merge_multimodal_embeddings
Roger Wang's avatar
Roger Wang committed
27
28
29
30

logger = init_logger(__name__)


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


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

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


PaliGemmaImageInputs = Union[PaliGemmaImagePixelInputs,
                             PaliGemmaImageEmbeddingInputs]


Roger Wang's avatar
Roger Wang committed
50
51
def get_max_paligemma_image_tokens(ctx: InputContext):
    hf_config = ctx.get_hf_config(PaliGemmaConfig)
52
    vision_config = hf_config.vision_config
Roger Wang's avatar
Roger Wang committed
53

54
    return get_max_siglip_image_tokens(vision_config)
Roger Wang's avatar
Roger Wang committed
55
56


57
58
def dummy_data_for_paligemma(ctx: InputContext, seq_len: int,
                             mm_counts: Mapping[str, int]):
Roger Wang's avatar
Roger Wang committed
59
60
    hf_config = ctx.get_hf_config(PaliGemmaConfig)
    vision_config = hf_config.vision_config
61
    num_images = mm_counts["image"]
Roger Wang's avatar
Roger Wang committed
62

63
64
    seq_data = dummy_seq_data_for_siglip(
        vision_config,
Roger Wang's avatar
Roger Wang committed
65
        seq_len,
66
        num_images,
Roger Wang's avatar
Roger Wang committed
67
68
69
        image_token_id=hf_config.image_token_index,
    )

70
    mm_data = dummy_image_for_siglip(vision_config, num_images)
Roger Wang's avatar
Roger Wang committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
    return seq_data, mm_data


def input_processor_for_paligemma(ctx: InputContext, llm_inputs: LLMInputs):

    """
    The correct prompt format needs to be:
    '<image>' * image_feature_size + '<bos>' + prompt + '\n'

    See https://github.com/huggingface/transformers/blob/25245ec26dc29bcf6102e1b4ddd0dfd02e720cf5/src/transformers/models/paligemma/processing_paligemma.py#L55
    """ # noqa

    multi_modal_data = llm_inputs.get("multi_modal_data")
    if multi_modal_data is None or "image" not in multi_modal_data:
        return llm_inputs

    model_config = ctx.model_config
    hf_config = ctx.get_hf_config(PaliGemmaConfig)

    tokenizer = cached_get_tokenizer(model_config.tokenizer)
    image_feature_size = hf_config.text_config.num_image_tokens
    image_token_str = tokenizer.decode(hf_config.image_token_index)
    bos_token = tokenizer.decode(hf_config.bos_token_id)
    image_token_str_pad = image_token_str * image_feature_size
    image_token_ids_pad = [hf_config.image_token_index] * image_feature_size

    orig_prompt = llm_inputs.get("prompt")
    orig_prompt_ids = llm_inputs.get("prompt_token_ids")

100
    if orig_prompt is not None and image_token_str in orig_prompt:
Roger Wang's avatar
Roger Wang committed
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
        logger.warning(
            "The image token '%s' was detected in the prompt and "
            "will be removed. Please follow the proper prompt format"
            " documented on HuggingFace.", image_token_str)
        orig_prompt = orig_prompt.replace(image_token_str, "")
        orig_prompt_ids.remove(hf_config.image_token_index)

    new_prompt = f"{image_token_str_pad}{bos_token}{orig_prompt}\n"
    new_token_ids = image_token_ids_pad + orig_prompt_ids + [108]  #newline

    # NOTE: Create a defensive copy of the original inputs
    return LLMInputs(prompt_token_ids=new_token_ids,
                     prompt=new_prompt,
                     multi_modal_data=multi_modal_data)


class PaliGemmaMultiModalProjector(nn.Module):

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

122
        self.linear = nn.Linear(vision_hidden_size, projection_dim, bias=True)
Roger Wang's avatar
Roger Wang committed
123
124

    def forward(self, image_features: torch.Tensor) -> torch.Tensor:
125
        hidden_states = self.linear(image_features)
Roger Wang's avatar
Roger Wang committed
126
127
128
129
130
131
132
        return hidden_states


@MULTIMODAL_REGISTRY.register_image_input_mapper()
@MULTIMODAL_REGISTRY.register_max_image_tokens(get_max_paligemma_image_tokens)
@INPUT_REGISTRY.register_dummy_data(dummy_data_for_paligemma)
@INPUT_REGISTRY.register_input_processor(input_processor_for_paligemma)
133
class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal):
Roger Wang's avatar
Roger Wang committed
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

    def __init__(self,
                 config: PaliGemmaConfig,
                 multimodal_config: MultiModalConfig,
                 cache_config: Optional[CacheConfig] = None,
                 quant_config: Optional[QuantizationConfig] = None) -> None:
        super().__init__()

        self.config = config
        self.multimodal_config = multimodal_config

        self.vision_tower = SiglipVisionModel(config.vision_config)
        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
151
152
        self.language_model = GemmaForCausalLM(config.text_config,
                                               cache_config, quant_config)
Roger Wang's avatar
Roger Wang committed
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
        self.unpadded_vocab_size = config.text_config.vocab_size
        logit_scale = getattr(config, "logit_scale", 1.0)
        self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
                                                config.vocab_size, logit_scale)
        self.sampler = Sampler()

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

177
        if pixel_values is None and image_embeds is None:
Roger Wang's avatar
Roger Wang committed
178
179
            return None

180
181
182
183
        if pixel_values is not None:
            if not isinstance(pixel_values, torch.Tensor):
                raise ValueError("Incorrect type of pixel values. "
                                 f"Got type: {type(pixel_values)}")
184
185
186
187

            # Remove the N dimension until multiple images are supported.
            pixel_values = pixel_values.squeeze(1)

188
189
190
191
192
193
194
195
196
            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)}")
197
198
199
200

            # Remove the N dimension until multiple images are supported.
            image_embeds = image_embeds.squeeze(1)

201
202
203
204
205
206
            return PaliGemmaImageEmbeddingInputs(
                type="image_embeds",
                data=image_embeds,
            )

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

208
209
210
211
212
    def _image_pixels_to_features(
        self,
        vision_tower: SiglipVisionModel,
        pixel_values: torch.Tensor,
    ) -> torch.Tensor:
Roger Wang's avatar
Roger Wang committed
213

214
        target_dtype = vision_tower.get_input_embeddings().weight.dtype
215
        image_features = vision_tower(pixel_values.to(dtype=target_dtype))
Roger Wang's avatar
Roger Wang committed
216

217
        return image_features
Roger Wang's avatar
Roger Wang committed
218

219
    def _process_image_input(
220
        self,
221
        image_input: PaliGemmaImageInputs,
222
    ) -> torch.Tensor:
Roger Wang's avatar
Roger Wang committed
223

224
225
        if image_input["type"] == "image_embeds":
            return image_input["data"]
Roger Wang's avatar
Roger Wang committed
226

227
228
229
        assert self.vision_tower is not None
        pixel_values = image_input["data"]
        image_features = self._image_pixels_to_features(
230
231
232
            self.vision_tower,
            pixel_values,
        )
Roger Wang's avatar
Roger Wang committed
233
234
235

        return self.multi_modal_projector(image_features)

236
237
238
    def forward(self,
                input_ids: torch.Tensor,
                positions: torch.Tensor,
Roger Wang's avatar
Roger Wang committed
239
240
                kv_caches: List[torch.Tensor],
                attn_metadata: AttentionMetadata,
241
                intermediate_tensors: Optional[IntermediateTensors] = None,
Roger Wang's avatar
Roger Wang committed
242
243
244
245
246
247
248
249
250
251
                **kwargs: object) -> SamplerOutput:

        parsed_image_input = self._parse_and_validate_image_input(**kwargs)

        if parsed_image_input is not None:
            vision_embeddings = self._process_image_input(parsed_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)

252
253
            inputs_embeds = self.language_model.model.get_input_embeddings(
                input_ids)
Roger Wang's avatar
Roger Wang committed
254

255
            inputs_embeds = merge_multimodal_embeddings(
Roger Wang's avatar
Roger Wang committed
256
257
258
259
260
261
262
                input_ids, inputs_embeds, vision_embeddings,
                self.config.image_token_index)

            input_ids = None
        else:
            inputs_embeds = None

263
264
265
266
267
268
        hidden_states = self.language_model.model(input_ids,
                                                  positions,
                                                  kv_caches,
                                                  attn_metadata,
                                                  None,
                                                  inputs_embeds=inputs_embeds)
Roger Wang's avatar
Roger Wang committed
269
270
271

        return hidden_states

272
273
274
275
276
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
277
278
        return self.language_model.compute_logits(hidden_states,
                                                  sampling_metadata)
Roger Wang's avatar
Roger Wang committed
279
280
281
282
283
284

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

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
        # prepare weight iterators for components
        vit_weights, mlp_weights, llm_weights = itertools.tee(weights, 3)

        # load vision tower
        vit_weights = filter_weights(vit_weights, "vision_tower")
        self.vision_tower.load_weights(vit_weights)

        # load mlp projector
        mlp_weights = filter_weights(mlp_weights, "multi_modal_projector")
        mlp_params_dict = dict(self.multi_modal_projector.named_parameters())
        for name, loaded_weight in mlp_weights:
            param = mlp_params_dict[name]
            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            weight_loader(param, loaded_weight)

        # load llm backbone
        llm_weights = filter_weights(llm_weights, "language_model")
        self.language_model.load_weights(llm_weights)