siglip.py 22.9 KB
Newer Older
1
2
3
4
"""Implementation of SiglipVisionModel intended to be only used
within a vision language model."""

import math
5
from typing import Iterable, List, Optional, Set, Tuple, Union
6

7
import numpy as np
8
9
10
import torch
from PIL import Image
from torch import nn
11
from transformers import SiglipVisionConfig
12

13
from vllm.attention.layer import MultiHeadAttention
14
from vllm.config import ModelConfig
15
from vllm.distributed import divide, get_tensor_model_parallel_world_size
16
from vllm.inputs import DecoderOnlyInputs, token_inputs
17
18
19
20
21
22
23
from vllm.model_executor.layers.activation import get_act_fn
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               QKVParallelLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.vocab_parallel_embedding import (
    VocabParallelEmbedding)
24
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
25
from vllm.multimodal.utils import (cached_get_tokenizer,
26
                                   consecutive_placeholder_ranges,
27
                                   repeat_and_pad_placeholder_tokens)
28
from vllm.sequence import SequenceData
29

30
from .vision import VisionEncoderInfo, resolve_visual_encoder_outputs
31

32
33

def get_siglip_patch_grid_length(*, image_size: int, patch_size: int) -> int:
34
35
    # Since interpolation is applied, the image size need not be divisible
    # assert image_size % patch_size == 0
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
    return image_size // patch_size


def get_siglip_num_patches(*, image_size: int, patch_size: int) -> int:
    grid_length = get_siglip_patch_grid_length(image_size=image_size,
                                               patch_size=patch_size)
    return grid_length * grid_length


def get_siglip_image_feature_size(hf_config: SiglipVisionConfig) -> int:
    return get_siglip_num_patches(image_size=hf_config.image_size,
                                  patch_size=hf_config.patch_size)


def get_max_siglip_image_tokens(hf_config: SiglipVisionConfig) -> int:
    return get_siglip_image_feature_size(hf_config)


def dummy_seq_data_for_siglip(
    hf_config: SiglipVisionConfig,
    seq_len: int,
57
    num_images: int,
58
59
60
    *,
    image_token_id: int,
    image_feature_size_override: Optional[int] = None,
61
    mm_key: str = "image",
62
63
64
65
66
67
):
    if image_feature_size_override is None:
        image_feature_size = get_siglip_image_feature_size(hf_config)
    else:
        image_feature_size = image_feature_size_override

68
    return SequenceData.from_prompt_token_counts(
69
70
        (image_token_id, image_feature_size * num_images),
        (0, seq_len - image_feature_size * num_images),
71
72
73
74
75
    ), {
        mm_key:
        consecutive_placeholder_ranges(num_items=num_images,
                                       item_size=image_feature_size)
    }
76
77
78
79


def dummy_image_for_siglip(
    hf_config: SiglipVisionConfig,
80
    num_images: int,
81
82
83
84
85
86
87
88
89
90
91
    *,
    image_width_override: Optional[int] = None,
    image_height_override: Optional[int] = None,
):
    width = height = hf_config.image_size
    if image_width_override is not None:
        width = image_width_override
    if image_height_override is not None:
        height = image_height_override

    image = Image.new("RGB", (width, height), color=0)
92
    return {"image": image if num_images == 1 else [image] * num_images}
93
94


95
96
97
def dummy_video_for_siglip(
    hf_config: SiglipVisionConfig,
    num_frames: int,
98
    num_videos: int = 1,
99
100
101
102
103
104
105
106
107
108
109
    *,
    image_width_override: Optional[int] = None,
    image_height_override: Optional[int] = None,
):
    pil_frame = dummy_image_for_siglip(
        hf_config,
        num_images=1,
        image_width_override=image_width_override,
        image_height_override=image_height_override)
    np_frame = np.array(pil_frame["image"])
    mm_data_per_video = np.repeat([np_frame], num_frames, axis=0)
110
111
    video_data = [mm_data_per_video] * num_videos
    mm_data = {"video": video_data}
112
113
114
    return mm_data


115
116
117
def input_processor_for_siglip(
    model_config: ModelConfig,
    hf_config: SiglipVisionConfig,
118
    inputs: DecoderOnlyInputs,
119
120
    *,
    image_token_id: int,
121
    image_feature_size_override: Optional[Union[int, List[int]]] = None,
122
):
123
    multi_modal_data = inputs.get("multi_modal_data")
