class_docstr: Window based multi-head self attention (W-MSA) module with relative position bias.
  It supports both of shifted and non-shifted window.
  Args:
      dim (`int`):
          Number of input channels.
      window_size (`tuple[int]`):
          The height and width of the window.
      num_heads (`int`):
          Number of attention heads.
      qkv_bias (`bool`):
          If True, add a learnable bias to query, key, value. Default: True
      attn_drop (`float`):
          Dropout ratio of attention weight. Default: 0.0
      proj_drop (`float`):
          Dropout ratio of output. Default: 0.0
      pretrained_window_size (`tuple[int]`):
          The height and width of the window in pre-training.
docstr: Args:
      x: input features with shape of (num_windows*B, N, C)
      mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
code: def forward(self, x, mask: Optional[torch.Tensor] = None):
      """
      Args:
          x: input features with shape of (num_windows*B, N, C)
          mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
      """
      b_, n, c = x.shape
      if self.is_v2:
          qkv_bias = None
          if self.q_bias is not None:
              qkv_bias = torch.cat(
                  (
                      self.q_bias,
                      torch.zeros_like(
                          self.v_bias, requires_grad=False
                      ),
                      self.v_bias,
                  )
              )
          qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
  
      else:
          qkv = self.qkv(x)
      qkv = qkv.reshape(b_, n, 3, self.num_heads, -1).permute(
          2, 0, 3, 1, 4
      )
      q, k, v = (
          qkv[0],
          qkv[1],
          qkv[2],
      )  # make torchscript happy (cannot use tensor as tuple)
      if self.is_v2:
          # cosine attention
          attn = F.normalize(q, dim=-1) @ F.normalize(
              k, dim=-1
          ).transpose(-2, -1)
          logit_scale = torch.clamp(
              self.logit_scale, max=torch.log(torch.tensor(1.0 / 0.01))
          ).exp()
          attn = attn * logit_scale
  
          relative_position_bias_table = self.cpb_mlp(
              self.relative_coords_table
          ).view(-1, self.num_heads)
          relative_position_bias = relative_position_bias_table[
              self.relative_position_index.view(-1)
          ].view(
              self.window_size[0] * self.window_size[1],
              self.window_size[0] * self.window_size[1],
              -1,
          )  # Wh*Ww,Wh*Ww,nH
          relative_position_bias = relative_position_bias.permute(
              2, 0, 1
          ).contiguous()  # nH, Wh*Ww, Wh*Ww
          relative_position_bias = 16 * torch.sigmoid(
              relative_position_bias
          )
  
      else:
          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.window_size[0] * self.window_size[1],
              self.window_size[0] * self.window_size[1],
              -1,
          )  # Wh*Ww,Wh*Ww,nH
          relative_position_bias = relative_position_bias.permute(
              2, 0, 1
          ).contiguous()  # nH, Wh*Ww, Wh*Ww
  
      attn = attn + relative_position_bias.unsqueeze(0)
  
      if mask is not None:
          nw = mask.shape[0]
          attn = attn.view(
              b_ // nw, nw, self.num_heads, n, n
          ) + mask.unsqueeze(1).unsqueeze(0)
          attn = attn.view(-1, self.num_heads, n, n)
          attn = self.softmax(attn)
      else:
          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
  
file_path: towhee/towhee/models/layers/window_attention.py
imports: import torch
  import torch.nn.functional
  from torch import nn
  import numpy
  from typing import Optional
  from towhee.models.utils.weight_init import trunc_normal_
label_desc: Use image node_id 2 for calling the WindowAttention obj's underlying member method code with example usage: obj.forward(x, mask) and returns: x