


import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor


def        _UpperCamelCase							(      lowercase__      ):
						if "cls_token" in name:
												__SCREAMING_SNAKE_CASE						:     Any					=		name.replace('''cls_token'''    ,       '''vit.embeddings.cls_token'''      )
						if "mask_token" in name:
												__SCREAMING_SNAKE_CASE						:     Any					=		name.replace('''mask_token'''    ,       '''decoder.mask_token'''      )
						if "decoder_pos_embed" in name:
												__SCREAMING_SNAKE_CASE						:     Any					=		name.replace('''decoder_pos_embed'''    ,       '''decoder.decoder_pos_embed'''      )
						if "pos_embed" in name and "decoder" not in name:
												__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		name.replace('''pos_embed'''    ,       '''vit.embeddings.position_embeddings'''      )
						if "patch_embed.proj" in name:
												__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		name.replace('''patch_embed.proj'''    ,       '''vit.embeddings.patch_embeddings.projection'''      )
						if "patch_embed.norm" in name:
												__SCREAMING_SNAKE_CASE						:     Any					=		name.replace('''patch_embed.norm'''    ,       '''vit.embeddings.norm'''      )
						if "decoder_blocks" in name:
												__SCREAMING_SNAKE_CASE						:     Any					=		name.replace('''decoder_blocks'''    ,       '''decoder.decoder_layers'''      )
						if "blocks" in name:
												__SCREAMING_SNAKE_CASE						:     Any					=		name.replace('''blocks'''    ,       '''vit.encoder.layer'''      )
						if "attn.proj" in name:
												__SCREAMING_SNAKE_CASE						:     List[str]					=		name.replace('''attn.proj'''    ,       '''attention.output.dense'''      )
						if "attn" in name:
												__SCREAMING_SNAKE_CASE						:     List[Any]					=		name.replace('''attn'''    ,       '''attention.self'''      )
						if "norm1" in name:
												__SCREAMING_SNAKE_CASE						:     Tuple					=		name.replace('''norm1'''    ,       '''layernorm_before'''      )
						if "norm2" in name:
												__SCREAMING_SNAKE_CASE						:     List[str]					=		name.replace('''norm2'''    ,       '''layernorm_after'''      )
						if "mlp.fc1" in name:
												__SCREAMING_SNAKE_CASE						:     Any					=		name.replace('''mlp.fc1'''    ,       '''intermediate.dense'''      )
						if "mlp.fc2" in name:
												__SCREAMING_SNAKE_CASE						:     Any					=		name.replace('''mlp.fc2'''    ,       '''output.dense'''      )
						if "decoder_embed" in name:
												__SCREAMING_SNAKE_CASE						:     List[Any]					=		name.replace('''decoder_embed'''    ,       '''decoder.decoder_embed'''      )
						if "decoder_norm" in name:
												__SCREAMING_SNAKE_CASE						:     List[Any]					=		name.replace('''decoder_norm'''    ,       '''decoder.decoder_norm'''      )
						if "decoder_pred" in name:
												__SCREAMING_SNAKE_CASE						:     Any					=		name.replace('''decoder_pred'''    ,       '''decoder.decoder_pred'''      )
						if "norm.weight" in name and "decoder" not in name:
												__SCREAMING_SNAKE_CASE						:     Dict					=		name.replace('''norm.weight'''    ,       '''vit.layernorm.weight'''      )
						if "norm.bias" in name and "decoder" not in name:
												__SCREAMING_SNAKE_CASE						:     Optional[int]					=		name.replace('''norm.bias'''    ,       '''vit.layernorm.bias'''      )

						return name


