fuyu.py 15.7 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# adapted from https://github.com/huggingface/transformers/blob/v4.39.3/src/transformers/models/fuyu/modeling_fuyu.py
# Copyright 2023 The vLLM team.
# Copyright 2023 HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" PyTorch Fuyu model."""
import math
20
from collections.abc import Iterable, Mapping, Sequence
21
from typing import Literal, Optional, Set, Tuple, TypedDict, Union
22
23
24

import torch
import torch.nn as nn
25
26
from transformers import (BatchFeature, FuyuConfig, FuyuImageProcessor,
                          FuyuProcessor)
27

28
from vllm.config import VllmConfig
29
from vllm.model_executor.layers.linear import ColumnParallelLinear
30
from vllm.model_executor.layers.sampler import SamplerOutput
31
32
from vllm.model_executor.models.persimmon import PersimmonForCausalLM
from vllm.model_executor.sampling_metadata import SamplingMetadata
33
from vllm.multimodal import MULTIMODAL_REGISTRY
34
from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargs
35
36
from vllm.multimodal.parse import (ImageProcessorItems, ImageSize,
                                   MultiModalDataItems)
37
from vllm.multimodal.processing import (BaseMultiModalProcessor,
38
                                        BaseProcessingInfo, PromptReplacement,
39
                                        PromptUpdate, PromptUpdateDetails)
40
from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
41
from vllm.sequence import IntermediateTensors
42
from vllm.utils import flatten_2d_lists
43

44
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
45
46
from .utils import (AutoWeightsLoader, flatten_bn, maybe_prefix,
                    merge_multimodal_embeddings)
47
from .vision import scatter_patch_features, select_patch_features
48
49
50
51
52
53

# Cannot find the following 2 numbers from hf config.
_IMAGE_TOKEN_ID = 71011
_NEWLINE_TOKEN_ID = 71019


54
55
class FuyuImagePatchInputs(TypedDict):
    type: Literal["image_patches"]
56
    flat_data: torch.Tensor
57
58
    """
    Shape: 
59
60
61
    `(batch_size * num_patches, patch_size_x * patch_size_y * num_channels)`
    """

62
    patches_per_image: list[int]
63
    """
64
65
66
67
    The number of total patches for each image in the batch.

    This is used to split the embeddings which has the first two dimensions
    flattened just like `flat_data`.
68
    """
69
70
71
72
73
    embed_is_patch: Union[torch.Tensor, list[torch.Tensor]]
    """
    A boolean mask indicating which image embeddings correspond
    to patch tokens.
    """
74
75


76
class FuyuProcessingInfo(BaseProcessingInfo):
77

78
    def get_hf_config(self):
79
        return self.ctx.get_hf_config(FuyuConfig)
80

81
82
    def get_hf_processor(self, **kwargs: object):
        return self.ctx.get_hf_processor(FuyuProcessor, **kwargs)
83

84
85
    def get_image_processor(self) -> FuyuImageProcessor:
        return self.get_hf_processor().image_processor
86

87
88
89
    def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
        return {"image": 1}

