






"""simple docstring"""


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



def 					lowerCAmelCase				(    __UpperCamelCase		,       __UpperCamelCase=False ):




  """simple docstring"""







  __A									= []
  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"
    __A									= [(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				(    __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase=False ):




  """simple docstring"""







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



def 					lowerCAmelCase				(    __UpperCamelCase ):




  """simple docstring"""







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



def 					lowerCAmelCase				(    __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase ):




  """simple docstring"""







  __A									= dct.pop(__UpperCamelCase )
  __A									= val



def 					lowerCAmelCase				(    ):




  """simple docstring"""







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



@torch.no_grad()
def 					lowerCAmelCase				(    __UpperCamelCase		,       __UpperCamelCase		,       __UpperCamelCase=True ):




  """simple docstring"""







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

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

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

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

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

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

  Path(__UpperCamelCase ).mkdir(exist_ok=__UpperCamelCase )
  print(f'Saving model {model_name} to {pytorch_dump_folder_path}' )
  model.save_pretrained(__UpperCamelCase )
  print(f'Saving image processor to {pytorch_dump_folder_path}' )
  image_processor.save_pretrained(__UpperCamelCase )


if __name__ == "__main__":
       lowercase_				     =     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)
       lowercase_				     =     parser.parse_args()
       convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)












"""simple docstring"""


import os
import tempfile
import unittest
from pathlib import Path

from transformers import AutoConfig, is_torch_available
from transformers.testing_utils import require_torch, torch_device


if is_torch_available():
       from transformers import PyTorchBenchmark, PyTorchBenchmarkArguments


@require_torch
class 					snake_case		(							unittest.TestCase						):


   '''simple docstring'''

   def 				_SCREAMING_SNAKE_CASE				(	self :     Any,       _lowerCamelCase :     Optional[int]  ):

     '''simple docstring'''
     for model_result in results.values():
       for batch_size, sequence_length in zip(model_result['''bs'''],       model_result['''ss''']  ):
         __A									= model_result['''result'''][batch_size][sequence_length]
         self.assertIsNotNone(_lowerCamelCase  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[Any]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[Any]  ):

     '''simple docstring'''
     __A									= '''sgugger/tiny-distilbert-classification'''
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       only_pretrain_model=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Any  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       torchscript=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   @unittest.skipIf(torch_device == '''cpu''',       '''Cant do half precision'''  )
   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[Any]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       fpaa=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Dict  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= AutoConfig.from_pretrained(_lowerCamelCase  )
     # set architectures equal to `None`
     __A									= None
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase,       configs=[config]  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     int  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result  )
     self.check_results_dict_not_empty(results.memory_train_result  )


   @unittest.skipIf(torch_device == '''cpu''',       '''Can\'t do half precision'''  )
   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[int]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       fpaa=_lowerCamelCase,       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result  )
     self.check_results_dict_not_empty(results.memory_train_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     str  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= AutoConfig.from_pretrained(_lowerCamelCase  )
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase,       configs=[config]  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     List[Any]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tinier_bart'''
     __A									= AutoConfig.from_pretrained(_lowerCamelCase  )
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase,       configs=[config]  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[Any]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     __A									= AutoConfig.from_pretrained(_lowerCamelCase  )
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase,       configs=[config]  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result  )
     self.check_results_dict_not_empty(results.memory_train_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Dict  ):

     '''simple docstring'''
     __A									= '''sshleifer/tinier_bart'''
     __A									= AutoConfig.from_pretrained(_lowerCamelCase  )
     __A									= PyTorchBenchmarkArguments(
         models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       multi_process=_lowerCamelCase,       )
     __A									= PyTorchBenchmark(_lowerCamelCase,       configs=[config]  )
     __A									= benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result  )
     self.check_results_dict_not_empty(results.memory_train_result  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[Any]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''
     with tempfile.TemporaryDirectory() as tmp_dir:
       __A									= PyTorchBenchmarkArguments(
           models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       save_to_csv=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       inference_time_csv_file=os.path.join(_lowerCamelCase,       '''inf_time.csv'''  ),       train_memory_csv_file=os.path.join(_lowerCamelCase,       '''train_mem.csv'''  ),       inference_memory_csv_file=os.path.join(_lowerCamelCase,       '''inf_mem.csv'''  ),       train_time_csv_file=os.path.join(_lowerCamelCase,       '''train_time.csv'''  ),       env_info_csv_file=os.path.join(_lowerCamelCase,       '''env.csv'''  ),       multi_process=_lowerCamelCase,       )
       __A									= PyTorchBenchmark(_lowerCamelCase  )
       benchmark.run()
       self.assertTrue(Path(os.path.join(_lowerCamelCase,       '''inf_time.csv'''  )  ).exists()  )
       self.assertTrue(Path(os.path.join(_lowerCamelCase,       '''train_time.csv'''  )  ).exists()  )
       self.assertTrue(Path(os.path.join(_lowerCamelCase,       '''inf_mem.csv'''  )  ).exists()  )
       self.assertTrue(Path(os.path.join(_lowerCamelCase,       '''train_mem.csv'''  )  ).exists()  )
       self.assertTrue(Path(os.path.join(_lowerCamelCase,       '''env.csv'''  )  ).exists()  )






   def 				_SCREAMING_SNAKE_CASE				(	self :     Union[str, Any]  ):

     '''simple docstring'''
     __A									= '''sshleifer/tiny-gpt2'''

     def _check_summary_is_not_empty(_lowerCamelCase :     List[Any]  ):
       self.assertTrue(hasattr(_lowerCamelCase,       '''sequential'''  )  )
       self.assertTrue(hasattr(_lowerCamelCase,       '''cumulative'''  )  )
       self.assertTrue(hasattr(_lowerCamelCase,       '''current'''  )  )
       self.assertTrue(hasattr(_lowerCamelCase,       '''total'''  )  )

     with tempfile.TemporaryDirectory() as tmp_dir:
       __A									= PyTorchBenchmarkArguments(
           models=[MODEL_ID],       training=_lowerCamelCase,       inference=_lowerCamelCase,       sequence_lengths=[8],       batch_sizes=[1],       log_filename=os.path.join(_lowerCamelCase,       '''log.txt'''  ),       log_print=_lowerCamelCase,       trace_memory_line_by_line=_lowerCamelCase,       multi_process=_lowerCamelCase,       )
       __A									= PyTorchBenchmark(_lowerCamelCase  )
       __A									= benchmark.run()
       _check_summary_is_not_empty(result.inference_summary  )
       _check_summary_is_not_empty(result.train_summary  )
       self.assertTrue(Path(os.path.join(_lowerCamelCase,       '''log.txt'''  )  ).exists()  )




