class_docstr: Vision Transformer
  
  A PyTorch impl of : `Space-Time Transformer` from Frozen-in-time  - by Max Bain.
      https://arxiv.org/abs/2104.00650
  
  Based off:
   - ViT implementation from the timm library [https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py]
  lucidrains timesformer implementation [https://github.com/lucidrains/TimeSformer-pytorch].
  
  Notable differences:
   - allows for variable length input frames (<= num_frames)
   - allows for variable length input resolution  (<= (img_size, img_size)) [UNTESTED]
   - different attention block mechanism
class_docstr_tok: ["Vision","Transformer","A","PyTorch","impl","of",":","`","Space-Time","Transformer","`","from","Frozen-in-time","-","by","Max","Bain",".","https",":","\/\/arxiv.org\/abs\/2104.00650","Based","off",":","-","ViT","implementation","from","the","timm","library","[","https",":","\/\/github.com\/rwightman\/pytorch-image-models\/blob\/master\/timm\/models\/vision_transformer.py","]","lucidrains","timesformer","implementation","[","https",":","\/\/github.com\/lucidrains\/TimeSformer-pytorch","]",".","Notable","differences",":","-","allows","for","variable","length","input","frames","(","<","=","num_frames",")","-","allows","for","variable","length","input","resolution","(","<","=","(","img_size",",","img_size",")",")","[","UNTESTED","]","-","different","attention","block","mechanism"]
docstr: Args:
      img_size (`int` or `tuple`):
          Input image size.
      patch_size (`int` or `tuple`):
          Patch size.
      in_chans (`int`):
          Number of input channels.
      num_classes (`int`):
          Number of classes for classification head.
      embed_dim (`int`):
          Embedding dimension.
      depth (`int`):
          Depth of transformer.
      num_heads (`int`):
          Number of attention heads.
      mlp_ratio (`int`):
          Ratio of mlp hidden dim to embedding dim.
      qkv_bias (`bool`):
          Enable bias for qkv if True.
      qk_scale (`float`):
          Override default qk scale of head_dim ** -0.5 if set.
      representation_size (Optional[int]):
          Enable and set representation layer (pre-logits) to this value if set.
      drop_rate (`float`):
          Dropout rate.
      attn_drop_rate (`float`):
          Attention dropout rate.
      drop_path_rate (`float`):
          Stochastic depth rate.
      norm_layer (`nn.Module`):
          Normalization layer.
      num_frames (`int`):
          maximum number of frames expected as input.
docstr_tok: ["Args",":","img_size","(","`","int","`","or","`","tuple","`",")",":","Input","image","size",".","patch_size","(","`","int","`","or","`","tuple","`",")",":","Patch","size",".","in_chans","(","`","int","`",")",":","Number","of","input","channels",".","num_classes","(","`","int","`",")",":","Number","of","classes","for","classification","head",".","embed_dim","(","`","int","`",")",":","Embedding","dimension",".","depth","(","`","int","`",")",":","Depth","of","transformer",".","num_heads","(","`","int","`",")",":","Number","of","attention","heads",".","mlp_ratio","(","`","int","`",")",":","Ratio","of","mlp","hidden","dim","to","embedding","dim",".","qkv_bias","(","`","bool","`",")",":","Enable","bias","for","qkv","if","True",".","qk_scale","(","`","float","`",")",":","Override","default","qk","scale","of","head_dim","*","*","-0.5","if","set",".","representation_size","(","Optional","[","int","]",")",":","Enable","and","set","representation","layer","(","pre-logits",")","to","this","value","if","set",".","drop_rate","(","`","float","`",")",":","Dropout","rate",".","attn_drop_rate","(","`","float","`",")",":","Attention","dropout","rate",".","drop_path_rate","(","`","float","`",")",":","Stochastic","depth","rate",".","norm_layer","(","`","nn.Module","`",")",":","Normalization","layer",".","num_frames","(","`","int","`",")",":","maximum","number","of","frames","expected","as","input","."]
code: def __init__(
      self,
      img_size=224,
      patch_size=16,
      in_chans=3,
      num_classes=1000,
      embed_dim=768,
      depth=12,
      num_heads=12,
      mlp_ratio=4.0,
      qkv_bias=True,
      qk_scale=None,
      representation_size=None,
      drop_rate=0.0,
      attn_drop_rate=0.0,
      drop_path_rate=0.0,
      norm_layer=None,
      num_frames=8,
  ):
      """
      Args:
          img_size (`int` or `tuple`):
              Input image size.
          patch_size (`int` or `tuple`):
              Patch size.
          in_chans (`int`):
              Number of input channels.
          num_classes (`int`):
              Number of classes for classification head.
          embed_dim (`int`):
              Embedding dimension.
          depth (`int`):
              Depth of transformer.
          num_heads (`int`):
              Number of attention heads.
          mlp_ratio (`int`):
              Ratio of mlp hidden dim to embedding dim.
          qkv_bias (`bool`):
              Enable bias for qkv if True.
          qk_scale (`float`):
              Override default qk scale of head_dim ** -0.5 if set.
          representation_size (Optional[int]):
              Enable and set representation layer (pre-logits) to this value if set.
          drop_rate (`float`):
              Dropout rate.
          attn_drop_rate (`float`):
              Attention dropout rate.
          drop_path_rate (`float`):
              Stochastic depth rate.
          norm_layer (`nn.Module`):
              Normalization layer.
          num_frames (`int`):
              maximum number of frames expected as input.
  
      """
      super().__init__()
      self.num_classes = num_classes
      self.num_features = (
          self.embed_dim
      ) = embed_dim  # num_features for consistency with other models
      self.num_frames = num_frames
      self.embed_dim = embed_dim
  
      norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
  
      self.patch_embed = PatchEmbed2D(
          img_size=img_size,
          patch_size=patch_size,
          in_chans=in_chans,
          embed_dim=embed_dim,
          flatten=False,
      )
  
      num_patches = self.patch_embed.num_patches * num_frames
      self.patches_per_frame = num_patches // num_frames
  
      self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
      self.pos_embed = nn.Parameter(
          torch.zeros(1, self.patches_per_frame + 1, embed_dim)
      )  # remember to take pos_embed[1:] for tiling over time
  
      self.pos_drop = nn.Dropout(p=drop_rate)
  
      dpr = [
          x.item() for x in torch.linspace(0, drop_path_rate, depth)
      ]  # stochastic depth decay rule
      self.blocks = nn.ModuleList(
          [
              VideoBridgeBlock(
                  dim=embed_dim,
                  num_heads=num_heads,
                  mlp_ratio=mlp_ratio,
                  qkv_bias=qkv_bias,
                  qk_scale=qk_scale,
                  drop=drop_rate,
                  attn_drop=attn_drop_rate,
                  drop_path=dpr[i],
                  norm_layer=norm_layer,
              )
              for i in range(depth)
          ]
      )
  
      self.norm1 = norm_layer(embed_dim)
      self.norm2 = norm_layer(embed_dim)
      # Representation layer
      if representation_size:
          self.num_features = representation_size
          self.pre_logits = nn.Sequential(
              OrderedDict(
                  [
                      ("fc", nn.Linear(embed_dim, representation_size)),
                      ("act", nn.Tanh()),
                  ]
              )
          )
      else:
          self.pre_logits = nn.Identity()
  
      # Classifier head
      self.head = (
          nn.Linear(self.num_features, num_classes)
          if num_classes > 0
          else nn.Identity()
      )
  
      trunc_normal_(self.pos_embed, std=0.02)
      trunc_normal_(self.cls_token, std=0.02)
  
      # if num_frames > 1, then we perform ViT inflation and initialise time attention to zero so not necessary.
      if num_frames == 1:
          self.apply(init_vit_weights)
  
      # einops transformations
      self.einops_from_space = "b (f n) d"
      self.einops_to_space = "(b f) n d"
  
imports: from collections import OrderedDict
  from functools import partial
  import torch
  from einops import rearrange, repeat
  from timm.models.layers import DropPath, trunc_normal_
  from torch import einsum, nn
  from towhee.models.layers.mlp import Mlp
  from towhee.models.layers.patch_embed2d import PatchEmbed2D
  from towhee.models.utils.init_vit_weights import init_vit_weights
  from towhee.models.frozen_in_time.frozen_video_transformer import attn
label_desc: Use image node_id 1 to create a new VideoBridgeFormer object from inherited base classes: nn with example: obj = VideoBridgeFormer(img_size, patch_size, in_chans, num_classes, embed_dim, depth, num_heads, mlp_ratio, qkv_bias, qk_scale, representation_size, drop_rate, attn_drop_rate, drop_path_rate, norm_layer, num_frames)