90
91
92
93
94
    def get_mm_max_tokens_per_item(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> Mapping[str, int]:
95
96
97
98
99
100
101
102
103
104
105
        target_width, target_height = self.get_image_size_with_most_features()

        max_ncols, max_nrows = self.get_image_feature_grid_size(
            image_width=target_width,
            image_height=target_height,
        )
        max_image_tokens = (max_ncols + 1) * max_nrows

        return {"image": max_image_tokens}

    def get_image_feature_grid_size(
106
107
108
109
110
        self,
        *,
        image_width: int,
        image_height: int,
    ) -> tuple[int, int]:
111
        image_processor = self.get_image_processor()
112
113
        target_width = image_processor.size["width"]
        target_height = image_processor.size["height"]
114
115
        patch_width = image_processor.patch_size["width"]
        patch_height = image_processor.patch_size["height"]
116
117
118
119
120
121
122
123
124

        if not (image_width <= target_width and image_height <= target_height):
            height_scale_factor = target_height / image_height
            width_scale_factor = target_width / image_width
            optimal_scale_factor = min(height_scale_factor, width_scale_factor)

            image_height = int(image_height * optimal_scale_factor)
            image_width = int(image_width * optimal_scale_factor)

125
126
        ncols = math.ceil(image_width / patch_width)
        nrows = math.ceil(image_height / patch_height)
127
128
        return ncols, nrows

129
130
    def get_image_size_with_most_features(self) -> ImageSize:
        image_processor = self.get_image_processor()
131
132
133
        return ImageSize(width=image_processor.size["width"],
                         height=image_processor.size["height"])

134
135
136

class FuyuDummyInputsBuilder(BaseDummyInputsBuilder[FuyuProcessingInfo]):

137
138
139
140
141
    def get_dummy_processor_inputs(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> ProcessorInputs:
142
143
        target_width, target_height = \
            self.info.get_image_size_with_most_features()
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
        num_images = mm_counts.get("image", 0)

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

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


159
class FuyuMultiModalProcessor(BaseMultiModalProcessor[FuyuProcessingInfo]):
160
161
162
163
164
165
166
167
168

    def _call_hf_processor(
        self,
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
    ) -> BatchFeature:
        if not mm_data:
            # Avoid warning from HF logger for text-only input
169
170
            prompt_ids = self.info.get_tokenizer().encode(prompt)
            prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
            return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")

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

        image_patches = processed_outputs.get("image_patches")
        if image_patches is not None:
            images = mm_data["images"]
            assert isinstance(images, list)

            # Original output: (1, num_images, Pn, Px * Py * C)
            # New output: (num_images, Pn, Px * Py * C)
            assert (isinstance(image_patches, list)
                    and len(image_patches) == 1)
            assert (isinstance(image_patches[0], torch.Tensor)
                    and len(image_patches[0]) == len(images))

            processed_outputs["image_patches"] = image_patches[0]

193
194
195
196
197
198
199
200
201
202
203
204
205
            # get patch grid size for each image
            embed_is_patch = []
            for image in images:
                ncols, nrows = self.info.get_image_feature_grid_size(
                    image_width=image.width,
                    image_height=image.height,
                )

                mask = torch.tensor(([True] * ncols + [False]) * nrows)
                embed_is_patch.append(mask)

            processed_outputs["embed_is_patch"] = embed_is_patch

206
207
        return processed_outputs

208
209
210
211
212
213
    def _apply_hf_processor_tokens_only(
        self,
        prompt_tokens: list[int],
    ) -> list[int]:
        # HF processor adds boa_token_id
        tokenizer = self.info.get_tokenizer()
214
215
216
        vocab = tokenizer.get_vocab()

        boa_token_id = vocab["<0x04>"]
217
218
219

        return prompt_tokens + [boa_token_id]

220
221
222
223
224
    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
225
226
        return dict(image_patches=MultiModalFieldConfig.batched("image"),
                    embed_is_patch=MultiModalFieldConfig.batched("image"))
227

228
    def _get_prompt_updates(
229
230
231
232
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
        out_mm_kwargs: MultiModalKwargs,
233
    ) -> Sequence[PromptUpdate]:
234
        hf_config = self.info.get_hf_config()
235
        bos_token_id = hf_config.bos_token_id
236
        assert isinstance(bos_token_id, int)
237

238
        tokenizer = self.info.get_tokenizer()
239
240
241
242
243
244
        eot_token_id = tokenizer.bos_token_id
        assert isinstance(eot_token_id, int)

        def get_replacement_fuyu(item_idx: int):
            images = mm_items.get_items("image", ImageProcessorItems)
            image_size = images.get_image_size(item_idx)
245

246
            ncols, nrows = self.info.get_image_feature_grid_size(
247
248
                image_width=image_size.width,
                image_height=image_size.height,
249
            )
250
251
            image_tokens = ([_IMAGE_TOKEN_ID] * ncols +
                            [_NEWLINE_TOKEN_ID]) * nrows
252

253
            return PromptUpdateDetails(
254
255
256
                full=image_tokens + [bos_token_id],
                features=image_tokens,
            )
257
258
259
260
261
262
263
264
265
266

        return [
            PromptReplacement(
                modality="image",
                target=[eot_token_id],
                replacement=get_replacement_fuyu,
            )
        ]


267
268
269
@MULTIMODAL_REGISTRY.register_processor(FuyuMultiModalProcessor,
                                        info=FuyuProcessingInfo,
                                        dummy_inputs=FuyuDummyInputsBuilder)
270
class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
271

272
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
273
        super().__init__()
274
275
276
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config
277
278
279
        self.config = config
        self.multimodal_config = multimodal_config

280
        self.vocab_size = config.text_config.vocab_size
281
282
283
284
285
286
287
        self.image_token_id = _IMAGE_TOKEN_ID
        self.image_feature_size = config.patch_size**2 * config.num_channels

        self.vision_embed_tokens = ColumnParallelLinear(
            self.image_feature_size,
            config.hidden_size,
            quant_config=quant_config,
288
            gather_output=True,
289
        )
290
        self.language_model = PersimmonForCausalLM(
291
292
293
            vllm_config=vllm_config.with_hf_config(config.text_config),
            prefix=maybe_prefix(prefix, "language_model"),
        )
294
295
296
297
298
299
        self.make_empty_intermediate_tensors = (
            self.language_model.make_empty_intermediate_tensors)

    @property
    def sampler(self):
        return self.language_model.sampler
300

301
302
303
304
305
306
307
308
309
310
311
312
313
    def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:

        h = w = self.config.patch_size
        num_channels = self.config.num_channels
        expected_dims = num_channels * h * w

        def _validate_shape(d: torch.Tensor):
            actual_dims = d.size(-1)

            if actual_dims != expected_dims:
                expected_expr = str(expected_dims)
                raise ValueError(
                    "The expected shape of pixel values per image per batch "
314
                    f"per patch is {expected_expr}. "
315
316
317
318
319
320
321
                    f"You supplied {tuple(d.shape)}.")

        for d in data:
            _validate_shape(d)

        return data.to(self.vision_embed_tokens.weight.dtype)

322
    def _parse_and_validate_image_input(
323
324
            self, **kwargs: object) -> Optional[FuyuImagePatchInputs]:
        image_patches = kwargs.pop("image_patches", None)
325
        embed_is_patch = kwargs.pop("embed_is_patch", None)
326
327
        if image_patches is not None:
            if not isinstance(image_patches, (torch.Tensor, list)):
328
                raise ValueError("Incorrect type of image patches. "
329
                                 f"Got type: {type(image_patches)}")
330

331
332
333
            if not isinstance(embed_is_patch, (torch.Tensor, list)):
                raise ValueError("Incorrect type of embed_is_patch. "
                                 f"Got type: {type(embed_is_patch)}")
334
335
336
337
            image_patches_flat = flatten_bn(image_patches)

            return FuyuImagePatchInputs(
                type="image_patches",
338
                flat_data=self._validate_pixel_values(
339
340
                    flatten_bn(image_patches_flat, concat=True)),
                patches_per_image=[x.size(0) for x in image_patches_flat],
341
                embed_is_patch=embed_is_patch,
342
            )
343

344
345
        return None

346
    def _process_image_input(
347
            self, image_input: FuyuImagePatchInputs) -> MultiModalEmbeddings:
348
        image_patches_flat = image_input["flat_data"]
349
        patches_per_image = image_input["patches_per_image"]
350
351

        assert self.vision_embed_tokens is not None
352
353
354
        vision_embeddings_flat, _ = self.vision_embed_tokens(
            image_patches_flat)
        return vision_embeddings_flat.split(patches_per_image, dim=0)
355

356
    def get_multimodal_embeddings(
357
            self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
358
359
360
361
        image_input = self._parse_and_validate_image_input(**kwargs)
        if image_input is None:
            return None
        vision_embeddings = self._process_image_input(image_input)
362
363
364
365
366
367
        #return vision_embeddings
        return flatten_2d_lists(
            scatter_patch_features(*args) for args in zip(
                vision_embeddings,
                image_input["embed_is_patch"],
            ))
368
369
370
371

    def get_input_embeddings(
        self,
        input_ids: torch.Tensor,
372
        multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
373
374
375
376
    ) -> torch.Tensor:
        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
        if multimodal_embeddings is not None:
            inputs_embeds = merge_multimodal_embeddings(
377
378
                input_ids, inputs_embeds,
                select_patch_features(multimodal_embeddings), _IMAGE_TOKEN_ID)
379
380
        return inputs_embeds

381
382
383
384
385
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        intermediate_tensors: Optional[IntermediateTensors] = None,
386
        inputs_embeds: Optional[torch.Tensor] = None,
387
388
        **kwargs: object,
    ):
389
390
        if intermediate_tensors is not None:
            inputs_embeds = None
391
392
393
394
395
396
397
398

        # 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
399
400
401
402

        hidden_states = self.language_model(
            input_ids=input_ids,
            positions=positions,
403
            intermediate_tensors=intermediate_tensors,
404
405
406
407
            inputs_embeds=inputs_embeds,
        )
        return hidden_states

408
409
410
411
412
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
413
414
415
416
417
418
419
420
421
422
423
424
        logits = self.language_model.logits_processor(
            self.language_model.lm_head, hidden_states, sampling_metadata)
        return logits

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

425
426
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
427
        loader = AutoWeightsLoader(self)
428
        return loader.load_weights(weights)