





import argparse
import math
import os
from copy import deepcopy

import torch
from audio_diffusion.models import DiffusionAttnUnetaD
from diffusion import sampling
from torch import nn

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel


_snake_case  :   str						   =  {
    'gwf-440k': {
        'url': 'https://model-server.zqevans2.workers.dev/gwf-440k.ckpt',
        'sample_rate': 48000,
        'sample_size': 65536,
    },
    'jmann-small-190k': {
        'url': 'https://model-server.zqevans2.workers.dev/jmann-small-190k.ckpt',
        'sample_rate': 48000,
        'sample_size': 65536,
    },
    'jmann-large-580k': {
        'url': 'https://model-server.zqevans2.workers.dev/jmann-large-580k.ckpt',
        'sample_rate': 48000,
        'sample_size': 131072,
    },
    'maestro-uncond-150k': {
        'url': 'https://model-server.zqevans2.workers.dev/maestro-uncond-150k.ckpt',
        'sample_rate': 16000,
        'sample_size': 65536,
    },
    'unlocked-uncond-250k': {
        'url': 'https://model-server.zqevans2.workers.dev/unlocked-uncond-250k.ckpt',
        'sample_rate': 16000,
        'sample_size': 65536,
    },
    'honk-140k': {
        'url': 'https://model-server.zqevans2.workers.dev/honk-140k.ckpt',
        'sample_rate': 16000,
        'sample_size': 65536,
    },
}

def   a_   (			lowerCAmelCase_				:      Any,     lowerCAmelCase_				:      Dict     ):
   return torch.atana(lowerCAmelCase_,     lowerCAmelCase_     ) / math.pi * 2

def   a_   (			lowerCAmelCase_				:      List[Any]     ):
   __lowerCAmelCase							     =     torch.sin(t * math.pi / 2     ) ** 2
   __lowerCAmelCase							     =     (1 - sigma**2) ** 0.5
   return alpha_sigma_to_t(lowerCAmelCase_,     lowerCAmelCase_     )



class     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


  pass





class     _UpperCAmelCase   (   nn.Module   ):

  """simple docstring"""


  def __init__(    self			:		Optional[int] ,  lowerCAmelCase_			:		int					)							->						List[str]:
     super().__init__()

     __lowerCAmelCase							     =     DiffusionAttnUnetaD(lowerCAmelCase_ ,  n_attn_layers=4					)
     __lowerCAmelCase							     =     deepcopy(self.diffusion					)
     __lowerCAmelCase							     =     torch.quasirandom.SobolEngine(1 ,  scramble=lowerCAmelCase_					)





def   a_   (			lowerCAmelCase_				:      str     ):
   __lowerCAmelCase							     =     MODELS_MAP[model_name]['url']
   os.system(F"""wget {url} ./"""     )

   return F"""./{model_name}.ckpt"""


_snake_case  :   int						   =  {
    '1': 'resnets.0',
    '2': 'attentions.0',
    '3': 'resnets.1',
    '4': 'attentions.1',
    '5': 'resnets.2',
    '6': 'attentions.2',
}
_snake_case  :   Optional[Any]						   =  {
    '8': 'resnets.0',
    '9': 'attentions.0',
    '10': 'resnets.1',
    '11': 'attentions.1',
    '12': 'resnets.2',
    '13': 'attentions.2',
}
_snake_case  :   int						   =  {
    '1': 'resnets.0',
    '2': 'attentions.0',
    '3': 'resnets.1',
    '4': 'attentions.1',
    '5': 'resnets.2',
    '6': 'attentions.2',
    '8': 'resnets.3',
    '9': 'attentions.3',
    '10': 'resnets.4',
    '11': 'attentions.4',
    '12': 'resnets.5',
    '13': 'attentions.5',
}
_snake_case  :   List[str]						   =  {
    '0': 'resnets.0',
    '1': 'resnets.1',
    '2': 'resnets.2',
    '4': 'resnets.0',
    '5': 'resnets.1',
    '6': 'resnets.2',
}

_snake_case  :   List[str]						   =  {
    'skip': 'conv_skip',
    'main.0': 'conv_1',
    'main.1': 'group_norm_1',
    'main.3': 'conv_2',
    'main.4': 'group_norm_2',
}