124
    if multi_modal_data is None or "image" not in multi_modal_data:
125
        return inputs
126

127
128
129
130
131
    if "multi_modal_placeholders" in inputs and "image" in inputs[
            "multi_modal_placeholders"]:
        # The inputs already have placeholders.
        return inputs

132
133
134
    tokenizer = cached_get_tokenizer(model_config.tokenizer)

    if image_feature_size_override is None:
135
136
137
138
        image_data = multi_modal_data["image"]
        if isinstance(image_data, Image.Image):
            image_feature_size = get_siglip_image_feature_size(hf_config)
        elif isinstance(image_data, torch.Tensor):
139
            num_images, image_feature_size, hidden_size = image_data.shape
140
141
        else:
            raise TypeError(f"Invalid image type: {type(image_data)}")
142
143
144
    else:
        image_feature_size = image_feature_size_override

145
    new_prompt, new_token_ids, ranges = repeat_and_pad_placeholder_tokens(
146
        tokenizer,
147
148
        inputs.get("prompt"),
        inputs["prompt_token_ids"],
149
        placeholder_token_id=image_token_id,
150
151
152
153
        repeat_count=image_feature_size,
    )

    # NOTE: Create a defensive copy of the original inputs
154
155
156
157
    return token_inputs(prompt_token_ids=new_token_ids,
                        prompt=new_prompt,
                        multi_modal_data=multi_modal_data,
                        multi_modal_placeholders={"image": ranges})
158
159


160
161
162
163
164
165
166
167
168
169
170
171
172
class SiglipEncoderInfo(VisionEncoderInfo[SiglipVisionConfig]):

    def get_num_image_tokens(
        self,
        *,
        image_width: int,
        image_height: int,
    ) -> int:
        return get_siglip_image_feature_size(self.vision_config)

    def get_max_image_tokens(self) -> int:
        return get_max_siglip_image_tokens(self.vision_config)

173
174
175
176
177
178
179
    def get_image_size(self) -> int:
        return self.vision_config.image_size

    def get_patch_size(self) -> int:
        return self.vision_config.patch_size

    def get_patch_grid_length(self) -> int:
180
181
182
183
184
185
        return get_siglip_patch_grid_length(
            image_size=self.vision_config.image_size,
            patch_size=self.vision_config.patch_size,
        )


186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
# Adapted from https://github.com/huggingface/transformers/blob/v4.43.3/src/transformers/models/siglip/modeling_siglip.py#L249 # noqa
class SiglipVisionEmbeddings(nn.Module):

    def __init__(self, config: SiglipVisionConfig):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.image_size = config.image_size
        self.patch_size = config.patch_size

        self.patch_embedding = nn.Conv2d(
            in_channels=config.num_channels,
            out_channels=self.embed_dim,
            kernel_size=self.patch_size,
            stride=self.patch_size,
            padding="valid",
        )

        self.num_patches = (self.image_size // self.patch_size)**2
        self.num_positions = self.num_patches
        self.position_embedding = VocabParallelEmbedding(
            self.num_positions, self.embed_dim)
        self.register_buffer(
            "position_ids",
            torch.arange(self.num_positions, dtype=torch.int64).expand(
                (1, -1)),
            persistent=False,
        )

    def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int,
                                 width: int) -> torch.Tensor:
        """
        This method is an adapted method for SigLIP (due to SigLIP not having
        class embedding unlike other ViTs) that allows the model to interpolate
        the pre-trained position encodings such that it can be usable on higher
        resolution images.

        Source:
        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
        """
        position_embeddings = self.position_embedding.weight.unsqueeze(0)
        num_patches = embeddings.shape[1]
        num_positions = position_embeddings.shape[1]
        if num_patches == num_positions and height == width:
            return position_embeddings

        dim = embeddings.shape[-1]
        height = height // self.patch_size
        width = width // self.patch_size
        # we add a small number to avoid floating point error
        # in the interpolation
        # see discussion at https://github.com/facebookresearch/dino/issues/8
        height, width = height + 0.1, width + 0.1

        patch_pos_embed = position_embeddings.reshape(
            1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)),
            dim)
        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
        patch_pos_embed = nn.functional.interpolate(
            patch_pos_embed,
            scale_factor=(
                height / math.sqrt(num_positions),
                width / math.sqrt(num_positions),
            ),
            mode="bicubic",
            align_corners=False,
        )
        if (int(height) != patch_pos_embed.shape[-2]
                or int(width) != patch_pos_embed.shape[-1]):
            raise ValueError("Width or height does not match with "
                             "the interpolated position embeddings")

        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return patch_pos_embed

    def forward(self,
                pixel_values: torch.Tensor,
                interpolate_pos_encoding: bool = False) -> torch.Tensor:
        _, _, height, width = pixel_values.shape
        target_dtype = self.patch_embedding.weight.dtype
        patch_embeds = self.patch_embedding(pixel_values.to(
            dtype=target_dtype))  # shape = [*, width, grid, grid]
        embeddings = patch_embeds.flatten(2).transpose(1, 2)

        if interpolate_pos_encoding:
            embeddings = embeddings + self.interpolate_pos_encoding(
                embeddings, height, width)
        else:
            embeddings = embeddings + self.position_embedding(
                self.position_ids)
        return embeddings


