Upload models_mamba_ecg.py with huggingface_hub

models_mamba_ecg.py

@@ -0,0 +1,795 @@
+# Copyright (c) 2015-present, Facebook, Inc.
+# All rights reserved.
+import torch
+import torch.nn as nn
+import torch.nn.functional as functional
+from functools import partial
+from torch import Tensor
+from typing import Optional
+
+from timm.models.vision_transformer import VisionTransformer, _cfg
+# from timm.models.registry import register_model
+# from timm.models.layers import trunc_normal_, lecun_normal_
+
+from timm.models import register_model
+from timm.layers import trunc_normal_, lecun_normal_
+
+from timm.layers import DropPath, to_2tuple
+
+# from timm.models.layers import DropPath, to_2tuple
+from timm.models.vision_transformer import _load_weights
+
+import math
+
+from collections import namedtuple
+
+from mamba_ssm.modules.mamba_simple import Mamba
+from mamba_ssm.utils.generation import GenerationMixin
+from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf
+
+from rope import *
+import random
+import sys
+
+try:
+    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
+except ImportError:
+    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
+
+
+# layer_norm_fn and rms_norm_fn both are normalization method
+
+__all__ = [
+    'vim_tiny_patch16_224', 'vim_small_patch16_224', 'vim_base_patch16_224',
+    'vim_tiny_patch16_384', 'vim_small_patch16_384', 'vim_base_patch16_384',
+]
+
+
+'''
+in the original script(ft-vim-s.sh)
+img_size = 224, 
+patch_size = 16, 
+stride = 8, 
+in_chans = 3, 
+embed_dim = 768
+------------------------------------
+self.img_size: (224, 224)
+self.patch_size: (16, 16)
+self.grid_size: (27, 27)
+self.num_patches: 729
+self.flatten: True
+self.norm: nn.Identity()
+'''
+class PatchEmbed(nn.Module):
+    """ 2D Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, stride=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = ((img_size[0] - patch_size[0]) // stride + 1, (img_size[1] - patch_size[1]) // stride + 1)
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride)
+        # if the norm_layer is not none or null, the self.norm = norm_layer(embed_dim)
+        # otherwise, self.norm = nn.Identity()
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+        
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        
+        x = self.proj(x)
+        print("This is the shape after the CNN", x.shape)
+
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+            print("This is the shape after the flatten:", x.shape)
+
+        x = self.norm(x)
+        return x
+
+class PatchEmbed_spectrogram(nn.Module):
+    """ 2D spectrogram to Patch Embedding
+    """
+    def __init__(self, img_size_f = 128, img_size_t = 64, patch_size=6, stride=3, in_chans=12, embed_dim=432, flatten=True):
+        super().__init__()
+        # img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+
+        self.img_size_f = img_size_f
+        self.img_size_t = img_size_t
+
+        self.patch_size = patch_size
+        self.grid_size = ((img_size_f - patch_size[0]) // stride + 1, (img_size_t - patch_size[1]) // stride + 1)
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride)
+        # if the norm_layer is not none or null, the self.norm = norm_layer(embed_dim)
+        # otherwise, self.norm = nn.Identity()
+        # self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+        
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert H == self.img_size_f and W == self.img_size_t, \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size_f}*{self.img_size_t})."
+        
+        x = self.proj(x)
+
+        # This is the shape after the CNN torch.Size([1, 432, 41, 20])
+        # print("This is the shape after the CNN", x.shape)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        # This is the shape after the flatten: torch.Size([1, 820, 432])
+        # print("This is the shape after the flatten:", x.shape)
+        return x
+
+
+class CNN_layers(nn.Module):
+
+    def __init__(self, embed_size = 384):
+        super().__init__()
+        self.multiple_cnn = nn.Sequential( 
+            nn.Conv1d(12, 128, kernel_size=14, stride=3, padding=2, bias=False),
+            nn.BatchNorm1d(128),
+            nn.ReLU(inplace=True),
+
+            nn.Conv1d(128, embed_size, kernel_size=15, stride=4, padding=100, bias=False),
+            nn.BatchNorm1d(embed_size),
+            nn.ReLU(inplace=True))
+        
+    def forward(self, x):
+        # print("This is the shape of enconder:(before)", x.shape)
+        # This is the shape of enconder:(before) torch.Size([44, 12, 8192])
+        x = self.multiple_cnn(x)
+        x = x.transpose(1, 2)
+        # print("This is the shape of enconder:(after)", x.shape)
+        # This is the shape of enconder:(after) torch.Size([44, 729, 384])
+
+        # sys.exit()
+        return x 
+
+'''
+dim: 384
+mixer_cls: mixer_cla is an instance of mamba.
+drop_path = 0.
+norm_cls = nn.LayerNorm
+fused_add_norm = True, 
+residual_in_fp32 = True
+'''
+class Block(nn.Module):
+    def __init__(
+        self, dim, mixer_cls, norm_cls=nn.LayerNorm, fused_add_norm=False, residual_in_fp32=False, drop_path=0.,
+    ):
+        """
+        Simple block wrapping a mixer class with LayerNorm/RMSNorm and residual connection"
+
+        This Block has a slightly different structure compared to a regular
+        prenorm Transformer block.
+        The standard block is: LN -> MHA/MLP -> Add.
+        [Ref: https://arxiv.org/abs/2002.04745]
+        Here we have: Add -> LN -> Mixer, returning both
+        the hidden_states (output of the mixer) and the residual.
+        This is purely for performance reasons, as we can fuse add and LayerNorm.
+        The residual needs to be provided (except for the very first block).
+        """
+        super().__init__()
+        
+        self.residual_in_fp32 = residual_in_fp32
+        self.fused_add_norm = fused_add_norm
+        self.mixer = mixer_cls(dim)
+        self.norm = norm_cls(dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        
+        # fused_add_norm true
+        if self.fused_add_norm:
+            assert RMSNorm is not None, "RMSNorm import fails"
+            assert isinstance(self.norm, (nn.LayerNorm, RMSNorm)), "Only LayerNorm and RMSNorm are supported for fused_add_norm"
+
+    def forward(self, hidden_states: Tensor, residual: Optional[Tensor] = None, inference_params=None):
+
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            hidden_states: the sequence to the encoder layer (required).
+            residual: hidden_states = Mixer(LN(residual))
+        """
+        
+        if not self.fused_add_norm:
+            if residual is None:
+                residual = hidden_states
+            else:
+                residual = residual + self.drop_path(hidden_states)
+            
+            hidden_states = self.norm(residual.to(dtype=self.norm.weight.dtype))
+            if self.residual_in_fp32:
+                residual = residual.to(torch.float32)
+        
+        # since the self.fused_add_norm is true, the code is going below
+        # fused_add_norm_fn = layer_norm_fn
+        ###########
+        # hidden_states: Tensor
+        # self.norm.weight = torch.nn.LayerNorm.weight
+        # self.norm.bias = torch.nn.LayerNorm.bias
+        # residual: Optional[Tensor] = None
+        # prenorm=True
+        # residual_in_fp32 = True
+        # eps = torch.nn.LayerNorm.eps
+        
+        else:
+            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm, RMSNorm) else layer_norm_fn
+            if residual is None:
+                hidden_states, residual = fused_add_norm_fn(
+                    hidden_states,
+                    self.norm.weight,
+                    self.norm.bias,
+                    residual=residual,
+                    prenorm=True,
+                    residual_in_fp32=self.residual_in_fp32,
+                    eps=self.norm.eps,
+                )
+            else:
+                hidden_states, residual = fused_add_norm_fn(
+                    self.drop_path(hidden_states),
+                    self.norm.weight,
+                    self.norm.bias,
+                    residual=residual,
+                    prenorm=True,
+                    residual_in_fp32=self.residual_in_fp32,
+                    eps=self.norm.eps,
+                )  
+
+        # inference_params=None
+        hidden_states = self.mixer(hidden_states, inference_params=inference_params)
+        return hidden_states, residual
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        print("the code is going through allocate_inference_cache in the block")
+
+        return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
+
+'''
+from torch.nn.ModuleList()
+
+embed_dim=384, (embedding dimension)
+device: None 
+dtype: None
+ssm_cfg: None
+norm_epsilon: float = 1e-5
+rms_norm: bool = False
+residual_in_fp32 = True
+fused_add_norm = True
+if_bimamba = False
+bimamba_type = "v2"
+inter_dpr: [0.0, 0.0, 0.004347826354205608, ...,0.09565217792987823, 0.10000000149011612]
+if_devide_out = True
+init_layer_scale = None
+layer_idx = i
+'''
+
+def create_block(
+    d_model,
+    ssm_cfg=None,
+    norm_epsilon=1e-5,
+    drop_path=0.,
+    rms_norm=False,
+    residual_in_fp32=False,
+    fused_add_norm=False,
+    layer_idx=None,
+    device=None,
+    dtype=None,
+    if_bimamba=False,
+    bimamba_type="none",
+    if_devide_out=False,
+    init_layer_scale=None,
+    block_name = "default_value"
+):
+    if if_bimamba:
+        bimamba_type = "v1"
+
+    if ssm_cfg is None:
+        ssm_cfg = {}
+
+    factory_kwargs = {"device": device, "dtype": dtype}
+
+    if block_name == "VisionMamba":
+        mixer_cls = partial(Mamba, layer_idx=layer_idx, bimamba_type=bimamba_type, if_devide_out=if_devide_out, init_layer_scale=init_layer_scale, **ssm_cfg, **factory_kwargs)
+    elif block_name == "OriginalMamba":
+        mixer_cls = partial(Mamba, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)
+    else:
+        raise ValueError(f"No matching condition for value: {block_name}")
+    
+    
+
+    # rms_norm = False
+    norm_cls = partial(nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs)
+    
+    block = Block(
+        d_model,
+        mixer_cls,
+        norm_cls=norm_cls,
+        drop_path=drop_path,
+        fused_add_norm=fused_add_norm,
+        residual_in_fp32=residual_in_fp32,
+    )
+    block.layer_idx = layer_idx
+    return block
+
+
+# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
+def _init_weights(
+    module,
+    n_layer,
+    initializer_range=0.02,  # Now only used for embedding layer.
+    rescale_prenorm_residual=True,
+    n_residuals_per_layer=1,  # Change to 2 if we have MLP
+):
+    if isinstance(module, nn.Linear):
+        if module.bias is not None:
+            if not getattr(module.bias, "_no_reinit", False):
+                nn.init.zeros_(module.bias)
+    elif isinstance(module, nn.Embedding):
+        nn.init.normal_(module.weight, std=initializer_range)
+
+    if rescale_prenorm_residual:
+        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+        #
+        # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+        for name, p in module.named_parameters():
+            if name in ["out_proj.weight", "fc2.weight"]:
+                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                # We need to reinit p since this code could be called multiple times
+                # Having just p *= scale would repeatedly scale it down
+                nn.init.kaiming_uniform_(p, a=math.sqrt(5))
+                with torch.no_grad():
+                    p /= math.sqrt(n_residuals_per_layer * n_layer)
+
+
+def segm_init_weights(m):
+    if isinstance(m, nn.Linear):
+        trunc_normal_(m.weight, std=0.02)
+        if isinstance(m, nn.Linear) and m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.Conv2d):
+        # NOTE conv was left to pytorch default in my original init
+        lecun_normal_(m.weight)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+    elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
+        nn.init.zeros_(m.bias)
+        nn.init.ones_(m.weight)
+
+'''
+below is the 'ft-vim-s.sh':
+
+patch_size=16, (just patch size) 
+stride=8,  (just stride)
+embed_dim=384, (embedding dimension)
+depth=24, (?) the number of block
+rms_norm = True, (?)
+residual_in_fp32 = True, (?) 
+fused_add_norm = True, (?)
+final_pool_type = 'mean', (?)
+if_abs_pos_embed = True, (?)
+if_rope = False, (?)
+if_rope_residual = False, (?) 
+bimamba_type = "v2", (?)
+if_cls_token = True, (?)
+if_devide_out = True, (?)
+use_middle_cls_token = True, 
+**kwargs
+'''
+
+class VisionMamba(nn.Module):
+    def __init__(self, 
+                 img_size=224, 
+                 patch_size=16, 
+                 stride=16,
+                 depth=24, 
+                 embed_dim=192, 
+                 channels=3, 
+                 num_classes=26,
+                 ssm_cfg=None, 
+                 drop_rate=0.,
+                 drop_path_rate=0,
+                 norm_epsilon: float = 1e-5, 
+                 rms_norm: bool = False, 
+                 initializer_cfg=None,
+                 fused_add_norm=False,
+                 residual_in_fp32=False,
+                 device=None,
+                 dtype=None,
+                 ft_seq_len=None,
+                 pt_hw_seq_len=14,
+                 if_bidirectional=False,
+                 final_pool_type='none',
+                 if_abs_pos_embed=False,
+                 if_rope=False,
+                 if_rope_residual=False,
+                 flip_img_sequences_ratio=-1.,
+                 if_bimamba=False,
+                 bimamba_type="none",
+                 if_cls_token=False,
+                 if_devide_out=False,
+                 init_layer_scale=None,
+                 use_double_cls_token=False,
+                 use_middle_cls_token=False,
+                 **kwargs):
+        # print("The program is coming the init")
+        factory_kwargs = {"device": device, "dtype": dtype}
+        # factory_kwargs: {'device': None, 'dtype': None}
+        # add factory_kwargs into kwargs
+        block_name = kwargs.get('block', 'default_value')
+
+        kwargs.update(factory_kwargs) 
+
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+        self.fused_add_norm = fused_add_norm
+        self.if_bidirectional = if_bidirectional
+        self.final_pool_type = final_pool_type
+        self.if_abs_pos_embed = if_abs_pos_embed
+        self.if_rope = if_rope
+        self.if_rope_residual = if_rope_residual
+        self.flip_img_sequences_ratio = flip_img_sequences_ratio
+        self.if_cls_token = if_cls_token
+        self.use_double_cls_token = use_double_cls_token
+        self.use_middle_cls_token = use_middle_cls_token
+        self.num_tokens = 1 if if_cls_token else 0
+
+        # pretrain parameters
+        self.num_classes = num_classes
+        self.d_model = self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+
+        # self.patch_embed = PatchEmbed(img_size=img_size, patch_size=patch_size, stride=stride, in_chans=channels, embed_dim=embed_dim)
+        # num_patches = self.patch_embed.num_patches
+
+        self.CNN_layers = CNN_layers()
+        num_patches = 729
+
+        # self.ECG_patch_embedding = ECG_Patch_embedding()
+        # num_patches = 1023
+        # self.LC = LeadCombiner(lead=6, out_ch=8)
+
+        # self.CNN_layers = D2_CNN_layers()
+        # num_patches = 828
+
+        # self.CNN_layers = CNN_layers()
+        # num_patches = 775
+        
+        # if_cls_token: True (in the original script)
+        if if_cls_token:
+            # use_double_cls_token: False (in the original script)
+            if use_double_cls_token:
+                self.cls_token_head = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+                self.cls_token_tail = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+                self.num_tokens = 2
+
+            else:
+                self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+                # self.num_tokens = 1
+        
+        # if_abs_pos_embed: True (in the original script)
+        if if_abs_pos_embed:
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, self.embed_dim))
+            self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # if_rope: False (in the original script)
+
+        if if_rope:
+            pass
+            # half_head_dim = embed_dim // 2
+            # self.rope = TimeSeriesRotaryEmbeddingFast(dim=embed_dim, seq_len= num_patches + self.num_tokens)
+        
+        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+        # self.head_LC = nn.Linear(16, num_classes)
+        # depth: 24; drop_path_rate: 0.1
+        # TODO: release this comment
+        # dpr: [0.0, 0.004347826354205608, 0.008695652708411217, ..., 0.09130434691905975, 0.09565217792987823, 0.10000000149011612]
+
+        if drop_path_rate == 0:
+            print("This is the drop_path_rate:", drop_path_rate)
+            dpr = [x.item() for x in torch.full((depth,), drop_path_rate)]            
+        else:
+            print("This is the drop_path_rate:", drop_path_rate)
+            print("follow the stochastic depth decay rule")
+            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+
+        # import ipdb;ipdb.set_trace()
+        # inter_dpr: [0.0, 0.0, 0.004347826354205608, ...,0.09565217792987823, 0.10000000149011612]
+        inter_dpr = [0.0] + dpr
+        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+
+        # transformer blocks
+        # depth: 
+        self.layers = nn.ModuleList(
+            [
+                create_block(
+                    embed_dim,
+                    ssm_cfg=ssm_cfg,
+                    norm_epsilon=norm_epsilon,
+                    rms_norm=rms_norm,
+                    residual_in_fp32=residual_in_fp32,
+                    fused_add_norm=fused_add_norm,
+                    layer_idx=i,
+                    if_bimamba=if_bimamba,
+                    bimamba_type=bimamba_type,
+                    drop_path=inter_dpr[i],
+                    if_devide_out=if_devide_out,
+                    init_layer_scale=init_layer_scale,
+                    block_name = block_name,
+                    **factory_kwargs,
+                )
+                for i in range(depth)
+            ]
+        )
+        
+        # output head
+        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(embed_dim, eps=norm_epsilon, **factory_kwargs)
+
+        # self.pre_logits = nn.Identity()
+
+        # original init
+        # self.patch_embed.apply(segm_init_weights)
+
+        self.CNN_layers.apply(segm_init_weights)
+
+        # self.ECG_patch_embedding.apply(segm_init_weights)
+
+        self.head.apply(segm_init_weights)
+
+        # self.head_LC.apply(segm_init_weights)
+        # self.LC.apply(segm_init_weights)
+
+
+        # if_abs_pos_embed: True (in the original script)
+        if if_abs_pos_embed:
+            trunc_normal_(self.pos_embed, std=.02)
+        # if_cls_token: True (in the original script)
+        if if_cls_token:
+            if use_double_cls_token:
+                trunc_normal_(self.cls_token_head, std=.02)
+                trunc_normal_(self.cls_token_tail, std=.02)
+            # the code is coming here
+            else:
+                trunc_normal_(self.cls_token, std=.02)
+
+        # mamba init
+        self.apply(partial(_init_weights, n_layer=depth, **(initializer_cfg if initializer_cfg is not None else {}),))
+
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        print("the code is going through allocate_inference_cache in the vision mamba")
+        return {
+            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
+            for i, layer in enumerate(self.layers)
+        }
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token", "dist_token", "cls_token_head", "cls_token_tail"}
+
+    @torch.jit.ignore()
+    def load_pretrained(self, checkpoint_path, prefix=""):
+        _load_weights(self, checkpoint_path, prefix)
+
+    # x: this is the input 224*224(tensor)
+    # inference_params: None
+    # if_random_cls_token_position：False
+    # if_random_token_rank: False
+
+    def forward_features(self, x, inference_params=None, if_random_cls_token_position=False, if_random_token_rank=False):
+        # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+        # with slight modifications to add the dist_token
+        
+        # x = self.patch_embed(x)
+
+        # x = self.CNN_layers(x)
+
+        x = self.CNN_layers(x)
+
+        # B: batch size 16
+        # M: 729 the number of patch,(27*27)
+        # D: the hidden state dimension, 384 (small-size variant), this is set by author of Vim, it is 768 in the convential Vit
+        # N: SSM dimension, SSM dimension N to 16.
+        # L: the number of blocks, we set the number of blocks L to 24
+
+        B, M, _ = x.shape
+        
+        # if_cls_token: True (in the original script)
+        if self.if_cls_token:
+
+            # self.use_double_cls_token: False (in the original script)
+            if self.use_double_cls_token:
+                cls_token_head = self.cls_token_head.expand(B, -1, -1)
+                cls_token_tail = self.cls_token_tail.expand(B, -1, -1)
+                token_position = [0, M + 1]
+                x = torch.cat((cls_token_head, x, cls_token_tail), dim=1)
+                M = x.shape[1]
+            else:
+                # self.use_middle_cls_token: True(in the original script)
+                if self.use_middle_cls_token:
+                    cls_token = self.cls_token.expand(B, -1, -1)
+                    token_position = M // 2
+                    # add cls token in the middle
+                    x = torch.cat((x[:, :token_position, :], cls_token, x[:, token_position:, :]), dim=1)
+                elif if_random_cls_token_position:
+                    cls_token = self.cls_token.expand(B, -1, -1)
+                    token_position = random.randint(0, M)
+                    x = torch.cat((x[:, :token_position, :], cls_token, x[:, token_position:, :]), dim=1)
+                    print("token_position: ", token_position)
+                else:
+                    cls_token = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+                    token_position = 0
+                    x = torch.cat((cls_token, x), dim=1)
+                M = x.shape[1]
+
+        # # if_abs_pos_embed: True (in the original script)
+        if self.if_abs_pos_embed:
+            # if new_grid_size[0] == self.patch_embed.grid_size[0] and new_grid_size[1] == self.patch_embed.grid_size[1]:
+            #     x = x + self.pos_embed
+            # else:
+            #     pos_embed = interpolate_pos_embed_online(
+            #                 self.pos_embed, self.patch_embed.grid_size, new_grid_size,0
+            #             )
+            x = x + self.pos_embed
+            x = self.pos_drop(x)
+
+
+
+        if_flip_img_sequences = False
+        if self.flip_img_sequences_ratio > 0 and (self.flip_img_sequences_ratio - random.random()) > 1e-5:
+            x = x.flip([1])
+            if_flip_img_sequences = True
+
+        # mamba impl
+        # if_bidirectional： false
+        # inference_params： None
+        residual = None
+        hidden_states = x
+        if not self.if_bidirectional:
+            for layer in self.layers:
+                
+                # here is false in the original script
+                if if_flip_img_sequences and self.if_rope:
+                    hidden_states = hidden_states.flip([1])
+                    if residual is not None:
+                        residual = residual.flip([1])
+
+                # rope about, defaule is false
+                if self.if_rope:
+                    hidden_states = self.rope(hidden_states)
+                    if residual is not None and self.if_rope_residual:
+                        residual = self.rope(residual)
+
+                # here is false in the original script
+                if if_flip_img_sequences and self.if_rope:
+                    hidden_states = hidden_states.flip([1])
+                    if residual is not None:
+                        residual = residual.flip([1])
+
+                hidden_states, residual = layer(hidden_states, residual, inference_params=inference_params)
+                # sys.exit()
+
+        else:
+            # get two layers in a single for-loop
+            for i in range(len(self.layers) // 2):
+                if self.if_rope:
+                    hidden_states = self.rope(hidden_states)
+                    if residual is not None and self.if_rope_residual:
+                        residual = self.rope(residual)
+
+                hidden_states_f, residual_f = self.layers[i * 2](
+                    hidden_states, residual, inference_params=inference_params
+                )
+                hidden_states_b, residual_b = self.layers[i * 2 + 1](
+                    hidden_states.flip([1]), None if residual == None else residual.flip([1]), inference_params=inference_params
+                )
+                hidden_states = hidden_states_f + hidden_states_b.flip([1])
+                residual = residual_f + residual_b.flip([1])
+        
+        # fused_add_norm: True
+
+        if not self.fused_add_norm:
+            if residual is None:
+                residual = hidden_states
+            else:
+                residual = residual + self.drop_path(hidden_states)
+            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
+        else:
+            # Set prenorm = False here since we don't need the residual
+            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn
+
+            hidden_states = fused_add_norm_fn(
+                self.drop_path(hidden_states),
+                self.norm_f.weight,
+                self.norm_f.bias,
+                eps=self.norm_f.eps,
+                residual=residual,
+                prenorm=False,
+                residual_in_fp32=self.residual_in_fp32,
+            )
+
+        # return only cls token if it exists
+        # if_cls_token: True (in the original script)
+        # self.use_middle_cls_token: True
+        if self.if_cls_token:
+            if self.use_double_cls_token:
+                return (hidden_states[:, token_position[0], :] + hidden_states[:, token_position[1], :]) / 2
+            else:
+                if self.use_middle_cls_token:
+                    return hidden_states[:, token_position, :]
+                elif if_random_cls_token_position:
+                    return hidden_states[:, token_position, :]
+                else:
+                    return hidden_states[:, token_position, :]
+
+        # self.final_pol_type = 'mean'
+        if self.final_pool_type == 'none':
+            return hidden_states[:, -1, :]
+        elif self.final_pool_type == 'mean':
+            return hidden_states.mean(dim=1)
+        elif self.final_pool_type == 'max':
+            return hidden_states
+        elif self.final_pool_type == 'all':
+            return hidden_states
+        else:
+            raise NotImplementedError
+
+    def forward(self, x, return_features=False, inference_params=None, if_random_cls_token_position=False, if_random_token_rank=False):
+        x = self.forward_features(x, inference_params, if_random_cls_token_position=if_random_cls_token_position, if_random_token_rank=if_random_token_rank)
+        # batch_number = x.shape[0]
+        # x = x.view(batch_number, 8, 6, 8)
+        # x = self.LC(x)
+        # x = self.head_LC(x)
+        # print("This is the shape of X (after the fully connected layer):", x.shape)
+        # return_features = False 
+        # print("This is the return feature:", return_features)
+        # if return_features:
+        #     return x
+        
+        # print("This is the shape of X (Before the fully connected layer):", x.shape)
+        x = self.head(x)
+        
+        # final_pool_type = 'mean' in original script
+        if self.final_pool_type == 'max':
+            x = x.max(dim=1)[0]
+        # sys.exit()
+        return x
+
+
+
+# below is for the vision in mamba
+@register_model
+def ecg_vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2(pretrained=False, depth=5, fused_add_norm = True, drop_path_rate = 0.1, if_divide_out = True, use_middle_cls_token = True, **kwargs):
+    model = VisionMamba(patch_size=16, stride=8, embed_dim=384, depth=depth, rms_norm=True, residual_in_fp32=True, drop_path_rate = drop_path_rate, fused_add_norm = fused_add_norm, final_pool_type='mean', if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type="v2", if_cls_token=True, if_devide_out=if_divide_out, use_middle_cls_token=use_middle_cls_token, **kwargs)
+    
+    # As a reminder:
+    print("This is whether the fused_add_norm:", fused_add_norm)
+    print("This is whether the if_divide_out:", if_divide_out)
+    print("This is whether the use_middle_cls_token:", use_middle_cls_token)
+
+
+    model.default_cfg = _cfg()
+    
+    return model
+
+# below is for the original mamba and 24 blocks
+# @register_model
+# def ecg_vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2(pretrained=False, **kwargs):
+#     model = VisionMamba(patch_size=16, stride=8, embed_dim=384, depth=24, rms_norm=True, residual_in_fp32=True, fused_add_norm=True, final_pool_type='mean', if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type="v2", if_cls_token=True, if_devide_out=True, use_middle_cls_token=True, **kwargs)
+#     model.default_cfg = _cfg()
+    
+#     return model