





'''simple docstring'''



import argparse
import json
import os

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

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
_SCREAMING_SNAKE_CASE    =			logging.get_logger(__name__)




def       _lowerCAmelCase  ( lowerCamelCase_	:     int ):
			__lowercase   =							ConvNextVaConfig()

			if "atto" in checkpoint_url:
						__lowercase   =							[2, 2, 6, 2]
						__lowercase   =							[4_0, 8_0, 1_6_0, 3_2_0]
			if "femto" in checkpoint_url:
						__lowercase   =							[2, 2, 6, 2]
						__lowercase   =							[4_8, 9_6, 1_9_2, 3_8_4]
			if "pico" in checkpoint_url:
						__lowercase   =							[2, 2, 6, 2]
						__lowercase   =							[6_4, 1_2_8, 2_5_6, 5_1_2]
			if "nano" in checkpoint_url:
						__lowercase   =							[2, 2, 8, 2]
						__lowercase   =							[8_0, 1_6_0, 3_2_0, 6_4_0]
			if "tiny" in checkpoint_url:
						__lowercase   =							[3, 3, 9, 3]
						__lowercase   =							[9_6, 1_9_2, 3_8_4, 7_6_8]
			if "base" in checkpoint_url:
						__lowercase   =							[3, 3, 2_7, 3]
						__lowercase   =							[1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
			if "large" in checkpoint_url:
						__lowercase   =							[3, 3, 2_7, 3]
						__lowercase   =							[1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
			if "huge" in checkpoint_url:
						__lowercase   =							[3, 3, 2_7, 3]
						__lowercase   =							[3_5_2, 7_0_4, 1_4_0_8, 2_8_1_6]

			__lowercase   =							1_0_0_0
			__lowercase   =							'''imagenet-1k-id2label.json'''
			__lowercase   =							(1, 1_0_0_0)

			__lowercase   =							'''huggingface/label-files'''
			__lowercase   =							num_labels
			__lowercase   =							json.load(open(hf_hub_download(lowerCamelCase_						,				lowerCamelCase_						,				repo_type='''dataset''' )						,				'''r''' ) )
			__lowercase   =							{int(lowerCamelCase_ ): v for k, v in idalabel.items()}

			__lowercase   =							idalabel
			__lowercase   =							{v: k for k, v in idalabel.items()}
			__lowercase   =							hidden_sizes
			__lowercase   =							depths

			return config, expected_shape




def       _lowerCAmelCase  ( lowerCamelCase_	:     Optional[int] ):
			if "downsample_layers.0.0" in name:
						__lowercase   =							name.replace('''downsample_layers.0.0'''						,				'''embeddings.patch_embeddings''' )
			if "downsample_layers.0.1" in name:
						__lowercase   =							name.replace('''downsample_layers.0.1'''						,				'''embeddings.norm''' )  # we rename to layernorm later on
			if "downsample_layers.1.0" in name:
						__lowercase   =							name.replace('''downsample_layers.1.0'''						,				'''stages.1.downsampling_layer.0''' )
			if "downsample_layers.1.1" in name:
						__lowercase   =							name.replace('''downsample_layers.1.1'''						,				'''stages.1.downsampling_layer.1''' )
			if "downsample_layers.2.0" in name:
						__lowercase   =							name.replace('''downsample_layers.2.0'''						,				'''stages.2.downsampling_layer.0''' )
			if "downsample_layers.2.1" in name:
						__lowercase   =							name.replace('''downsample_layers.2.1'''						,				'''stages.2.downsampling_layer.1''' )
			if "downsample_layers.3.0" in name:
						__lowercase   =							name.replace('''downsample_layers.3.0'''						,				'''stages.3.downsampling_layer.0''' )
			if "downsample_layers.3.1" in name:
						__lowercase   =							name.replace('''downsample_layers.3.1'''						,				'''stages.3.downsampling_layer.1''' )
			if "stages" in name and "downsampling_layer" not in name:
						# stages.0.0. for instance should be renamed to stages.0.layers.0.
						__lowercase   =							name[: len('''stages.0''' )] + '''.layers''' + name[len('''stages.0''' ) :]
			if "gamma" in name:
						__lowercase   =							name.replace('''gamma'''						,				'''weight''' )
			if "beta" in name:
						__lowercase   =							name.replace('''beta'''						,				'''bias''' )
			if "stages" in name:
						__lowercase   =							name.replace('''stages'''						,				'''encoder.stages''' )
			if "norm" in name:
						__lowercase   =							name.replace('''norm'''						,				'''layernorm''' )
			if "head" in name:
						__lowercase   =							name.replace('''head'''						,				'''classifier''' )

			return name




def       _lowerCAmelCase  ( ):
			__lowercase   =							'''http://images.cocodataset.org/val2017/000000039769.jpg'''
			__lowercase   =							Image.open(requests.get(lowerCamelCase_						,				stream=lowerCamelCase_ ).raw )
			return im




def       _lowerCAmelCase  ( lowerCamelCase_	:     Dict ):
			if "224" in checkpoint_url:
						__lowercase   =							2_2_4
						__lowercase   =							2_2_4 / 2_5_6
			elif "384" in checkpoint_url:
						__lowercase   =							3_8_4
						__lowercase   =							None
			else:
						__lowercase   =							5_1_2
						__lowercase   =							None

			return ConvNextImageProcessor(
			    size=lowerCamelCase_						,				crop_pct=lowerCamelCase_						,				image_mean=[0.4_85, 0.4_56, 0.4_06]						,				image_std=[0.2_29, 0.2_24, 0.2_25]						,				resample=PILImageResampling.BICUBIC						,				)




@torch.no_grad()
def       _lowerCAmelCase  ( lowerCamelCase_	:     Any						,				lowerCamelCase_	:     List[str]						,				lowerCamelCase_	:     Dict						,				lowerCamelCase_	:     Tuple ):
			print('''Downloading original model from checkpoint...''' )
			# define ConvNeXTV2 configuration based on URL
			__lowercase  ,						__lowercase   =							get_convnextva_config(lowerCamelCase_ )
			# load original state_dict from URL
			__lowercase   =							torch.hub.load_state_dict_from_url(lowerCamelCase_ )['''model''']

			print('''Converting model parameters...''' )
			# rename keys
			for key in state_dict.copy().keys():
						__lowercase   =							state_dict.pop(lowerCamelCase_ )
						__lowercase   =							val
			# add prefix to all keys expect classifier head
			for key in state_dict.copy().keys():
						__lowercase   =							state_dict.pop(lowerCamelCase_ )
						if not key.startswith('''classifier''' ):
									__lowercase   =							'''convnextv2.''' + key
						__lowercase   =							val

			# load HuggingFace model
			__lowercase   =							ConvNextVaForImageClassification(lowerCamelCase_ )
			model.load_state_dict(lowerCamelCase_ )
			model.eval()

			# Check outputs on an image, prepared by ConvNextImageProcessor
			__lowercase   =							convert_preprocessor(lowerCamelCase_ )
			__lowercase   =							preprocessor(images=prepare_img()						,				return_tensors='''pt''' )
			__lowercase   =							model(**lowerCamelCase_ ).logits

			# note: the logits below were obtained without center cropping
			if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
						__lowercase   =							torch.tensor([-0.39_30, 0.17_47, -0.52_46, 0.41_77, 0.42_95] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
						__lowercase   =							torch.tensor([-0.17_27, -0.53_41, -0.78_18, -0.47_45, -0.65_66] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
						__lowercase   =							torch.tensor([-0.03_33, 0.15_63, -0.91_37, 0.10_54, 0.03_81] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
						__lowercase   =							torch.tensor([-0.17_44, -0.15_55, -0.07_13, 0.09_50, -0.14_31] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
						__lowercase   =							torch.tensor([0.99_96, 0.19_66, -0.43_86, -0.34_72, 0.66_61] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
						__lowercase   =							torch.tensor([-0.25_53, -0.67_08, -0.13_59, 0.25_18, -0.24_88] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
						__lowercase   =							torch.tensor([-0.06_73, -0.56_27, -0.37_53, -0.27_22, 0.01_78] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
						__lowercase   =							torch.tensor([-0.63_77, -0.74_58, -0.21_50, 0.11_84, -0.05_97] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
						__lowercase   =							torch.tensor([1.07_99, 0.23_22, -0.88_60, 1.02_19, 0.62_31] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
						__lowercase   =							torch.tensor([0.37_66, 0.49_17, -1.14_26, 0.99_42, 0.60_24] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
						__lowercase   =							torch.tensor([0.42_20, -0.69_19, -0.43_17, -0.28_81, -0.66_09] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
						__lowercase   =							torch.tensor([0.10_82, -0.82_86, -0.50_95, 0.46_81, -0.80_85] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
						__lowercase   =							torch.tensor([-0.24_19, -0.62_21, 0.21_76, -0.09_80, -0.75_27] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
						__lowercase   =							torch.tensor([0.03_91, -0.43_71, 0.37_86, 0.12_51, -0.27_84] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
						__lowercase   =							torch.tensor([-0.05_04, 0.56_36, -0.17_29, -0.65_07, -0.39_49] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
						__lowercase   =							torch.tensor([0.35_60, 0.94_86, 0.31_49, -0.26_67, -0.51_38] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
						__lowercase   =							torch.tensor([-0.24_69, -0.45_50, -0.58_53, -0.08_10, 0.03_09] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
						__lowercase   =							torch.tensor([-0.30_90, 0.08_02, -0.06_82, -0.19_79, -0.28_26] )
			else:
						raise ValueError(f"Unknown URL: {checkpoint_url}" )

			assert torch.allclose(logits[0, :5]						,				lowerCamelCase_						,				atol=1E-3 )
			assert logits.shape == expected_shape
			print('''Model outputs match the original results!''' )

			if save_model:
						print('''Saving model to local...''' )
						# Create folder to save model
						if not os.path.isdir(lowerCamelCase_ ):
									os.mkdir(lowerCamelCase_ )

						model.save_pretrained(lowerCamelCase_ )
						preprocessor.save_pretrained(lowerCamelCase_ )

			__lowercase   =							'''convnextv2'''
			if "atto" in checkpoint_url:
						model_name += "-atto"
			if "femto" in checkpoint_url:
						model_name += "-femto"
			if "pico" in checkpoint_url:
						model_name += "-pico"
			if "nano" in checkpoint_url:
						model_name += "-nano"
			elif "tiny" in checkpoint_url:
						model_name += "-tiny"
			elif "base" in checkpoint_url:
						model_name += "-base"
			elif "large" in checkpoint_url:
						model_name += "-large"
			elif "huge" in checkpoint_url:
						model_name += "-huge"
			if "22k" in checkpoint_url and "1k" not in checkpoint_url:
						model_name += "-22k"
			elif "22k" in checkpoint_url and "1k" in checkpoint_url:
						model_name += "-22k-1k"
			elif "1k" in checkpoint_url:
						model_name += "-1k"
			if "224" in checkpoint_url:
						model_name += "-224"
			elif "384" in checkpoint_url:
						model_name += "-384"
			elif "512" in checkpoint_url:
						model_name += "-512"

			if push_to_hub:
						print(f"Pushing {model_name} to the hub..." )
						model.push_to_hub(lowerCamelCase_ )
						preprocessor.push_to_hub(lowerCamelCase_ )


if __name__ == "__main__":
	_SCREAMING_SNAKE_CASE    =			argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '''--checkpoint_url''',
	    default='''https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt''',
	    type=str,
	    help='''URL of the original ConvNeXTV2 checkpoint you\'d like to convert.''',
	)
	parser.add_argument(
	    '''--pytorch_dump_folder_path''',
	    default='''model''',
	    type=str,
	    help='''Path to the output PyTorch model directory.''',
	)
	parser.add_argument('''--save_model''', action='''store_true''', help='''Save model to local''')
	parser.add_argument('''--push_to_hub''', action='''store_true''', help='''Push model and image preprocessor to the hub''')

	_SCREAMING_SNAKE_CASE    =			parser.parse_args()
	convert_convnextva_checkpoint(
	    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
	)









'''simple docstring'''






def 					_snake_case   (     A       )						->    str:
							return " ".join(sorted(set(sentence.split()       )       )       )


if __name__ == "__main__":
		import doctest

		doctest.testmod()