
from typing import List, Optional, Union

import torch
from transformers import (
    XLMRobertaTokenizer,
)

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from .text_encoder import MultilingualCLIP


lowercase__       =logging.get_logger(__name__)  # pylint: disable=invalid-name

lowercase__       ='\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/Kandinsky-2-1-prior")\n        >>> pipe_prior.to("cuda")\n\n        >>> prompt = "red cat, 4k photo"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> negative_image_emb = out.negative_image_embeds\n\n        >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1")\n        >>> pipe.to("cuda")\n\n        >>> image = pipe(\n        ...     prompt,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=negative_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ... ).images\n\n        >>> image[0].save("cat.png")\n        ```\n'




def    __UpperCamelCase     (				lowerCAmelCase__							:				Optional[int]  ,   lowerCAmelCase__							:				Union[str, Any]  ,   lowerCAmelCase__							:				Optional[Any]=8  ):
 __a :  str						=  h // scale_factor**2
 if h % scale_factor**2 != 0:
  new_h += 1
 __a :  List[str]						=  w // scale_factor**2
 if w % scale_factor**2 != 0:
  new_w += 1
 return new_h * scale_factor, new_w * scale_factor



class 		UpperCamelCase__  (       snake_case__    ):
   def __init__(self :					Union[str, Any]					,				snake_case_ :					Any					,				snake_case_ :					int					,				snake_case_ :					Optional[Any]					,				snake_case_ :					Any					,				snake_case_ :					Union[str, Any]					,				):
    super().__init__()

    self.register_modules(
        text_encoder=snake_case_					,				tokenizer=snake_case_					,				unet=snake_case_					,				scheduler=snake_case_					,				movq=snake_case_					,				)
    __a :  Optional[Any]						=  2 ** (len(self.movq.config.block_out_channels   ) - 1)
   def       lowerCAmelCase			(self :					Optional[Any]					,				snake_case_ :					Tuple					,				snake_case_ :					int					,				snake_case_ :					int					,				snake_case_ :					Tuple					,				snake_case_ :					List[str]					,				snake_case_ :					str   ):







    if latents is None:
     __a :  Optional[int]						=  randn_tensor(snake_case_					,				generator=snake_case_					,				device=snake_case_					,				dtype=snake_case_   )
    else:
     if latents.shape != shape:
      raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}"   )
     __a :  Union[str, Any]						=  latents.to(snake_case_   )

    __a :  List[Any]						=  latents * scheduler.init_noise_sigma
    return latents
   def       lowerCAmelCase			(self :					str					,				snake_case_ :					List[str]					,				snake_case_ :					str					,				snake_case_ :					Dict					,				snake_case_ :					str					,				snake_case_ :					Optional[int]=None					,				):
    __a :  Any						=  len(snake_case_   ) if isinstance(snake_case_					,				snake_case_   ) else 1
    # get prompt text embeddings
    __a :  Tuple						=  self.tokenizer(
        snake_case_					,				padding='''max_length'''					,				truncation=snake_case_					,				max_length=7_7					,				return_attention_mask=snake_case_					,				add_special_tokens=snake_case_					,				return_tensors='''pt'''					,				)

    __a :  Tuple						=  text_inputs.input_ids
    __a :  List[str]						=  self.tokenizer(snake_case_					,				padding='''longest'''					,				return_tensors='''pt'''   ).input_ids

    if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(snake_case_					,				snake_case_   ):
     __a :  Dict						=  self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]   )
     logger.warning(
         '''The following part of your input was truncated because CLIP can only handle sequences up to'''
         f" {self.tokenizer.model_max_length} tokens: {removed_text}"   )

    __a :  Tuple						=  text_input_ids.to(snake_case_   )
    __a :  str						=  text_inputs.attention_mask.to(snake_case_   )

    __a				,       __a :  Tuple						=  self.text_encoder(
        input_ids=snake_case_					,				attention_mask=snake_case_   )

    __a :  Optional[Any]						=  prompt_embeds.repeat_interleave(snake_case_					,				dim=0   )
    __a :  str						=  text_encoder_hidden_states.repeat_interleave(snake_case_					,				dim=0   )
    __a :  Dict						=  text_mask.repeat_interleave(snake_case_					,				dim=0   )

    if do_classifier_free_guidance:
     __a :  Tuple						=  4_2
     if negative_prompt is None:
      __a :  int						=  [''''''] * batch_size
     elif type(snake_case_   ) is not type(snake_case_   ):
      raise TypeError(
          f"`negative_prompt` should be the same type to `prompt`, but got {type(snake_case_   )} !="
          f" {type(snake_case_   )}."   )
     elif isinstance(snake_case_					,				snake_case_   ):
      __a :  str						=  [negative_prompt]
     elif batch_size != len(snake_case_   ):
      raise ValueError(
          f"`negative_prompt`: {negative_prompt} has batch size {len(snake_case_   )}, but `prompt`:"
          f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
          ''' the batch size of `prompt`.'''   )
     else:
      __a :  Dict						=  negative_prompt

     __a :  Dict						=  self.tokenizer(
         snake_case_					,				padding='''max_length'''					,				max_length=7_7					,				truncation=snake_case_					,				return_attention_mask=snake_case_					,				add_special_tokens=snake_case_					,				return_tensors='''pt'''					,				)
     __a :  Optional[Any]						=  uncond_input.input_ids.to(snake_case_   )
     __a :  Dict						=  uncond_input.attention_mask.to(snake_case_   )

     __a				,       __a :  int						=  self.text_encoder(
         input_ids=snake_case_					,				attention_mask=snake_case_   )

     # duplicate unconditional embeddings for each generation per prompt, using mps friendly method

     __a :  List[Any]						=  negative_prompt_embeds.shape[1]
     __a :  str						=  negative_prompt_embeds.repeat(1					,				snake_case_   )
     __a :  Union[str, Any]						=  negative_prompt_embeds.view(batch_size * num_images_per_prompt					,				snake_case_   )

     __a :  str						=  uncond_text_encoder_hidden_states.shape[1]
     __a :  Any						=  uncond_text_encoder_hidden_states.repeat(1					,				snake_case_					,				1   )
     __a :  List[str]						=  uncond_text_encoder_hidden_states.view(
         batch_size * num_images_per_prompt					,				snake_case_					,				-1   )
     __a :  List[str]						=  uncond_text_mask.repeat_interleave(snake_case_					,				dim=0   )

     # done duplicates

     # For classifier free guidance, we need to do two forward passes.
     # Here we concatenate the unconditional and text embeddings into a single batch
     # to avoid doing two forward passes
     __a :  Optional[int]						=  torch.cat([negative_prompt_embeds, prompt_embeds]   )
     __a :  List[Any]						=  torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]   )

     __a :  List[str]						=  torch.cat([uncond_text_mask, text_mask]   )

    return prompt_embeds, text_encoder_hidden_states, text_mask
   def       lowerCAmelCase			(self :					List[str]					,				snake_case_ :					int=0   ):







    if is_accelerate_available():
     from accelerate import cpu_offload
    else:
     raise ImportError('''Please install accelerate via `pip install accelerate`'''   )

    __a :  Dict						=  torch.device(f"cuda:{gpu_id}"   )

    __a :  Optional[Any]						=  [
        self.unet,
        self.text_encoder,
        self.movq,
    ]
    for cpu_offloaded_model in models:
     if cpu_offloaded_model is not None:
      cpu_offload(snake_case_					,				snake_case_   )
   def       lowerCAmelCase			(self :					Optional[Any]					,				snake_case_ :					Dict=0   ):







    if is_accelerate_available() and is_accelerate_version('''>='''					,				'''0.17.0.dev0'''   ):
     from accelerate import cpu_offload_with_hook
    else:
     raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''   )

    __a :  List[Any]						=  torch.device(f"cuda:{gpu_id}"   )

    if self.device.type != "cpu":
     self.to('''cpu'''					,				silence_dtype_warnings=snake_case_   )
     torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

    __a :  List[Any]						=  None
    for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
     __a				,       __a :  List[str]						=  cpu_offload_with_hook(snake_case_					,				snake_case_					,				prev_module_hook=snake_case_   )

    if self.safety_checker is not None:
     __a				,       __a :  Tuple						=  cpu_offload_with_hook(self.safety_checker					,				snake_case_					,				prev_module_hook=snake_case_   )

    # We'll offload the last model manually.
    __a :  Optional[Any]						=  hook
   @property
   # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
   def       lowerCAmelCase			(self :					List[Any]   ):







    if not hasattr(self.unet					,				'''_hf_hook'''   ):
     return self.device
    for module in self.unet.modules():
     if (
         hasattr(snake_case_					,				'''_hf_hook'''   )
         and hasattr(module._hf_hook					,				'''execution_device'''   )
         and module._hf_hook.execution_device is not None
     ):
      return torch.device(module._hf_hook.execution_device   )
    return self.device







   @torch.no_grad()
   @replace_example_docstring(snake_case_   )
   def __call__(self :					Optional[Any]					,				snake_case_ :					str					,				snake_case_ :					int					,				snake_case_ :					str					,				snake_case_ :					Any = None					,				snake_case_ :					List[Any] = 5_1_2					,				snake_case_ :					Union[str, Any] = 5_1_2					,				snake_case_ :					str = 1_0_0					,				snake_case_ :					Union[str, Any] = 4.0					,				snake_case_ :					Tuple = 1					,				snake_case_ :					Dict = None					,				snake_case_ :					Optional[Any] = None					,				snake_case_ :					Optional[int] = "pil"					,				snake_case_ :					List[str] = True					,				):







    if isinstance(snake_case_					,				snake_case_   ):
     __a :  Any						=  1
    elif isinstance(snake_case_					,				snake_case_   ):
     __a :  List[Any]						=  len(snake_case_   )
    else:
     raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(snake_case_   )}"   )

    __a :  int						=  self._execution_device

    __a :  Union[str, Any]						=  batch_size * num_images_per_prompt
    __a :  Dict						=  guidance_scale > 1.0

    __a				,       __a				,       __a :  int						=  self._encode_prompt(
        snake_case_					,				snake_case_					,				snake_case_					,				snake_case_					,				snake_case_   )

    if isinstance(snake_case_					,				snake_case_   ):
     __a :  Optional[Any]						=  torch.cat(snake_case_					,				dim=0   )
    if isinstance(snake_case_					,				snake_case_   ):
     __a :  List[str]						=  torch.cat(snake_case_					,				dim=0   )

    if do_classifier_free_guidance:
     __a :  int						=  image_embeds.repeat_interleave(snake_case_					,				dim=0   )
     __a :  List[Any]						=  negative_image_embeds.repeat_interleave(snake_case_					,				dim=0   )

    __a :  List[str]						=  torch.cat([negative_image_embeds, image_embeds]					,				dim=0   ).to(
        dtype=prompt_embeds.dtype					,				device=snake_case_   )

    self.scheduler.set_timesteps(snake_case_					,				device=snake_case_   )
    __a :  Union[str, Any]						=  self.scheduler.timesteps

    __a :  Tuple						=  self.unet.config.in_channels

    __a				,       __a :  Any						=  get_new_h_w(snake_case_					,				snake_case_					,				self.movq_scale_factor   )

    # create initial latent
    __a :  Union[str, Any]						=  self.prepare_latents(
        (batch_size, num_channels_latents, height, width)					,				text_encoder_hidden_states.dtype					,				snake_case_					,				snake_case_					,				snake_case_					,				self.scheduler					,				)

    for i, t in enumerate(self.progress_bar(snake_case_   )   ):
     # expand the latents if we are doing classifier free guidance
     __a :  List[Any]						=  torch.cat([latents] * 2   ) if do_classifier_free_guidance else latents

     __a :  int						=  {'''text_embeds''': prompt_embeds, '''image_embeds''': image_embeds}
     __a :  List[Any]						=  self.unet(
         sample=snake_case_					,				timestep=snake_case_					,				encoder_hidden_states=snake_case_					,				added_cond_kwargs=snake_case_					,				return_dict=snake_case_					,				)[0]

     if do_classifier_free_guidance:
      __a				,       __a :  List[str]						=  noise_pred.split(latents.shape[1]					,				dim=1   )
      __a				,       __a :  List[Any]						=  noise_pred.chunk(2   )
      __a				,       __a :  List[Any]						=  variance_pred.chunk(2   )
      __a :  str						=  noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
      __a :  str						=  torch.cat([noise_pred, variance_pred_text]					,				dim=1   )

     if not (
         hasattr(self.scheduler.config					,				'''variance_type'''   )
         and self.scheduler.config.variance_type in ["learned", "learned_range"]
     ):
      __a				,       __a :  str						=  noise_pred.split(latents.shape[1]					,				dim=1   )

     # compute the previous noisy sample x_t -> x_t-1
     __a :  Dict						=  self.scheduler.step(
         snake_case_					,				snake_case_					,				snake_case_					,				generator=snake_case_					,				).prev_sample
    # post-processing
    __a :  List[Any]						=  self.movq.decode(snake_case_					,				force_not_quantize=snake_case_   )['''sample''']

    if output_type not in ["pt", "np", "pil"]:
     raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"   )

    if output_type in ["np", "pil"]:
     __a :  Optional[int]						=  image * 0.5 + 0.5
     __a :  List[Any]						=  image.clamp(0					,				1   )
     __a :  int						=  image.cpu().permute(0					,				2					,				3					,				1   ).float().numpy()

    if output_type == "pil":
     __a :  Optional[int]						=  self.numpy_to_pil(snake_case_   )

    if not return_dict:
     return (image,)

    return ImagePipelineOutput(images=snake_case_   )








