    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., qkv_bias=True, qk_scale=None, representation_size=None,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=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=.02)
        trunc_normal_(self.cls_token, std=.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'
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.
towhee/models/bridgeformer/bridge_former_training_block.py