_snake_case  :   List[Any]						   =  {
    'norm': 'group_norm',
    'qkv_proj': ['query', 'key', 'value'],
    'out_proj': ['proj_attn'],
}

def   a_   (			lowerCAmelCase_				:      Tuple     ):
   if name.startswith('skip'     ):
      return name.replace('skip',     RES_CONV_MAP['skip']     )

   # name has to be of format main.{digit}
   if not name.startswith('main.'     ):
      raise ValueError(F"""ResConvBlock error with {name}"""     )

   return name.replace(name[:6],     RES_CONV_MAP[name[:6]]     )

def   a_   (			lowerCAmelCase_				:      Union[str, Any]     ):
   for key, value in ATTN_MAP.items():
      if name.startswith(lowerCAmelCase_     ) and not isinstance(lowerCAmelCase_,     lowerCAmelCase_     ):
         return name.replace(lowerCAmelCase_,     lowerCAmelCase_     )
      elif name.startswith(lowerCAmelCase_     ):
         return [name.replace(lowerCAmelCase_,     lowerCAmelCase_     ) for v in value]
   raise ValueError(F"""Attn error with {name}"""     )

def   a_   (			lowerCAmelCase_				:      List[Any],     lowerCAmelCase_				:      Optional[int]=13     ):
   __lowerCAmelCase							     =     input_string

   if string.split('.'     )[0] == "timestep_embed":
      return string.replace('timestep_embed',     'time_proj'     )

   __lowerCAmelCase							     =     0
   if string.startswith('net.3.'     ):
      depth += 1
      __lowerCAmelCase							     =     string[6:]
   elif string.startswith('net.'     ):
      __lowerCAmelCase							     =     string[4:]

   while string.startswith('main.7.'     ):
      depth += 1
      __lowerCAmelCase							     =     string[7:]

   if string.startswith('main.'     ):
      __lowerCAmelCase							     =     string[5:]

   # mid block
   if string[:2].isdigit():
      __lowerCAmelCase							     =     string[:2]
      __lowerCAmelCase							     =     string[2:]
   else:
      __lowerCAmelCase							     =     string[0]
      __lowerCAmelCase							     =     string[1:]

   if depth == max_depth:
      __lowerCAmelCase							     =     MID_NUM_TO_LAYER[layer_num]
      __lowerCAmelCase							     =     'mid_block'
   elif depth > 0 and int(lowerCAmelCase_     ) < 7:
      __lowerCAmelCase							     =     DOWN_NUM_TO_LAYER[layer_num]
      __lowerCAmelCase							     =     F"""down_blocks.{depth}"""
   elif depth > 0 and int(lowerCAmelCase_     ) > 7:
      __lowerCAmelCase							     =     UP_NUM_TO_LAYER[layer_num]
      __lowerCAmelCase							     =     F"""up_blocks.{max_depth - depth - 1}"""
   elif depth == 0:
      __lowerCAmelCase							     =     DEPTH_0_TO_LAYER[layer_num]
      __lowerCAmelCase							     =     F"""up_blocks.{max_depth - 1}""" if int(lowerCAmelCase_     ) > 3 else 'down_blocks.0'

   if not string_left.startswith('.'     ):
      raise ValueError(F"""Naming error with {input_string} and string_left: {string_left}."""     )

   __lowerCAmelCase							     =     string_left[1:]

   if "resnets" in new_layer:
      __lowerCAmelCase							     =     convert_resconv_naming(lowerCAmelCase_     )
   elif "attentions" in new_layer:
      __lowerCAmelCase							     =     convert_attn_naming(lowerCAmelCase_     )
      __lowerCAmelCase							     =     new_string_left

   if not isinstance(lowerCAmelCase_,     lowerCAmelCase_     ):
      __lowerCAmelCase							     =     prefix + '.' + new_layer + '.' + string_left
   else:
      __lowerCAmelCase							     =     [prefix + '.' + new_layer + '.' + s for s in string_left]
   return new_string

