





import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin

class   A							(						__lowercase		,   __lowercase    ):



      @register_to_config
      def __init__(						self:    List[Any]		,	_lowerCAmelCase:    int = 128		,	_lowerCAmelCase:    int = 256		,	_lowerCAmelCase:    float = 20_00.0		,	_lowerCAmelCase:    int = 768		,	_lowerCAmelCase:    int = 12		,	_lowerCAmelCase:    int = 12		,	_lowerCAmelCase:    int = 64		,	_lowerCAmelCase:    int = 2048		,	_lowerCAmelCase:    float = 0.1		,	)    ->     Dict:
         '''simple docstring'''

         super().__init__()

         UpperCAmelCase_           =nn.Sequential(
             nn.Linear(_lowerCAmelCase		,	d_model * 4		,	bias=_lowerCAmelCase							)		,	nn.SiLU()		,	nn.Linear(d_model * 4		,	d_model * 4		,	bias=_lowerCAmelCase							)		,	nn.SiLU()		,	)

         UpperCAmelCase_           =nn.Embedding(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =False

         UpperCAmelCase_           =nn.Linear(_lowerCAmelCase		,	_lowerCAmelCase		,	bias=_lowerCAmelCase							)

         UpperCAmelCase_           =nn.Dropout(p=_lowerCAmelCase							)

         UpperCAmelCase_           =nn.ModuleList()
         for lyr_num in range(_lowerCAmelCase							):
            # FiLM conditional T5 decoder
            UpperCAmelCase_           =DecoderLayer(d_model=_lowerCAmelCase		,	d_kv=_lowerCAmelCase		,	num_heads=_lowerCAmelCase		,	d_ff=_lowerCAmelCase		,	dropout_rate=_lowerCAmelCase							)
            self.decoders.append(_lowerCAmelCase							)

         UpperCAmelCase_           =TaLayerNorm(_lowerCAmelCase							)

         UpperCAmelCase_           =nn.Dropout(p=_lowerCAmelCase							)
         UpperCAmelCase_           =nn.Linear(_lowerCAmelCase		,	_lowerCAmelCase		,	bias=_lowerCAmelCase							)



      def 						lowerCAmelCase__     (						self:    Any		,	_lowerCAmelCase:    Optional[int]		,	_lowerCAmelCase:    Optional[Any]							)    ->     Optional[Any]:
         '''simple docstring'''

         UpperCAmelCase_           =torch.mul(query_input.unsqueeze(-1							)		,	key_input.unsqueeze(-2							)							)
         return mask.unsqueeze(-3							)






      def 						lowerCAmelCase__     (						self:    List[str]		,	_lowerCAmelCase:    List[Any]		,	_lowerCAmelCase:    Optional[Any]		,	_lowerCAmelCase:    Optional[Any]							)    ->     Optional[Any]:
         '''simple docstring'''

         UpperCAmelCase_    ,							UpperCAmelCase_    ,							UpperCAmelCase_           =decoder_input_tokens.shape
         assert decoder_noise_time.shape == (batch,)

         # decoder_noise_time is in [0, 1), so rescale to expected timing range.
         UpperCAmelCase_           =get_timestep_embedding(
             decoder_noise_time * self.config.max_decoder_noise_time		,	embedding_dim=self.config.d_model		,	max_period=self.config.max_decoder_noise_time		,	).to(dtype=self.dtype							)

         UpperCAmelCase_           =self.conditioning_emb(_lowerCAmelCase							).unsqueeze(1							)

         assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

         UpperCAmelCase_           =decoder_input_tokens.shape[1]

         # If we want to use relative positions for audio context, we can just offset
         # this sequence by the length of encodings_and_masks.
         UpperCAmelCase_           =torch.broadcast_to(
             torch.arange(_lowerCAmelCase		,	device=decoder_input_tokens.device							)		,	(batch, seq_length)		,	)

         UpperCAmelCase_           =self.position_encoding(_lowerCAmelCase							)

         UpperCAmelCase_           =self.continuous_inputs_projection(_lowerCAmelCase							)
         inputs += position_encodings
         UpperCAmelCase_           =self.dropout(_lowerCAmelCase							)

         # decoder: No padding present.
         UpperCAmelCase_           =torch.ones(
             decoder_input_tokens.shape[:2]		,	device=decoder_input_tokens.device		,	dtype=inputs.dtype							)

         # Translate encoding masks to encoder-decoder masks.
         UpperCAmelCase_           =[(x, self.encoder_decoder_mask(_lowerCAmelCase		,	_lowerCAmelCase							)) for x, y in encodings_and_masks]

         # cross attend style: concat encodings
         UpperCAmelCase_           =torch.cat([x[0] for x in encodings_and_encdec_masks]		,	dim=1							)
         UpperCAmelCase_           =torch.cat([x[1] for x in encodings_and_encdec_masks]		,	dim=-1							)

         for lyr in self.decoders:
            UpperCAmelCase_           =lyr(
                _lowerCAmelCase		,	conditioning_emb=_lowerCAmelCase		,	encoder_hidden_states=_lowerCAmelCase		,	encoder_attention_mask=_lowerCAmelCase		,	)[0]

         UpperCAmelCase_           =self.decoder_norm(_lowerCAmelCase							)
         UpperCAmelCase_           =self.post_dropout(_lowerCAmelCase							)

         UpperCAmelCase_           =self.spec_out(_lowerCAmelCase							)
         return spec_out

class   A							(						nn.Module    ):



      def __init__(						self:    Optional[int]		,	_lowerCAmelCase:    Any		,	_lowerCAmelCase:    List[Any]		,	_lowerCAmelCase:    List[str]		,	_lowerCAmelCase:    Any		,	_lowerCAmelCase:    List[Any]		,	_lowerCAmelCase:    Optional[Any]=1e-6							)    ->     List[Any]:
         '''simple docstring'''

         super().__init__()
         UpperCAmelCase_           =nn.ModuleList()

         # cond self attention: layer 0
         self.layer.append(
             TaLayerSelfAttentionCond(d_model=_lowerCAmelCase		,	d_kv=_lowerCAmelCase		,	num_heads=_lowerCAmelCase		,	dropout_rate=_lowerCAmelCase							)							)

         # cross attention: layer 1
         self.layer.append(
             TaLayerCrossAttention(
                 d_model=_lowerCAmelCase		,	d_kv=_lowerCAmelCase		,	num_heads=_lowerCAmelCase		,	dropout_rate=_lowerCAmelCase		,	layer_norm_epsilon=_lowerCAmelCase		,	)							)

         # Film Cond MLP + dropout: last layer
         self.layer.append(
             TaLayerFFCond(d_model=_lowerCAmelCase		,	d_ff=_lowerCAmelCase		,	dropout_rate=_lowerCAmelCase		,	layer_norm_epsilon=_lowerCAmelCase							)							)






      def 						lowerCAmelCase__     (						self:    int		,	_lowerCAmelCase:    str		,	_lowerCAmelCase:    Dict=None		,	_lowerCAmelCase:    List[str]=None		,	_lowerCAmelCase:    str=None		,	_lowerCAmelCase:    List[str]=None		,	_lowerCAmelCase:    Optional[Any]=None		,	)    ->     Optional[Any]:
         '''simple docstring'''

         UpperCAmelCase_           =self.layer[0](
             _lowerCAmelCase		,	conditioning_emb=_lowerCAmelCase		,	attention_mask=_lowerCAmelCase		,	)

         if encoder_hidden_states is not None:
            UpperCAmelCase_           =torch.where(encoder_attention_mask > 0		,	0		,	-1e10							).to(
                encoder_hidden_states.dtype							)

            UpperCAmelCase_           =self.layer[1](
                _lowerCAmelCase		,	key_value_states=_lowerCAmelCase		,	attention_mask=_lowerCAmelCase		,	)

         # Apply Film Conditional Feed Forward layer
         UpperCAmelCase_           =self.layer[-1](_lowerCAmelCase		,	_lowerCAmelCase							)

         return (hidden_states,)

class   A							(						nn.Module    ):



      def __init__(						self:    List[str]		,	_lowerCAmelCase:    Tuple		,	_lowerCAmelCase:    str		,	_lowerCAmelCase:    Any		,	_lowerCAmelCase:    List[Any]							)    ->     List[Any]:
         '''simple docstring'''

         super().__init__()
         UpperCAmelCase_           =TaLayerNorm(_lowerCAmelCase							)
         UpperCAmelCase_           =TaFiLMLayer(in_features=d_model * 4		,	out_features=_lowerCAmelCase							)
         UpperCAmelCase_           =Attention(query_dim=_lowerCAmelCase		,	heads=_lowerCAmelCase		,	dim_head=_lowerCAmelCase		,	out_bias=_lowerCAmelCase		,	scale_qk=_lowerCAmelCase							)
         UpperCAmelCase_           =nn.Dropout(_lowerCAmelCase							)






      def 						lowerCAmelCase__     (						self:    int		,	_lowerCAmelCase:    Dict		,	_lowerCAmelCase:    str=None		,	_lowerCAmelCase:    Tuple=None		,	)    ->     str:
         '''simple docstring'''

         UpperCAmelCase_           =self.layer_norm(_lowerCAmelCase							)

         if conditioning_emb is not None:
            UpperCAmelCase_           =self.FiLMLayer(_lowerCAmelCase		,	_lowerCAmelCase							)

         # Self-attention block
         UpperCAmelCase_           =self.attention(_lowerCAmelCase							)

         UpperCAmelCase_           =hidden_states + self.dropout(_lowerCAmelCase							)

         return hidden_states

class   A							(						nn.Module    ):



      def __init__(						self:    Optional[int]		,	_lowerCAmelCase:    Tuple		,	_lowerCAmelCase:    Tuple		,	_lowerCAmelCase:    Any		,	_lowerCAmelCase:    Union[str, Any]		,	_lowerCAmelCase:    Dict							)    ->     Optional[int]:
         '''simple docstring'''

         super().__init__()
         UpperCAmelCase_           =Attention(query_dim=_lowerCAmelCase		,	heads=_lowerCAmelCase		,	dim_head=_lowerCAmelCase		,	out_bias=_lowerCAmelCase		,	scale_qk=_lowerCAmelCase							)
         UpperCAmelCase_           =TaLayerNorm(_lowerCAmelCase		,	eps=_lowerCAmelCase							)
         UpperCAmelCase_           =nn.Dropout(_lowerCAmelCase							)






      def 						lowerCAmelCase__     (						self:    str		,	_lowerCAmelCase:    Dict		,	_lowerCAmelCase:    Optional[int]=None		,	_lowerCAmelCase:    List[Any]=None		,	)    ->     Optional[Any]:
         '''simple docstring'''

         UpperCAmelCase_           =self.layer_norm(_lowerCAmelCase							)
         UpperCAmelCase_           =self.attention(
             _lowerCAmelCase		,	encoder_hidden_states=_lowerCAmelCase		,	attention_mask=attention_mask.squeeze(1							)		,	)
         UpperCAmelCase_           =hidden_states + self.dropout(_lowerCAmelCase							)
         return layer_output

class   A							(						nn.Module    ):



      def __init__(						self:    List[Any]		,	_lowerCAmelCase:    List[Any]		,	_lowerCAmelCase:    int		,	_lowerCAmelCase:    int		,	_lowerCAmelCase:    int							)    ->     Optional[int]:
         '''simple docstring'''

         super().__init__()
         UpperCAmelCase_           =TaDenseGatedActDense(d_model=_lowerCAmelCase		,	d_ff=_lowerCAmelCase		,	dropout_rate=_lowerCAmelCase							)
         UpperCAmelCase_           =TaFiLMLayer(in_features=d_model * 4		,	out_features=_lowerCAmelCase							)
         UpperCAmelCase_           =TaLayerNorm(_lowerCAmelCase		,	eps=_lowerCAmelCase							)
         UpperCAmelCase_           =nn.Dropout(_lowerCAmelCase							)






      def 						lowerCAmelCase__     (						self:    Union[str, Any]		,	_lowerCAmelCase:    List[Any]		,	_lowerCAmelCase:    Optional[int]=None							)    ->     List[Any]:
         '''simple docstring'''

         UpperCAmelCase_           =self.layer_norm(_lowerCAmelCase							)
         if conditioning_emb is not None:
            UpperCAmelCase_           =self.film(_lowerCAmelCase		,	_lowerCAmelCase							)

         UpperCAmelCase_           =self.DenseReluDense(_lowerCAmelCase							)
         UpperCAmelCase_           =hidden_states + self.dropout(_lowerCAmelCase							)
         return hidden_states

class   A							(						nn.Module    ):



      def __init__(						self:    Dict		,	_lowerCAmelCase:    int		,	_lowerCAmelCase:    Union[str, Any]		,	_lowerCAmelCase:    Optional[Any]							)    ->     Optional[Any]:
         '''simple docstring'''

         super().__init__()
         UpperCAmelCase_           =nn.Linear(_lowerCAmelCase		,	_lowerCAmelCase		,	bias=_lowerCAmelCase							)
         UpperCAmelCase_           =nn.Linear(_lowerCAmelCase		,	_lowerCAmelCase		,	bias=_lowerCAmelCase							)
         UpperCAmelCase_           =nn.Linear(_lowerCAmelCase		,	_lowerCAmelCase		,	bias=_lowerCAmelCase							)
         UpperCAmelCase_           =nn.Dropout(_lowerCAmelCase							)
         UpperCAmelCase_           =NewGELUActivation()






      def 						lowerCAmelCase__     (						self:    int		,	_lowerCAmelCase:    Optional[int]							)    ->     Optional[int]:
         '''simple docstring'''

         UpperCAmelCase_           =self.act(self.wi_a(_lowerCAmelCase							)							)
         UpperCAmelCase_           =self.wi_a(_lowerCAmelCase							)
         UpperCAmelCase_           =hidden_gelu * hidden_linear
         UpperCAmelCase_           =self.dropout(_lowerCAmelCase							)

         UpperCAmelCase_           =self.wo(_lowerCAmelCase							)
         return hidden_states

class   A							(						nn.Module    ):



      def __init__(						self:    Any		,	_lowerCAmelCase:    List[Any]		,	_lowerCAmelCase:    Optional[Any]=1e-6							)    ->     Dict:
         '''simple docstring'''

         super().__init__()
         UpperCAmelCase_           =nn.Parameter(torch.ones(_lowerCAmelCase							)							)
         UpperCAmelCase_           =eps






      def 						lowerCAmelCase__     (						self:    List[Any]		,	_lowerCAmelCase:    Optional[Any]							)    ->     Any:
         '''simple docstring'''

         UpperCAmelCase_           =hidden_states.to(torch.floataa							).pow(2							).mean(-1		,	keepdim=_lowerCAmelCase							)
         UpperCAmelCase_           =hidden_states * torch.rsqrt(variance + self.variance_epsilon							)

         # convert into half-precision if necessary
         if self.weight.dtype in [torch.floataa, torch.bfloataa]:
            UpperCAmelCase_           =hidden_states.to(self.weight.dtype							)

         return self.weight * hidden_states

class   A							(						nn.Module    ):



      def 						lowerCAmelCase__     (						self:    List[str]		,	_lowerCAmelCase:    torch.Tensor							)    ->     torch.Tensor:
         '''simple docstring'''

         return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi							) * (input + 0.04_47_15 * torch.pow(_lowerCAmelCase		,	3.0							))							))

