

'''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()
__snake_case					:      Optional[int] 		=     logging.get_logger(__name__)
def 		__lowerCamelCase			(  __snake_case	:						int,						__snake_case	:						Tuple=False			)			->				Optional[Any]:






						"""simple docstring"""





						A__	:  str           =[]
						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__	:  Optional[int]           =[(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			(  __snake_case	:						List[str],						__snake_case	:						Optional[Any],						__snake_case	:						Any=False			)			->				int:






						"""simple docstring"""





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






						"""simple docstring"""





						A__	:  str           =["""head.weight""", """head.bias"""]
						for k in ignore_keys:
												state_dict.pop(UpperCAmelCase__,						UpperCAmelCase__			)
def 		__lowerCamelCase			(  __snake_case	:						Dict,						__snake_case	:						int,						__snake_case	:						int			)			->				Tuple:






						"""simple docstring"""





						A__	:  Optional[int]           =dct.pop(UpperCAmelCase__			)
						A__	:  Tuple           =val
def 		__lowerCamelCase			(  )			->				Any:






						"""simple docstring"""





						A__	:  str           ="""http://images.cocodataset.org/val2017/000000039769.jpg"""
						A__	:  str           =Image.open(requests.get(UpperCAmelCase__,						stream=UpperCAmelCase__			).raw			)
						return im
@torch.no_grad()
def 		__lowerCamelCase			(  __snake_case	:						str,						__snake_case	:						Optional[int],						__snake_case	:						Tuple=True			)			->				Dict:






						"""simple docstring"""





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

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

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

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

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



import importlib.metadata
import operator
import re
import sys
from typing import Optional

from packaging import version


lowercase   =	{
    '''<''': operator.lt,
    '''<=''': operator.le,
    '''==''': operator.eq,
    '''!=''': operator.ne,
    '''>=''': operator.ge,
    '''>''': operator.gt,
}



def   __lowerCAmelCase    (				UpperCAmelCase__					:   Tuple      ,					UpperCAmelCase__					:   int      ,					UpperCAmelCase__					:   Optional[int]      ,					UpperCAmelCase__					:   Optional[int]      ,					UpperCAmelCase__					:   Optional[int]      ,					UpperCAmelCase__					:   int							)	->     Union[str, Any]:
							if got_ver is None or want_ver is None:
														raise ValueError(
														    F'''Unable to compare versions for {requirement}: need={want_ver} found={got_ver}. This is unusual. Consider'''
														    F''' reinstalling {pkg}.'''							)
							if not ops[op](version.parse(UpperCAmelCase__							)      ,					version.parse(UpperCAmelCase__							)							):
														raise ImportError(
														    F'''{requirement} is required for a normal functioning of this module, but found {pkg}=={got_ver}.{hint}'''							)



def   __lowerCAmelCase    (				UpperCAmelCase__					:   str      ,					UpperCAmelCase__					:   Optional[str] = None							)	->     None:
							lowerCamelCase_					=       F'''\n{hint}''' if hint is not None else """"""

							# non-versioned check
							if re.match(R"""^[\w_\-\d]+$"""      ,					UpperCAmelCase__							):
														lowerCamelCase_	, lowerCamelCase_	, lowerCamelCase_					=       requirement, None, None
							else:
														lowerCamelCase_					=       re.findall(R"""^([^!=<>\s]+)([\s!=<>]{1,2}.+)"""      ,					UpperCAmelCase__							)
														if not match:
																					raise ValueError(
																					    """requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23, but"""
																					    F''' got {requirement}'''							)
														lowerCamelCase_	, lowerCamelCase_					=       match[0]
														lowerCamelCase_					=       want_full.split(""","""							)  # there could be multiple requirements
														lowerCamelCase_					=       {}
														for w in want_range:
																					lowerCamelCase_					=       re.findall(R"""^([\s!=<>]{1,2})(.+)"""      ,					UpperCAmelCase__							)
																					if not match:
																												raise ValueError(
																												    """requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23,"""
																												    F''' but got {requirement}'''							)
																					lowerCamelCase_	, lowerCamelCase_					=       match[0]
																					lowerCamelCase_					=       want_ver
																					if op not in ops:
																												raise ValueError(F'''{requirement}: need one of {list(ops.keys()							)}, but got {op}'''							)

    # special case
							if pkg == "python":
														lowerCamelCase_					=       """.""".join([str(UpperCAmelCase__							) for x in sys.version_info[:3]]							)
														for op, want_ver in wanted.items():
																					_compare_versions(UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__							)
														return

							# check if any version is installed
							try:
														lowerCamelCase_					=       importlib.metadata.version(UpperCAmelCase__							)
							except importlib.metadata.PackageNotFoundError:
														raise importlib.metadata.PackageNotFoundError(
														    F'''The \'{requirement}\' distribution was not found and is required by this application. {hint}'''							)

							# check that the right version is installed if version number or a range was provided
							if want_ver is not None:
														for op, want_ver in wanted.items():
																					_compare_versions(UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__							)



def   __lowerCAmelCase    (				UpperCAmelCase__					:   List[str]							)	->     Dict:
							lowerCamelCase_					=       """Try: pip install transformers -U or pip install -e '.[dev]' if you're working with git main"""
							return require_version(UpperCAmelCase__      ,					UpperCAmelCase__							)





