





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





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





def       lowercase_      (		lowercase__						)   ->Tuple:
						_snake_case:   Optional[int]  					=		['head.weight', 'head.bias']
						for k in ignore_keys:
												state_dict.pop(lowercase__				,  lowercase__						)





def       lowercase_      (		lowercase__				,  lowercase__				,  lowercase__						)   ->List[Any]:
						_snake_case:   List[str]  					=		dct.pop(lowercase__						)
						_snake_case:   List[str]  					=		val





def       lowercase_      (		)   ->Union[str, Any]:
						_snake_case:   Tuple  					=		'http://images.cocodataset.org/val2017/000000039769.jpg'
						_snake_case:   Optional[Any]  					=		Image.open(requests.get(lowercase__				,  stream=lowercase__						).raw						)
						return im





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

						# load original model from torch hub
						_snake_case:   List[str]  					=		torch.hub.load('facebookresearch/dino:main'				,  lowercase__						)
						original_model.eval()

						# load state_dict of original model, remove and rename some keys
						_snake_case:   List[Any]  					=		original_model.state_dict()
						if base_model:
												remove_classification_head_(lowercase__						)
						_snake_case:   Any  					=		create_rename_keys(lowercase__				,  base_model=lowercase__						)
						for src, dest in rename_keys:
												rename_key(lowercase__				,  lowercase__				,  lowercase__						)
						read_in_q_k_v(lowercase__				,  lowercase__				,  lowercase__						)

						# load HuggingFace model
						if base_model:
												_snake_case:   Union[str, Any]  					=		ViTModel(lowercase__				,  add_pooling_layer=lowercase__						).eval()
						else:
												_snake_case:   Optional[Any]  					=		ViTForImageClassification(lowercase__						).eval()
						model.load_state_dict(lowercase__						)

						# Check outputs on an image, prepared by ViTImageProcessor
						_snake_case:   Tuple  					=		ViTImageProcessor()
						_snake_case:   List[Any]  					=		image_processor(images=prepare_img()				,  return_tensors='pt'						)
						_snake_case:   List[Any]  					=		encoding['pixel_values']
						_snake_case:   Optional[Any]  					=		model(lowercase__						)

						if base_model:
												_snake_case:   List[Any]  					=		original_model(lowercase__						)
												assert torch.allclose(lowercase__				,  outputs.last_hidden_state[:, 0, :]				,  atol=1e-1						)
						else:
												_snake_case:   str  					=		original_model(lowercase__						)
												assert logits.shape == outputs.logits.shape
												assert torch.allclose(lowercase__				,  outputs.logits				,  atol=1e-3						)

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


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












'''simple docstring'''






A   : List[str]             =			'\n# Transformers 설치 방법\n! pip install transformers datasets\n# 마지막 릴리스 대신 소스에서 설치하려면, 위 명령을 주석으로 바꾸고 아래 명령을 해제하세요.\n# ! pip install git+https://github.com/huggingface/transformers.git\n'

A   : List[str]             =			[{'type': 'code', 'content': INSTALL_CONTENT}]
A   : Dict             =			{
    '{processor_class}': 'FakeProcessorClass',
    '{model_class}': 'FakeModelClass',
    '{object_class}': 'FakeObjectClass',
}