def        _UpperCamelCase							(      lowercase__    ,       lowercase__      ):
						for key in orig_state_dict.copy().keys():
												__SCREAMING_SNAKE_CASE						:     List[str]					=		orig_state_dict.pop(lowercase__      )

												if "qkv" in key:
																		__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		key.split('''.'''      )
																		__SCREAMING_SNAKE_CASE						:     str					=		int(key_split[1]      )
																		if "decoder_blocks" in key:
																								__SCREAMING_SNAKE_CASE						:     List[Any]					=		config.decoder_hidden_size
																								__SCREAMING_SNAKE_CASE						:     str					=		'''decoder.decoder_layers.'''
																								if "weight" in key:
																														__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		val[:dim, :]
																														__SCREAMING_SNAKE_CASE						:     List[Any]					=		val[dim : dim * 2, :]
																														__SCREAMING_SNAKE_CASE						:     str					=		val[-dim:, :]
																								elif "bias" in key:
																														__SCREAMING_SNAKE_CASE						:     List[str]					=		val[:dim]
																														__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		val[dim : dim * 2]
																														__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		val[-dim:]
																		else:
																								__SCREAMING_SNAKE_CASE						:     Optional[int]					=		config.hidden_size
																								__SCREAMING_SNAKE_CASE						:     int					=		'''vit.encoder.layer.'''
																								if "weight" in key:
																														__SCREAMING_SNAKE_CASE						:     Dict					=		val[:dim, :]
																														__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		val[dim : dim * 2, :]
																														__SCREAMING_SNAKE_CASE						:     List[str]					=		val[-dim:, :]
																								elif "bias" in key:
																														__SCREAMING_SNAKE_CASE						:     Dict					=		val[:dim]
																														__SCREAMING_SNAKE_CASE						:     List[str]					=		val[dim : dim * 2]
																														__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		val[-dim:]

												else:
																		__SCREAMING_SNAKE_CASE						:     int					=		val

						return orig_state_dict


