Iambackup/Unlimited-OCR
012
1import torch.nn as nn2import torch3import torch.nn.functional as F4import copy5 6from contextlib import nullcontext7import math8from typing import Optional, Tuple9# from megatron.model import LayerNorm10 11from einops import rearrange12from easydict import EasyDict as adict13 14 15from typing import Optional, Tuple, Type16from functools import partial17 18 19 20class MlpProjector(nn.Module):21 22 def __init__(self, cfg):23 24 super().__init__()25 26 self.cfg = cfg27 28 if cfg.projector_type == "identity":29 modules = nn.Identity()30 31 elif cfg.projector_type == "linear":32 modules = nn.Linear(cfg.input_dim, cfg.n_embed)33 34 elif cfg.projector_type == "mlp_gelu":35 mlp_depth = cfg.get("depth", 1)36 modules = [nn.Linear(cfg.input_dim, cfg.n_embed)]37 for _ in range(1, mlp_depth):38 modules.append(nn.GELU())39 modules.append(nn.Linear(cfg.n_embed, cfg.n_embed))40 modules = nn.Sequential(*modules)41 42 elif cfg.projector_type == "normlayer_downsample_mlp_gelu":43 mlp_depth = cfg.get("depth", 1)44 mlp_ratio = cfg.get("mlp_ratio", 1)45 modules = [46 nn.LayerNorm(cfg.input_dim * cfg.downsample_ratio * cfg.downsample_ratio),47 nn.Linear(cfg.input_dim * cfg.downsample_ratio * cfg.downsample_ratio, cfg.n_embed * mlp_ratio)48 ]49 for _ in range(1, mlp_depth - 1):50 modules.append(nn.GELU())51 modules.append(nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed * mlp_ratio))52 modules.append(nn.GELU())53 modules.append(nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed))54 modules = nn.Sequential(*modules)55 56 elif cfg.projector_type == "downsample_mlp_gelu":57 mlp_depth = cfg.get("depth", 1)58 mlp_ratio = cfg.get("mlp_ratio", 1)59 modules = [nn.Linear(cfg.input_dim * cfg.downsample_ratio * cfg.downsample_ratio, cfg.n_embed * mlp_ratio)]60 for _ in range(1, mlp_depth - 1):61 modules.append(nn.GELU())62 modules.append(nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed * mlp_ratio))63 modules.append(nn.GELU())64 modules.append(nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed))65 modules = nn.Sequential(*modules)66 67 elif cfg.projector_type == "low_high_hybrid_split_mlp_gelu":68 mlp_depth = cfg.get("depth", 1)69 self.high_up_proj = nn.Linear(cfg.input_dim, cfg.n_embed // 2)70 self.low_up_proj = nn.Linear(cfg.input_dim, cfg.n_embed // 2)71 72 modules = []73 for _ in range(1, mlp_depth):74 modules.append(nn.GELU())75 modules.append(nn.Linear(cfg.n_embed, cfg.n_embed))76 modules = nn.Sequential(*modules)77 78 elif cfg.projector_type == "hybrid_split_feature_mlp_gelu":79 mlp_depth = cfg.get("depth", 1)80 channel_div = cfg.get("channel_div", 0.5)81 self.high_up_proj = nn.Linear(cfg.input_dim[0], int(cfg.n_embed * channel_div))82 self.low_up_proj = nn.Linear(cfg.input_dim[1], cfg.n_embed - int(cfg.n_embed * channel_div))83 84 modules = []85 for _ in range(1, mlp_depth):86 modules.append(nn.GELU())87 modules.append(nn.Linear(cfg.n_embed, cfg.n_embed))88 modules = nn.Sequential(*modules)89 90 elif cfg.projector_type == "low_high_split_mlp_gelu":91 mlp_depth = cfg.get("depth", 1)92 modules = []93 for _ in range(1, mlp_depth):94 modules.append(nn.GELU())95 modules.append(nn.Linear(cfg.n_embed // 2, cfg.n_embed // 2))96 modules = nn.Sequential(*modules)97 self.high_layers = nn.Sequential(*modules)98 self.low_layers = copy.deepcopy(modules)99 100 else:101 raise ValueError(f"Unknown projector type: {cfg.projector_type}")102 103 if cfg.get("token_pooling", False):104 self.token_pooling_layer = nn.Linear(cfg.input_dim * 4, cfg.input_dim)105 106 if cfg.get("conv_fusion_high_low_features", False):107 self.fusion_layer = nn.Linear(cfg.input_dim, cfg.input_dim)108 self.layers = modules109 110 def forward(self, x):111 if self.cfg.get("token_pooling", False):112 batch_size, wxh, channels = x.shape113 w = h = int(wxh**0.5)114 x = x.view(batch_size, w, h, channels)115 x = x.permute(0, 3, 1, 2)116 # import ipdb; ipdb.set_trace()117 patches = x.unfold(2, 2, 2).unfold(3, 2, 2)118 batch_size, channels, h_patches, w_patches, _, _ = patches.size()119 # 在通道维度上拼接120 patches = patches.contiguous().view(batch_size, channels, h_patches * w_patches, -1)121 122 # 通过线性层123 patches = patches.permute(0, 2, 1, 3).contiguous()124 patches = patches.view(batch_size, h_patches * w_patches, channels * 4)125 126 x = self.token_pooling_layer(patches)127 128 if self.cfg.get("conv_fusion_high_low_features", False):129 x = self.fusion_layer(x[:, 0]) + x[:, 1]130 131 if self.cfg.projector_type == 'low_high_hybrid_split_mlp_gelu':132 high_x, low_x = x[0], x[1]133 high_x = self.high_up_proj(high_x)134 low_x = self.low_up_proj(low_x)135 x = torch.concat([high_x, low_x], dim=-1)136 137 if self.cfg.projector_type == 'hybrid_split_feature_mlp_gelu':138 high_x = x[...,:self.cfg.input_dim[0]]139 low_x = x[...,self.cfg.input_dim[0]:]140 high_x = self.high_up_proj(high_x)141 low_x = self.low_up_proj(low_x)142 x = torch.concat([high_x, low_x], dim=-1)143 144 if self.cfg.projector_type == 'low_high_split_mlp_gelu':145 high_x, low_x = x[0], x[1]146 high_x = self.high_layers(high_x)147 low_x = self.low_layers(low_x)148 x = torch.concat([high_x, low_x], dim=-1)149 return x150 151 if self.cfg.projector_type == 'downsample_mlp_gelu' or self.cfg.projector_type == 'normlayer_downsample_mlp_gelu':152 bs, hw, input_dim = x.shape153 h = w = int((hw) ** 0.5)154 155 """compute padding"""156 if h % self.cfg.downsample_ratio:157 pad = self.cfg.downsample_ratio - h % self.cfg.downsample_ratio158 else:159 pad = 0160 x = x.reshape(bs, h, w, input_dim)161 if pad > 0:162 x = F.pad(x, (0, 0, 0, pad, 0, pad), "constant", 0)163 164 """4 to 1 concat"""165 x = x.permute(0, 3, 1, 2) # B, C, H, W166 x = F.unfold(x, kernel_size=self.cfg.downsample_ratio, stride=self.cfg.downsample_ratio, padding=0) # B, C*4, HW // 4167 x = x.permute(0, 2, 1)168 169 return self.layers(x)170 171 @staticmethod172 def get_flops_per_sample(cfg):173 if cfg.projector_type == "linear":174 fwd = 2 * cfg.input_dim * cfg.n_embed175 176 elif "mlp_gelu" in cfg.projector_type :177 mlp_depth = cfg.get("depth", 1)178 downsample_ratio = cfg.get("downsample_ratio", 1)179 input_dim = sum(cfg.input_dim) if isinstance(cfg.input_dim, list) else cfg.input_dim180 input_dim = input_dim * downsample_ratio * downsample_ratio181 fwd = 2 * input_dim * cfg.n_embed + (mlp_depth - 1) * 2 * cfg.n_embed * cfg.n_embed182 else:183 fwd = 0184 185 return fwd * 3186 187 188#===================clip============================================================189 190class LayerNormfp32(torch.nn.LayerNorm):191 """Subclass torch's LayerNorm to handle fp16."""192 193 def forward(self, x: torch.Tensor):194 orig_type = x.dtype195 ret = super().forward(x.type(torch.float32))196 return ret.type(orig_type)197 198 199def get_abs_pos(abs_pos, tgt_size):200 # abs_pos: L, C201 # tgt_size: M202 # return: M, C203 204 # print(tgt_size)205 # print(abs_pos.shape)206 # exit()207 dim = abs_pos.size(-1)208 # print(dim)209 abs_pos_new = abs_pos.squeeze(0)210 cls_token, old_pos_embed = abs_pos_new[:1], abs_pos_new[1:]211 212 213 214 src_size = int(math.sqrt(abs_pos_new.shape[0] - 1))215 tgt_size = int(math.sqrt(tgt_size))216 dtype = abs_pos.dtype217 218 if src_size != tgt_size:219 old_pos_embed = old_pos_embed.view(1, src_size, src_size, dim).permute(0, 3, 1,220 2).contiguous()221 old_pos_embed = old_pos_embed.to(torch.float32)222 new_pos_embed = F.interpolate(223 old_pos_embed,224 size=(tgt_size, tgt_size),225 mode='bicubic',226 antialias=True,227 align_corners=False,228 ).to(dtype)229 new_pos_embed = new_pos_embed.permute(0, 2, 3, 1)230 new_pos_embed = new_pos_embed.view(tgt_size * tgt_size, dim)231 vision_pos_embed = torch.cat([cls_token, new_pos_embed], dim=0)232 vision_pos_embed = vision_pos_embed.view(1, tgt_size * tgt_size + 1, dim)233 return vision_pos_embed234 else:235 return abs_pos236 237@torch.jit.script238def quick_gelu(x):239 return x * torch.sigmoid(1.702 * x)240 241 242 243class CLIPVisionEmbeddings(nn.Module):244 def __init__(self, hidden_size=1024, image_size=224, patch_size=14, num_channels=3):245 super().__init__()246 self.embed_dim = hidden_size247 self.image_size = image_size248 self.patch_size = patch_size249 250 self.class_embedding = torch.nn.Parameter(torch.randn(self.embed_dim))251 252 self.patch_embedding = torch.nn.Conv2d(253 in_channels=num_channels,254 out_channels=self.embed_dim,255 kernel_size=self.patch_size,256 stride=self.patch_size,257 bias=False,258 )259 260 self.num_patches = (self.image_size // self.patch_size) ** 2261 self.num_positions = self.num_patches + 1262 self.position_embedding = torch.nn.Embedding(self.num_positions, self.embed_dim)263 self.register_buffer(264 "position_ids", torch.arange(self.num_positions).expand((1, -1))265 )266 267 def forward(self, pixel_values, patch_embeds):268 batch_size = pixel_values.shape[0]269 # patch_embeds = self.patch_embedding(270 # pixel_values271 # ) # shape = [*, width, grid, grid]272 273 274 if patch_embeds is not None:275 patch_embeds = patch_embeds276 # print(patch_embeds.shape)277 else:278 patch_embeds = self.patch_embedding(pixel_values) 279 # print(111111)280 # shape = [*, width, grid, grid]281 # patch_embeds = patch_embeds.flatten(2).transpose(1, 2)282 283 patch_embeds = patch_embeds.flatten(2).transpose(1, 2)284 285 286 class_embeds = self.class_embedding.expand(batch_size, 1, -1)287 embeddings = torch.cat([class_embeds, patch_embeds], dim=1)288 289 # x = torch.cat([cls_token, x], dim=1)290 embeddings = embeddings + get_abs_pos(self.position_embedding(self.position_ids), embeddings.size(1))291 # embeddings = embeddings + self.position_embedding(self.position_ids)292 return embeddings293 294 295class NoTPFeedForward(nn.Module):296 def __init__(297 self,298 cfg,299 dim: int,300 hidden_dim: int,301 ):302 super().__init__()303 304 self.fc1 = torch.nn.Linear(dim, hidden_dim, bias=True)305 self.fc2 = torch.nn.Linear(hidden_dim, dim, bias=True)306 307 def forward(self, x):308 output = self.fc2(quick_gelu(self.fc1(x)))309 return output310 311 312 313 314class NoTPAttention(torch.nn.Module):315 def __init__(self, cfg):316 super().__init__()317 self.num_heads = cfg.num_attention_heads318 self.n_local_heads = cfg.num_attention_heads319 self.head_dim = cfg.hidden_size // cfg.num_attention_heads320 self.max_seq_len = cfg.seq_length321 self.use_flash_attention = cfg.use_flash_attn322 323 self.qkv_proj = torch.nn.Linear(cfg.hidden_size, cfg.hidden_size * 3, bias=True)324 self.out_proj = torch.nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=True)325 326 # self.core_attention = CoreAttention(cfg, AttnType.self_attn)327 328 self.attn_drop = cfg.attention_dropout329 330 def forward(331 self,332 x: torch.Tensor,333 ):334 bsz, seqlen, _ = x.shape335 xqkv = self.qkv_proj(x)336 xqkv = xqkv.view(bsz, seqlen, 3, self.num_heads, self.head_dim)337 338 if self.use_flash_attention:339 340 xq, xk, xv = torch.split(xqkv, 1, dim=2)341 xq = xq.squeeze(2)342 xk = xk.squeeze(2)343 xv = xv.squeeze(2)344 # xq, xk, xv = xqkv[:, :, 0, ...], xqkv[:, :, 1, ...], xqkv[:, :, 2, ...]345 346 # （B, num_head, S, head_size)347 xq = xq.permute(0, 2, 1, 3)348 xk = xk.permute(0, 2, 1, 3)349 xv = xv.permute(0, 2, 1, 3)350 # with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):351 output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None)352 output = output.permute(0, 2, 1, 3).reshape(bsz, seqlen, -1)353 # output = output.permute(0, 2, 1, 3).contiguous().view(bsz, seqlen, -1)354 else:355 # print(22222)356 xq, xk, xv = torch.split(xqkv, 1, dim=2)357 xq = xq.squeeze(2)358 xk = xk.squeeze(2)359 xv = xv.squeeze(2)360 # xq, xk, xv = xqkv[:, :, 0, ...], xqkv[:, :, 1, ...], xqkv[:, :, 2, ...]361 362 # （B, num_head, S, head_size)363 xq = xq.permute(0, 2, 1, 3)364 xk = xk.permute(0, 2, 1, 3)365 xv = xv.permute(0, 2, 1, 3)366 # with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):367 output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None)368 output = output.permute(0, 2, 1, 3).reshape(bsz, seqlen, -1)369 # output = output.permute(0, 2, 1, 3).contiguous().view(bsz, seqlen, -1)370 output = self.out_proj(output)371 return output372 373class NoTPTransformerBlock(nn.Module):374 def __init__(self, cfg, layer_id: int, multiple_of=256):375 super().__init__()376 377 self.n_heads = cfg.num_attention_heads378 self.dim = cfg.hidden_size379 self.head_dim = cfg.hidden_size // cfg.num_attention_heads380 self.self_attn = NoTPAttention(cfg)381 self.mlp = NoTPFeedForward(382 cfg, dim=cfg.hidden_size, hidden_dim=cfg.ffn_hidden_size383 )384 self.layer_id = layer_id385 self.layer_norm1 = torch.nn.LayerNorm(386 cfg.hidden_size, eps=cfg.layernorm_epsilon387 )388 self.layer_norm2 = torch.nn.LayerNorm(389 cfg.hidden_size, eps=cfg.layernorm_epsilon390 )391 392 def forward(self, x: torch.Tensor):393 residual = self.self_attn.forward(self.layer_norm1(x))394 h = x + residual395 out = h + self.mlp.forward(self.layer_norm2(h))396 return out397 398 399class NoTPTransformer(nn.Module):400 def __init__(self, cfg):401 super().__init__()402 403 self.cfg = cfg404 # self.recompute_list = self.cfg.get("recompute_list", [])405 self.num_layers = cfg.num_layers # _get_num_layers(cfg)406 407 self.layers = torch.nn.ModuleList()408 for layer_id in range(self.num_layers):409 self.layers.append(410 NoTPTransformerBlock(411 cfg,412 layer_id + 1,413 )414 )415 416 def forward(417 self,418 hidden_states,419 ):420 421 for lid, layer in enumerate(self.layers):422 # if lid in self.recompute_list:423 # def custom(layer_id):424 # def custom_forward(*args, **kwargs):425 # x_ = self.layers[layer_id](*args, **kwargs)426 # return x_427 428 # return custom_forward429 430 # assert hidden_states.requires_grad == True, logger.warning(431 # "When using recalculation, the input must have grad fn"432 # )433 # hidden_states = tensor_parallel.checkpoint(434 # custom(lid),435 # False,436 # hidden_states.contiguous()437 # )438 # else:439 hidden_states = layer(hidden_states)440 441 return hidden_states442 443 444# from megatron.core.tensor_parallel.layers import non_tensor_paralleled, local_dp_reduce, local_dp_scatter445 446class VitModel(nn.Module):447 def __init__(448 self,449 cfg,450 freeze_embed=False,451 freeze_pre_norm=False452 ) -> None:453 super().__init__()454 455 self.embeddings = CLIPVisionEmbeddings(hidden_size=cfg.hidden_size, image_size=cfg.image_size, patch_size=cfg.patch_size)456 457 if freeze_embed:458 for name, param in self.embeddings.named_parameters():459 param.requires_grad = False460 461 self.transformer = NoTPTransformer(cfg=cfg)462 463 if cfg.get("fp32norm", False):464 logger.info("Load fp32 layernorm for ViT.")465 self.pre_layrnorm = LayerNormfp32(466 cfg.hidden_size,467 eps=cfg.get("pre_layernorm_epsilon", 1e-5),468 )469 else:470 self.pre_layrnorm = torch.nn.LayerNorm(471 cfg.hidden_size,472 eps=cfg.get("pre_layernorm_epsilon", 1e-5),473 )474 475 # self.pre_layrnorm = RMSNorm(476 # cfg.hidden_size,477 # eps=cfg.get("pre_layernorm_epsilon", 1e-5),478 # sequence_parallel=False,479 # use_fp32=True,480 # use_optimus=True,481 # )482 483 if freeze_pre_norm:484 for name, param in self.pre_layrnorm.named_parameters():485 param.requires_grad = False486 487 for p in self.parameters():488 p.micro_dp = True489 490 def set_input_tensor(self, input_tensor):491 if not isinstance(input_tensor, list):492 input_tensor = [input_tensor]493 self.transformer.set_input_tensor(input_tensor[0])494 495 def __str__(self) -> str:496 return "open_clip"497 498 def forward(499 self,500 x,501 patch_embeds502 ):503 x = self.embeddings(x, patch_embeds)504 hidden_states = self.pre_layrnorm(x)505 506 # hidden_states, dis = local_dp_scatter(hidden_states)507 output = self.transformer(hidden_states)508 509 # output = local_dp_reduce(output, dis)510 511 return output512 513 514vit_model_cfg = adict(515 num_layers=24,516 hidden_size=1024,517 num_heads = 16,518 num_attention_heads=16,519 ffn_hidden_size=4096,520 seq_length=256,521 max_position_embeddings=256,522 use_flash_attn=False,523 understand_projector_stride=2,524 hidden_dropout = 0.0,525 attention_dropout = 0.0,526 no_persist_layer_norm = False,527 layernorm_epsilon = 1e-5,528 pre_layernorm_epsilon = 1e-5,529 image_size = 224,530 patch_size = 14,531 recompute_list = []532)533 534def build_clip_l():535 return VitModel(536 cfg=vit_model_cfg,537 freeze_embed=False,538 freeze_pre_norm=False,539 )540 541 542 543 544 545#=========================Sam-Vary=================================546 547 548def get_abs_pos_sam(abs_pos, tgt_size):549 550 dtype = abs_pos.dtype551 552 src_size = abs_pos.size(1)553 554 if src_size != tgt_size:555 old_pos_embed = abs_pos.permute(0, 3, 1, 2)556 old_pos_embed = old_pos_embed.to(torch.float32)557 new_pos_embed = F.interpolate(558 old_pos_embed,559 size=(tgt_size, tgt_size),560 mode='bicubic',561 antialias=True,562 align_corners=False,563 ).to(dtype)564 new_pos_embed = new_pos_embed.permute(0, 2, 3, 1)565 return new_pos_embed566 else:567 return abs_pos568 569 570 571 572class MLPBlock(nn.Module):573 def __init__(574 self,575 embedding_dim: int,576 mlp_dim: int,577 act: Type[nn.Module] = nn.GELU,578 ) -> None:579 super().__init__()580 self.lin1 = nn.Linear(embedding_dim, mlp_dim)581 self.lin2 = nn.Linear(mlp_dim, embedding_dim)582 self.act = act()583 584 def forward(self, x: torch.Tensor) -> torch.Tensor:585 return self.lin2(self.act(self.lin1(x)))586 587 588# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa589# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa590class LayerNorm2d(nn.Module):591 def __init__(self, num_channels: int, eps: float = 1e-6) -> None:592 super().__init__()593 self.weight = nn.Parameter(torch.ones(num_channels))594 self.bias = nn.Parameter(torch.zeros(num_channels))595 self.eps = eps596 597 def forward(self, x: torch.Tensor) -> torch.Tensor:598 u = x.mean(1, keepdim=True)599 s = (x - u).pow(2).mean(1, keepdim=True)600 x = (x - u) / torch.sqrt(s + self.eps)601 x = self.weight[:, None, None] * x + self.bias[:, None, None]602 return x603 604 605# 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 # noqa606class ImageEncoderViT(nn.Module):607 def __init__(608 self,609 img_size: int = 1024,610 patch_size: int = 16,611 in_chans: int = 3,612 embed_dim: int = 768,613 depth: int = 12,614 num_heads: int = 12,615 mlp_ratio: float = 4.0,616 out_chans: int = 256,617 qkv_bias: bool = True,618 norm_layer: Type[nn.Module] = nn.LayerNorm,619 act_layer: Type[nn.Module] = nn.GELU,620 use_abs_pos: bool = True,621 use_rel_pos: bool = False,622 rel_pos_zero_init: bool = True,623 window_size: int = 0,624 global_attn_indexes: Tuple[int, ...] = (),625 ) -> None:626 """627 Args:628 img_size (int): Input image size.629 patch_size (int): Patch size.630 in_chans (int): Number of input image channels.631 embed_dim (int): Patch embedding dimension.632 depth (int): Depth of ViT.633 num_heads (int): Number of attention heads in each ViT block.634 mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.635 qkv_bias (bool): If True, add a learnable bias to query, key, value.636 norm_layer (nn.Module): Normalization layer.637 act_layer (nn.Module): Activation layer.638 use_abs_pos (bool): If True, use absolute positional embeddings.639 use_rel_pos (bool): If True, add relative positional embeddings to the attention map.640 rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.641 window_size (int): Window size for window attention blocks.642 global_attn_indexes (list): Indexes for blocks using global attention.643 """644 super().__init__()645 self.img_size = img_size646 647 self.patch_embed = PatchEmbed(648 kernel_size=(patch_size, patch_size),649 stride=(patch_size, patch_size),650 in_chans=in_chans,651 embed_dim=embed_dim,652 )653 654 self.pos_embed: Optional[nn.Parameter] = None655 if use_abs_pos:656 # Initialize absolute positional embedding with pretrain image size.657 self.pos_embed = nn.Parameter(658 torch.zeros(1, img_size // patch_size, img_size // patch_size, embed_dim)659 )660 661 self.blocks = nn.ModuleList()662 for i in range(depth):663 block = Block(664 dim=embed_dim,665 num_heads=num_heads,666 mlp_ratio=mlp_ratio,667 qkv_bias=qkv_bias,668 norm_layer=norm_layer,669 act_layer=act_layer,670 use_rel_pos=use_rel_pos,671 rel_pos_zero_init=rel_pos_zero_init,672 window_size=window_size if i not in global_attn_indexes else 0,673 input_size=(img_size // patch_size, img_size // patch_size),674 )675 self.blocks.append(block)676 677 self.neck = nn.Sequential(678 nn.Conv2d(679 embed_dim,680 out_chans,681 kernel_size=1,682 bias=False,683 ),684 LayerNorm2d(out_chans),685 nn.Conv2d(686 out_chans,687 out_chans,688 kernel_size=3,689 padding=1,690 bias=False,691 ),692 LayerNorm2d(out_chans),693 )694 695 self.net_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1, bias=False)696 self.net_3 = nn.Conv2d(512, 1024, kernel_size=3, stride=2, padding=1, bias=False)697 698 def forward(self, x: torch.Tensor) -> torch.Tensor:699 x = self.patch_embed(x)700 if self.pos_embed is not None:701 # x = x + self.pos_embed702 x = x + get_abs_pos_sam(self.pos_embed, x.size(1))703 704 for blk in self.blocks:705 x = blk(x)706 707 x = self.neck(x.permute(0, 3, 1, 2))708 x2 = self.net_2(x)709 x3 = self.net_3(x2.clone())710 711 return x3712 713 714class Block(nn.Module):715 """Transformer blocks with support of window attention and residual propagation blocks"""716 717 def __init__(718 self,719 dim: int,720 num_heads: int,721 mlp_ratio: float = 4.0,722 qkv_bias: bool = True,723 norm_layer: Type[nn.Module] = nn.LayerNorm,724 act_layer: Type[nn.Module] = nn.GELU,725 use_rel_pos: bool = False,726 rel_pos_zero_init: bool = True,727 window_size: int = 0,728 input_size: Optional[Tuple[int, int]] = None,729 ) -> None:730 """731 Args:732 dim (int): Number of input channels.733 num_heads (int): Number of attention heads in each ViT block.734 mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.735 qkv_bias (bool): If True, add a learnable bias to query, key, value.736 norm_layer (nn.Module): Normalization layer.737 act_layer (nn.Module): Activation layer.738 use_rel_pos (bool): If True, add relative positional embeddings to the attention map.739 rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.740 window_size (int): Window size for window attention blocks. If it equals 0, then741 use global attention.742 input_size (tuple(int, int) or None): Input resolution for calculating the relative743 positional parameter size.744 """745 super().__init__()746 self.norm1 = norm_layer(dim)747 self.attn = Attention(748 dim,749 num_heads=num_heads,750 qkv_bias=qkv_bias,751 use_rel_pos=use_rel_pos,752 rel_pos_zero_init=rel_pos_zero_init,753 input_size=input_size if window_size == 0 else (window_size, window_size),754 )755 756 self.norm2 = norm_layer(dim)757 self.mlp = MLPBlock(embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer)758 759 self.window_size = window_size760 761 def forward(self, x: torch.Tensor) -> torch.Tensor:762 shortcut = x763 x = self.norm1(x)764 # Window partition765 if self.window_size > 0:766 H, W = x.shape[1], x.shape[2]767 x, pad_hw = window_partition(x, self.window_size)768 769 x = self.attn(x)770 # Reverse window partition771 if self.window_size > 0:772 x = window_unpartition(x, self.window_size, pad_hw, (H, W))773 774 x = shortcut + x775 x = x + self.mlp(self.norm2(x))776 777 return x778 779 780class Attention(nn.Module):781 """Multi-head Attention block with relative position embeddings."""782 783 def __init__(784 self,785 dim: int,786 num_heads: int = 8,787 qkv_bias: bool = True,788 use_rel_pos: bool = False,789 rel_pos_zero_init: bool = True,790 input_size: Optional[Tuple[int, int]] = None,791 ) -> None:792 """793 Args:794 dim (int): Number of input channels.795 num_heads (int): Number of attention heads.796 qkv_bias (bool): If True, add a learnable bias to query, key, value.797 rel_pos (bool): If True, add relative positional embeddings to the attention map.798 rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.799 input_size (tuple(int, int) or None): Input resolution for calculating the relative800 positional parameter size.801 """802 super().__init__()803 self.num_heads = num_heads804 head_dim = dim // num_heads805 self.scale = head_dim**-0.5806 807 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)808 self.proj = nn.Linear(dim, dim)809 810 self.use_rel_pos = use_rel_pos811 if self.use_rel_pos:812 assert (813 input_size is not None814 ), "Input size must be provided if using relative positional encoding."815 # initialize relative positional embeddings816 self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))817 self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))818 819 def forward(self, x: torch.Tensor) -> torch.Tensor:820 B, H, W, _ = x.shape821 # qkv with shape (3, B, nHead, H * W, C)822 qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)823 # q, k, v with shape (B * nHead, H * W, C)824 q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)825 826 rel_h, rel_w = None, None827 if self.use_rel_pos:828 rel_h, rel_w = add_decomposed_rel_pos(q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W))829 830 q = q.view(B, self.num_heads, H * W, -1)831 k = k.view(B, self.num_heads, H * W, -1)832 v = v.view(B, self.num_heads, H * W, -1)833 834 if self.use_rel_pos:835 rel_h = rel_h.view(B, self.num_heads, rel_h.size(1), rel_h.size(2), rel_h.size(3))836 rel_w = rel_w.view(B, self.num_heads, rel_w.size(1), rel_w.size(2), rel_w.size(3))837 attn_bias = (rel_h + rel_w).view(B, self.num_heads, rel_h.size(2), rel_h.size(3) * rel_w.size(4))838 x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_bias)839 # x = _attention_rel_h_rel_w(q, k, v, rel_h, rel_w)840 else:841 x = torch.nn.functional.scaled_dot_product_attention(q, k, v)842 843 x = x.view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1)844 845 x = self.proj(x)846 847 return x848 849 850def window_partition(x: torch.Tensor, window_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]:851 """852 Partition into non-overlapping windows with padding if needed.853 Args:854 x (tensor): input tokens with [B, H, W, C].855 window_size (int): window size.856 857 Returns:858 windows: windows after partition with [B * num_windows, window_size, window_size, C].859 (Hp, Wp): padded height and width before partition860 """861 B, H, W, C = x.shape862 863 pad_h = (window_size - H % window_size) % window_size864 pad_w = (window_size - W % window_size) % window_size865 if pad_h > 0 or pad_w > 0:866 x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))867 Hp, Wp = H + pad_h, W + pad_w868 869 x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)870 windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)871 return windows, (Hp, Wp)872 873 874def window_unpartition(875 windows: torch.Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int]876) -> torch.Tensor:877 """878 Window unpartition into original sequences and removing padding.879 Args:880 windows (tensor): input tokens with [B * num_windows, window_size, window_size, C].881 window_size (int): window size.882 pad_hw (Tuple): padded height and width (Hp, Wp).883 hw (Tuple): original height and width (H, W) before padding.884 885 Returns:886 x: unpartitioned sequences with [B, H, W, C].887 """888 Hp, Wp = pad_hw889 H, W = hw890 B = windows.shape[0] // (Hp * Wp // window_size // window_size)891 x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)892 x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1)893 894 if Hp > H or Wp > W:895 x = x[:, :H, :W, :].contiguous()896 return x897 898 899def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:900 """901 Get relative positional embeddings according to the relative positions of902 query and key sizes.903 Args:904 q_size (int): size of query q.905 k_size (int): size of key k.906 rel_pos (Tensor): relative position embeddings (L, C).907 908 Returns:909 Extracted positional embeddings according to relative positions.910 """911 max_rel_dist = int(2 * max(q_size, k_size) - 1)912 # Interpolate rel pos if needed.913 if rel_pos.shape[0] != max_rel_dist:914 # Interpolate rel pos.915 dtype = rel_pos.dtype916 rel_pos = rel_pos.to(torch.float32)917 rel_pos_resized = F.interpolate(918 rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),919 size=max_rel_dist,920 mode="linear",921 ).to(dtype)922 rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)923 else:924 rel_pos_resized = rel_pos925 926 # Scale the coords with short length if shapes for q and k are different.927 q_coords = torch.arange(q_size, device=rel_pos.device)[:, None] * max(k_size / q_size, 1.0)928 k_coords = torch.arange(k_size, device=rel_pos.device)[None, :] * max(q_size / k_size, 1.0)929 relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)930 931 return rel_pos_resized[relative_coords.long()]932 933 934def add_decomposed_rel_pos(935 q: torch.Tensor,936 rel_pos_h: torch.Tensor,937 rel_pos_w: torch.Tensor,938 q_size: Tuple[int, int],939 k_size: Tuple[int, int],940) -> torch.Tensor:941 """942 Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.943 https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py # noqa B950944 Args:945 q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).946 rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis.947 rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis.948 q_size (Tuple): spatial sequence size of query q with (q_h, q_w).949 k_size (Tuple): spatial sequence size of key k with (k_h, k_w).950 951 Returns:952 attn (Tensor): attention map with added relative positional embeddings.953 """954 q_h, q_w = q_size955 k_h, k_w = k_size956 Rh = get_rel_pos(q_h, k_h, rel_pos_h)957 Rw = get_rel_pos(q_w, k_w, rel_pos_w)958 959 B, _, dim = q.shape960 r_q = q.reshape(B, q_h, q_w, dim)961 rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh)962 rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw)963 rel_h = rel_h.unsqueeze(-1)964 rel_w = rel_w.unsqueeze(-2)965 rel_h = rel_h.reshape(B, q_h * q_w, k_h, 1)966 rel_w = rel_w.reshape(B, q_h * q_w, 1, k_w)967 968 return rel_h, rel_w969 970 971class PatchEmbed(nn.Module):972 """973 Image to Patch Embedding.974 """975 976 def __init__(977 self,978 kernel_size: Tuple[int, int] = (16, 16),979 stride: Tuple[int, int] = (16, 16),980 padding: Tuple[int, int] = (0, 0),981 in_chans: int = 3,982 embed_dim: int = 768,983 ) -> None:984 """985 Args:986 kernel_size (Tuple): kernel size of the projection layer.987 stride (Tuple): stride of the projection layer.988 padding (Tuple): padding size of the projection layer.989 in_chans (int): Number of input image channels.990 embed_dim (int): Patch embedding dimension.991 """992 super().__init__()993 994 self.proj = nn.Conv2d(995 in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding996 )997 998 def forward(self, x: torch.Tensor) -> torch.Tensor:999 x = self.proj(x)1000 # B C H W -> B H W C1001 x = x.permute(0, 2, 3, 1)1002 return x1003 1004 1005def build_sam_vit_b(checkpoint=None):1006 return _build_sam(1007 encoder_embed_dim=768,1008 encoder_depth=12,1009 encoder_num_heads=12,1010 encoder_global_attn_indexes=[2, 5, 8, 11],1011 checkpoint=checkpoint,1012 )1013 1014def build_sam_fast_vit_b(checkpoint=None, compile_mode='max-autotune', dtype=torch.bfloat16):1015 image_encoder = build_sam_vit_b(checkpoint).eval().to(dtype)1016 # sam = _apply_eval_dtype_sam(sam, dtype)1017 image_encoder = torch.compile(image_encoder, mode=compile_mode)1018 return image_encoder1019 1020 1021def _build_sam(1022 encoder_embed_dim,1023 encoder_depth,1024 encoder_num_heads,1025 encoder_global_attn_indexes,1026 checkpoint=None,1027):1028 prompt_embed_dim = 2561029 image_size = 10241030 vit_patch_size = 161031 image_embedding_size = image_size // vit_patch_size1032 image_encoder=ImageEncoderViT(1033 depth=encoder_depth,1034 embed_dim=encoder_embed_dim,1035 img_size=image_size,1036 mlp_ratio=4,1037 norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),1038 num_heads=encoder_num_heads,1039 patch_size=vit_patch_size,1040 qkv_bias=True,1041 use_rel_pos=True,1042 global_attn_indexes=encoder_global_attn_indexes,1043 window_size=14,1044 out_chans=prompt_embed_dim,1045 )1046 image_encoder.eval()1047 if checkpoint is not None:1048 # with open(checkpoint, "rb") as f:1049 state_dict = torch.load(checkpoint)1050 # print(state_dict.keys())1051 # for key in state_dict:1052 # image_encoder.load_state_dict({k[14:]: v for k, v in state_dict.items() if 'image_encoder' in k}, strict=False)1053 # ocr-anyting1054 # image_encoder.load_state_dict(state_dict, strict=True)1055 # tob1056 image_encoder.load_state_dict({k[30:]: v for k, v in state_dict.items() if 'vision_tower_high' in k}, strict=True)1057 print(checkpoint)1058 return image_encoder