





'''simple docstring'''

import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)




def    _a    (							lowerCAmelCase_			,      lowerCAmelCase_=False	):





						"""simple docstring"""



						_snake_case		:	List[str]					     =     []
						for i in range(config.num_hidden_layers	):
												# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
												rename_keys.append((f'''module.blocks.{i}.norm1.weight''', f'''vit.encoder.layer.{i}.layernorm_before.weight''')	)
												rename_keys.append((f'''module.blocks.{i}.norm1.bias''', f'''vit.encoder.layer.{i}.layernorm_before.bias''')	)
												rename_keys.append(
												    (f'''module.blocks.{i}.attn.proj.weight''', f'''vit.encoder.layer.{i}.attention.output.dense.weight''')	)
												rename_keys.append((f'''module.blocks.{i}.attn.proj.bias''', f'''vit.encoder.layer.{i}.attention.output.dense.bias''')	)
												rename_keys.append((f'''module.blocks.{i}.norm2.weight''', f'''vit.encoder.layer.{i}.layernorm_after.weight''')	)
												rename_keys.append((f'''module.blocks.{i}.norm2.bias''', f'''vit.encoder.layer.{i}.layernorm_after.bias''')	)
												rename_keys.append((f'''module.blocks.{i}.mlp.fc1.weight''', f'''vit.encoder.layer.{i}.intermediate.dense.weight''')	)
												rename_keys.append((f'''module.blocks.{i}.mlp.fc1.bias''', f'''vit.encoder.layer.{i}.intermediate.dense.bias''')	)
												rename_keys.append((f'''module.blocks.{i}.mlp.fc2.weight''', f'''vit.encoder.layer.{i}.output.dense.weight''')	)
												rename_keys.append((f'''module.blocks.{i}.mlp.fc2.bias''', f'''vit.encoder.layer.{i}.output.dense.bias''')	)

						# projection layer + position embeddings
						rename_keys.extend(
						    [
						        ('''module.cls_token''', '''vit.embeddings.cls_token'''),
						        ('''module.patch_embed.proj.weight''', '''vit.embeddings.patch_embeddings.projection.weight'''),
						        ('''module.patch_embed.proj.bias''', '''vit.embeddings.patch_embeddings.projection.bias'''),
						        ('''module.pos_embed''', '''vit.embeddings.position_embeddings'''),
						    ]	)

						if base_model:
												# layernorm + pooler
												rename_keys.extend(
												    [
												        ('''module.norm.weight''', '''layernorm.weight'''),
												        ('''module.norm.bias''', '''layernorm.bias'''),
												    ]	)

												# if just the base model, we should remove "vit" from all keys that start with "vit"
												_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    _a    (							lowerCAmelCase_			,      lowerCAmelCase_			,      lowerCAmelCase_=False	):





						"""simple docstring"""



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




def    _a    (							lowerCAmelCase_	):





						"""simple docstring"""



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




def    _a    (							lowerCAmelCase_	):





						"""simple docstring"""



						_snake_case		:	List[str]					     =     [
						    '''module.fc.fc1.weight''',
						    '''module.fc.fc1.bias''',
						    '''module.fc.bn1.weight''',
						    '''module.fc.bn1.bias''',
						    '''module.fc.bn1.running_mean''',
						    '''module.fc.bn1.running_var''',
						    '''module.fc.bn1.num_batches_tracked''',
						    '''module.fc.fc2.weight''',
						    '''module.fc.fc2.bias''',
						    '''module.fc.bn2.weight''',
						    '''module.fc.bn2.bias''',
						    '''module.fc.bn2.running_mean''',
						    '''module.fc.bn2.running_var''',
						    '''module.fc.bn2.num_batches_tracked''',
						    '''module.fc.fc3.weight''',
						    '''module.fc.fc3.bias''',
						]
						for k in ignore_keys:
												state_dict.pop(lowerCAmelCase__			,      lowerCAmelCase__	)




def    _a    (							lowerCAmelCase_			,      lowerCAmelCase_			,      lowerCAmelCase_	):





						"""simple docstring"""



						_snake_case		:	Dict					     =     dct.pop(lowerCAmelCase__	)
						_snake_case		:	Any					     =     val