279
class SiglipAttention(nn.Module):
280
281
282

    def __init__(
        self,
283
        config: SiglipVisionConfig,
284
        quant_config: Optional[QuantizationConfig] = None,
285
286
        prefix: str = "",
    ) -> None:
287
        super().__init__()
288

289
290
        self.config = config
        self.embed_dim = config.hidden_size
291
292
293
        self.num_heads = config.num_attention_heads
        self.head_dim = self.embed_dim // self.num_heads
        if self.head_dim * self.num_heads != self.embed_dim:
294
295
296
            raise ValueError(f"embed_dim must be divisible by num_heads (got "
                             "`embed_dim`: {self.embed_dim} and `num_heads`:"
                             f" {self.num_heads}).")
297

298
299
300
301
302
        self.scale = self.head_dim**-0.5
        self.dropout = config.attention_dropout
        self.qkv_proj = QKVParallelLinear(
            hidden_size=self.embed_dim,
            head_size=self.head_dim,
303
            total_num_heads=self.num_heads,
304
            quant_config=quant_config,
305
            prefix=f"{prefix}.qkv_proj",
306
        )
307

308
309
310
311
        self.out_proj = RowParallelLinear(
            input_size=self.embed_dim,
            output_size=self.embed_dim,
            quant_config=quant_config,
312
            prefix=f"{prefix}.out_proj",
313
314
        )

315
316
        self.tp_size = get_tensor_model_parallel_world_size()
        self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
317

318
319
        self.attn = MultiHeadAttention(self.num_heads_per_partition,
                                       self.head_dim, self.scale)
320