def   a_   (			lowerCAmelCase_				:      Dict     ):
   __lowerCAmelCase							     =     {}
   for k, v in state_dict.items():
      if k.endswith('kernel'     ):
         # up- and downsample layers, don't have trainable weights
         continue

      __lowerCAmelCase							     =     rename(lowerCAmelCase_     )

      # check if we need to transform from Conv => Linear for attention
      if isinstance(lowerCAmelCase_,     lowerCAmelCase_     ):
         __lowerCAmelCase							     =     transform_conv_attns(lowerCAmelCase_,     lowerCAmelCase_,     lowerCAmelCase_     )
      else:
         __lowerCAmelCase							     =     v

   return new_state_dict

def   a_   (			lowerCAmelCase_				:      List[Any],     lowerCAmelCase_				:      Union[str, Any],     lowerCAmelCase_				:      Tuple     ):
   if len(lowerCAmelCase_     ) == 1:
      if len(v.shape     ) == 3:
         # weight
         __lowerCAmelCase							     =     v[:, :, 0]
      else:
         # bias
         __lowerCAmelCase							     =     v
   else:
      # qkv matrices
      __lowerCAmelCase							     =     v.shape[0]
      __lowerCAmelCase							     =     trippled_shape // 3
      for i in range(3     ):
         if len(v.shape     ) == 3:
            __lowerCAmelCase							     =     v[i * single_shape : (i + 1) * single_shape, :, 0]
         else:
            __lowerCAmelCase							     =     v[i * single_shape : (i + 1) * single_shape]
   return new_state_dict

def   a_   (			lowerCAmelCase_				:      str     ):
   __lowerCAmelCase							     =     torch.device('cuda' if torch.cuda.is_available() else 'cpu'     )

   __lowerCAmelCase							     =     args.model_path.split('/'     )[-1].split('.'     )[0]
   if not os.path.isfile(args.model_path     ):
      assert (
          model_name == args.model_path
      ), F"""Make sure to provide one of the official model names {MODELS_MAP.keys()}"""
      __lowerCAmelCase							     =     download(lowerCAmelCase_     )

   __lowerCAmelCase							     =     MODELS_MAP[model_name]['sample_rate']
   __lowerCAmelCase							     =     MODELS_MAP[model_name]['sample_size']

   __lowerCAmelCase							     =     Object()
   __lowerCAmelCase							     =     sample_size
   __lowerCAmelCase							     =     sample_rate
   __lowerCAmelCase							     =     0

   __lowerCAmelCase							     =     UNetaDModel(sample_size=lowerCAmelCase_,     sample_rate=lowerCAmelCase_     )
   __lowerCAmelCase							     =     diffusers_model.state_dict()

   __lowerCAmelCase							     =     DiffusionUncond(lowerCAmelCase_     )
   orig_model.load_state_dict(torch.load(args.model_path,     map_location=lowerCAmelCase_     )['state_dict']     )
   __lowerCAmelCase							     =     orig_model.diffusion_ema.eval()
   __lowerCAmelCase							     =     orig_model.state_dict()
   __lowerCAmelCase							     =     rename_orig_weights(lowerCAmelCase_     )

   __lowerCAmelCase							     =     set(renamed_state_dict.keys()     ) - set(diffusers_state_dict.keys()     )
   __lowerCAmelCase							     =     set(diffusers_state_dict.keys()     ) - set(renamed_state_dict.keys()     )

   assert len(lowerCAmelCase_     ) == 0, F"""Problem with {renamed_minus_diffusers}"""
   assert all(k.endswith('kernel'     ) for k in list(lowerCAmelCase_     )     ), F"""Problem with {diffusers_minus_renamed}"""

   for key, value in renamed_state_dict.items():
      assert (
          diffusers_state_dict[key].squeeze().shape == value.squeeze().shape
      ), F"""Shape for {key} doesn't match. Diffusers: {diffusers_state_dict[key].shape} vs. {value.shape}"""
      if key == "time_proj.weight":
         __lowerCAmelCase							     =     value.squeeze()

      __lowerCAmelCase							     =     value

   diffusers_model.load_state_dict(lowerCAmelCase_     )

   __lowerCAmelCase							     =     100
   __lowerCAmelCase							     =     33

   __lowerCAmelCase							     =     IPNDMScheduler(num_train_timesteps=lowerCAmelCase_     )

   __lowerCAmelCase							     =     torch.manual_seed(lowerCAmelCase_     )
   __lowerCAmelCase							     =     torch.randn([1, 2, config.sample_size],     generator=lowerCAmelCase_     ).to(lowerCAmelCase_     )

   __lowerCAmelCase							     =     torch.linspace(1,     0,     steps + 1,     device=lowerCAmelCase_     )[:-1]
   __lowerCAmelCase							     =     get_crash_schedule(lowerCAmelCase_     )

   __lowerCAmelCase							     =     DanceDiffusionPipeline(unet=lowerCAmelCase_,     scheduler=lowerCAmelCase_     )

   __lowerCAmelCase							     =     torch.manual_seed(33     )
   __lowerCAmelCase							     =     pipe(num_inference_steps=lowerCAmelCase_,     generator=lowerCAmelCase_     ).audios

   __lowerCAmelCase							     =     sampling.iplms_sample(lowerCAmelCase_,     lowerCAmelCase_,     lowerCAmelCase_,     {}     )
   __lowerCAmelCase							     =     generated.clamp(-1,     1     )

   __lowerCAmelCase							     =     (generated - audio).abs().sum()
   __lowerCAmelCase							     =     (generated - audio).abs().max()

   if args.save:
      pipe.save_pretrained(args.checkpoint_path     )

   print('Diff sum',     lowerCAmelCase_     )
   print('Diff max',     lowerCAmelCase_     )

   assert diff_max < 1E-3, F"""Diff max: {diff_max} is too much :-/"""

   print(F"""Conversion for {model_name} successful!"""     )


