content: import os
  import torch
  import torch.nn.functional as F
  from timm.models.layers import DropPath, to_2tuple
  from torch import nn
  
  import torch.utils.checkpoint as checkpoint
  from functools import partial
  from einops import rearrange
  
  from .flash_attention_class import FlashAttention
  from flash_attn.modules.mlp import FusedMLP
  from flash_attn.ops.rms_norm import DropoutAddRMSNorm
  
  
  MODEL_PATH = 'your_model_path/internvl'
  _MODELS = {
      # see InternVL
      "internvl_c_13b_224px": os.path.join(MODEL_PATH, "internvl_c_13b_224px.pth"),
  }
  
  
  class CrossAttention(nn.Module):
      def __init__(
              self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
              proj_drop=0., attn_head_dim=None, out_dim=None):
          super().__init__()
          if out_dim is None:
              out_dim = dim
          self.num_heads = num_heads
          head_dim = dim // num_heads
          if attn_head_dim is not None:
              head_dim = attn_head_dim
          all_head_dim = head_dim * self.num_heads
          self.scale = qk_scale or head_dim ** -0.5
          assert all_head_dim == dim
          
          self.q = nn.Linear(dim, all_head_dim, bias=False)
          self.k = nn.Linear(dim, all_head_dim, bias=False)
          self.v = nn.Linear(dim, all_head_dim, bias=False)
          
          if qkv_bias:
              self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
              self.k_bias = nn.Parameter(torch.zeros(all_head_dim))
              self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
          else:
              self.q_bias = None
              self.k_bias = None
              self.v_bias = None
          
          self.attn_drop = nn.Dropout(attn_drop)
          self.proj = nn.Linear(all_head_dim, out_dim)
          self.proj_drop = nn.Dropout(proj_drop)
      
      def forward(self, x, k=None, v=None, return_attn=False):
          B, N, C = x.shape
          N_k = k.shape[1]
          N_v = v.shape[1]
          
          q_bias, k_bias, v_bias = None, None, None
          if self.q_bias is not None:
              q_bias = self.q_bias
              k_bias = self.k_bias
              v_bias = self.v_bias
          
          q = F.linear(input=x, weight=self.q.weight, bias=q_bias)
          q = q.reshape(B, N, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)  # (B, N_head, N_q, dim)
          
          k = F.linear(input=k, weight=self.k.weight, bias=k_bias)
          k = k.reshape(B, N_k, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)
          
          v = F.linear(input=v, weight=self.v.weight, bias=v_bias)
          v = v.reshape(B, N_v, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)
          
          q = q * self.scale
          attn = (q @ k.transpose(-2, -1))  # (B, N_head, N_q, N_k)
          
          attn = attn.softmax(dim=-1)
          attn = self.attn_drop(attn)
          
          x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
          x = self.proj(x)
          x = self.proj_drop(x)
          
          if return_attn:
              return x, attn.mean(1) # (B, n_head, n_q, C) => (B, n_q, C)
          else:
              return x, None
  
  
  class AttentiveBlock(nn.Module):
      
      def __init__(self, dim, num_heads, qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                   drop_path=0., norm_layer=nn.LayerNorm, attn_head_dim=None, out_dim=None):
          super().__init__()
          
          self.norm1_q = norm_layer(dim)
          self.norm1_k = norm_layer(dim)
          self.norm1_v = norm_layer(dim)
          self.cross_attn = CrossAttention(
              dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop,
              proj_drop=drop, attn_head_dim=attn_head_dim, out_dim=out_dim)
          
          self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
      
      def forward(self, x_q, x_kv, pos_q, pos_k, bool_masked_pos, rel_pos_bias=None, return_attn=False):
          x_q = self.norm1_q(x_q + pos_q)
          x_k = self.norm1_k(x_kv + pos_k)
          x_v = self.norm1_v(x_kv)
          x, attn = self.cross_attn(x_q, k=x_k, v=x_v, return_attn=return_attn)
          return x, attn
  
  
  class AttentionPoolingBlock(AttentiveBlock):
      
      def forward(self, x, return_attn=False):
          x_q = x.mean(1, keepdim=True)
          x_kv, pos_q, pos_k = x, 0, 0
          x, attn = super().forward(x_q, x_kv, pos_q, pos_k, bool_masked_pos=None, rel_pos_bias=None, return_attn=return_attn)
          x = x.squeeze(1)
          return x, attn
  
  
  class RMSNorm(nn.Module):
      def __init__(self, hidden_size, eps=1e-6):
          super().__init__()
          self.weight = nn.Parameter(torch.ones(hidden_size))
          self.variance_epsilon = eps
      
      def forward(self, hidden_states):
          input_dtype = hidden_states.dtype
          hidden_states = hidden_states.to(torch.float32)
          variance = hidden_states.pow(2).mean(-1, keepdim=True)
          hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
          return self.weight * hidden_states.to(input_dtype)
  
  
  class LayerScale(nn.Module):
      def __init__(self, dim, init_values=1e-5, inplace=False, force_fp32=False):
          super().__init__()
          self.inplace = inplace
          self.gamma = nn.Parameter(init_values * torch.ones(dim))
          self.force_fp32 = force_fp32
      
      @torch.cuda.amp.autocast(enabled=False)
      def forward(self, x):
          if self.force_fp32:
              output_type = x.dtype
              out = x.float().mul_(self.gamma.float()) if self.inplace else x.float() * self.gamma.float()
              return out.to(dtype=output_type)
          else:
              out = x.mul_(self.gamma) if self.inplace else x * self.gamma
              return out
  
  
  class Attention(nn.Module):
      def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., use_flash_attn=False,
                   causal=False, norm_layer=nn.LayerNorm, qk_normalization=False, use_fused_rmsnorm=False):
          super().__init__()
          assert dim % num_heads == 0, 'dim should be divisible by num_heads'
          self.num_heads = num_heads
          head_dim = dim // num_heads
          self.scale = head_dim ** -0.5
          
          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)
          
          self.use_flash_attn = use_flash_attn
          if use_flash_attn:
              self.causal = causal
              self.inner_attn = FlashAttention(attention_dropout=attn_drop)
          
          self.qk_normalization = qk_normalization
          self.q_norm = norm_layer(dim) if qk_normalization else nn.Identity()
          self.k_norm = norm_layer(dim) if qk_normalization else nn.Identity()
          self.use_fused_rmsnorm = use_fused_rmsnorm
      
      def _naive_attn(self, x):
          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.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
          
          if self.qk_normalization:
              B_, H_, N_, D_ = q.shape
              q = self.q_norm(q.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)
              k = self.k_norm(k.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)
          
          attn = ((q * self.scale) @ k.transpose(-2, -1))
          # attn = attn - attn.max(-1)[0].unsqueeze(-1)  # in case of overflow for fp16
          attn = attn.softmax(dim=-1)
          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
      
      def _flash_attn(self, x, key_padding_mask=None, need_weights=False):
          
          qkv = self.qkv(x)
          qkv = rearrange(qkv, "b s (three h d) -> b s three h d", three=3, h=self.num_heads)
          
          if self.qk_normalization:
              q, k, v = qkv.unbind(2)
              if self.use_fused_rmsnorm:
                  q = self.q_norm(q.flatten(-2, -1))[0].view(q.shape)
                  k = self.k_norm(k.flatten(-2, -1))[0].view(k.shape)
              else:
                  q = self.q_norm(q.flatten(-2, -1)).view(q.shape)
                  k = self.k_norm(k.flatten(-2, -1)).view(k.shape)
              qkv = torch.stack([q, k, v], dim=2)
          
          context, _ = self.inner_attn(
              qkv, key_padding_mask=key_padding_mask, need_weights=need_weights, causal=self.causal
          )
          outs = self.proj(rearrange(context, "b s h d -> b s (h d)"))
          outs = self.proj_drop(outs)
          return outs
      
      def forward(self, x):
          x = self._naive_attn(x) if not self.use_flash_attn else self._flash_attn(x)
          return x
  
  
  class Mlp(nn.Module):
      """ MLP as used in Vision Transformer, MLP-Mixer and related networks
      """
      
      def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU,
                   bias=True, drop=0.):
          super().__init__()
          out_features = out_features or in_features
          hidden_features = hidden_features or in_features
          bias = to_2tuple(bias)
          drop_probs = to_2tuple(drop)
          
          self.fc1 = nn.Linear(in_features, hidden_features, bias=bias[0])
          self.act = act_layer()
          self.drop1 = nn.Dropout(drop_probs[0])
          self.fc2 = nn.Linear(hidden_features, out_features, bias=bias[1])
          self.drop2 = nn.Dropout(drop_probs[1])
      
      def forward(self, x):
          x = self.fc1(x)
          x = self.act(x)
          x = self.drop1(x)
          x = self.fc2(x)
          x = self.drop2(x)
          return x
  
  
  class Block(nn.Module):
      
      def __init__(
              self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., init_values=None,
              drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_flash_attn=False, use_fused_mlp=False,
              fused_mlp_heuristic=1, with_cp=False, qk_normalization=False, layerscale_no_force_fp32=False,
              use_fused_rmsnorm=False):
          super().__init__()
          
          self.norm1 = norm_layer(dim)
          self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop,
                                use_flash_attn=use_flash_attn, causal=False, norm_layer=norm_layer,
                                qk_normalization=qk_normalization,
                                use_fused_rmsnorm=use_fused_rmsnorm)
          self.ls1 = LayerScale(dim, init_values=init_values,
                                force_fp32=(not layerscale_no_force_fp32)) if init_values else nn.Identity()
          # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
          self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
          
          self.norm2 = norm_layer(dim)
          mlp_hidden_dim = int(dim * mlp_ratio)
          if use_fused_mlp:
              self.mlp = FusedMLP(in_features=dim, hidden_features=mlp_hidden_dim, heuristic=fused_mlp_heuristic)
          else:
              self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
          self.ls2 = LayerScale(dim, init_values=init_values,
                                force_fp32=(not layerscale_no_force_fp32)) if init_values else nn.Identity()
          self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
          
          self.with_cp = with_cp
          self.use_fused_rmsnorm = use_fused_rmsnorm
      
      def forward(self, x, residual=None):
          
          def _inner_forward(x, residual=None):
              if self.use_fused_rmsnorm:
                  x, residual = self.norm1(x, residual)
                  x = self.drop_path1(self.ls1(self.attn(x)))
                  x, residual = self.norm2(x, residual)
                  x = self.drop_path2(self.ls2(self.mlp(x)))
                  return x, residual
              else:
                  assert residual is None
                  x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))
                  x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
                  return x
          
          if self.with_cp:
              return checkpoint.checkpoint(_inner_forward, x, residual)
          else:
              return _inner_forward(x, residual=residual)
  
  
  class PatchEmbed(nn.Module):
      """ 3D Image to Patch Embedding
      """
      
      def __init__(self, img_size=224, patch_size=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], img_size[1] // patch_size[1])
          self.num_patches = self.grid_size[0] * self.grid_size[1]
          self.flatten = flatten
          
          self.proj = nn.Conv3d(
              in_chans, embed_dim, 
              kernel_size=(1, patch_size[0], patch_size[1]), 
              stride=(1, patch_size[0], patch_size[1]), 
          )
          self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
      
      def forward(self, x):
          x = self.proj(x)
          if self.flatten:
              x = x.flatten(3).permute(0, 2, 3, 1)  # (N, C, T, H, W) => (N, T, H * W, C)
          x = self.norm(x)
          return x
  
  
  class InternVL_CLIP(nn.Module):
      def __init__(
              self,
              in_chans: int = 3,
              patch_size: int = 14,
              img_size: int = 224,
              qkv_bias: bool = False,
              drop_path_rate: float = 0.2,
              embed_dim: int = 3200,
              num_heads: int = 25,
              mlp_ratio: int = 4,
              init_values: float = 0.1,
              qk_normalization: bool = True,
              depth: int = 48,
              use_flash_attn: bool = True,
              use_fused_rmsnorm: bool = True,
              use_fused_mlp: bool = True,
              fused_mlp_heuristic: int = 1,
              with_cp: bool = False,
              attn_pool_num_heads: int = 16,
              clip_embed_dim: int = 768,
              layerscale_no_force_fp32: bool = True,
              # for unmasked teacher
              clip_norm_type: str = 'l2',
              return_attn: bool = True,
              clip_return_layer: int = 1,
              clip_return_interval: int = 1,
              ):
          super().__init__()
          
          assert use_flash_attn == use_fused_rmsnorm == use_fused_mlp, print(
              'use_flash_attn, use_fused_rmsnorm and use_fused_mlp should be consistent')
          
          self.use_flash_attn = use_flash_attn
          self.embed_dim = embed_dim
  
          self.clip_norm_type = clip_norm_type
          self.return_attn = return_attn
          self.return_index = []
          for i in range(clip_return_layer):
              self.return_index.append(depth - int(i * clip_return_interval) - 1)
          print(f'Normalization Type: {clip_norm_type}')
          print(f'Return Attention: {return_attn}')
          print(f'Teacher Return Interval: {self.return_index}')
          
          """ only use image encoder of InternVL """
          if use_fused_rmsnorm:
              norm_layer_for_blocks = partial(DropoutAddRMSNorm, eps=1e-6, prenorm=True)
          else:
              norm_layer_for_blocks = partial(RMSNorm, eps=1e-6)
          self.norm_layer_for_blocks = norm_layer_for_blocks
          self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
          num_patches = self.patch_embed.num_patches
          self.num_patches = num_patches
          self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)
          self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
          
          dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
          
          self.blocks = nn.ModuleList([
              Block(embed_dim, num_heads, mlp_ratio, qkv_bias=qkv_bias,
                    norm_layer=norm_layer_for_blocks,
                    drop_path=dpr[i], init_values=init_values, attn_drop=0.,
                    use_flash_attn=use_flash_attn, use_fused_mlp=use_fused_mlp,
                    fused_mlp_heuristic=fused_mlp_heuristic,
                    with_cp=with_cp,
                    qk_normalization=qk_normalization,
                    layerscale_no_force_fp32=layerscale_no_force_fp32,
                    use_fused_rmsnorm=use_fused_rmsnorm)
              for i in range(depth)])
          self.clip_projector = AttentionPoolingBlock(
              dim=embed_dim, num_heads=attn_pool_num_heads, qkv_bias=True, qk_scale=None,
              drop=0., attn_drop=0., norm_layer=partial(nn.LayerNorm, eps=1e-5), out_dim=clip_embed_dim)
      
      @property
      def dtype(self):
          return self.patch_embed.proj.weight.dtype
      
      def forward(self, image):
          x = self.patch_embed(image.type(self.dtype))
          B, T, HW, C = x.size()
          x = x.reshape(B * T, HW, C)
          
          cls_tokens = self.cls_token.expand(B * T, -1, -1)
          x = torch.cat((cls_tokens, x), dim=1)
          x = x + self.pos_embed
  
          residual = None
          z = []
          for idx, blk in enumerate(self.blocks):
              if isinstance(x, tuple) and len(x) == 2:
                  x, residual = x
              x = blk(x, residual=residual)
              # return intermediate features
              if idx in self.return_index:
                  if isinstance(x, tuple) and len(x) == 2:
                      tmp_x, tmp_residual = x
                      if residual is not None:
                          z.append(tmp_x + tmp_residual)
                  else:
                      z.append(x)
  
          if isinstance(x, tuple) and len(x) == 2:
              x, residual = x
              if residual is not None:
                  x = x + residual
  
          x, attn = self.clip_projector(x, return_attn=self.return_attn)
  
          if self.clip_norm_type == 'l2':
              # normalization of intermediate features
              z = torch.stack(z) # (K, BT, HW+1, C)
              K = z.shape[0]
              cls_tokens, z = z[:, :, :1, :], z[:, :, 1:, :]
              cls_tokens = cls_tokens.view(K, B, T, 1, C).mean(2) # (K, BT, 1, C) => (K, B, 1, C)
              z = z.reshape(K, B, T * HW, C)
              z = torch.cat((cls_tokens, z), dim=2) # (K, B, HWT+1, C)
              z = z / z.norm(dim=-1, keepdim=True)
              # normalization of final features
              x = x.view(B, T, -1).mean(1) # (BT, C) => (B, C)
              x = x / x.norm(dim=-1, keepdim=True)
          elif self.clip_norm_type == 'none':
              pass
          else:
              raise NotImplementedError
  
          if self.return_attn:
              return z, x, attn[:, 0, 1:] # (B * T, HW)
          else:
              return z, x
  
  
  def inflate_weight(weight_2d, time_dim, center=True):
      print(f'Init center: {center}')
      if center:
          weight_3d = torch.zeros(*weight_2d.shape)
          weight_3d = weight_3d.unsqueeze(2).repeat(1, 1, time_dim, 1, 1)
          middle_idx = time_dim // 2
          weight_3d[:, :, middle_idx, :, :] = weight_2d
      else:
          weight_3d = weight_2d.unsqueeze(2).repeat(1, 1, time_dim, 1, 1)
          weight_3d = weight_3d / time_dim
      return weight_3d
  
  
  def process_checkpoint(ckpt, model):
      new_ckpt = {}
      state_dict_3d = model.state_dict()
      for k, v in ckpt['module'].items():
          new_k = k
          if 'patch_embed' in new_k and new_k in state_dict_3d.keys() and v.shape != state_dict_3d[new_k].shape:
              print(new_k)
              print(f'Inflate: {k}, {v.shape} => {state_dict_3d[new_k].shape}')
              time_dim = state_dict_3d[new_k].shape[2]
              v = inflate_weight(v, time_dim)
          new_ckpt[new_k] = v
  
      # interpolate position embedding
      pos_embed_checkpoint = new_ckpt['pos_embed']
      embedding_size = pos_embed_checkpoint.shape[-1]
      num_patches = model.num_patches
      orig_size = int((pos_embed_checkpoint.shape[-2] - 1) ** 0.5)
      new_size = int(num_patches ** 0.5)
      if orig_size != new_size:
          print(f'pos_embed from {orig_size} to {new_size}')
          extra_tokens = pos_embed_checkpoint[:, :1]
          pos_tokens = pos_embed_checkpoint[:, 1:]
          pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
          pos_tokens = torch.nn.functional.interpolate(
              pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
          pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(0, 2).unsqueeze(0)
          new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
          new_ckpt['pos_embed'] = new_pos_embed
  
      return new_ckpt
  
  
  def internvl_clip_6b(
          img_size, 
          clip_norm_type='l2',
          return_attn=True,
          clip_return_layer=1,
          clip_return_interval=1
      ):
      model = InternVL_CLIP(
          img_size=img_size, 
          layerscale_no_force_fp32=False,
          clip_norm_type=clip_norm_type,
          return_attn=return_attn,
          clip_return_layer=clip_return_layer,
          clip_return_interval=clip_return_interval,
      )
      
      ckpt = torch.load(_MODELS["internvl_c_13b_224px"], map_location='cpu')
      new_ckpt = process_checkpoint(ckpt, model)
      message = model.load_state_dict(new_ckpt, strict=False)
      print(message)
      return model.eval()
  
  
  if __name__ == '__main__':
      import time
      from fvcore.nn import FlopCountAnalysis
      from fvcore.nn import flop_count_table
      import numpy as np
  
      seed = 4217
      np.random.seed(seed)
      torch.manual_seed(seed)
      torch.cuda.manual_seed(seed)
      torch.cuda.manual_seed_all(seed)
  
      num_frames = 8
      img_size = 224
      video = torch.rand(1, 3, num_frames, img_size, img_size).cuda().half()
  
      model = internvl_clip_6b(img_size).cuda().half()
      # flops = FlopCountAnalysis(model, video)
      model(video)
  
  