321
322
323
324
325
326
    def forward(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        """Input shape: Batch x Time x Channel"""
        qkv_states, _ = self.qkv_proj(hidden_states)
327
328
        query_states, key_states, value_states = qkv_states.chunk(3, dim=-1)

329
        out = self.attn(query_states, key_states, value_states)
330
        attn_output, _ = self.out_proj(out)
331

332
        return attn_output, None
333
334
335
336
337
338


class SiglipMLP(nn.Module):

    def __init__(
        self,
339
        config: SiglipVisionConfig,
340
        quant_config: Optional[QuantizationConfig] = None,
341
342
        prefix: str = "",
    ) -> None:
343
        super().__init__()
344

345
346
        self.config = config
        self.activation_fn = get_act_fn(config.hidden_act)
347
348
349
350
351
352
353
354
355
356
        # Special handling for BNB quantization
        if quant_config and quant_config.get_name() == "bitsandbytes":
            quantizable = True
        else:
            # For other quantization, we require the hidden size to be a 
            # multiple of 64
            quantizable = (
                config.hidden_size % 64 == 0
                and config.intermediate_size % 64 == 0
            )
357
358
359
360
        self.fc1 = ColumnParallelLinear(
            config.hidden_size,
            config.intermediate_size,
            quant_config=quant_config if quantizable else None,
361
            prefix=f"{prefix}.fc1",
362
363
364
365
366
        )
        self.fc2 = RowParallelLinear(
            config.intermediate_size,
            config.hidden_size,
            quant_config=quant_config if quantizable else None,
367
            prefix=f"{prefix}.fc2",
368
369
370
371
372
373
374
375
376
377
378
379
380
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states, _ = self.fc1(hidden_states)
        hidden_states = self.activation_fn(hidden_states)
        hidden_states, _ = self.fc2(hidden_states)
        return hidden_states


class SiglipEncoderLayer(nn.Module):

    def __init__(
        self,
381
        config: SiglipVisionConfig,
382
        quant_config: Optional[QuantizationConfig] = None,
383
384
        prefix: str = "",
    ) -> None:
385
        super().__init__()
386

387
388
        self.embed_dim = config.hidden_size

389
390
391
392
393
        self.self_attn = SiglipAttention(
            config,
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
        )
394
395
396
397
398
        self.layer_norm1 = nn.LayerNorm(self.embed_dim,
                                        eps=config.layer_norm_eps)
        self.mlp = SiglipMLP(
            config,
            quant_config=quant_config,
399
            prefix=f"{prefix}.mlp",
400
401
402
403
404
405
406
        )
        self.layer_norm2 = nn.LayerNorm(self.embed_dim,
                                        eps=config.layer_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
407
    ) -> Tuple[torch.Tensor, None]:
408
409
410
        residual = hidden_states

        hidden_states = self.layer_norm1(hidden_states)
411
        hidden_states, _ = self.self_attn(hidden_states=hidden_states)
412
413
414
415
416
417
418
419
420
421
422
423
424
425
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.layer_norm2(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states, None


class SiglipEncoder(nn.Module):

    def __init__(
        self,
426
        config: SiglipVisionConfig,
427
        quant_config: Optional[QuantizationConfig] = None,
428
        num_hidden_layers_override: Optional[int] = None,
429
430
        prefix: str = "",
    ) -> None:
431
        super().__init__()
432

433
        self.config = config
434
435
436
437
438
439

        if num_hidden_layers_override is None:
            num_hidden_layers = config.num_hidden_layers
        else:
            num_hidden_layers = num_hidden_layers_override

440
        self.layers = nn.ModuleList([
441
442
443
444
            SiglipEncoderLayer(config,
                               quant_config=quant_config,
                               prefix=f"{prefix}.layers.{layer_idx}")
            for layer_idx in range(num_hidden_layers)
445
446
447
448
449
        ])

    def forward(
        self,
        inputs_embeds: torch.Tensor,
450
451
452
        return_all_hidden_states: bool,
    ) -> Union[torch.Tensor, list[torch.Tensor]]:
        hidden_states_pool = []
453
        hidden_states = inputs_embeds
454

455
456
        for encoder_layer in self.layers:
            hidden_states, _ = encoder_layer(hidden_states)
457
458
459
460
461
462
            if return_all_hidden_states:
                hidden_states_pool.append(hidden_states)
        # If we have multiple feature sample layers, we return all hidden
        # states in order and grab the ones we need by index.
        if return_all_hidden_states:
            return hidden_states_pool
463
464
465
466
467
468
469
470
471
472
        return hidden_states


class SiglipMultiheadAttentionPoolingHead(nn.Module):
    """Multihead Attention Pooling."""

    def __init__(
        self,
        config: SiglipVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
473
474
        prefix: str = "",
    ) -> None:
475
476
477
478
479
480
481
482
        super().__init__()

        self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
        # TODO(ChristopherCho): Implement vLLM version of MultiheadAttention
        self.attention = torch.nn.MultiheadAttention(
            config.hidden_size, config.num_attention_heads, batch_first=True)
        self.layernorm = nn.LayerNorm(config.hidden_size,
                                      eps=config.layer_norm_eps)
483
484
485
        self.mlp = SiglipMLP(config=config,
                             quant_config=quant_config,
                             prefix=f"{prefix}.mlp")
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        batch_size = hidden_state.shape[0]
        probe = self.probe.repeat(batch_size, 1, 1)

        hidden_state = self.attention(probe, hidden_state, hidden_state)[0]

        residual = hidden_state
        hidden_state = self.layernorm(hidden_state)
        hidden_state = residual + self.mlp(hidden_state)

        return hidden_state[:, 0]


class SiglipVisionTransformer(nn.Module):

    def __init__(
        self,
        config: SiglipVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
506
        *,
507
        num_hidden_layers_override: Optional[int] = None,
508
509
510
        require_post_norm: Optional[bool] = None,
        prefix: str = "",
    ) -> None:
511
        super().__init__()
512

513
514
515
516
        self.config = config
        embed_dim = config.hidden_size

        self.embeddings = SiglipVisionEmbeddings(config)
517

518
519
520
        self.encoder = SiglipEncoder(
            config,
            quant_config=quant_config,
521
            num_hidden_layers_override=num_hidden_layers_override,
522
            prefix=f"{prefix}.encoder",
523
        )
524

525
        num_hidden_layers = config.num_hidden_layers
526
527
        if len(self.encoder.layers) > config.num_hidden_layers:
            raise ValueError(
528
                f"The original encoder only has {num_hidden_layers} "
529
530
                f"layers, but you requested {len(self.encoder.layers)} layers."
            )
531
532
533
534
535
536

        # If possible, skip post_layernorm to conserve memory
        if require_post_norm is None:
            require_post_norm = len(self.encoder.layers) == num_hidden_layers

        if require_post_norm:
537
538
539
            self.post_layernorm = nn.LayerNorm(embed_dim,
                                               eps=config.layer_norm_eps)
        else:
540
541
            self.post_layernorm = None

542
543
544
545
        self.use_head = (True if not hasattr(config, "vision_use_head") else
                         config.vision_use_head)
        if self.use_head:
            self.head = SiglipMultiheadAttentionPoolingHead(
546
547
548
549
                config=config,
                quant_config=quant_config,
                prefix=f"{prefix}.head",
            )
550
551
552
553
554

    def forward(
        self,
        pixel_values: torch.Tensor,
        interpolate_pos_encoding: bool = True,
555
        feature_sample_layers: Optional[list[int]] = None,
556
    ) -> torch.Tensor:
557

558
559
560
561
562
        hidden_states = self.embeddings(
            pixel_values,
            interpolate_pos_encoding=interpolate_pos_encoding,
        )

563
564
565
566
567
568
569
570
        return_all_hidden_states = feature_sample_layers is not None

        # Produces either the last layer output or all of the hidden states,
        # depending on if we have feature_sample_layers or not
        encoder_outputs = self.encoder(
            inputs_embeds=hidden_states,
            return_all_hidden_states=return_all_hidden_states,
        )
571

572
573
574
575
        # Handle post-norm (if applicable) and stacks feature layers if needed
        encoder_outputs = resolve_visual_encoder_outputs(
            encoder_outputs, feature_sample_layers, self.post_layernorm,
            self.config.num_hidden_layers)
576

577
        # TODO: add this back when pooled_output is used in inference.
578
        # if self.use_head:
579
        # pooled_output = self.head(encoder_outputs)
580

581
        return encoder_outputs
582
583
584
585
586
587
588
589
590
591


class SiglipVisionModel(nn.Module):
    config_class = SiglipVisionConfig
    main_input_name = "pixel_values"

    def __init__(
        self,
        config: SiglipVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
592
        *,
593
        num_hidden_layers_override: Optional[int] = None,
594
595
596
        require_post_norm: Optional[bool] = None,
        prefix: str = "",
    ) -> None:
597
        super().__init__()
598

599
600
601
        self.vision_model = SiglipVisionTransformer(
            config,
            quant_config,
602
            num_hidden_layers_override=num_hidden_layers_override,
603
604
            require_post_norm=require_post_norm,
            prefix=f"{prefix}.vision_model",
605
606
607
608
609
610
611
612
613
        )

    def get_input_embeddings(self) -> nn.Module:
        return self.vision_model.embeddings.patch_embedding

    def forward(
        self,
        pixel_values: torch.Tensor,
        interpolate_pos_encoding: bool = False,
614
        feature_sample_layers: Optional[list[int]] = None,
615
616
617
618
    ) -> torch.Tensor:
        return self.vision_model(
            pixel_values=pixel_values,
            interpolate_pos_encoding=interpolate_pos_encoding,
619
            feature_sample_layers=feature_sample_layers,
620
        )
621

622
623
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
624
625
626
627
628
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
629
        ]
630
        params_dict = dict(self.named_parameters())
631
        loaded_params: Set[str] = set()
632
633
634
        layer_count = len(self.vision_model.encoder.layers)

        for name, loaded_weight in weights:
635
            # post_layernorm is optional in SiglipVisionModel
636
637
            if (name.startswith("vision_model.post_layernorm")
                    and self.vision_model.post_layernorm is None):
638
639
                continue

640
            # omit layers when num_hidden_layers_override is set
641
            if name.startswith("vision_model.encoder.layers"):
642
643
644
645
                layer_idx = int(name.split(".")[3])
                if layer_idx >= layer_count:
                    continue

646
647
648
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
                if weight_name not in name:
                    continue
649
                name = name.replace(weight_name, param_name)
650

651
                param = params_dict[name]
652
653
654
655
656
657
658
659
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
660
661
            loaded_params.add(name)
        return loaded_params