



"""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_						:    Dict				=      logging.get_logger(__name__)
def  __snake_case							(	__A      :						Dict							,		__A      :						Dict=False     )     ->       List[Any]:






				'''simple docstring'''
				SCREAMING_SNAKE_CASE						:							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"
								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  __snake_case							(	__A      :						Optional[Any]							,		__A      :						str							,		__A      :						Union[str, Any]=False     )     ->       Any:






				'''simple docstring'''
				for i in range(config.num_hidden_layers     ):
								if base_model:
												SCREAMING_SNAKE_CASE						:							Tuple							=     ''
								else:
												SCREAMING_SNAKE_CASE						:							List[Any]							=     'vit.'
								# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
								SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     state_dict.pop(F"""blocks.{i}.attn.qkv.weight"""     )
								SCREAMING_SNAKE_CASE						:							int							=     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						:							int							=     in_proj_weight[
								    : config.hidden_size, :
								]
								SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     in_proj_bias[: config.hidden_size]
								SCREAMING_SNAKE_CASE						:							Dict							=     in_proj_weight[
								    config.hidden_size : config.hidden_size * 2, :
								]
								SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     in_proj_bias[
								    config.hidden_size : config.hidden_size * 2
								]
								SCREAMING_SNAKE_CASE						:							Optional[int]							=     in_proj_weight[
								    -config.hidden_size :, :
								]
								SCREAMING_SNAKE_CASE						:							Optional[Any]							=     in_proj_bias[-config.hidden_size :]
def  __snake_case							(	__A      :						Union[str, Any]     )     ->       Union[str, Any]:






				'''simple docstring'''
				SCREAMING_SNAKE_CASE						:							Optional[int]							=     ['head.weight', 'head.bias']
				for k in ignore_keys:
								state_dict.pop(__A							,		__A     )
def  __snake_case							(	__A      :						Optional[int]							,		__A      :						Any							,		__A      :						Optional[int]     )     ->       Any:






				'''simple docstring'''
				SCREAMING_SNAKE_CASE						:							Optional[Any]							=     dct.pop(__A     )
				SCREAMING_SNAKE_CASE						:							int							=     val
def  __snake_case							(	)     ->       Any:






				'''simple docstring'''
				SCREAMING_SNAKE_CASE						:							int							=     'http://images.cocodataset.org/val2017/000000039769.jpg'
				SCREAMING_SNAKE_CASE						:							Any							=     Image.open(requests.get(__A							,		stream=__A     ).raw     )
				return im
@torch.no_grad()
def  __snake_case							(	__A      :						int							,		__A      :						List[Any]							,		__A      :						Union[str, Any]=True     )     ->       List[Any]:






				'''simple docstring'''
				SCREAMING_SNAKE_CASE						:							Union[str, 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						:							Optional[int]							=     1000
								SCREAMING_SNAKE_CASE						:							Dict							=     'huggingface/label-files'
								SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     'imagenet-1k-id2label.json'
								SCREAMING_SNAKE_CASE						:							Optional[Any]							=     json.load(open(hf_hub_download(__A							,		__A							,		repo_type='dataset'     )							,		'r'     )     )
								SCREAMING_SNAKE_CASE						:							List[Any]							=     {int(__A     ): v for k, v in idalabel.items()}
								SCREAMING_SNAKE_CASE						:							List[Any]							=     idalabel
								SCREAMING_SNAKE_CASE						:							Any							=     {v: k for k, v in idalabel.items()}
				# size of the architecture
				if model_name in ["dino_vits8", "dino_vits16"]:
								SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     384
								SCREAMING_SNAKE_CASE						:							str							=     1536
								SCREAMING_SNAKE_CASE						:							str							=     12
								SCREAMING_SNAKE_CASE						:							List[str]							=     6

				# load original model from torch hub
				SCREAMING_SNAKE_CASE						:							Optional[int]							=     torch.hub.load('facebookresearch/dino:main'							,		__A     )
				original_model.eval()

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

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

				if base_model:
								SCREAMING_SNAKE_CASE						:							Dict							=     original_model(__A     )
								assert torch.allclose(__A							,		outputs.last_hidden_state[:, 0, :]							,		atol=1E-1     )
				else:
								SCREAMING_SNAKE_CASE						:							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__":
							A_						:    Tuple				=      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_						:    Union[str, Any]				=      parser.parse_args()
							convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)










"""simple docstring"""

import argparse

import torch

from transformers import BertConfig, BertForPreTraining, load_tf_weights_in_bert
from transformers.utils import logging


logging.set_verbosity_info()
def  __snake_case							(	__A      :						Union[str, Any]							,		__A      :						Any							,		__A      :						Dict     )     ->       Tuple:






				'''simple docstring'''
				# Initialise PyTorch model
				SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     BertConfig.from_json_file(__A     )
				print(F"""Building PyTorch model from configuration: {config}"""     )
				SCREAMING_SNAKE_CASE						:							List[Any]							=     BertForPreTraining(__A     )

				# Load weights from tf checkpoint
				load_tf_weights_in_bert(__A							,		__A							,		__A     )

				# Save pytorch-model
				print(F"""Save PyTorch model to {pytorch_dump_path}"""     )
				torch.save(model.state_dict()							,		__A     )


if __name__ == "__main__":
							A_						:    int				=      argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
							)
							parser.add_argument(
							    '--bert_config_file',
							    default=None,
							    type=str,
							    required=True,
							    help=(
							        'The config json file corresponding to the pre-trained BERT model. \n'
							        'This specifies the model architecture.'
							    ),
							)
							parser.add_argument(
							    '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
							)
							A_						:    str				=      parser.parse_args()
							convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)