if __name__ == "__main__":
     _snake_case  :   Union[str, Any]						   =  argparse.ArgumentParser()

     parser.add_argument('--model_path', default=None, type=str, required=True, help='Path to the model to convert.')
     parser.add_argument(
         '--save', default=True, type=bool, required=False, help='Whether to save the converted model or not.'
     )
     parser.add_argument('--checkpoint_path', default=None, type=str, required=True, help='Path to the output model.')
     _snake_case  :   str						   =  parser.parse_args()

     main(args)








import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DeiTImageProcessor, ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
_snake_case  :   List[str]						   =  logging.get_logger(__name__)

def   a_   (			lowerCAmelCase_				:      Optional[Any],     lowerCAmelCase_				:      Union[str, Any]=False     ):
   __lowerCAmelCase							     =     []
   for i in range(config.num_hidden_layers     ):
      # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
      rename_keys.append((F"""blocks.{i}.norm1.weight""", F"""vit.encoder.layer.{i}.layernorm_before.weight""")     )
      rename_keys.append((F"""blocks.{i}.norm1.bias""", F"""vit.encoder.layer.{i}.layernorm_before.bias""")     )
      rename_keys.append((F"""blocks.{i}.attn.proj.weight""", F"""vit.encoder.layer.{i}.attention.output.dense.weight""")     )
      rename_keys.append((F"""blocks.{i}.attn.proj.bias""", F"""vit.encoder.layer.{i}.attention.output.dense.bias""")     )
      rename_keys.append((F"""blocks.{i}.norm2.weight""", F"""vit.encoder.layer.{i}.layernorm_after.weight""")     )
      rename_keys.append((F"""blocks.{i}.norm2.bias""", F"""vit.encoder.layer.{i}.layernorm_after.bias""")     )
      rename_keys.append((F"""blocks.{i}.mlp.fc1.weight""", F"""vit.encoder.layer.{i}.intermediate.dense.weight""")     )
      rename_keys.append((F"""blocks.{i}.mlp.fc1.bias""", F"""vit.encoder.layer.{i}.intermediate.dense.bias""")     )
      rename_keys.append((F"""blocks.{i}.mlp.fc2.weight""", F"""vit.encoder.layer.{i}.output.dense.weight""")     )
      rename_keys.append((F"""blocks.{i}.mlp.fc2.bias""", F"""vit.encoder.layer.{i}.output.dense.bias""")     )

   # projection layer + position embeddings
   rename_keys.extend(
       [
           ('cls_token', 'vit.embeddings.cls_token'),
           ('patch_embed.proj.weight', 'vit.embeddings.patch_embeddings.projection.weight'),
           ('patch_embed.proj.bias', 'vit.embeddings.patch_embeddings.projection.bias'),
           ('pos_embed', 'vit.embeddings.position_embeddings'),
       ]     )

   if base_model:
      # layernorm + pooler
      rename_keys.extend(
          [
              ('norm.weight', 'layernorm.weight'),
              ('norm.bias', 'layernorm.bias'),
              ('pre_logits.fc.weight', 'pooler.dense.weight'),
              ('pre_logits.fc.bias', 'pooler.dense.bias'),
          ]     )

      # if just the base model, we should remove "vit" from all keys that start with "vit"
      __lowerCAmelCase							     =     [(pair[0], pair[1][4:]) if pair[1].startswith('vit'     ) else pair for pair in rename_keys]
   else:
      # layernorm + classification head
      rename_keys.extend(
          [
              ('norm.weight', 'vit.layernorm.weight'),
              ('norm.bias', 'vit.layernorm.bias'),
              ('head.weight', 'classifier.weight'),
              ('head.bias', 'classifier.bias'),
          ]     )

   return rename_keys

