



'''simple docstring'''




import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor



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

							return name



def       __lowerCamelCase  (					_lowercase	,				_lowercase		)      ->			Optional[int]:
							for key in orig_state_dict.copy().keys():
														UpperCAmelCase     :		List[str]							     =		orig_state_dict.pop(_lowercase		)

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

														else:
																					UpperCAmelCase     :		Tuple							     =		val

							return orig_state_dict



def       __lowerCamelCase  (					_lowercase	,				_lowercase		)      ->			Dict:
							UpperCAmelCase     :		int							     =		ViTMAEConfig()
							if "large" in checkpoint_url:
														UpperCAmelCase     :		List[str]							     =		1_0_2_4
														UpperCAmelCase     :		Tuple							     =		4_0_9_6
														UpperCAmelCase     :		Tuple							     =		2_4
														UpperCAmelCase     :		Dict							     =		1_6
							elif "huge" in checkpoint_url:
														UpperCAmelCase     :		int							     =		1_4
														UpperCAmelCase     :		Optional[Any]							     =		1_2_8_0
														UpperCAmelCase     :		Tuple							     =		5_1_2_0
														UpperCAmelCase     :		Optional[Any]							     =		3_2
														UpperCAmelCase     :		List[str]							     =		1_6

							UpperCAmelCase     :		List[str]							     =		ViTMAEForPreTraining(_lowercase		)

							UpperCAmelCase     :		List[str]							     =		torch.hub.load_state_dict_from_url(_lowercase	,				map_location="""cpu"""		)["""model"""]

							UpperCAmelCase     :		Dict							     =		ViTMAEImageProcessor(size=config.image_size		)

							UpperCAmelCase     :		Optional[Any]							     =		convert_state_dict(_lowercase	,				_lowercase		)

							model.load_state_dict(_lowercase		)
							model.eval()

							UpperCAmelCase     :		Tuple							     =		"""https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg"""

							UpperCAmelCase     :		str							     =		Image.open(requests.get(_lowercase	,				stream=_lowercase		).raw		)
							UpperCAmelCase     :		List[str]							     =		ViTMAEImageProcessor(size=config.image_size		)
							UpperCAmelCase     :		Any							     =		image_processor(images=_lowercase	,				return_tensors="""pt"""		)

							# forward pass
							torch.manual_seed(2		)
							UpperCAmelCase     :		List[Any]							     =		model(**_lowercase		)
							UpperCAmelCase     :		int							     =		outputs.logits

							if "large" in checkpoint_url:
														UpperCAmelCase     :		List[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:
														UpperCAmelCase     :		Optional[int]							     =		torch.tensor(
														    [[-1.1599, -0.9199, -1.2221], [-1.1952, -0.9269, -1.2307], [-1.2143, -0.9337, -1.2262]]		)
							else:
														UpperCAmelCase     :		List[str]							     =		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__":
							a   :			List[str]					  =    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."""
							)

							a   :			List[str]					  =    parser.parse_args()
							convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)






'''simple docstring'''




import contextlib
import importlib
import io
import unittest

import transformers

# Try to import everything from transformers to ensure every object can be loaded.
from transformers import *  # noqa F406
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available


if is_torch_available():
							from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification

if is_tf_available():
							from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification

if is_flax_available():
							from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification


a   :			str					  =    DUMMY_UNKNOWN_IDENTIFIER
# An actual model hosted on huggingface.co

a   :			int					  =    """main"""
# Default branch name
a   :			Any					  =    """f2c752cfc5c0ab6f4bdec59acea69eefbee381c2"""
# One particular commit (not the top of `main`)
a   :			str					  =    """aaaaaaa"""
# This commit does not exist, so we should 404.

a   :			int					  =    """d9e9f15bc825e4b2c9249e9578f884bbcb5e3684"""
# Sha-1 of config.json on the top of `main`, for checking purposes
a   :			Any					  =    """4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3"""



@contextlib.contextmanager
def       __lowerCamelCase  (					)      ->			List[str]:
							print("""Welcome!"""		)
							yield
							print("""Bye!"""		)



@contextlib.contextmanager
def       __lowerCamelCase  (					)      ->			Optional[int]:
							print("""Bonjour!"""		)
							yield
							print("""Au revoir!"""		)

class 			UpperCamelCase_ (							unittest.TestCase			):





	def  _lowercase(  self    )			->			List[Any]:
								# If the spec is missing, importlib would not be able to import the module dynamically.
								assert transformers.__spec__ is not None
								assert importlib.util.find_spec("""transformers"""    ) is not None




class 			UpperCamelCase_ (							unittest.TestCase			):





	@unittest.mock.patch("""sys.stdout"""		,		new_callable=io.StringIO    )
	def  _lowercase(  self		,		A    )			->			Tuple:
								with ContextManagers([]    ):
															print("""Transformers are awesome!"""    )
								# The print statement adds a new line at the end of the output
								self.assertEqual(mock_stdout.getvalue()		,		"""Transformers are awesome!\n"""    )





	@unittest.mock.patch("""sys.stdout"""		,		new_callable=io.StringIO    )
	def  _lowercase(  self		,		A    )			->			Dict:
								with ContextManagers([context_en()]    ):
															print("""Transformers are awesome!"""    )
								# The output should be wrapped with an English welcome and goodbye
								self.assertEqual(mock_stdout.getvalue()		,		"""Welcome!\nTransformers are awesome!\nBye!\n"""    )





	@unittest.mock.patch("""sys.stdout"""		,		new_callable=io.StringIO    )
	def  _lowercase(  self		,		A    )			->			Union[str, Any]:
								with ContextManagers([context_fr(), context_en()]    ):
															print("""Transformers are awesome!"""    )
								# The output should be wrapped with an English and French welcome and goodbye
								self.assertEqual(mock_stdout.getvalue()		,		"""Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n"""    )





	@require_torch
	def  _lowercase(  self    )			->			Optional[int]:
								self.assertEqual(find_labels(A    )		,		["""labels"""]    )
								self.assertEqual(find_labels(A    )		,		["""labels""", """next_sentence_label"""]    )
								self.assertEqual(find_labels(A    )		,		["""start_positions""", """end_positions"""]    )

								class 			UpperCamelCase_ (							__magic_name__			):
									pass




								self.assertEqual(find_labels(A    )		,		["""labels"""]    )





	@require_tf
	def  _lowercase(  self    )			->			int:
								self.assertEqual(find_labels(A    )		,		["""labels"""]    )
								self.assertEqual(find_labels(A    )		,		["""labels""", """next_sentence_label"""]    )
								self.assertEqual(find_labels(A    )		,		["""start_positions""", """end_positions"""]    )

								class 			UpperCamelCase_ (							__magic_name__			):
									pass




								self.assertEqual(find_labels(A    )		,		["""labels"""]    )






	@require_flax
	def  _lowercase(  self    )			->			Any:
								# Flax models don't have labels
								self.assertEqual(find_labels(A    )		,		[]    )
								self.assertEqual(find_labels(A    )		,		[]    )
								self.assertEqual(find_labels(A    )		,		[]    )

								class 			UpperCamelCase_ (							__magic_name__			):
									pass




								self.assertEqual(find_labels(A    )		,		[]    )