def    _a    (							lowerCAmelCase_			,      lowerCAmelCase_	):





						"""simple docstring"""



						_snake_case		:	Any					     =     ViTMSNConfig()
						_snake_case		:	Optional[int]					     =     1_000

						_snake_case		:	str					     =     '''datasets/huggingface/label-files'''
						_snake_case		:	Optional[int]					     =     '''imagenet-1k-id2label.json'''
						_snake_case		:	str					     =     json.load(open(hf_hub_download(lowerCAmelCase__			,      lowerCAmelCase__	)			,      '''r'''	)	)
						_snake_case		:	int					     =     {int(lowerCAmelCase__	): v for k, v in idalabel.items()}
						_snake_case		:	int					     =     idalabel
						_snake_case		:	Dict					     =     {v: k for k, v in idalabel.items()}

						if "s16" in checkpoint_url:
												_snake_case		:	List[str]					     =     384
												_snake_case		:	List[str]					     =     1_536
												_snake_case		:	Optional[Any]					     =     6
						elif "l16" in checkpoint_url:
												_snake_case		:	List[str]					     =     1_024
												_snake_case		:	Any					     =     4_096
												_snake_case		:	Union[str, Any]					     =     24
												_snake_case		:	List[str]					     =     16
												_snake_case		:	List[str]					     =     0.1
						elif "b4" in checkpoint_url:
												_snake_case		:	List[Any]					     =     4
						elif "l7" in checkpoint_url:
												_snake_case		:	List[Any]					     =     7
												_snake_case		:	Union[str, Any]					     =     1_024
												_snake_case		:	List[Any]					     =     4_096
												_snake_case		:	Optional[Any]					     =     24
												_snake_case		:	Union[str, Any]					     =     16
												_snake_case		:	List[Any]					     =     0.1

						_snake_case		:	Union[str, Any]					     =     ViTMSNModel(lowerCAmelCase__	)

						_snake_case		:	List[Any]					     =     torch.hub.load_state_dict_from_url(lowerCAmelCase__			,      map_location='''cpu'''	)['''target_encoder''']

						_snake_case		:	List[Any]					     =     ViTImageProcessor(size=config.image_size	)

						remove_projection_head(lowerCAmelCase__	)
						_snake_case		:	Any					     =     create_rename_keys(lowerCAmelCase__			,      base_model=lowerCAmelCase__	)

						for src, dest in rename_keys:
												rename_key(lowerCAmelCase__			,      lowerCAmelCase__			,      lowerCAmelCase__	)
						read_in_q_k_v(lowerCAmelCase__			,      lowerCAmelCase__			,      base_model=lowerCAmelCase__	)

						model.load_state_dict(lowerCAmelCase__	)
						model.eval()

						_snake_case		:	List[Any]					     =     '''http://images.cocodataset.org/val2017/000000039769.jpg'''

						_snake_case		:	Optional[int]					     =     Image.open(requests.get(lowerCAmelCase__			,      stream=lowerCAmelCase__	).raw	)
						_snake_case		:	Tuple					     =     ViTImageProcessor(
						    size=config.image_size			,      image_mean=lowerCAmelCase__			,      image_std=lowerCAmelCase__	)
						_snake_case		:	Tuple					     =     image_processor(images=lowerCAmelCase__			,      return_tensors='''pt'''	)

						# forward pass
						torch.manual_seed(2	)
						_snake_case		:	Tuple					     =     model(**lowerCAmelCase__	)
						_snake_case		:	Optional[Any]					     =     outputs.last_hidden_state

						# The following Colab Notebook was used to generate these outputs:
						# https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
						if "s16" in checkpoint_url:
												_snake_case		:	Dict					     =     torch.tensor([[-1.0_915, -1.4_876, -1.1_809]]	)
						elif "b16" in checkpoint_url:
												_snake_case		:	str					     =     torch.tensor([[14.2_889, -18.9_045, 11.7_281]]	)
						elif "l16" in checkpoint_url:
												_snake_case		:	Optional[Any]					     =     torch.tensor([[41.5_028, -22.8_681, 45.6_475]]	)
						elif "b4" in checkpoint_url:
												_snake_case		:	Union[str, Any]					     =     torch.tensor([[-4.3_868, 5.2_932, -0.4_137]]	)
						else:
												_snake_case		:	Optional[int]					     =     torch.tensor([[-0.1_792, -0.6_465, 2.4_263]]	)

						# verify logits
						assert torch.allclose(last_hidden_state[:, 0, :3]			,      lowerCAmelCase__			,      atol=1E-4	)

						print(f'''Saving model to {pytorch_dump_folder_path}'''	)
						model.save_pretrained(lowerCAmelCase__	)

						print(f'''Saving image processor to {pytorch_dump_folder_path}'''	)
						image_processor.save_pretrained(lowerCAmelCase__	)


if __name__ == "__main__":
			UpperCAmelCase					:		List[Any]		=							argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    '--checkpoint_url',
			    default='https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar',
			    type=str,
			    help='URL of the checkpoint you\'d like to convert.',
			)
			parser.add_argument(
			    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
			)

			UpperCAmelCase					:		Any		=							parser.parse_args()
			convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









'''simple docstring'''

# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCAmelCase					:		Dict		=							{'configuration_timm_backbone': ['TimmBackboneConfig']}


try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase					:		Union[str, Any]		=							['TimmBackbone']


if TYPE_CHECKING:
   from .configuration_timm_backbone import TimmBackboneConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_timm_backbone import TimmBackbone

else:
   import sys

   UpperCAmelCase					:		Dict		=							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