def   a_   (			lowerCAmelCase_				:      List[Any],     lowerCAmelCase_				:      Dict,     lowerCAmelCase_				:      int=False     ):
   for i in range(config.num_hidden_layers     ):
      if base_model:
         __lowerCAmelCase							     =     ''
      else:
         __lowerCAmelCase							     =     'vit.'
      # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
      __lowerCAmelCase							     =     state_dict.pop(F"""blocks.{i}.attn.qkv.weight"""     )
      __lowerCAmelCase							     =     state_dict.pop(F"""blocks.{i}.attn.qkv.bias"""     )
      # next, add query, keys and values (in that order) to the state dict
      __lowerCAmelCase							     =     in_proj_weight[
          : config.hidden_size, :
      ]
      __lowerCAmelCase							     =     in_proj_bias[: config.hidden_size]
      __lowerCAmelCase							     =     in_proj_weight[
          config.hidden_size : config.hidden_size * 2, :
      ]
      __lowerCAmelCase							     =     in_proj_bias[
          config.hidden_size : config.hidden_size * 2
      ]
      __lowerCAmelCase							     =     in_proj_weight[
          -config.hidden_size :, :
      ]
      __lowerCAmelCase							     =     in_proj_bias[-config.hidden_size :]

def   a_   (			lowerCAmelCase_				:      Optional[int]     ):
   __lowerCAmelCase							     =     ['head.weight', 'head.bias']
   for k in ignore_keys:
      state_dict.pop(lowerCAmelCase_,     lowerCAmelCase_     )

def   a_   (			lowerCAmelCase_				:      Optional[Any],     lowerCAmelCase_				:      Tuple,     lowerCAmelCase_				:      Optional[int]     ):
   __lowerCAmelCase							     =     dct.pop(lowerCAmelCase_     )
   __lowerCAmelCase							     =     val

def   a_   (			):
   __lowerCAmelCase							     =     'http://images.cocodataset.org/val2017/000000039769.jpg'
   __lowerCAmelCase							     =     Image.open(requests.get(lowerCAmelCase_,     stream=lowerCAmelCase_     ).raw     )
   return im

