





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





def 						lowerCamelCase__     (     UpperCamelCase__			:       Optional[Any]       ,    UpperCamelCase__			:       List[str]=False						)			-> Any:



					'''simple docstring'''




					_snake_case									=				[]
					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									=				[(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__     (     UpperCamelCase__			:       List[str]       ,    UpperCamelCase__			:       Tuple       ,    UpperCamelCase__			:       List[Any]=False						)			-> str:



					'''simple docstring'''




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





def 						lowerCamelCase__     (     UpperCamelCase__			:       Any						)			-> Optional[int]:



					'''simple docstring'''




					_snake_case									=				['head.weight', 'head.bias']
					for k in ignore_keys:
										state_dict.pop(lowercase__       ,    lowercase__						)





def 						lowerCamelCase__     (     UpperCamelCase__			:       Tuple       ,    UpperCamelCase__			:       Union[str, Any]       ,    UpperCamelCase__			:       Tuple						)			-> List[Any]:



					'''simple docstring'''




					_snake_case									=				dct.pop(lowercase__						)
					_snake_case									=				val





def 						lowerCamelCase__     (     )			-> List[Any]:



					'''simple docstring'''




					_snake_case									=				'http://images.cocodataset.org/val2017/000000039769.jpg'
					_snake_case									=				Image.open(requests.get(lowercase__       ,    stream=lowercase__						).raw						)
					return im





@torch.no_grad()
def 						lowerCamelCase__     (     UpperCamelCase__			:       Optional[int]       ,    UpperCamelCase__			:       int       ,    UpperCamelCase__			:       Optional[int]=True						)			-> Dict:



					'''simple docstring'''




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

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

					# load state_dict of original model, remove and rename some keys
					_snake_case									=				original_model.state_dict()
					if base_model:
										remove_classification_head_(lowercase__						)
					_snake_case									=				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									=				ViTModel(lowercase__       ,    add_pooling_layer=lowercase__						).eval()
					else:
										_snake_case									=				ViTForImageClassification(lowercase__						).eval()
					model.load_state_dict(lowercase__						)

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

					if base_model:
										_snake_case									=				original_model(lowercase__						)
										assert torch.allclose(lowercase__       ,    outputs.last_hidden_state[:, 0, :]       ,    atol=1e-1						)
					else:
										_snake_case									=				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__":
					UpperCAmelCase_      							=     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)
					UpperCAmelCase_      							=     parser.parse_args()
					convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)










from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


UpperCAmelCase_      							=     {"""configuration_swin""": ["""SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP""", """SwinConfig""", """SwinOnnxConfig"""]}


try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCAmelCase_      							=     [
					    """SWIN_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """SwinForImageClassification""",
					    """SwinForMaskedImageModeling""",
					    """SwinModel""",
					    """SwinPreTrainedModel""",
					    """SwinBackbone""",
					]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCAmelCase_      							=     [
					    """TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """TFSwinForImageClassification""",
					    """TFSwinForMaskedImageModeling""",
					    """TFSwinModel""",
					    """TFSwinPreTrainedModel""",
					]

if TYPE_CHECKING:
					from .configuration_swin import SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP, SwinConfig, SwinOnnxConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_swin import (
										    SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
										    SwinBackbone,
										    SwinForImageClassification,
										    SwinForMaskedImageModeling,
										    SwinModel,
										    SwinPreTrainedModel,
										)

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_swin import (
										    TF_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST,
										    TFSwinForImageClassification,
										    TFSwinForMaskedImageModeling,
										    TFSwinModel,
										    TFSwinPreTrainedModel,
										)

else:
					import sys

					UpperCAmelCase_      							=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



