

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


def 	_A  (					A__  ,       A__=False   ):






 """simple docstring"""



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






 """simple docstring"""



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


def 	_A  (					A__   ):






 """simple docstring"""



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


def 	_A  (					A__  ,       A__  ,       A__   ):






 """simple docstring"""



 __lowercase													=				dct.pop(A__   )
 __lowercase													=				val


def 	_A  (					):






 """simple docstring"""



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


@torch.no_grad()
def 	_A  (					A__  ,       A__  ,       A__=True   ):






 """simple docstring"""



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

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

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

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

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

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

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


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




'''simple docstring'''






import copy
import json
import os
import tempfile

from transformers import is_torch_available

from .test_configuration_utils import config_common_kwargs






class     lowercase_			(lowerCamelCase__   ):




  """simple docstring"""




  def __init__(   self     :  Optional[Any]      ,lowercase__     :  int      ,lowercase__     :  List[str]=None      ,lowercase__     :  Any=True      ,lowercase__     :  Union[str, Any]=None      ,**lowercase__     :  Dict    ):
   __lowercase													=				parent
   __lowercase													=				config_class
   __lowercase													=				has_text_modality
   __lowercase													=				kwargs
   __lowercase													=				common_properties






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )
   __lowercase													=				(
       ['''hidden_size''', '''num_attention_heads''', '''num_hidden_layers''']
       if self.common_properties is None
       else self.common_properties
   )

   # Add common fields for text models
   if self.has_text_modality:
    common_properties.extend(['''vocab_size''']    )

   # Test that config has the common properties as getters
   for prop in common_properties:
    self.parent.assertTrue(hasattr(lowercase__      ,lowercase__    )      ,msg=F"`{prop}` does not exist"    )

   # Test that config has the common properties as setter
   for idx, name in enumerate(lowercase__    ):
    try:
     setattr(lowercase__      ,lowercase__      ,lowercase__    )
     self.parent.assertEqual(
         getattr(lowercase__      ,lowercase__    )      ,lowercase__      ,msg=F"`{name} value {idx} expected, but was {getattr(lowercase__      ,lowercase__    )}"    )
    except NotImplementedError:
     # Some models might not be able to implement setters for common_properties
     # In that case, a NotImplementedError is raised
     pass

        # Test if config class can be called with Config(prop_name=..)
   for idx, name in enumerate(lowercase__    ):
    try:
     __lowercase													=				self.config_class(**{name: idx}    )
     self.parent.assertEqual(
         getattr(lowercase__      ,lowercase__    )      ,lowercase__      ,msg=F"`{name} value {idx} expected, but was {getattr(lowercase__      ,lowercase__    )}"    )
    except NotImplementedError:
     # Some models might not be able to implement setters for common_properties
     # In that case, a NotImplementedError is raised
     pass






  def       SCREAMING_SNAKE_CASE   (   self     :  str    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )
   __lowercase													=				json.loads(config.to_json_string()    )
   for key, value in self.inputs_dict.items():
    self.parent.assertEqual(obj[key]      ,lowercase__    )






  def       SCREAMING_SNAKE_CASE   (   self     :  int    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )

   with tempfile.TemporaryDirectory() as tmpdirname:
    __lowercase													=				os.path.join(lowercase__      ,'''config.json'''    )
    config_first.to_json_file(lowercase__    )
    __lowercase													=				self.config_class.from_json_file(lowercase__    )

   self.parent.assertEqual(config_second.to_dict()      ,config_first.to_dict()    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[Any]    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )

   with tempfile.TemporaryDirectory() as tmpdirname:
    config_first.save_pretrained(lowercase__    )
    __lowercase													=				self.config_class.from_pretrained(lowercase__    )

   self.parent.assertEqual(config_second.to_dict()      ,config_first.to_dict()    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   __lowercase													=				self.config_class(**self.inputs_dict    )

   __lowercase													=				'''test'''
   with tempfile.TemporaryDirectory() as tmpdirname:
    __lowercase													=				os.path.join(lowercase__      ,lowercase__    )
    config_first.save_pretrained(lowercase__    )
    __lowercase													=				self.config_class.from_pretrained(lowercase__      ,subfolder=lowercase__    )

   self.parent.assertEqual(config_second.to_dict()      ,config_first.to_dict()    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[Any]    ):
   __lowercase													=				self.config_class(**self.inputs_dict      ,num_labels=5    )
   self.parent.assertEqual(len(config.idalabel    )      ,5    )
   self.parent.assertEqual(len(config.labelaid    )      ,5    )

   __lowercase													=				3
   self.parent.assertEqual(len(config.idalabel    )      ,3    )
   self.parent.assertEqual(len(config.labelaid    )      ,3    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Tuple    ):
   if self.config_class.is_composition:
    return
   __lowercase													=				self.config_class()
   self.parent.assertIsNotNone(lowercase__    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Any    ):
   __lowercase													=				copy.deepcopy(lowercase__    )
   __lowercase													=				self.config_class(**lowercase__    )
   __lowercase													=				[]
   for key, value in config_common_kwargs.items():
    if key == "torch_dtype":
     if not is_torch_available():
      continue
     else:
      import torch

      if config.torch_dtype != torch.floataa:
       wrong_values.append(('''torch_dtype''', config.torch_dtype, torch.floataa)    )
    elif getattr(lowercase__      ,lowercase__    ) != value:
     wrong_values.append((key, getattr(lowercase__      ,lowercase__    ), value)    )

   if len(lowercase__    ) > 0:
    __lowercase													=				'''\n'''.join([F"- {v[0]}: got {v[1]} instead of {v[2]}" for v in wrong_values]    )
    raise ValueError(F"The following keys were not properly set in the config:\n{errors}"    )






  def       SCREAMING_SNAKE_CASE   (   self     :  Optional[int]    ):
   self.create_and_test_config_common_properties()
   self.create_and_test_config_to_json_string()
   self.create_and_test_config_to_json_file()
   self.create_and_test_config_from_and_save_pretrained()
   self.create_and_test_config_from_and_save_pretrained_subfolder()
   self.create_and_test_config_with_num_labels()
   self.check_config_can_be_init_without_params()
   self.check_config_arguments_init()