import argparse
import json
from pathlib import Path

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

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


logging.set_verbosity_info()
_a      	=  logging.get_logger(__name__)


def __A    (      __lowerCAmelCase ,		__lowerCAmelCase=False	)->  Union[str, Any]:





						"""simple docstring"""



						_UpperCAmelCase									=			[]
						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'),
												    ]	)

												# if just the base model, we should remove "vit" from all keys that start with "vit"
												_UpperCAmelCase									=			[(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 ,		__lowerCAmelCase ,		__lowerCAmelCase=False	)->  List[str]:





						"""simple docstring"""



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


def __A    (      __lowerCAmelCase	)->  Optional[Any]:





						"""simple docstring"""



						_UpperCAmelCase									=			['head.weight', 'head.bias']
						for k in ignore_keys:
												state_dict.pop(__lowerCAmelCase ,		__lowerCAmelCase	)


def __A    (      __lowerCAmelCase ,		__lowerCAmelCase ,		__lowerCAmelCase	)->  int:





						"""simple docstring"""



						_UpperCAmelCase									=			dct.pop(__lowerCAmelCase	)
						_UpperCAmelCase									=			val


def __A    (      )->  str:





						"""simple docstring"""



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


@torch.no_grad()
def __A    (      __lowerCAmelCase ,		__lowerCAmelCase ,		__lowerCAmelCase=True	)->  List[str]:





						"""simple docstring"""



						_UpperCAmelCase									=			ViTConfig()
						# patch_size
						if model_name[-1] == "8":
												_UpperCAmelCase									=			8
						# set labels if required
						if not base_model:
												_UpperCAmelCase									=			1_000
												_UpperCAmelCase									=			'huggingface/label-files'
												_UpperCAmelCase									=			'imagenet-1k-id2label.json'
												_UpperCAmelCase									=			json.load(open(hf_hub_download(__lowerCAmelCase ,		__lowerCAmelCase ,		repo_type='dataset'	) ,		'r'	)	)
												_UpperCAmelCase									=			{int(__lowerCAmelCase	): v for k, v in idalabel.items()}
												_UpperCAmelCase									=			idalabel
												_UpperCAmelCase									=			{v: k for k, v in idalabel.items()}
						# size of the architecture
						if model_name in ["dino_vits8", "dino_vits16"]:
												_UpperCAmelCase									=			384
												_UpperCAmelCase									=			1_536
												_UpperCAmelCase									=			12
												_UpperCAmelCase									=			6

						# load original model from torch hub
						_UpperCAmelCase									=			torch.hub.load('facebookresearch/dino:main' ,		__lowerCAmelCase	)
						original_model.eval()

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

						# load HuggingFace model
						if base_model:
												_UpperCAmelCase									=			ViTModel(__lowerCAmelCase ,		add_pooling_layer=__lowerCAmelCase	).eval()
						else:
												_UpperCAmelCase									=			ViTForImageClassification(__lowerCAmelCase	).eval()
						model.load_state_dict(__lowerCAmelCase	)

						# Check outputs on an image, prepared by ViTImageProcessor
						_UpperCAmelCase									=			ViTImageProcessor()
						_UpperCAmelCase									=			image_processor(images=prepare_img() ,		return_tensors='pt'	)
						_UpperCAmelCase									=			encoding['pixel_values']
						_UpperCAmelCase									=			model(__lowerCAmelCase	)

						if base_model:
												_UpperCAmelCase									=			original_model(__lowerCAmelCase	)
												assert torch.allclose(__lowerCAmelCase ,		outputs.last_hidden_state[:, 0, :] ,		atol=1E-1	)
						else:
												_UpperCAmelCase									=			original_model(__lowerCAmelCase	)
												assert logits.shape == outputs.logits.shape
												assert torch.allclose(__lowerCAmelCase ,		outputs.logits ,		atol=1E-3	)

						Path(__lowerCAmelCase	).mkdir(exist_ok=__lowerCAmelCase	)
						print(F"""Saving model {model_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__":
		_a      	=  argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--model_name''',
		    default='''dino_vitb16''',
		    type=str,
		    help='''Name of the model trained with DINO you\'d like to convert.''',
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
		)
		parser.add_argument(
		    '''--base_model''',
		    action='''store_true''',
		    help='''Whether to only convert the base model (no projection head weights).''',
		)

		parser.set_defaults(base_model=True)
		_a      	=  parser.parse_args()
		convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)
