



import argparse
import json
from collections import OrderedDict
from pathlib import Path

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

from transformers import PoolFormerConfig, PoolFormerForImageClassification, PoolFormerImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
__snake_case     :str			       =logging.get_logger(__name__)






def  lowerCamelCase_	(      lowerCAmelCase__	:	List[Any]			,	lowerCAmelCase__	:	Any			,	lowerCAmelCase__	:	str			,	lowerCAmelCase__	:	Tuple				)       ->							Optional[Any]:

							'''simple docstring'''






							A      						=    original_name.split('.'				)[0]
							A      						=    key.split('.'				)
							A      						=    int(key_list[key_list.index(lowerCAmelCase__				) - 2]				)
							A      						=    int(key_list[key_list.index(lowerCAmelCase__				) - 1]				)
							A      						=    orig_block_num - offset

							A      						=    key.replace(F'''{orig_block_num}.{layer_num}.{original_name}'''			,	F'''block.{new_block_num}.{layer_num}.{new_name}'''				)
							return key






def  lowerCamelCase_	(      lowerCAmelCase__	:	Optional[int]				)       ->							Tuple:

							'''simple docstring'''






							A      						=    OrderedDict()
							A						,   A      						=    0, 0
							for key, value in state_dict.items():
														if key.startswith('network'				):
																					A      						=    key.replace('network'			,	'poolformer.encoder'				)
														if "proj" in key:
																					# Works for the first embedding as well as the internal embedding layers
																					if key.endswith('bias'				) and "patch_embed" not in key:
																												patch_emb_offset += 1
																					A      						=    key[: key.find('proj'				)]
																					A      						=    key.replace(lowerCAmelCase__			,	F'''patch_embeddings.{total_embed_found}.'''				)
																					A      						=    key.replace('proj'			,	'projection'				)
																					if key.endswith('bias'				):
																												total_embed_found += 1
														if "patch_embeddings" in key:
																					A      						=    'poolformer.encoder.' + key
														if "mlp.fc1" in key:
																					A      						=    replace_key_with_offset(lowerCAmelCase__			,	lowerCAmelCase__			,	'mlp.fc1'			,	'output.conv1'				)
														if "mlp.fc2" in key:
																					A      						=    replace_key_with_offset(lowerCAmelCase__			,	lowerCAmelCase__			,	'mlp.fc2'			,	'output.conv2'				)
														if "norm1" in key:
																					A      						=    replace_key_with_offset(lowerCAmelCase__			,	lowerCAmelCase__			,	'norm1'			,	'before_norm'				)
														if "norm2" in key:
																					A      						=    replace_key_with_offset(lowerCAmelCase__			,	lowerCAmelCase__			,	'norm2'			,	'after_norm'				)
														if "layer_scale_1" in key:
																					A      						=    replace_key_with_offset(lowerCAmelCase__			,	lowerCAmelCase__			,	'layer_scale_1'			,	'layer_scale_1'				)
														if "layer_scale_2" in key:
																					A      						=    replace_key_with_offset(lowerCAmelCase__			,	lowerCAmelCase__			,	'layer_scale_2'			,	'layer_scale_2'				)
														if "head" in key:
																					A      						=    key.replace('head'			,	'classifier'				)
														A      						=    value
							return new_state_dict






def  lowerCamelCase_	(      )       ->							Any:

							'''simple docstring'''






							A      						=    'http://images.cocodataset.org/val2017/000000039769.jpg'
							A      						=    Image.open(requests.get(lowerCAmelCase__			,	stream=lowerCAmelCase__				).raw				)

							return image