@torch.no_grad()
def   a_   (			lowerCAmelCase_				:      List[str],     lowerCAmelCase_				:      Union[str, Any]     ):
   __lowerCAmelCase							     =     ViTConfig()
   __lowerCAmelCase							     =     False
   # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
   if vit_name[-5:] == "in21k":
      __lowerCAmelCase							     =     True
      __lowerCAmelCase							     =     int(vit_name[-12:-10]     )
      __lowerCAmelCase							     =     int(vit_name[-9:-6]     )
   else:
      __lowerCAmelCase							     =     1000
      __lowerCAmelCase							     =     'huggingface/label-files'
      __lowerCAmelCase							     =     'imagenet-1k-id2label.json'
      __lowerCAmelCase							     =     json.load(open(hf_hub_download(lowerCAmelCase_,     lowerCAmelCase_,     repo_type='dataset'     ),     'r'     )     )
      __lowerCAmelCase							     =     {int(lowerCAmelCase_     ): v for k, v in idalabel.items()}
      __lowerCAmelCase							     =     idalabel
      __lowerCAmelCase							     =     {v: k for k, v in idalabel.items()}
      __lowerCAmelCase							     =     int(vit_name[-6:-4]     )
      __lowerCAmelCase							     =     int(vit_name[-3:]     )
   # size of the architecture
   if "deit" in vit_name:
      if vit_name[9:].startswith('tiny'     ):
         __lowerCAmelCase							     =     192
         __lowerCAmelCase							     =     768
         __lowerCAmelCase							     =     12
         __lowerCAmelCase							     =     3
      elif vit_name[9:].startswith('small'     ):
         __lowerCAmelCase							     =     384
         __lowerCAmelCase							     =     1536
         __lowerCAmelCase							     =     12
         __lowerCAmelCase							     =     6
      else:
         pass
   else:
      if vit_name[4:].startswith('small'     ):
         __lowerCAmelCase							     =     768
         __lowerCAmelCase							     =     2304
         __lowerCAmelCase							     =     8
         __lowerCAmelCase							     =     8
      elif vit_name[4:].startswith('base'     ):
         pass
      elif vit_name[4:].startswith('large'     ):
         __lowerCAmelCase							     =     1024
         __lowerCAmelCase							     =     4096
         __lowerCAmelCase							     =     24
         __lowerCAmelCase							     =     16
      elif vit_name[4:].startswith('huge'     ):
         __lowerCAmelCase							     =     1280
         __lowerCAmelCase							     =     5120
         __lowerCAmelCase							     =     32
         __lowerCAmelCase							     =     16

    # load original model from timm
   __lowerCAmelCase							     =     timm.create_model(lowerCAmelCase_,     pretrained=lowerCAmelCase_     )
   timm_model.eval()

   # load state_dict of original model, remove and rename some keys
   __lowerCAmelCase							     =     timm_model.state_dict()
   if base_model:
      remove_classification_head_(lowerCAmelCase_     )
   __lowerCAmelCase							     =     create_rename_keys(lowerCAmelCase_,     lowerCAmelCase_     )
   for src, dest in rename_keys:
      rename_key(lowerCAmelCase_,     lowerCAmelCase_,     lowerCAmelCase_     )
   read_in_q_k_v(lowerCAmelCase_,     lowerCAmelCase_,     lowerCAmelCase_     )

   # load HuggingFace model
   if vit_name[-5:] == "in21k":
      __lowerCAmelCase							     =     ViTModel(lowerCAmelCase_     ).eval()
   else:
      __lowerCAmelCase							     =     ViTForImageClassification(lowerCAmelCase_     ).eval()
   model.load_state_dict(lowerCAmelCase_     )

   # Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
   if "deit" in vit_name:
      __lowerCAmelCase							     =     DeiTImageProcessor(size=config.image_size     )
   else:
      __lowerCAmelCase							     =     ViTImageProcessor(size=config.image_size     )
   __lowerCAmelCase							     =     image_processor(images=prepare_img(),     return_tensors='pt'     )
   __lowerCAmelCase							     =     encoding['pixel_values']
   __lowerCAmelCase							     =     model(lowerCAmelCase_     )

   if base_model:
      __lowerCAmelCase							     =     timm_model.forward_features(lowerCAmelCase_     )
      assert timm_pooled_output.shape == outputs.pooler_output.shape
      assert torch.allclose(lowerCAmelCase_,     outputs.pooler_output,     atol=1E-3     )
   else:
      __lowerCAmelCase							     =     timm_model(lowerCAmelCase_     )
      assert timm_logits.shape == outputs.logits.shape
      assert torch.allclose(lowerCAmelCase_,     outputs.logits,     atol=1E-3     )

   Path(lowerCAmelCase_     ).mkdir(exist_ok=lowerCAmelCase_     )
   print(F"""Saving model {vit_name} to {pytorch_dump_folder_path}"""     )
   model.save_pretrained(lowerCAmelCase_     )
   print(F"""Saving image processor to {pytorch_dump_folder_path}"""     )
   image_processor.save_pretrained(lowerCAmelCase_     )


if __name__ == "__main__":
     _snake_case  :   List[Any]						   =  argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '--vit_name',
         default='vit_base_patch16_224',
         type=str,
         help='Name of the ViT timm model you\'d like to convert.',
     )
     parser.add_argument(
         '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
     )

     _snake_case  :   Optional[int]						   =  parser.parse_args()
     convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)