class   A							(						nn.Module    ):



      def __init__(						self:    Union[str, Any]		,	_lowerCAmelCase:    Tuple		,	_lowerCAmelCase:    Union[str, Any]							)    ->     Dict:
         '''simple docstring'''

         super().__init__()
         UpperCAmelCase_           =nn.Linear(_lowerCAmelCase		,	out_features * 2		,	bias=_lowerCAmelCase							)






      def 						lowerCAmelCase__     (						self:    Any		,	_lowerCAmelCase:    str		,	_lowerCAmelCase:    int							)    ->     str:
         '''simple docstring'''

         UpperCAmelCase_           =self.scale_bias(_lowerCAmelCase							)
         UpperCAmelCase_    ,							UpperCAmelCase_           =torch.chunk(_lowerCAmelCase		,	2		,	-1							)
         UpperCAmelCase_           =x * (1 + scale) + shift
         return x








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				:       int      =  logging.get_logger(__name__)





def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase=False			)    ->						Optional[int]:


						'''simple docstring'''






						__snake_case				=		[]
						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"
												__snake_case				=		[(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 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase=False			)    ->						int:


						'''simple docstring'''






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





def 						_lowerCAmelCase							(     _lowerCAmelCase			)    ->						List[str]:


						'''simple docstring'''






						__snake_case				=		["head.weight", "head.bias"]
						for k in ignore_keys:
												state_dict.pop(_lowerCAmelCase   ,						_lowerCAmelCase			)





def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase			)    ->						Dict:


						'''simple docstring'''






						__snake_case				=		dct.pop(_lowerCAmelCase			)
						__snake_case				=		val





