text: import math
  import torch
  import torch.nn as nn
  from torch.nn import functional as F
  import numpy as np
  from timm.models.layers import DropPath, to_2tuple, trunc_normal_
  
  class Mlp(nn.Module):
      def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
          super().__init__()
          out_features = out_features or in_features
          hidden_features = hidden_features or in_features
          self.fc1 = nn.Linear(in_features, hidden_features)
          self.act = act_layer()
          self.fc2 = nn.Linear(hidden_features, out_features)
          self.drop = nn.Dropout(drop)
  
      def forward(self, x, H=None, W=None):
          x = self.fc1(x)
          x = self.act(x)
          x = self.drop(x)
          x = self.fc2(x)
          x = self.drop(x)
          return x      
  class DWConv(nn.Module):
      def __init__(self, dim=768):
          super(DWConv, self).__init__()
          self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)
  
      def forward(self, x, H, W):
          B, N, C = x.shape
          x = x.transpose(1, 2).view(B, C, H, W)
          x = self.dwconv(x)
          x = x.flatten(2).transpose(1, 2)
  
          return x
  class Relative_Attention(nn.Module):
      def __init__(self,dim,img_size,extra_token_num=1,num_heads=8,qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
          super().__init__()
          self.num_heads = num_heads
          self.extra_token_num = extra_token_num
          head_dim = dim // num_heads
          self.img_size = img_size # h,w
          self.scale = qk_scale or head_dim ** -0.5
           # define a parameter table of relative position bias,add cls_token bias
          self.relative_position_bias_table = nn.Parameter(
              torch.zeros((2 * img_size[0] - 1) * (2 * img_size[1] - 1) + 1, num_heads))  # 2*h-1 * 2*w-1 + 1, nH
          
          # get pair-wise relative position index for each token
          coords_h = torch.arange(self.img_size[0])
          coords_w = torch.arange(self.img_size[1])
          coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, h, w
          coords_flatten = torch.flatten(coords, 1)  # 2, h*w
          relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, h*w, h*w
          relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # h*w, h*w, 2
          relative_coords[:, :, 0] += self.img_size[0] - 1  # shift to start from 0
          relative_coords[:, :, 1] += self.img_size[1] - 1
          relative_coords[:, :, 0] *= 2 * self.img_size[1] - 1
          relative_position_index = relative_coords.sum(-1)  # h*w, h*w
          relative_position_index = F.pad(relative_position_index,(extra_token_num,0,extra_token_num,0))
          relative_position_index = relative_position_index.long()
          self.register_buffer("relative_position_index", relative_position_index)
          self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
          self.attn_drop = nn.Dropout(attn_drop)
          self.proj = nn.Linear(dim, dim)
          self.proj_drop = nn.Dropout(proj_drop)
          trunc_normal_(self.relative_position_bias_table, std=.02)
          self.softmax = nn.Softmax(dim=-1)
      def forward(self, x,):
          """
          Args:
              x: input features with shape of (B, N, C)
          """
          B_, N, C = x.shape
          qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
          q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
  
          q = q * self.scale
          attn = (q @ k.transpose(-2, -1))
  
          relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
              self.img_size[0] * self.img_size[1] + self.extra_token_num, self.img_size[0] * self.img_size[1] + self.extra_token_num, -1)  # h*w+1,h*w+1,nH
          
          relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, h*w+1, h*w+1
          attn = attn + relative_position_bias.unsqueeze(0)
  
          attn = self.softmax(attn)
          
          attn = self.attn_drop(attn)
  
          x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
          x = self.proj(x)
          x = self.proj_drop(x)
          return x
  class OverlapPatchEmbed(nn.Module):
      """ Image to Patch Embedding
      """
  
      def __init__(self, patch_size=7, stride=4, in_chans=3, embed_dim=768):
          super().__init__()
          patch_size = to_2tuple(patch_size)
          self.patch_size = patch_size
          self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride,
                                padding=(patch_size[0] // 2, patch_size[1] // 2))
          self.norm = nn.LayerNorm(embed_dim)
  
          self.apply(self._init_weights)
  
      def _init_weights(self, m):
          if isinstance(m, nn.Linear):
              trunc_normal_(m.weight, std=.02)
              if isinstance(m, nn.Linear) and m.bias is not None:
                  nn.init.constant_(m.bias, 0)
          elif isinstance(m, nn.LayerNorm):
              nn.init.constant_(m.bias, 0)
              nn.init.constant_(m.weight, 1.0)
          elif isinstance(m, nn.Conv2d):
              fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
              fan_out //= m.groups
              m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
              if m.bias is not None:
                  m.bias.data.zero_()
  
      def forward(self, x):
          x = self.proj(x)
          _, _, H, W = x.shape
          x = x.flatten(2).transpose(1, 2)
          x = self.norm(x)
  
          return x, H, W        
  class MHSABlock(nn.Module):
      def __init__(self, input_dim, output_dim,image_size, stride, num_heads,extra_token_num=1,mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                   drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
          super().__init__()
          if stride != 1:
              self.patch_embed = OverlapPatchEmbed(patch_size=3,stride=stride,in_chans=input_dim,embed_dim=output_dim)
              self.img_size = image_size//2
          else:
              self.patch_embed = None
              self.img_size = image_size
          self.img_size = to_2tuple(self.img_size)
          
          self.norm1 = norm_layer(output_dim)
          self.attn = Relative_Attention(
              output_dim,self.img_size, extra_token_num=extra_token_num,num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
          self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
          self.norm2 = norm_layer(output_dim)
          mlp_hidden_dim = int(output_dim * mlp_ratio)
          self.mlp = Mlp(in_features=output_dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
  
      def forward(self, x, H, W, extra_tokens=None):
          if self.patch_embed is not None:
              x,_,_ = self.patch_embed(x)
  
              extra_tokens = [token.expand(x.shape[0],-1,-1) for token in extra_tokens]
              extra_tokens.append(x)
              x = torch.cat(extra_tokens,dim=1)
          x = x + self.drop_path(self.attn(self.norm1(x)))
          x = x + self.drop_path(self.mlp(self.norm2(x),H//2,W//2))
          return x
  
  