



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 	_lowerCAmelCase						(				__lowerCamelCase   :     Dict		,	__lowerCamelCase   :     Any=False     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Tuple       =	[]
 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"
  __SCREAMING_SNAKE_CASE       :    str       =	[(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   :     Union[str, Any]		,	__lowerCamelCase   :     List[str]		,	__lowerCamelCase   :     Optional[Any]=False     ):



 """simple docstring"""





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






def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[str]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Dict       =	["head.weight", "head.bias"]
 for k in ignore_keys:
  state_dict.pop(__lowerCamelCase		,	__lowerCamelCase     )






def 	_lowerCAmelCase						(				__lowerCamelCase   :     str		,	__lowerCamelCase   :     List[Any]		,	__lowerCamelCase   :     Optional[int]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	dct.pop(__lowerCamelCase     )
 __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	val






def 	_lowerCAmelCase						(				):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    List[Any]       =	"http://images.cocodataset.org/val2017/000000039769.jpg"
 __SCREAMING_SNAKE_CASE       :    Tuple       =	Image.open(requests.get(__lowerCamelCase		,	stream=__lowerCamelCase     ).raw     )
 return im






@torch.no_grad()
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Any		,	__lowerCamelCase   :     int		,	__lowerCamelCase   :     Optional[Any]=True     ):



 """simple docstring"""





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

 # load original model from torch hub
 __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	torch.hub.load("facebookresearch/dino:main"		,	__lowerCamelCase     )
 original_model.eval()

 # load state_dict of original model, remove and rename some keys
 __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	original_model.state_dict()
 if base_model:
  remove_classification_head_(__lowerCamelCase     )
 __SCREAMING_SNAKE_CASE       :    List[Any]       =	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:
  __SCREAMING_SNAKE_CASE       :    str       =	ViTModel(__lowerCamelCase		,	add_pooling_layer=__lowerCamelCase     ).eval()
 else:
  __SCREAMING_SNAKE_CASE       :    int       =	ViTForImageClassification(__lowerCamelCase     ).eval()
 model.load_state_dict(__lowerCamelCase     )

 # Check outputs on an image, prepared by ViTImageProcessor
 __SCREAMING_SNAKE_CASE       :    List[Any]       =	ViTImageProcessor()
 __SCREAMING_SNAKE_CASE       :    int       =	image_processor(images=prepare_img()		,	return_tensors="pt"     )
 __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	encoding["pixel_values"]
 __SCREAMING_SNAKE_CASE       :    Dict       =	model(__lowerCamelCase     )

 if base_model:
  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	original_model(__lowerCamelCase     )
  assert torch.allclose(__lowerCamelCase		,	outputs.last_hidden_state[:, 0, :]		,	atol=1E-1     )
 else:
  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	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__":
      _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)











import inspect
import os
import unittest

import torch

import accelerate
from accelerate import debug_launcher
from accelerate.test_utils import (
    execute_subprocess_async,
    require_cpu,
    require_huggingface_suite,
    require_multi_gpu,
    require_single_gpu,
)
from accelerate.utils import patch_environment

@require_huggingface_suite
class       _SCREAMING_SNAKE_CASE	(unittest.TestCase ):





     def 	__snake_case    (       self :							Optional[Any]				)->Union[str, Any]:
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	inspect.getfile(accelerate.test_utils				)
      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	os.path.sep.join(
          mod_file.split(os.path.sep				)[:-1] + ["scripts", "external_deps", "test_metrics.py"]				)

      from accelerate.test_utils.scripts.external_deps import test_metrics  # noqa: F401

      __SCREAMING_SNAKE_CASE       :    Tuple       =	test_metrics





     @require_cpu
     def 	__snake_case    (       self :							Union[str, Any]				)->List[Any]:
      debug_launcher(self.test_metrics.main     ,  num_processes=1				)





     @require_cpu
     def 	__snake_case    (       self :							str				)->List[str]:
      debug_launcher(self.test_metrics.main				)





     @require_single_gpu
     def 	__snake_case    (       self :							Union[str, Any]				)->str:
      self.test_metrics.main()






     @require_multi_gpu
     def 	__snake_case    (       self :							List[str]				)->int:
      print(F"""Found {torch.cuda.device_count()} devices."""				)
      __SCREAMING_SNAKE_CASE       :    Any       =	["torchrun", F"""--nproc_per_node={torch.cuda.device_count()}""", self.test_file_path]
      with patch_environment(omp_num_threads=1				):
       execute_subprocess_async(UpperCamelCase     ,  env=os.environ.copy()				)






