






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





def 		_UpperCAmelCase (				__A				: str   ):
				a_		:	Optional[Any]									=   ['''head.weight''', '''head.bias''']
				for k in ignore_keys:
								state_dict.pop(__A    ,		__A   )





def 		_UpperCAmelCase (				__A				: str    ,		__A				: List[str]    ,		__A				: Tuple   ):
				a_		:	Union[str, Any]									=   dct.pop(__A   )
				a_		:	str									=   val





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





@torch.no_grad()
def 		_UpperCAmelCase (				__A				: Optional[Any]    ,		__A				: Tuple    ,		__A				: Optional[Any]=True   ):
				a_		:	Dict									=   ViTConfig()
				# patch_size
				if model_name[-1] == "8":
								a_		:	List[Any]									=   8
				# set labels if required
				if not base_model:
								a_		:	Union[str, Any]									=   10_00
								a_		:	List[str]									=   '''huggingface/label-files'''
								a_		:	Optional[Any]									=   '''imagenet-1k-id2label.json'''
								a_		:	Tuple									=   json.load(open(hf_hub_download(__A    ,		__A    ,		repo_type='''dataset'''   )    ,		'''r'''   )   )
								a_		:	Any									=   {int(__A   ): v for k, v in idalabel.items()}
								a_		:	Any									=   idalabel
								a_		:	str									=   {v: k for k, v in idalabel.items()}
				# size of the architecture
				if model_name in ["dino_vits8", "dino_vits16"]:
								a_		:	Optional[int]									=   3_84
								a_		:	Union[str, Any]									=   15_36
								a_		:	Optional[Any]									=   12
								a_		:	List[str]									=   6

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

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

				# Check outputs on an image, prepared by ViTImageProcessor
				a_		:	Optional[int]									=   ViTImageProcessor()
				a_		:	str									=   image_processor(images=prepare_img()    ,		return_tensors='''pt'''   )
				a_		:	Tuple									=   encoding['''pixel_values''']
				a_		:	Dict									=   model(__A   )

				if base_model:
								a_		:	str									=   original_model(__A   )
								assert torch.allclose(__A    ,		outputs.last_hidden_state[:, 0, :]    ,		atol=1E-1   )
				else:
								a_		:	List[str]									=   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 os
import time

import pytest

from datasets.utils.filelock import FileLock, Timeout





def 		_UpperCAmelCase (				__A				: int   ):
				a_		:	Optional[Any]									=   FileLock(str(tmpdir / '''foo.lock'''   )   )
				a_		:	Union[str, Any]									=   FileLock(str(tmpdir / '''foo.lock'''   )   )
				a_		:	List[Any]									=   0.01
				with locka.acquire():
								with pytest.raises(__A   ):
												a_		:	List[Any]									=   time.time()
												locka.acquire(__A   )
												assert time.time() - _start > timeout





def 		_UpperCAmelCase (				__A				: Tuple   ):
				a_		:	Tuple									=   '''a''' * 10_00 + '''.lock'''
				a_		:	List[Any]									=   FileLock(str(tmpdir / filename   )   )
				assert locka._lock_file.endswith('''.lock'''   )
				assert not locka._lock_file.endswith(__A   )
				assert len(os.path.basename(locka._lock_file   )   ) <= 2_55
				a_		:	Optional[int]									=   FileLock(tmpdir / filename   )
				with locka.acquire():
								with pytest.raises(__A   ):
												locka.acquire(0   )