@torch.no_grad()
def  lowerCamelCase_	(      lowerCAmelCase__	:	Optional[int]			,	lowerCAmelCase__	:	Dict			,	lowerCAmelCase__	:	Optional[int]				)       ->							Union[str, Any]:

							'''simple docstring'''






							A      						=    PoolFormerConfig()

							# set attributes based on model_name
							A      						=    'huggingface/label-files'
							A      						=    model_name[-3:]
							A      						=    1000
							A      						=    'imagenet-1k-id2label.json'
							A      						=    (1, 1000)

							# set config attributes
							A      						=    json.load(open(hf_hub_download(lowerCAmelCase__			,	lowerCAmelCase__			,	repo_type='dataset'				)			,	'r'				)				)
							A      						=    {int(lowerCAmelCase__				): v for k, v in idalabel.items()}
							A      						=    idalabel
							A      						=    {v: k for k, v in idalabel.items()}
							if size == "s12":
														A      						=    [2, 2, 6, 2]
														A      						=    [64, 128, 320, 512]
														A      						=    4.0
														A      						=    0.9
							elif size == "s24":
														A      						=    [4, 4, 12, 4]
														A      						=    [64, 128, 320, 512]
														A      						=    4.0
														A      						=    0.9
							elif size == "s36":
														A      						=    [6, 6, 18, 6]
														A      						=    [64, 128, 320, 512]
														A      						=    4.0
														A      						=    1E-6
														A      						=    0.9
							elif size == "m36":
														A      						=    [6, 6, 18, 6]
														A      						=    [96, 192, 384, 768]
														A      						=    4.0
														A      						=    1E-6
														A      						=    0.95
							elif size == "m48":
														A      						=    [8, 8, 24, 8]
														A      						=    [96, 192, 384, 768]
														A      						=    4.0
														A      						=    1E-6
														A      						=    0.95
							else:
														raise ValueError(F'''Size {size} not supported'''				)

							# load image processor
							A      						=    PoolFormerImageProcessor(crop_pct=lowerCAmelCase__				)

							# Prepare image
							A      						=    prepare_img()
							A      						=    image_processor(images=lowerCAmelCase__			,	return_tensors='pt'				).pixel_values

							logger.info(F'''Converting model {model_name}...'''				)

							# load original state dict
							A      						=    torch.load(lowerCAmelCase__			,	map_location=torch.device('cpu'				)				)

							# rename keys
							A      						=    rename_keys(lowerCAmelCase__				)

							# create HuggingFace model and load state dict
							A      						=    PoolFormerForImageClassification(lowerCAmelCase__				)
							model.load_state_dict(lowerCAmelCase__				)
							model.eval()

							# Define image processor
							A      						=    PoolFormerImageProcessor(crop_pct=lowerCAmelCase__				)
							A      						=    image_processor(images=prepare_img()			,	return_tensors='pt'				).pixel_values

							# forward pass
							A      						=    model(lowerCAmelCase__				)
							A      						=    outputs.logits

							# define expected logit slices for different models
							if size == "s12":
														A      						=    torch.tensor([-0.3045, -0.6758, -0.4869]				)
							elif size == "s24":
														A      						=    torch.tensor([0.4402, -0.1374, -0.8045]				)
							elif size == "s36":
														A      						=    torch.tensor([-0.6080, -0.5133, -0.5898]				)
							elif size == "m36":
														A      						=    torch.tensor([0.3952, 0.2263, -1.2668]				)
							elif size == "m48":
														A      						=    torch.tensor([0.1167, -0.0656, -0.3423]				)
							else:
														raise ValueError(F'''Size {size} not supported'''				)

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

							# finally, save model and image processor
							logger.info(F'''Saving PyTorch model and image processor to {pytorch_dump_folder_path}...'''				)
							Path(lowerCAmelCase__				).mkdir(exist_ok=lowerCAmelCase__				)
							model.save_pretrained(lowerCAmelCase__				)
							print(F'''Saving image processor to {pytorch_dump_folder_path}'''				)
							image_processor.save_pretrained(lowerCAmelCase__				)