def 						_lowerCAmelCase							(     )    ->						Any:


						'''simple docstring'''






						__snake_case				=		"http://images.cocodataset.org/val2017/000000039769.jpg"
						__snake_case				=		Image.open(requests.get(_lowerCAmelCase   ,						stream=_lowerCAmelCase			).raw			)
						return im





@torch.no_grad()
def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase=True			)    ->						int:


						'''simple docstring'''






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

						# load original model from torch hub
						__snake_case				=		torch.hub.load("facebookresearch/dino:main"   ,						_lowerCAmelCase			)
						original_model.eval()

						# load state_dict of original model, remove and rename some keys
						__snake_case				=		original_model.state_dict()
						if base_model:
												remove_classification_head_(_lowerCAmelCase			)
						__snake_case				=		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:
												__snake_case				=		ViTModel(_lowerCAmelCase   ,						add_pooling_layer=_lowerCAmelCase			).eval()
						else:
												__snake_case				=		ViTForImageClassification(_lowerCAmelCase			).eval()
						model.load_state_dict(_lowerCAmelCase			)

						# Check outputs on an image, prepared by ViTImageProcessor
						__snake_case				=		ViTImageProcessor()
						__snake_case				=		image_processor(images=prepare_img()   ,						return_tensors="pt"			)
						__snake_case				=		encoding["pixel_values"]
						__snake_case				=		model(_lowerCAmelCase			)

						if base_model:
												__snake_case				=		original_model(_lowerCAmelCase			)
												assert torch.allclose(_lowerCAmelCase   ,						outputs.last_hidden_state[:, 0, :]   ,						atol=1E-1			)
						else:
												__snake_case				=		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				:       str      =  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				:       Any      =  parser.parse_args()
	convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)





