class_docstr: FrozenInTime model
  Args:
      img_size (`int`):
          The image or video height.
      patch_size (`int`):
          Patch height (equal to width).
      in_chans (`int`):
          Number of image channel.
      num_frames (`int`):
          Maximum number of frames expected as input.
      is_pretrained (`bool`):
          If true, then load the frozen pretrained model.
      weights_path (`str`):
          The frozen pretrained model path, if is_pretrained is true, and this not None, then load from this path,
          otherwise, will load the frozen model from an url.
      attention_style (`str`):
          Default is "frozen_in_time", if is "bridge_former", then is the BridgeFormer model.
      projection_dim (`int`):
          The output dim.
      video_pretrained_model (`str`):
          Default is "vit_base_16x224", if is_pretrained is false, then must load the pretrained vit parameters.
      video_is_load_pretrained (`bool`):
          If true, then will load vit pretrained weight.
      video_model_type (`str`) :
          Default is "SpaceTimeTransformer", the frozen model name, no need to change.
      text_pretrained_model (`str`):
          Default is "distilbert-base-uncased", the pretrained model for text encoder.
      text_is_load_pretrained (`bool`):
          If false, then will not load pretrained model for text encoder.
      projection (`str`):
          A layer.
      load_temporal_fix (`str`):
          If is_pretrained is true, this will use in function "_inflate_positional_embeds".
      device (`str`):
          "cpu" or "cuda".
      pretrained_url (`str`):
          If is_pretrained is true and weights_path is None, will load frozen pretrained model from this url.
