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
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
towhee/models/vit/vit.py