code: def create_model(
      model_name: str = None,
      pretrained: bool = False,
      weights_path: str = None,
      device: str = None,
      **kwargs,
  ) -> VitModel:
      """
      Create ViT model.
      Args:
          model_name (`str`):
              ViT model name.
          pretrained (`bool`):
              Whether it is a pretrained model.
          weights_path (`str`):
              If pretrained, you can use your local weights.
          device (`str`):
              Cuda or cpu.
  
      Returns:
          VitModel
      """
      if device is None:
          device = "cuda" if torch.cuda.is_available() else "cpu"
  
      if model_name is None:
          if pretrained:
              raise AssertionError(
                  "Fail to load pretrained model: no model name is specified."
              )
          model = VitModel(**kwargs)
      else:
          configs = get_configs(model_name)
          if "url" in configs:
              url = configs["url"]
              configs.pop("url")
  
          model = VitModel(**configs)
  
          if pretrained:
              if weights_path:
                  state_dict = torch.load(weights_path)
              elif url:
                  state_dict = model_zoo.load_url(
                      url, map_location=torch.device(device)
                  )
              else:
                  raise AssertionError(
                      "No model weights url or path is provided."
                  )
              model.load_state_dict(state_dict, strict=False)
  
      model.eval()
      return model
  
file_path: towhee/towhee/models/vit/vit.py
imports: import torch
  from torch import nn
  from torch.utils import model_zoo
  from functools import partial
  from collections import OrderedDict
  from towhee.models.utils.init_vit_weights import init_vit_weights
  from towhee.models.layers.patch_embed2d import PatchEmbed2D
  from .vit_utils import get_configs
  from .vit_block import Block
  from towhee.models.layers.layers_with_relprop import LayerNorm, GELU, Linear, IndexSelect, Add
label_desc: Use image node_id 2 for calling a global function with example usage: create_model(model_name, pretrained, weights_path, device) and returns: model