content: import numpy as np
  import torch
  import logging
  
  logger = logging.getLogger(__name__)
  
  # --------------------------------------------------------
  # 3D sine-cosine position embedding
  # References:
  # MVD: https://github.com/ruiwang2021/mvd/blob/main/modeling_finetune.py
  # --------------------------------------------------------
  def get_3d_sincos_pos_embed(embed_dim, grid_size, t_size, cls_token=False):
      """
      grid_size: int of the grid height and width
      t_size: int of the temporal size
      return:
      pos_embed: [t_size*grid_size*grid_size, embed_dim] or [1+t_size*grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
      """
      assert embed_dim % 4 == 0
      embed_dim_spatial = embed_dim // 4 * 3
      embed_dim_temporal = embed_dim // 4
  
      # spatial
      grid_h = np.arange(grid_size, dtype=np.float32)
      grid_w = np.arange(grid_size, dtype=np.float32)
      grid = np.meshgrid(grid_w, grid_h)  # here w goes first
      grid = np.stack(grid, axis=0)
  
      grid = grid.reshape([2, 1, grid_size, grid_size])
      pos_embed_spatial = get_2d_sincos_pos_embed_from_grid(
          embed_dim_spatial, grid
      )
  
      # temporal
      grid_t = np.arange(t_size, dtype=np.float32)
      pos_embed_temporal = get_1d_sincos_pos_embed_from_grid(
          embed_dim_temporal, grid_t
      )
  
      # concate: [T, H, W] order
      pos_embed_temporal = pos_embed_temporal[:, np.newaxis, :]
      pos_embed_temporal = np.repeat(
          pos_embed_temporal, grid_size**2, axis=1
      )  # [T, H*W, D // 4]
      pos_embed_spatial = pos_embed_spatial[np.newaxis, :, :]
      pos_embed_spatial = np.repeat(
          pos_embed_spatial, t_size, axis=0
      )  # [T, H*W, D // 4 * 3]
  
      pos_embed = np.concatenate([pos_embed_temporal, pos_embed_spatial], axis=-1)
      pos_embed = pos_embed.reshape([-1, embed_dim])  # [T*H*W, D]
  
      if cls_token:
          pos_embed = np.concatenate(
              [np.zeros([1, embed_dim]), pos_embed], axis=0
          )
      return pos_embed
  
  
  # --------------------------------------------------------
  # 2D sine-cosine position embedding
  # References:
  # Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
  # MoCo v3: https://github.com/facebookresearch/moco-v3
  # --------------------------------------------------------
  def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
      """
      grid_size: int of the grid height and width
      return:
      pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
      """
      grid_h = np.arange(grid_size, dtype=np.float32)
      grid_w = np.arange(grid_size, dtype=np.float32)
      grid = np.meshgrid(grid_w, grid_h)  # here w goes first
      grid = np.stack(grid, axis=0)
  
      grid = grid.reshape([2, 1, grid_size, grid_size])
      pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
      if cls_token:
          pos_embed = np.concatenate(
              [np.zeros([1, embed_dim]), pos_embed], axis=0
          )
      return pos_embed
  
  
  def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False):
      """
      t_size: int of the temporal size
      return:
      pos_embed: [t_size, embed_dim] or [1+t_size, embed_dim] (w/ or w/o cls_token)
      """
      grid_t = np.arange(t_size, dtype=np.float32)
      pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t)
      if cls_token:
          pos_embed = np.concatenate(
              [np.zeros([1, embed_dim]), pos_embed], axis=0
          )
      return pos_embed
  
  
  def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
      assert embed_dim % 2 == 0
  
      # use half of dimensions to encode grid_h
      emb_h = get_1d_sincos_pos_embed_from_grid(
          embed_dim // 2, grid[0]
      )  # (H*W, D/2)
      emb_w = get_1d_sincos_pos_embed_from_grid(
          embed_dim // 2, grid[1]
      )  # (H*W, D/2)
  
      emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
      return emb
  
  
  def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
      """
      embed_dim: output dimension for each position
      pos: a list of positions to be encoded: size (M,)
      out: (M, D)
      """
      assert embed_dim % 2 == 0
      omega = np.arange(embed_dim // 2, dtype=np.float32)
      omega /= embed_dim / 2.0
      omega = 1.0 / 10000**omega  # (D/2,)
  
      pos = pos.reshape(-1)  # (M,)
      out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
  
      emb_sin = np.sin(out)  # (M, D/2)
      emb_cos = np.cos(out)  # (M, D/2)
  
      emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
      return emb
  
  
  def interpolate_pos_embed(checkpoint_model, model, orig_t_size=4, pos_name='vision_encoder.pos_embed'):
      if pos_name in checkpoint_model:
          pos_embed_checkpoint = checkpoint_model[pos_name]
          embedding_size = pos_embed_checkpoint.shape[-1] # channel dim
          num_patches = model.patch_embed.num_patches # 
          num_extra_tokens = model.pos_embed.shape[-2] - num_patches # 0/1
  
          # we use 4 frames for pretraining
          new_t_size = model.T
          # height (== width) for the checkpoint position embedding
          orig_size = int(((pos_embed_checkpoint.shape[-2] - num_extra_tokens)//(orig_t_size)) ** 0.5)
          # height (== width) for the new position embedding
          new_size = int((num_patches // (new_t_size))** 0.5)
          
          # class_token and dist_token are kept unchanged
          if orig_t_size != new_t_size:
              logger.info(f"Temporal interpolate from {orig_t_size} to {new_t_size} ({pos_name})")
              extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
              # only the position tokens are interpolated
              pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
              # B, L, C -> B， T, HW, C -> BHW, C, T  (B = 1)
              pos_tokens = pos_tokens.view(1, orig_t_size, -1, embedding_size)
              pos_tokens = pos_tokens.permute(0, 2, 3, 1).reshape(-1, embedding_size, orig_t_size)
              pos_tokens = torch.nn.functional.interpolate(pos_tokens, size=new_t_size, mode='linear')
              pos_tokens = pos_tokens.view(1, -1, embedding_size, new_t_size)
              pos_tokens = pos_tokens.permute(0, 3, 1, 2).reshape(1, -1, embedding_size)
              new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
              checkpoint_model[pos_name] = new_pos_embed
              pos_embed_checkpoint = new_pos_embed
  
          # class_token and dist_token are kept unchanged
          if orig_size != new_size:
              logger.info(f"Position interpolate from {orig_size}x{orig_size} to {new_size}x{new_size} ({pos_name})")
              extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
              # only the position tokens are interpolated
              pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
              # B, L, C -> BT, H, W, C -> BT, C, H, W
              pos_tokens = pos_tokens.reshape(-1, new_t_size, orig_size, orig_size, embedding_size)
              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)
              # BT, C, H, W -> BT, H, W, C ->  B, T, H, W, C
              pos_tokens = pos_tokens.permute(0, 2, 3, 1).reshape(-1, new_t_size, new_size, new_size, embedding_size) 
              pos_tokens = pos_tokens.flatten(1, 3) # B, L, C
              new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
              checkpoint_model[pos_name] = new_pos_embed
  
  
  def interpolate_pos_embed_internvideo2(checkpoint_model, model, orig_t_size = 8):
      # interpolate position embedding
      for pos_name in ['pos_embed', 'clip_pos_embed']:
          if pos_name in checkpoint_model:
              pos_embed_checkpoint = checkpoint_model[pos_name]
              embedding_size = pos_embed_checkpoint.shape[-1] # channel dim
              num_patches = model.patch_embed.num_patches # 
              num_extra_tokens = model.pos_embed.shape[-2] - num_patches # 0/1
  
              # we use 8 frames for pretraining
              # new_t_size = args.num_frames * args.num_segments // model.patch_embed.tubelet_size
              new_t_size = model.num_frames // model.tubelet_size
              # height (== width) for the checkpoint position embedding
              orig_size = int(((pos_embed_checkpoint.shape[-2] - num_extra_tokens)//(orig_t_size)) ** 0.5)
              # height (== width) for the new position embedding
              new_size = int((num_patches // (new_t_size))** 0.5)
              
              # class_token and dist_token are kept unchanged
              if orig_t_size != new_t_size:
                  logger.info(f"Temporal interpolate from {orig_t_size} to {new_t_size} ({pos_name})")
                  extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
                  # only the position tokens are interpolated
                  pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
                  # B, L, C -> B， T, HW, C -> BHW, C, T  (B = 1)
                  pos_tokens = pos_tokens.view(1, orig_t_size, -1, embedding_size)
                  pos_tokens = pos_tokens.permute(0, 2, 3, 1).reshape(-1, embedding_size, orig_t_size)
                  pos_tokens = torch.nn.functional.interpolate(pos_tokens, size=new_t_size, mode='linear')
                  pos_tokens = pos_tokens.view(1, -1, embedding_size, new_t_size)
                  pos_tokens = pos_tokens.permute(0, 3, 1, 2).reshape(1, -1, embedding_size)
                  new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
                  checkpoint_model[pos_name] = new_pos_embed
                  pos_embed_checkpoint = new_pos_embed
  
              # class_token and dist_token are kept unchanged
              if orig_size != new_size:
                  logger.info(f"Position interpolate from {orig_size}x{orig_size} to {new_size}x{new_size} ({pos_name})")
                  extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
                  # only the position tokens are interpolated
                  pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
                  # B, L, C -> BT, H, W, C -> BT, C, H, W
                  pos_tokens = pos_tokens.reshape(-1, new_t_size, orig_size, orig_size, embedding_size)
                  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)
                  # BT, C, H, W -> BT, H, W, C ->  B, T, H, W, C
                  pos_tokens = pos_tokens.permute(0, 2, 3, 1).reshape(-1, new_t_size, new_size, new_size, embedding_size) 
                  pos_tokens = pos_tokens.flatten(1, 3) # B, L, C
                  new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
                  checkpoint_model[pos_name] = new_pos_embed
      
      if 'pos_embed_spatial' in checkpoint_model or 'pos_embed_temporal' in checkpoint_model:
          raise NotImplementedError
  
  
  def interpolate_pos_embed_internvideo2_new(checkpoint_model, model, orig_t_size = 8):
      pos_names = []
      for k in checkpoint_model.keys():
          if ('pos_embed' in k or 'clip_pos_embed' in k) and 'img_pos_embed' not in k:
              pos_names.append(k)
      
      logger.info(f"pos names list for interpolating: {pos_names}")
  
      assert len(pos_names) > 0, checkpoint_model.keys()
  
      if 'pos_embed_spatial' in checkpoint_model.keys() or 'pos_embed_temporal' in checkpoint_model.keys():
          raise NotImplementedError
      
      # interpolate position embedding
      for pos_name in pos_names:
  
          pos_embed_checkpoint = checkpoint_model[pos_name]
          embedding_size = pos_embed_checkpoint.shape[-1] # channel dim
          num_patches = model.patch_embed.num_patches # 
          num_extra_tokens = model.pos_embed.shape[-2] - num_patches # 0/1
  
          # we use 8 frames for pretraining
          # new_t_size = args.num_frames * args.num_segments // model.patch_embed.tubelet_size
          new_t_size = model.num_frames // model.tubelet_size
          # height (== width) for the checkpoint position embedding
          orig_size = int(((pos_embed_checkpoint.shape[-2] - num_extra_tokens)//(orig_t_size)) ** 0.5)
          # height (== width) for the new position embedding
          new_size = int((num_patches // (new_t_size))** 0.5)
          
          # class_token and dist_token are kept unchanged
          if orig_t_size != new_t_size:
              logger.info(f"Temporal interpolate from {orig_t_size} to {new_t_size} ({pos_name})")
              extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
              # only the position tokens are interpolated
              pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
              # B, L, C -> B， T, HW, C -> BHW, C, T  (B = 1)
              pos_tokens = pos_tokens.view(1, orig_t_size, -1, embedding_size)
              pos_tokens = pos_tokens.permute(0, 2, 3, 1).reshape(-1, embedding_size, orig_t_size)
              pos_tokens = torch.nn.functional.interpolate(pos_tokens, size=new_t_size, mode='linear')
              pos_tokens = pos_tokens.view(1, -1, embedding_size, new_t_size)
              pos_tokens = pos_tokens.permute(0, 3, 1, 2).reshape(1, -1, embedding_size)
              new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
              checkpoint_model[pos_name] = new_pos_embed
              pos_embed_checkpoint = new_pos_embed
  
          # class_token and dist_token are kept unchanged
          if orig_size != new_size:
              logger.info(f"Position interpolate from {orig_size}x{orig_size} to {new_size}x{new_size} ({pos_name})")
              extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
              # only the position tokens are interpolated
              pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
              # B, L, C -> BT, H, W, C -> BT, C, H, W
              pos_tokens = pos_tokens.reshape(-1, new_t_size, orig_size, orig_size, embedding_size)
              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)
              # BT, C, H, W -> BT, H, W, C ->  B, T, H, W, C
              pos_tokens = pos_tokens.permute(0, 2, 3, 1).reshape(-1, new_t_size, new_size, new_size, embedding_size) 
              pos_tokens = pos_tokens.flatten(1, 3) # B, L, C
              new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
              checkpoint_model[pos_name] = new_pos_embed