class_docstr_tok: ["FrozenInTime","model","Args",":","img_size","(","`","int","`",")",":","The","image","or","video","height",".","patch_size","(","`","int","`",")",":","Patch","height","(","equal","to","width",")",".","in_chans","(","`","int","`",")",":","Number","of","image","channel",".","num_frames","(","`","int","`",")",":","Maximum","number","of","frames","expected","as","input",".","is_pretrained","(","`","bool","`",")",":","If","true",",","then","load","the","frozen","pretrained","model",".","weights_path","(","`","str","`",")",":","The","frozen","pretrained","model","path",",","if","is_pretrained","is","true",",","and","this","not","None",",","then","load","from","this","path",",","otherwise",",","will","load","the","frozen","model","from","an","url",".","attention_style","(","`","str","`",")",":","Default","is","``","frozen_in_time","''",",","if","is","``","bridge_former","''",",","then","is","the","BridgeFormer","model",".","projection_dim","(","`","int","`",")",":","The","output","dim",".","video_pretrained_model","(","`","str","`",")",":","Default","is","``","vit_base_16x224","''",",","if","is_pretrained","is","false",",","then","must","load","the","pretrained","vit","parameters",".","video_is_load_pretrained","(","`","bool","`",")",":","If","true",",","then","will","load","vit","pretrained","weight",".","video_model_type","(","`","str","`",")",":","Default","is","``","SpaceTimeTransformer","''",",","the","frozen","model","name",",","no","need","to","change",".","text_pretrained_model","(","`","str","`",")",":","Default","is","``","distilbert-base-uncased","''",",","the","pretrained","model","for","text","encoder",".","text_is_load_pretrained","(","`","bool","`",")",":","If","false",",","then","will","not","load","pretrained","model","for","text","encoder",".","projection","(","`","str","`",")",":","A","layer",".","load_temporal_fix","(","`","str","`",")",":","If","is_pretrained","is","true",",","this","will","use","in","function","``","_inflate_positional_embeds","''",".","device","(","`","str","`",")",":","``","cpu","''","or","``","cuda","''",".","pretrained_url","(","`","str","`",")",":","If","is_pretrained","is","true","and","weights_path","is","None",",","will","load","frozen","pretrained","model","from","this","url","."]
code: def __init__(
      self,
      img_size: int = 224,
      patch_size: int = 16,
      in_chans: int = 3,
      num_frames: int = 4,
      is_pretrained: bool = False,
      weights_path: str = "",
      attention_style: str = "frozen_in_time",
      projection_dim: int = 256,
      video_pretrained_model: str = "vit_base_16x224",
      video_is_load_pretrained: bool = False,
      video_model_type: str = "SpaceTimeTransformer",
      text_pretrained_model: str = "distilbert-base-uncased",
      text_is_load_pretrained: bool = False,
      projection: str = "minimal",
      load_temporal_fix: str = "zeros",
      device: str = "cpu",
      pretrained_url: str = "https://www.robots.ox.ac.uk/~maxbain/frozen-in-time/models/cc-webvid2m-4f_stformer_b_16_224.pth.tar",
  ):
      super().__init__()
  
      self.text_pretrained_model = text_pretrained_model
      self.video_model_type = video_model_type
      self.num_frames = num_frames
      self.load_temporal_fix = load_temporal_fix
      self.attention_style = attention_style
      if text_is_load_pretrained:
          # load text pretrained model
          self.text_model = AutoModel.from_pretrained(
              self.text_pretrained_model
          )
      else:
          self.config = AutoConfig.from_pretrained(
              self.text_pretrained_model
          )
          self.text_model = AutoModel.from_config(config=self.config)
  
      self.text_model.train()
  
      # load video pretrained model
      if video_model_type == "SpaceTimeTransformer":
          model = SpaceTimeTransformer(
              img_size=img_size,
              patch_size=patch_size,
              in_chans=in_chans,
              num_frames=num_frames,
              time_init="zeros",
              attention_style=attention_style,
          )
          model.head = nn.Identity()
          model.pre_logits = nn.Identity()
          ftr_dim = model.embed_dim
          # if not load pretrained model of frozen in time, then must load pretrained vit model to initial
          if (not is_pretrained) and video_is_load_pretrained:
              if video_pretrained_model == "vit_base_16x224":
                  vit_model = vit.create_model(
                      model_name=video_pretrained_model, pretrained=True
                  )
              else:
                  raise NotImplementedError
              vit_checkpoint = vit_model.state_dict()
              model.load_state_dict(vit_checkpoint, strict=False)
          self.video_model = model
      else:
          raise NotImplementedError(
              f"{video_model_type} not implemented"
          )
      if self.attention_style == "bridge_former":
          self.video_model.fc = nn.Identity()
      # Project to a common embedding
      if projection == "minimal":
          txt_proj = nn.Sequential(
              nn.ReLU(),
              nn.Linear(
                  self.text_model.config.hidden_size, projection_dim
              ),
          )
  
          vid_proj = nn.Sequential(nn.Linear(ftr_dim, projection_dim))
      elif projection == "":
          txt_proj = nn.Identity()
          vid_proj = nn.Identity()
      else:
          raise NotImplementedError
      if self.attention_style == "bridge_former":
          self.text_proj = txt_proj
      else:
          self.txt_proj = txt_proj
      self.vid_proj = vid_proj
  
      if is_pretrained:
          if weights_path not in ["", None]:
              checkpoint = torch.load(weights_path, map_location=device)
              state_dict = checkpoint["state_dict"]
          else:
              if self.attention_style == "bridge_former":
                  raise NotImplementedError(
                      f"{self.attention_style} not have the weights url, "
                      f"need give the weights path"
                  )
  
              checkpoint = model_zoo.load_url(
                  pretrained_url, map_location=torch.device(device)
              )
              state_dict = checkpoint["state_dict"]
  
          new_state_dict = state_dict_data_parallel_fix(
              state_dict, self.state_dict()
          )
          if self.attention_style == "bridge_former":
              new_state_dict = remove_bridge_module_state_dic(
                  new_state_dict, self.state_dict()
              )
          new_state_dict = self._inflate_positional_embeds(
              new_state_dict
          )
          self.load_state_dict(new_state_dict, strict=True)
  
file_path: towhee/towhee/models/frozen_in_time/frozen_in_time.py
imports: import os
  import torch
  from torch import nn
  import torch.nn.functional
  from torch.utils import model_zoo
  from transformers import AutoModel, AutoConfig
  from towhee.models.frozen_in_time.frozen_video_transformer import SpaceTimeTransformer
  from towhee.models import vit
  from towhee.models.frozen_in_time.frozen_utils import sim_matrix, state_dict_data_parallel_fix, remove_bridge_module_state_dic
  import logging
label_desc: Use image node_id 1 to create a new FrozenInTime object from inherited base classes: nn with example: obj = FrozenInTime(img_size, patch_size, in_chans, num_frames, is_pretrained, weights_path, attention_style, projection_dim, video_pretrained_model, video_is_load_pretrained, video_model_type, text_pretrained_model, text_is_load_pretrained, projection, load_temporal_fix, device, pretrained_url)