| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # adapted from |
| 5 | # https://github.com/deepseek-ai/DeepSeek-OCR/blob/main/DeepSeek-OCR-master/DeepSeek-OCR-vllm/deepencoder/sam_vary_sdpa.py |
| 6 | |
| 7 | # Copyright (c) Meta Platforms, Inc. and affiliates. |
| 8 | # All rights reserved. |
| 9 | |
| 10 | # This source code is licensed under the license found in the |
| 11 | # LICENSE file in the root directory of this source tree. |
| 12 | import math |
| 13 | from collections.abc import Callable, Iterable |
| 14 | from functools import partial |
| 15 | |
| 16 | import torch |
| 17 | import torch.nn as nn |
| 18 | import torch.nn.functional as F |
| 19 | from transformers import CLIPVisionConfig |
| 20 | |
| 21 | from vllm.model_executor.custom_op import PluggableLayer |
| 22 | from vllm.model_executor.layers.attention import MMEncoderAttention |
| 23 | from vllm.model_executor.layers.conv import Conv2dLayer |
| 24 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 25 | from vllm.platforms import current_platform |
| 26 | from vllm.triton_utils import LOG2E, tl, triton |
| 27 | |
| 28 | from .clip import CLIPEncoder, CLIPVisionEmbeddings |
| 29 | from .utils import AutoWeightsLoader |
| 30 | |
| 31 | |
| 32 | @triton.jit |
| 33 | def _deepencoder_rel_pos_attention_kernel( |
| 34 | q_ptr, |
| 35 | k_ptr, |
| 36 | v_ptr, |
| 37 | rel_h_ptr, |
| 38 | rel_w_ptr, |
| 39 | out_ptr, |
| 40 | scale, |
| 41 | stride_qb, |
| 42 | stride_qh, |
| 43 | stride_qm, |
| 44 | stride_qd, |
| 45 | stride_kb, |
| 46 | stride_kh, |
| 47 | stride_kn, |
| 48 | stride_kd, |
| 49 | stride_vb, |
| 50 | stride_vh, |
| 51 | stride_vn, |
| 52 | stride_vd, |
| 53 | stride_rhb, |
| 54 | stride_rhh, |
| 55 | stride_rhm, |
| 56 | stride_rhk, |
| 57 | stride_rwb, |
| 58 | stride_rwh, |
| 59 | stride_rwm, |
| 60 | stride_rwk, |
| 61 | stride_ob, |
| 62 | stride_oh, |
| 63 | stride_om, |
| 64 | stride_od, |
| 65 | num_tokens: tl.constexpr, |
| 66 | num_heads: tl.constexpr, |
| 67 | key_width: tl.constexpr, |
| 68 | BLOCK_M: tl.constexpr, |
| 69 | BLOCK_N: tl.constexpr, |
| 70 | BLOCK_D: tl.constexpr, |
| 71 | ): |
| 72 | query_block = tl.program_id(0) |
| 73 | batch_head = tl.program_id(1) |
| 74 | batch = batch_head // num_heads |
| 75 | head = batch_head % num_heads |
| 76 | |
| 77 | offs_m = query_block * BLOCK_M + tl.arange(0, BLOCK_M) |
| 78 | offs_n = tl.arange(0, BLOCK_N) |
| 79 | offs_d = tl.arange(0, BLOCK_D) |
| 80 | mask_m = offs_m < num_tokens |
| 81 | |
| 82 | q_offsets = ( |
| 83 | batch * stride_qb |
| 84 | + head * stride_qh |
| 85 | + offs_m[:, None] * stride_qm |
| 86 | + offs_d[None, :] * stride_qd |
| 87 | ) |
| 88 | q = tl.load(q_ptr + q_offsets, mask=mask_m[:, None], other=0.0) |
| 89 | |
| 90 | row_max = tl.full([BLOCK_M], -float("inf"), tl.float32) |
| 91 | row_sum = tl.zeros([BLOCK_M], tl.float32) |
| 92 | accumulator = tl.zeros([BLOCK_M, BLOCK_D], tl.float32) |
| 93 | qk_scale = scale * LOG2E |
| 94 | |
| 95 | for start_n in range(0, num_tokens, BLOCK_N): |
| 96 | key_indices = start_n + offs_n |
| 97 | mask_n = key_indices < num_tokens |
| 98 | k_offsets = ( |
| 99 | batch * stride_kb |
| 100 | + head * stride_kh |
| 101 | + key_indices[None, :] * stride_kn |
| 102 | + offs_d[:, None] * stride_kd |
| 103 | ) |
| 104 | k = tl.load(k_ptr + k_offsets, mask=mask_n[None, :], other=0.0) |
| 105 | scores = tl.dot(q, k) * qk_scale |
| 106 | |
| 107 | key_h = key_indices // key_width |
| 108 | key_w = key_indices % key_width |
| 109 | rel_h_offsets = ( |
| 110 | batch * stride_rhb |
| 111 | + head * stride_rhh |
| 112 | + offs_m[:, None] * stride_rhm |
| 113 | + key_h[None, :] * stride_rhk |
| 114 | ) |
| 115 | rel_w_offsets = ( |
| 116 | batch * stride_rwb |
| 117 | + head * stride_rwh |
| 118 | + offs_m[:, None] * stride_rwm |
| 119 | + key_w[None, :] * stride_rwk |
| 120 | ) |
| 121 | score_mask = mask_m[:, None] & mask_n[None, :] |
| 122 | rel_h = tl.load(rel_h_ptr + rel_h_offsets, mask=score_mask, other=0.0) |
| 123 | rel_w = tl.load(rel_w_ptr + rel_w_offsets, mask=score_mask, other=0.0) |
| 124 | scores += (rel_h + rel_w) * LOG2E |
| 125 | scores = tl.where(score_mask, scores, -float("inf")) |
| 126 | |
| 127 | next_max = tl.maximum(row_max, tl.max(scores, axis=1)) |
| 128 | probabilities = tl.math.exp2(scores - next_max[:, None]) |
| 129 | alpha = tl.math.exp2(row_max - next_max) |
| 130 | next_sum = row_sum * alpha + tl.sum(probabilities, axis=1) |
| 131 | accumulator *= alpha[:, None] |
| 132 | |
| 133 | v_offsets = ( |
| 134 | batch * stride_vb |
| 135 | + head * stride_vh |
| 136 | + key_indices[:, None] * stride_vn |
| 137 | + offs_d[None, :] * stride_vd |
| 138 | ) |
| 139 | value = tl.load(v_ptr + v_offsets, mask=mask_n[:, None], other=0.0) |
| 140 | accumulator = tl.dot(probabilities.to(value.dtype), value, accumulator) |
| 141 | row_max = next_max |
| 142 | row_sum = next_sum |
| 143 | |
| 144 | output = accumulator / row_sum[:, None] |
| 145 | out_offsets = ( |
| 146 | batch * stride_ob |
| 147 | + head * stride_oh |
| 148 | + offs_m[:, None] * stride_om |
| 149 | + offs_d[None, :] * stride_od |
| 150 | ) |
| 151 | tl.store(out_ptr + out_offsets, output, mask=mask_m[:, None]) |
| 152 | |
| 153 | |
| 154 | def deepencoder_rel_pos_attention( |
| 155 | q: torch.Tensor, |
| 156 | k: torch.Tensor, |
| 157 | v: torch.Tensor, |
| 158 | rel_h: torch.Tensor, |
| 159 | rel_w: torch.Tensor, |
| 160 | width: int, |
| 161 | scale: float, |
| 162 | ) -> torch.Tensor: |
| 163 | """Run DeepEncoder global attention with fused relative-position bias.""" |
| 164 | batch, heads, num_tokens, head_dim = q.shape |
| 165 | |
| 166 | output = torch.empty_like(q) |
| 167 | block_m = 64 |
| 168 | block_n = 64 |
| 169 | grid = (triton.cdiv(num_tokens, block_m), batch * heads) |
| 170 | _deepencoder_rel_pos_attention_kernel[grid]( |
| 171 | q, |
| 172 | k, |
| 173 | v, |
| 174 | rel_h, |
| 175 | rel_w, |
| 176 | output, |
| 177 | scale, |
| 178 | *q.stride(), |
| 179 | *k.stride(), |
| 180 | *v.stride(), |
| 181 | *rel_h.stride(), |
| 182 | *rel_w.stride(), |
| 183 | *output.stride(), |
| 184 | num_tokens=num_tokens, |
| 185 | num_heads=heads, |
| 186 | key_width=width, |
| 187 | BLOCK_M=block_m, |
| 188 | BLOCK_N=block_n, |
| 189 | BLOCK_D=head_dim, |
| 190 | num_warps=4, |
| 191 | num_stages=2, |
| 192 | ) |
| 193 | return output |
| 194 | |
| 195 | |
| 196 | class MLPBlock(nn.Module): |
| 197 | def __init__( |
| 198 | self, |
| 199 | embedding_dim: int, |
| 200 | mlp_dim: int, |
| 201 | act: type[nn.Module] = nn.GELU, |
| 202 | ) -> None: |
| 203 | super().__init__() |
| 204 | self.lin1 = nn.Linear(embedding_dim, mlp_dim) |
| 205 | self.lin2 = nn.Linear(mlp_dim, embedding_dim) |
| 206 | self.act = act() |
| 207 | |
| 208 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 209 | return self.lin2(self.act(self.lin1(x))) |
| 210 | |
| 211 | |
| 212 | # From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa |
| 213 | # Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa |
| 214 | class LayerNorm2d(nn.Module): |
| 215 | def __init__(self, num_channels: int, eps: float = 1e-6) -> None: |
| 216 | super().__init__() |
| 217 | self.weight = nn.Parameter(torch.ones(num_channels)) |
| 218 | self.bias = nn.Parameter(torch.zeros(num_channels)) |
| 219 | self.eps = eps |
| 220 | |
| 221 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 222 | u = x.mean(1, keepdim=True) |
| 223 | s = (x - u).pow(2).mean(1, keepdim=True) |
| 224 | x = (x - u) / torch.sqrt(s + self.eps) |
| 225 | x = self.weight[:, None, None] * x + self.bias[:, None, None] |
| 226 | return x |
| 227 | |
| 228 | |
| 229 | # This class and its supporting functions below lightly adapted from the ViTDet backbone available at: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py # noqa |
| 230 | class ImageEncoderViT(nn.Module): |
| 231 | def __init__( |
| 232 | self, |
| 233 | img_size: int = 1024, |
| 234 | patch_size: int = 16, |
| 235 | in_chans: int = 3, |
| 236 | embed_dim: int = 768, |
| 237 | depth: int = 12, |
| 238 | num_heads: int = 12, |
| 239 | mlp_ratio: float = 4.0, |
| 240 | out_chans: int = 256, |
| 241 | qkv_bias: bool = True, |
| 242 | norm_layer: Callable[[int], nn.Module] = nn.LayerNorm, |
| 243 | act_layer: type[nn.Module] = nn.GELU, |
| 244 | use_abs_pos: bool = True, |
| 245 | use_rel_pos: bool = False, |
| 246 | rel_pos_zero_init: bool = True, |
| 247 | window_size: int = 0, |
| 248 | global_attn_indexes: tuple[int, ...] = (), |
| 249 | last_conv_output: int = 1024, |
| 250 | ) -> None: |
| 251 | """Args: |
| 252 | img_size (int): Input image size. |
| 253 | patch_size (int): Patch size. |
| 254 | in_chans (int): Number of input image channels. |
| 255 | embed_dim (int): Patch embedding dimension. |
| 256 | depth (int): Depth of ViT. |
| 257 | num_heads (int): Number of attention heads in each ViT block. |
| 258 | mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. |
| 259 | qkv_bias (bool): If True, add a learnable bias to query, key, value. |
| 260 | norm_layer (nn.Module): Normalization layer. |
| 261 | act_layer (nn.Module): Activation layer. |
| 262 | use_abs_pos (bool): If True, use absolute positional embeddings. |
| 263 | use_rel_pos (bool): If True, add relative positional embeddings to the attention map. |
| 264 | rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. |
| 265 | window_size (int): Window size for window attention blocks. |
| 266 | global_attn_indexes (list): Indexes for blocks using global attention. |
| 267 | |
| 268 | """ # noqa: E501 |
| 269 | super().__init__() |
| 270 | self.img_size = img_size |
| 271 | |
| 272 | self.patch_embed = PatchEmbed( |
| 273 | kernel_size=(patch_size, patch_size), |
| 274 | stride=(patch_size, patch_size), |
| 275 | in_chans=in_chans, |
| 276 | embed_dim=embed_dim, |
| 277 | ) |
| 278 | |
| 279 | self.pos_embed: nn.Parameter | None = None |
| 280 | if use_abs_pos: |
| 281 | # Initialize absolute positional embedding with pretrain image size. |
| 282 | self.pos_embed = nn.Parameter( |
| 283 | torch.zeros( |
| 284 | 1, img_size // patch_size, img_size // patch_size, embed_dim |
| 285 | ) |
| 286 | ) |
| 287 | |
| 288 | self.blocks = nn.ModuleList() |
| 289 | for i in range(depth): |
| 290 | block = Block( |
| 291 | dim=embed_dim, |
| 292 | num_heads=num_heads, |
| 293 | mlp_ratio=mlp_ratio, |
| 294 | qkv_bias=qkv_bias, |
| 295 | norm_layer=norm_layer, |
| 296 | act_layer=act_layer, |
| 297 | use_rel_pos=use_rel_pos, |
| 298 | rel_pos_zero_init=rel_pos_zero_init, |
| 299 | window_size=window_size if i not in global_attn_indexes else 0, |
| 300 | input_size=(img_size // patch_size, img_size // patch_size), |
| 301 | ) |
| 302 | self.blocks.append(block) |
| 303 | |
| 304 | self.neck = nn.Sequential( |
| 305 | Conv2dLayer( |
| 306 | embed_dim, |
| 307 | out_chans, |
| 308 | kernel_size=1, |
| 309 | bias=False, |
| 310 | ), |
| 311 | LayerNorm2d(out_chans), |
| 312 | Conv2dLayer( |
| 313 | out_chans, |
| 314 | out_chans, |
| 315 | kernel_size=3, |
| 316 | padding=1, |
| 317 | bias=False, |
| 318 | ), |
| 319 | LayerNorm2d(out_chans), |
| 320 | ) |
| 321 | |
| 322 | self.net_2 = Conv2dLayer( |
| 323 | 256, 512, kernel_size=3, stride=2, padding=1, bias=False |
| 324 | ) |
| 325 | self.net_3 = Conv2dLayer( |
| 326 | 512, last_conv_output, kernel_size=3, stride=2, padding=1, bias=False |
| 327 | ) |
| 328 | |
| 329 | def get_abs_pos(self, abs_pos: torch.Tensor, tgt_size: int): |
| 330 | dtype = abs_pos.dtype |
| 331 | |
| 332 | src_size = abs_pos.size(1) |
| 333 | |
| 334 | if src_size != tgt_size: |
| 335 | old_pos_embed = abs_pos.permute(0, 3, 1, 2) |
| 336 | old_pos_embed = old_pos_embed.to(torch.float32) |
| 337 | new_pos_embed = F.interpolate( |
| 338 | old_pos_embed, |
| 339 | size=(tgt_size, tgt_size), |
| 340 | mode="bicubic", |
| 341 | antialias=True, |
| 342 | align_corners=False, |
| 343 | ).to(dtype) |
| 344 | new_pos_embed = new_pos_embed.permute(0, 2, 3, 1) |
| 345 | return new_pos_embed |
| 346 | else: |
| 347 | return abs_pos |
| 348 | |
| 349 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 350 | x = self.patch_embed(x) |
| 351 | if self.pos_embed is not None: |
| 352 | x = x + self.get_abs_pos(self.pos_embed, x.size(1)) |
| 353 | |
| 354 | for blk in self.blocks: |
| 355 | x = blk(x) |
| 356 | |
| 357 | neck_output = self.neck(x.permute(0, 3, 1, 2)) |
| 358 | conv2_output = self.net_2(neck_output) |
| 359 | conv3_output = self.net_3(conv2_output) |
| 360 | |
| 361 | return conv3_output |
| 362 | |
| 363 | |
| 364 | class Block(nn.Module): |
| 365 | """Transformer blocks with support of window attention and residual propagation |
| 366 | blocks""" |
| 367 | |
| 368 | def __init__( |
| 369 | self, |
| 370 | dim: int, |
| 371 | num_heads: int, |
| 372 | mlp_ratio: float = 4.0, |
| 373 | qkv_bias: bool = True, |
| 374 | norm_layer: Callable[[int], nn.Module] = nn.LayerNorm, |
| 375 | act_layer: type[nn.Module] = nn.GELU, |
| 376 | use_rel_pos: bool = False, |
| 377 | rel_pos_zero_init: bool = True, |
| 378 | window_size: int = 0, |
| 379 | input_size: tuple[int, int] | None = None, |
| 380 | ) -> None: |
| 381 | """Args: |
| 382 | dim (int): Number of input channels. |
| 383 | num_heads (int): Number of attention heads in each ViT block. |
| 384 | mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. |
| 385 | qkv_bias (bool): If True, add a learnable bias to query, key, value. |
| 386 | norm_layer (nn.Module): Normalization layer. |
| 387 | act_layer (nn.Module): Activation layer. |
| 388 | use_rel_pos (bool): If True, add relative positional embeddings to the attention map. |
| 389 | rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. |
| 390 | window_size (int): Window size for window attention blocks. If it equals 0, then |
| 391 | use global attention. |
| 392 | input_size (tuple(int, int) or None): Input resolution for calculating the relative |
| 393 | positional parameter size. |
| 394 | |
| 395 | """ # noqa: E501 |
| 396 | super().__init__() |
| 397 | self.norm1 = norm_layer(dim) |
| 398 | self.attn = RelPosAttention( |
| 399 | dim, |
| 400 | num_heads=num_heads, |
| 401 | qkv_bias=qkv_bias, |
| 402 | use_rel_pos=use_rel_pos, |
| 403 | use_triton_attention=window_size == 0, |
| 404 | rel_pos_zero_init=rel_pos_zero_init, |
| 405 | input_size=input_size if window_size == 0 else (window_size, window_size), |
| 406 | ) |
| 407 | |
| 408 | self.norm2 = norm_layer(dim) |
| 409 | self.mlp = MLPBlock( |
| 410 | embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer |
| 411 | ) |
| 412 | |
| 413 | self.window_size = window_size |
| 414 | |
| 415 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 416 | shortcut = x |
| 417 | x = self.norm1(x) |
| 418 | # Window partition |
| 419 | if self.window_size > 0: |
| 420 | H, W = x.shape[1], x.shape[2] |
| 421 | x, pad_hw = window_partition(x, self.window_size) |
| 422 | |
| 423 | x = self.attn(x) |
| 424 | # Reverse window partition |
| 425 | if self.window_size > 0: |
| 426 | x = window_unpartition(x, self.window_size, pad_hw, (H, W)) |
| 427 | |
| 428 | x = shortcut + x |
| 429 | x = x + self.mlp(self.norm2(x)) |
| 430 | |
| 431 | return x |
| 432 | |
| 433 | |
| 434 | # --8<-- [start:rel_pos_attention] |
| 435 | @PluggableLayer.register("rel_pos_attention") |
| 436 | class RelPosAttention(PluggableLayer): |
| 437 | """Multi-head Attention block with relative position embeddings.""" |
| 438 | |
| 439 | # --8<-- [end:rel_pos_attention] |
| 440 | |
| 441 | def __init__( |
| 442 | self, |
| 443 | dim: int, |
| 444 | num_heads: int = 8, |
| 445 | qkv_bias: bool = True, |
| 446 | use_rel_pos: bool = False, |
| 447 | use_triton_attention: bool = False, |
| 448 | rel_pos_zero_init: bool = True, |
| 449 | input_size: tuple[int, int] | None = None, |
| 450 | ) -> None: |
| 451 | """Args: |
| 452 | dim (int): Number of input channels. |
| 453 | num_heads (int): Number of attention heads. |
| 454 | qkv_bias (bool): If True, add a learnable bias to query, key, value. |
| 455 | use_triton_attention (bool): If True, fuse relative position bias in |
| 456 | a Triton attention kernel on CUDA-alike platforms. |
| 457 | rel_pos_zero_init (bool): If True, zero initialize relative positional parameters. |
| 458 | input_size (tuple(int, int) or None): Input resolution for calculating the relative |
| 459 | positional parameter size. |
| 460 | |
| 461 | """ # noqa: E501 |
| 462 | super().__init__() |
| 463 | self.num_heads = num_heads |
| 464 | head_dim = dim // num_heads |
| 465 | self.scale = head_dim**-0.5 |
| 466 | |
| 467 | self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) |
| 468 | self.proj = nn.Linear(dim, dim) |
| 469 | |
| 470 | self.use_rel_pos = use_rel_pos |
| 471 | self.use_triton_attention = use_triton_attention |
| 472 | if self.use_rel_pos: |
| 473 | assert input_size is not None, ( |
| 474 | "Input size must be provided if using relative positional encoding." |
| 475 | ) |
| 476 | # initialize relative positional embeddings |
| 477 | self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim)) |
| 478 | self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim)) |
| 479 | |
| 480 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 481 | B, H, W, _ = x.shape |
| 482 | # qkv with shape (3, B, nHead, H * W, C) |
| 483 | qkv = ( |
| 484 | self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) |
| 485 | ) |
| 486 | # q, k, v with shape (B * nHead, H * W, C) |
| 487 | q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0) |
| 488 | |
| 489 | rel_h, rel_w = None, None |
| 490 | if self.use_rel_pos: |
| 491 | rel_h, rel_w = add_decomposed_rel_pos( |
| 492 | q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W) |
| 493 | ) |
| 494 | |
| 495 | q = q.view(B, self.num_heads, H * W, -1) |
| 496 | k = k.view(B, self.num_heads, H * W, -1) |
| 497 | v = v.view(B, self.num_heads, H * W, -1) |
| 498 | |
| 499 | if self.use_rel_pos: |
| 500 | assert rel_h is not None and rel_w is not None |
| 501 | rel_h = rel_h.view( |
| 502 | B, self.num_heads, rel_h.size(1), rel_h.size(2), rel_h.size(3) |
| 503 | ) |
| 504 | rel_w = rel_w.view( |
| 505 | B, self.num_heads, rel_w.size(1), rel_w.size(2), rel_w.size(3) |
| 506 | ) |
| 507 | use_triton = self.use_triton_attention and current_platform.is_cuda_alike() |
| 508 | if use_triton: |
| 509 | x = deepencoder_rel_pos_attention( |
| 510 | q, |
| 511 | k, |
| 512 | v, |
| 513 | rel_h.squeeze(-1), |
| 514 | rel_w.squeeze(-2), |
| 515 | W, |
| 516 | self.scale, |
| 517 | ) |
| 518 | else: |
| 519 | attn_bias = (rel_h + rel_w).view( |
| 520 | B, |
| 521 | self.num_heads, |
| 522 | rel_h.size(2), |
| 523 | rel_h.size(3) * rel_w.size(4), |
| 524 | ) |
| 525 | x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias) |
| 526 | else: |
| 527 | x = F.scaled_dot_product_attention(q, k, v) |
| 528 | |
| 529 | x = ( |
| 530 | x.view(B, self.num_heads, H, W, -1) |
| 531 | .permute(0, 2, 3, 1, 4) |
| 532 | .reshape(B, H, W, -1) |
| 533 | ) |
| 534 | |
| 535 | x = self.proj(x) |
| 536 | |
| 537 | return x |
| 538 | |
| 539 | |
| 540 | def window_partition( |
| 541 | x: torch.Tensor, window_size: int |
| 542 | ) -> tuple[torch.Tensor, tuple[int, int]]: |
| 543 | """Partition into non-overlapping windows with padding if needed. |
| 544 | |
| 545 | Args: |
| 546 | x (tensor): input tokens with [B, H, W, C]. |
| 547 | window_size (int): window size. |
| 548 | |
| 549 | Returns: |
| 550 | windows: windows after partition with [B * num_windows, window_size, window_size, C]. |
| 551 | (Hp, Wp): padded height and width before partition |
| 552 | |
| 553 | """ # noqa: E501 |
| 554 | B, H, W, C = x.shape |
| 555 | |
| 556 | pad_h = (window_size - H % window_size) % window_size |
| 557 | pad_w = (window_size - W % window_size) % window_size |
| 558 | if pad_h > 0 or pad_w > 0: |
| 559 | x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) |
| 560 | Hp, Wp = H + pad_h, W + pad_w |
| 561 | |
| 562 | x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C) |
| 563 | windows = ( |
| 564 | x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) |
| 565 | ) |
| 566 | return windows, (Hp, Wp) |
| 567 | |
| 568 | |
| 569 | def window_unpartition( |
| 570 | windows: torch.Tensor, |
| 571 | window_size: int, |
| 572 | pad_hw: tuple[int, int], |
| 573 | hw: tuple[int, int], |
| 574 | ) -> torch.Tensor: |
| 575 | """Window unpartition into original sequences and removing padding. |
| 576 | |
| 577 | Args: |
| 578 | windows (tensor): input tokens with [B * num_windows, window_size, window_size, C]. |
| 579 | window_size (int): window size. |
| 580 | pad_hw (Tuple): padded height and width (Hp, Wp). |
| 581 | hw (Tuple): original height and width (H, W) before padding. |
| 582 | |
| 583 | Returns: |
| 584 | x: unpartitioned sequences with [B, H, W, C]. |
| 585 | |
| 586 | """ # noqa: E501 |
| 587 | Hp, Wp = pad_hw |
| 588 | H, W = hw |
| 589 | B = windows.shape[0] // (Hp * Wp // window_size // window_size) |
| 590 | x = windows.view( |
| 591 | B, Hp // window_size, Wp // window_size, window_size, window_size, -1 |
| 592 | ) |
| 593 | x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1) |
| 594 | |
| 595 | if Hp > H or Wp > W: |
| 596 | x = x[:, :H, :W, :].contiguous() |
| 597 | return x |
| 598 | |
| 599 | |
| 600 | def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor: |
| 601 | """Get relative positional embeddings according to the relative positions of |
| 602 | query and key sizes. |
| 603 | |
| 604 | Args: |
| 605 | q_size (int): size of query q. |
| 606 | k_size (int): size of key k. |
| 607 | rel_pos (Tensor): relative position embeddings (L, C). |
| 608 | |
| 609 | Returns: |
| 610 | Extracted positional embeddings according to relative positions. |
| 611 | |
| 612 | """ |
| 613 | max_rel_dist = int(2 * max(q_size, k_size) - 1) |
| 614 | # Interpolate rel pos if needed. |
| 615 | if rel_pos.shape[0] != max_rel_dist: |
| 616 | # Interpolate rel pos. |
| 617 | dtype = rel_pos.dtype |
| 618 | rel_pos = rel_pos.to(torch.float32) |
| 619 | rel_pos_resized = F.interpolate( |
| 620 | rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1), |
| 621 | size=max_rel_dist, |
| 622 | mode="linear", |
| 623 | ).to(dtype) |
| 624 | rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0) |
| 625 | else: |
| 626 | rel_pos_resized = rel_pos |
| 627 | |
| 628 | # Scale the coords with short length if shapes for q and k are different. |
| 629 | q_coords = torch.arange(q_size, device=rel_pos.device)[:, None] * max( |
| 630 | k_size / q_size, 1.0 |
| 631 | ) |
| 632 | k_coords = torch.arange(k_size, device=rel_pos.device)[None, :] * max( |
| 633 | q_size / k_size, 1.0 |
| 634 | ) |
| 635 | relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0) |
| 636 | |
| 637 | return rel_pos_resized[relative_coords.long()] |
| 638 | |
| 639 | |
| 640 | def add_decomposed_rel_pos( |
| 641 | q: torch.Tensor, |
| 642 | rel_pos_h: torch.Tensor, |
| 643 | rel_pos_w: torch.Tensor, |
| 644 | q_size: tuple[int, int], |
| 645 | k_size: tuple[int, int], |
| 646 | ) -> torch.Tensor: |
| 647 | """Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`. |
| 648 | https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py |
| 649 | Args: |
| 650 | q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C). |
| 651 | rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis. |
| 652 | rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis. |
| 653 | q_size (Tuple): spatial sequence size of query q with (q_h, q_w). |
| 654 | k_size (Tuple): spatial sequence size of key k with (k_h, k_w). |
| 655 | |
| 656 | Returns: |
| 657 | attn (Tensor): attention map with added relative positional embeddings. |
| 658 | |
| 659 | """ # noqa: E501 |
| 660 | q_h, q_w = q_size |
| 661 | k_h, k_w = k_size |
| 662 | Rh = get_rel_pos(q_h, k_h, rel_pos_h) |
| 663 | Rw = get_rel_pos(q_w, k_w, rel_pos_w) |
| 664 | |
| 665 | B, _, dim = q.shape |
| 666 | r_q = q.reshape(B, q_h, q_w, dim) |
| 667 | rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh) |
| 668 | rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw) |
| 669 | rel_h = rel_h.unsqueeze(-1) |
| 670 | rel_w = rel_w.unsqueeze(-2) |
| 671 | rel_h = rel_h.reshape(B, q_h * q_w, k_h, 1) |
| 672 | rel_w = rel_w.reshape(B, q_h * q_w, 1, k_w) |
| 673 | |
| 674 | return rel_h, rel_w |
| 675 | |
| 676 | |
| 677 | class PatchEmbed(nn.Module): |
| 678 | """Image to Patch Embedding.""" |
| 679 | |
| 680 | def __init__( |
| 681 | self, |
| 682 | kernel_size: tuple[int, int] = (16, 16), |
| 683 | stride: tuple[int, int] = (16, 16), |
| 684 | padding: tuple[int, int] = (0, 0), |
| 685 | in_chans: int = 3, |
| 686 | embed_dim: int = 768, |
| 687 | ) -> None: |
| 688 | """Args: |
| 689 | kernel_size (Tuple): kernel size of the projection layer. |
| 690 | stride (Tuple): stride of the projection layer. |
| 691 | padding (Tuple): padding size of the projection layer. |
| 692 | in_chans (int): Number of input image channels. |
| 693 | embed_dim (int): Patch embedding dimension. |
| 694 | |
| 695 | """ |
| 696 | super().__init__() |
| 697 | |
| 698 | self.proj = Conv2dLayer( |
| 699 | in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding |
| 700 | ) |
| 701 | |
| 702 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 703 | x = self.proj(x) |
| 704 | # B C H W -> B H W C |
| 705 | x = x.permute(0, 2, 3, 1) |
| 706 | return x |
| 707 | |
| 708 | |
| 709 | # TODO(Isotr0py): use vision_config to build sam model |
| 710 | def build_sam_vit_b(): |
| 711 | return _build_sam( |
| 712 | encoder_embed_dim=768, |
| 713 | encoder_depth=12, |
| 714 | encoder_num_heads=12, |
| 715 | encoder_global_attn_indexes=[2, 5, 8, 11], |
| 716 | ) |
| 717 | |
| 718 | |
| 719 | def _build_sam( |
| 720 | encoder_embed_dim, |
| 721 | encoder_depth, |
| 722 | encoder_num_heads, |
| 723 | encoder_global_attn_indexes, |
| 724 | ): |
| 725 | prompt_embed_dim = 256 |
| 726 | image_size = 1024 |
| 727 | vit_patch_size = 16 |
| 728 | image_encoder = ImageEncoderViT( |
| 729 | depth=encoder_depth, |
| 730 | embed_dim=encoder_embed_dim, |
| 731 | img_size=image_size, |
| 732 | mlp_ratio=4, |
| 733 | norm_layer=partial(torch.nn.LayerNorm, eps=1e-6), |
| 734 | num_heads=encoder_num_heads, |
| 735 | patch_size=vit_patch_size, |
| 736 | qkv_bias=True, |
| 737 | use_rel_pos=True, |
| 738 | global_attn_indexes=encoder_global_attn_indexes, |
| 739 | window_size=14, |
| 740 | out_chans=prompt_embed_dim, |
| 741 | ) |
| 742 | return image_encoder |
| 743 | |
| 744 | |
| 745 | class DeepCLIPVisionEmbeddings(CLIPVisionEmbeddings): |
| 746 | def get_abs_pos(self, abs_pos: torch.Tensor, tgt_size: int): |
| 747 | # abs_pos: L, C |
| 748 | # tgt_size: M |
| 749 | # return: M, C |
| 750 | dim = abs_pos.size(-1) |
| 751 | abs_pos_new = abs_pos.squeeze(0) |
| 752 | cls_token, old_pos_embed = abs_pos_new[:1], abs_pos_new[1:] |
| 753 | |
| 754 | src_size = int(math.sqrt(abs_pos_new.shape[0] - 1)) |
| 755 | tgt_size = int(math.sqrt(tgt_size)) |
| 756 | dtype = abs_pos.dtype |
| 757 | |
| 758 | if src_size != tgt_size: |
| 759 | old_pos_embed = ( |
| 760 | old_pos_embed.view(1, src_size, src_size, dim) |
| 761 | .permute(0, 3, 1, 2) |
| 762 | .contiguous() |
| 763 | ) |
| 764 | old_pos_embed = old_pos_embed.to(torch.float32) |
| 765 | new_pos_embed = F.interpolate( |
| 766 | old_pos_embed, |
| 767 | size=(tgt_size, tgt_size), |
| 768 | mode="bicubic", |
| 769 | antialias=True, |
| 770 | align_corners=False, |
| 771 | ).to(dtype) |
| 772 | new_pos_embed = new_pos_embed.permute(0, 2, 3, 1) |
| 773 | new_pos_embed = new_pos_embed.view(tgt_size * tgt_size, dim) |
| 774 | vision_pos_embed = torch.cat([cls_token, new_pos_embed], dim=0) |
| 775 | vision_pos_embed = vision_pos_embed.view(1, tgt_size * tgt_size + 1, dim) |
| 776 | return vision_pos_embed |
| 777 | else: |
| 778 | return abs_pos |
| 779 | |
| 780 | def forward( |
| 781 | self, pixel_values: torch.Tensor, patch_embeds: torch.Tensor | None = None |
| 782 | ) -> torch.Tensor: |
| 783 | batch_size = pixel_values.shape[0] |
| 784 | if patch_embeds is None: |
| 785 | patch_embeds = self.patch_embedding(pixel_values) |
| 786 | patch_embeds = patch_embeds.flatten(2).transpose(1, 2) |
| 787 | |
| 788 | class_embeds = self.class_embedding.expand(batch_size, 1, -1) |
| 789 | embeddings = torch.cat([class_embeds, patch_embeds], dim=1) |
| 790 | embeddings = embeddings + self.get_abs_pos( |
| 791 | self.position_embedding(self.position_ids), embeddings.size(1) |
| 792 | ) |
| 793 | return embeddings |
| 794 | |
| 795 | |
| 796 | class DeepCLIPVisionTransformer(nn.Module): |
| 797 | def __init__( |
| 798 | self, |
| 799 | config: CLIPVisionConfig, |
| 800 | quant_config: QuantizationConfig | None = None, |
| 801 | *, |
| 802 | num_hidden_layers_override: int | None = None, |
| 803 | prefix: str = "", |
| 804 | ) -> None: |
| 805 | super().__init__() |
| 806 | |
| 807 | self.config = config |
| 808 | embed_dim = config.hidden_size |
| 809 | |
| 810 | self.embeddings = DeepCLIPVisionEmbeddings(config) |
| 811 | |
| 812 | # NOTE: This typo of "layrnorm" is not fixed on purpose to match |
| 813 | # the original transformers code and name of the model weights. |
| 814 | self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) |
| 815 | |
| 816 | self.transformer = CLIPEncoder( |
| 817 | config=config, |
| 818 | quant_config=quant_config, |
| 819 | num_hidden_layers_override=num_hidden_layers_override, |
| 820 | prefix=f"{prefix}.encoder", |
| 821 | attn_cls=MMEncoderAttention, |
| 822 | ) |
| 823 | |
| 824 | num_hidden_layers = config.num_hidden_layers |
| 825 | if len(self.transformer.layers) > config.num_hidden_layers: |
| 826 | raise ValueError( |
| 827 | f"The original encoder only has {num_hidden_layers} " |
| 828 | f"layers, but you requested {len(self.transformer.layers)} layers." |
| 829 | ) |
| 830 | |
| 831 | @property |
| 832 | def dtype(self): |
| 833 | return next(self.parameters()).dtype |
| 834 | |
| 835 | @property |
| 836 | def device(self): |
| 837 | return next(self.parameters()).device |
| 838 | |
| 839 | def forward( |
| 840 | self, |
| 841 | pixel_values: torch.Tensor, |
| 842 | patch_embeds: torch.Tensor | None = None, |
| 843 | *, |
| 844 | select_layers: list[int] | None = None, |
| 845 | ) -> torch.Tensor: |
| 846 | hidden_states = self.embeddings(pixel_values, patch_embeds) |
| 847 | hidden_states = self.pre_layrnorm(hidden_states) |
| 848 | |
| 849 | # Produces either the last layer output or all of the hidden states, |
| 850 | # depending on if we have select_layers or not |
| 851 | encoder_outputs = self.transformer( |
| 852 | inputs_embeds=hidden_states, |
| 853 | return_all_hidden_states=select_layers is not None, |
| 854 | ) |
| 855 | return encoder_outputs |
| 856 | |
| 857 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 858 | loader = AutoWeightsLoader(self) |
| 859 | return loader.load_weights(weights) |