def        _UpperCamelCase							(      lowercase__    ,       lowercase__      ):
						__SCREAMING_SNAKE_CASE						:     str					=		ViTMAEConfig()
						if "large" in checkpoint_url:
												__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		1024
												__SCREAMING_SNAKE_CASE						:     List[str]					=		4096
												__SCREAMING_SNAKE_CASE						:     Tuple					=		24
												__SCREAMING_SNAKE_CASE						:     List[str]					=		16
						elif "huge" in checkpoint_url:
												__SCREAMING_SNAKE_CASE						:     int					=		14
												__SCREAMING_SNAKE_CASE						:     Any					=		1280
												__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		5120
												__SCREAMING_SNAKE_CASE						:     Any					=		32
												__SCREAMING_SNAKE_CASE						:     List[str]					=		16

						__SCREAMING_SNAKE_CASE						:     Optional[int]					=		ViTMAEForPreTraining(lowercase__      )

						__SCREAMING_SNAKE_CASE						:     List[str]					=		torch.hub.load_state_dict_from_url(lowercase__    ,       map_location='''cpu'''      )['''model''']

						__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		ViTMAEImageProcessor(size=config.image_size      )

						__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		convert_state_dict(lowercase__    ,       lowercase__      )

						model.load_state_dict(lowercase__      )
						model.eval()

						__SCREAMING_SNAKE_CASE						:     Any					=		'''https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg'''

						__SCREAMING_SNAKE_CASE						:     Dict					=		Image.open(requests.get(lowercase__    ,       stream=lowercase__      ).raw      )
						__SCREAMING_SNAKE_CASE						:     str					=		ViTMAEImageProcessor(size=config.image_size      )
						__SCREAMING_SNAKE_CASE						:     str					=		image_processor(images=lowercase__    ,       return_tensors='''pt'''      )

						# forward pass
						torch.manual_seed(2      )
						__SCREAMING_SNAKE_CASE						:     Tuple					=		model(**lowercase__      )
						__SCREAMING_SNAKE_CASE						:     List[Any]					=		outputs.logits

						if "large" in checkpoint_url:
												__SCREAMING_SNAKE_CASE						:     str					=		torch.tensor(
												    [[-0.7309, -0.7128, -1.0169], [-1.0161, -0.9058, -1.1878], [-1.0478, -0.9411, -1.1911]]      )
						elif "huge" in checkpoint_url:
												__SCREAMING_SNAKE_CASE						:     List[Any]					=		torch.tensor(
												    [[-1.1599, -0.9199, -1.2221], [-1.1952, -0.9269, -1.2307], [-1.2143, -0.9337, -1.2262]]      )
						else:
												__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		torch.tensor(
												    [[-0.9192, -0.8481, -1.1259], [-1.1349, -1.0034, -1.2599], [-1.1757, -1.0429, -1.2726]]      )

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

						print(F'''Saving model 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__":
						__lowerCAmelCase	: Dict  				=argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--checkpoint_url',
						    default='https://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_base.pth',
						    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.'
						)

						__lowerCAmelCase	: Dict  				=parser.parse_args()
						convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)




import importlib.metadata
import operator
import re
import sys
from typing import Optional

from packaging import version


__lowerCAmelCase	: Union[str, Any]  				={
    '<': operator.lt,
    '<=': operator.le,
    '==': operator.eq,
    '!=': operator.ne,
    '>=': operator.ge,
    '>': operator.gt,
}


def        _UpperCamelCase							(      lowercase__    ,       lowercase__    ,       lowercase__    ,       lowercase__    ,       lowercase__    ,       lowercase__      ):
						if got_ver is None or want_ver is None:
												raise ValueError(
												    F'''Unable to compare versions for {requirement}: need={want_ver} found={got_ver}. This is unusual. Consider'''
												    F''' reinstalling {pkg}.'''      )
						if not ops[op](version.parse(lowercase__      )    ,       version.parse(lowercase__      )      ):
												raise ImportError(
												    F'''{requirement} is required for a normal functioning of this module, but found {pkg}=={got_ver}.{hint}'''      )


def        _UpperCamelCase							(      lowercase__    ,       lowercase__ = None      ):
						__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		F'''\n{hint}''' if hint is not None else ''''''

						# non-versioned check
						if re.match(R'''^[\w_\-\d]+$'''    ,       lowercase__      ):
												__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     str					=		requirement, None, None
						else:
												__SCREAMING_SNAKE_CASE						:     List[Any]					=		re.findall(R'''^([^!=<>\s]+)([\s!=<>]{1,2}.+)'''    ,       lowercase__      )
												if not match:
																		raise ValueError(
																		    '''requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23, but'''
																		    F''' got {requirement}'''      )
												__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     Tuple					=		match[0]
												__SCREAMING_SNAKE_CASE						:     Optional[int]					=		want_full.split(''','''      )  # there could be multiple requirements
												__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		{}
												for w in want_range:
																		__SCREAMING_SNAKE_CASE						:     Any					=		re.findall(R'''^([\s!=<>]{1,2})(.+)'''    ,       lowercase__      )
																		if not match:
																								raise ValueError(
																								    '''requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23,'''
																								    F''' but got {requirement}'''      )
																		__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     Tuple					=		match[0]
																		__SCREAMING_SNAKE_CASE						:     List[Any]					=		want_ver
																		if op not in ops:
																								raise ValueError(F'''{requirement}: need one of {list(ops.keys()      )}, but got {op}'''      )

    # special case
						if pkg == "python":
												__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		'''.'''.join([str(lowercase__      ) for x in sys.version_info[:3]]      )
												for op, want_ver in wanted.items():
																		_compare_versions(lowercase__    ,       lowercase__    ,       lowercase__    ,       lowercase__    ,       lowercase__    ,       lowercase__      )
												return

						# check if any version is installed
						try:
												__SCREAMING_SNAKE_CASE						:     Optional[int]					=		importlib.metadata.version(lowercase__      )
						except importlib.metadata.PackageNotFoundError:
												raise importlib.metadata.PackageNotFoundError(
												    F'''The \'{requirement}\' distribution was not found and is required by this application. {hint}'''      )

						# check that the right version is installed if version number or a range was provided
						if want_ver is not None:
												for op, want_ver in wanted.items():
																		_compare_versions(lowercase__    ,       lowercase__    ,       lowercase__    ,       lowercase__    ,       lowercase__    ,       lowercase__      )


def        _UpperCamelCase							(      lowercase__      ):
						__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		'''Try: pip install transformers -U or pip install -e \'.[dev]\' if you\'re working with git main'''
						return require_version(lowercase__    ,       lowercase__      )