if __name__ == "__main__":
				__snake_case     :Any			       =argparse.ArgumentParser()

				parser.add_argument(
				    '--model_name',
				    default='poolformer_s12',
				    type=str,
				    help='Name of the model you\'d like to convert.',
				)
				parser.add_argument(
				    '--checkpoint_path', default=None, type=str, help='Path to the original PyTorch checkpoint (.pth file).'
				)
				parser.add_argument(
				    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the folder to output PyTorch model.'
				)
				__snake_case     :Any			       =parser.parse_args()
				convert_poolformer_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path)






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      _UpperCamelCase  (a__ :Tuple ,      a__ :List[str]=False							):






							"""simple docstring"""





							UpperCamelCase__					  =      []
							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"
														UpperCamelCase__					  =      [(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__ :int ,      a__ :Tuple ,      a__ :Union[str, Any]=False							):






							"""simple docstring"""





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





def      _UpperCamelCase  (a__ :int							):






							"""simple docstring"""





							UpperCamelCase__					  =      ["""head.weight""", """head.bias"""]
							for k in ignore_keys:
														state_dict.pop(a__ ,      a__							)





def      _UpperCamelCase  (a__ :Optional[int]							):






							"""simple docstring"""





							UpperCamelCase__					  =      [
							    """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(a__ ,      a__							)





def      _UpperCamelCase  (a__ :List[str] ,      a__ :Any ,      a__ :str							):






							"""simple docstring"""





							UpperCamelCase__					  =      dct.pop(a__							)
							UpperCamelCase__					  =      val





def      _UpperCamelCase  (a__ :Dict ,      a__ :Dict							):






							"""simple docstring"""





							UpperCamelCase__					  =      ViTMSNConfig()
							UpperCamelCase__					  =      1000

							UpperCamelCase__					  =      """datasets/huggingface/label-files"""
							UpperCamelCase__					  =      """imagenet-1k-id2label.json"""
							UpperCamelCase__					  =      json.load(open(hf_hub_download(a__ ,      a__							) ,      """r"""							)							)
							UpperCamelCase__					  =      {int(a__							): v for k, v in idalabel.items()}
							UpperCamelCase__					  =      idalabel
							UpperCamelCase__					  =      {v: k for k, v in idalabel.items()}

							if "s16" in checkpoint_url:
														UpperCamelCase__					  =      384
														UpperCamelCase__					  =      1536
														UpperCamelCase__					  =      6
							elif "l16" in checkpoint_url:
														UpperCamelCase__					  =      1024
														UpperCamelCase__					  =      4096
														UpperCamelCase__					  =      24
														UpperCamelCase__					  =      16
														UpperCamelCase__					  =      0.1
							elif "b4" in checkpoint_url:
														UpperCamelCase__					  =      4
							elif "l7" in checkpoint_url:
														UpperCamelCase__					  =      7
														UpperCamelCase__					  =      1024
														UpperCamelCase__					  =      4096
														UpperCamelCase__					  =      24
														UpperCamelCase__					  =      16
														UpperCamelCase__					  =      0.1

							UpperCamelCase__					  =      ViTMSNModel(a__							)

							UpperCamelCase__					  =      torch.hub.load_state_dict_from_url(a__ ,      map_location="""cpu"""							)["""target_encoder"""]

							UpperCamelCase__					  =      ViTImageProcessor(size=config.image_size							)

							remove_projection_head(a__							)
							UpperCamelCase__					  =      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__ ,      base_model=a__							)

							model.load_state_dict(a__							)
							model.eval()

							UpperCamelCase__					  =      """http://images.cocodataset.org/val2017/000000039769.jpg"""

							UpperCamelCase__					  =      Image.open(requests.get(a__ ,      stream=a__							).raw							)
							UpperCamelCase__					  =      ViTImageProcessor(
							    size=config.image_size ,      image_mean=a__ ,      image_std=a__							)
							UpperCamelCase__					  =      image_processor(images=a__ ,      return_tensors="""pt"""							)

							# forward pass
							torch.manual_seed(2							)
							UpperCamelCase__					  =      model(**a__							)
							UpperCamelCase__					  =      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:
														UpperCamelCase__					  =      torch.tensor([[-1.0915, -1.4876, -1.1809]]							)
							elif "b16" in checkpoint_url:
														UpperCamelCase__					  =      torch.tensor([[14.2889, -18.9045, 11.7281]]							)
							elif "l16" in checkpoint_url:
														UpperCamelCase__					  =      torch.tensor([[41.5028, -22.8681, 45.6475]]							)
							elif "b4" in checkpoint_url:
														UpperCamelCase__					  =      torch.tensor([[-4.3868, 5.2932, -0.4137]]							)
							else:
														UpperCamelCase__					  =      torch.tensor([[-0.1792, -0.6465, 2.4263]]							)

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

							print(f"""Saving model 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__":
					UpperCamelCase__           =			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__           =			parser.parse_args()
					convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)




