



import argparse
import json
from pathlib import Path

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

from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


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


def   lowercase_	(   lowerCAmelCase__	:  Dict ):

			"""simple docstring"""





			__UpperCAmelCase   :      Any										=     ConvNextConfig()

			if "tiny" in checkpoint_url:
						__UpperCAmelCase   :      Any										=     [3, 3, 9, 3]
						__UpperCAmelCase   :      List[Any]										=     [96, 192, 384, 768]
			if "small" in checkpoint_url:
						__UpperCAmelCase   :      List[Any]										=     [3, 3, 27, 3]
						__UpperCAmelCase   :      int										=     [96, 192, 384, 768]
			if "base" in checkpoint_url:
						__UpperCAmelCase   :      int										=     [3, 3, 27, 3]
						__UpperCAmelCase   :      Tuple										=     [128, 256, 512, 1024]
			if "large" in checkpoint_url:
						__UpperCAmelCase   :      Tuple										=     [3, 3, 27, 3]
						__UpperCAmelCase   :      str										=     [192, 384, 768, 1536]
			if "xlarge" in checkpoint_url:
						__UpperCAmelCase   :      str										=     [3, 3, 27, 3]
						__UpperCAmelCase   :      Tuple										=     [256, 512, 1024, 2048]

			if "1k" in checkpoint_url:
						__UpperCAmelCase   :      Union[str, Any]										=     1000
						__UpperCAmelCase   :      List[Any]										=     '''imagenet-1k-id2label.json'''
						__UpperCAmelCase   :      Union[str, Any]										=     (1, 1000)
			else:
						__UpperCAmelCase   :      Optional[Any]										=     21841
						__UpperCAmelCase   :      str										=     '''imagenet-22k-id2label.json'''
						__UpperCAmelCase   :      Optional[int]										=     (1, 21841)

			__UpperCAmelCase   :      Any										=     '''huggingface/label-files'''
			__UpperCAmelCase   :      Any										=     num_labels
			__UpperCAmelCase   :      Tuple										=     json.load(open(hf_hub_download(lowerCAmelCase__			,							lowerCAmelCase__			,							repo_type="""dataset""" )			,							"""r""" ) )
			__UpperCAmelCase   :      int										=     {int(lowerCAmelCase__ ): v for k, v in idalabel.items()}
			if "1k" not in checkpoint_url:
						# this dataset contains 21843 labels but the model only has 21841
						# we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
						del idalabel[9205]
						del idalabel[15027]
			__UpperCAmelCase   :      str										=     idalabel
			__UpperCAmelCase   :      List[str]										=     {v: k for k, v in idalabel.items()}
			__UpperCAmelCase   :      Tuple										=     hidden_sizes
			__UpperCAmelCase   :      List[str]										=     depths

			return config, expected_shape


def   lowercase_	(   lowerCAmelCase__	:  Optional[Any] ):

			"""simple docstring"""





			if "downsample_layers.0.0" in name:
						__UpperCAmelCase   :      Any										=     name.replace("""downsample_layers.0.0"""			,							"""embeddings.patch_embeddings""" )
			if "downsample_layers.0.1" in name:
						__UpperCAmelCase   :      Any										=     name.replace("""downsample_layers.0.1"""			,							"""embeddings.norm""" )  # we rename to layernorm later on
			if "downsample_layers.1.0" in name:
						__UpperCAmelCase   :      int										=     name.replace("""downsample_layers.1.0"""			,							"""stages.1.downsampling_layer.0""" )
			if "downsample_layers.1.1" in name:
						__UpperCAmelCase   :      Optional[int]										=     name.replace("""downsample_layers.1.1"""			,							"""stages.1.downsampling_layer.1""" )
			if "downsample_layers.2.0" in name:
						__UpperCAmelCase   :      Optional[int]										=     name.replace("""downsample_layers.2.0"""			,							"""stages.2.downsampling_layer.0""" )
			if "downsample_layers.2.1" in name:
						__UpperCAmelCase   :      int										=     name.replace("""downsample_layers.2.1"""			,							"""stages.2.downsampling_layer.1""" )
			if "downsample_layers.3.0" in name:
						__UpperCAmelCase   :      Tuple										=     name.replace("""downsample_layers.3.0"""			,							"""stages.3.downsampling_layer.0""" )
			if "downsample_layers.3.1" in name:
						__UpperCAmelCase   :      Any										=     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.
						__UpperCAmelCase   :      Dict										=     name[: len("""stages.0""" )] + '''.layers''' + name[len("""stages.0""" ) :]
			if "stages" in name:
						__UpperCAmelCase   :      List[str]										=     name.replace("""stages"""			,							"""encoder.stages""" )
			if "norm" in name:
						__UpperCAmelCase   :      Union[str, Any]										=     name.replace("""norm"""			,							"""layernorm""" )
			if "gamma" in name:
						__UpperCAmelCase   :      Any										=     name.replace("""gamma"""			,							"""layer_scale_parameter""" )
			if "head" in name:
						__UpperCAmelCase   :      Optional[int]										=     name.replace("""head"""			,							"""classifier""" )

			return name


def   lowercase_	(   ):

			"""simple docstring"""





			__UpperCAmelCase   :      Union[str, Any]										=     '''http://images.cocodataset.org/val2017/000000039769.jpg'''
			__UpperCAmelCase   :      Union[str, Any]										=     Image.open(requests.get(lowerCAmelCase__			,							stream=lowerCAmelCase__ ).raw )
			return im


@torch.no_grad()
def   lowercase_	(   lowerCAmelCase__	:  Optional[Any]			,							lowerCAmelCase__	:  Optional[Any] ):

			"""simple docstring"""





			__UpperCAmelCase   :      int										=     get_convnext_config(lowerCAmelCase__ )
			# load original state_dict from URL
			__UpperCAmelCase   :      Any										=     torch.hub.load_state_dict_from_url(lowerCAmelCase__ )['''model''']
			# rename keys
			for key in state_dict.copy().keys():
						__UpperCAmelCase   :      List[str]										=     state_dict.pop(lowerCAmelCase__ )
						__UpperCAmelCase   :      Tuple										=     val
			# add prefix to all keys expect classifier head
			for key in state_dict.copy().keys():
						__UpperCAmelCase   :      List[Any]										=     state_dict.pop(lowerCAmelCase__ )
						if not key.startswith("""classifier""" ):
									__UpperCAmelCase   :      Dict										=     '''convnext.''' + key
						__UpperCAmelCase   :      Optional[int]										=     val

			# load HuggingFace model
			__UpperCAmelCase   :      Dict										=     ConvNextForImageClassification(lowerCAmelCase__ )
			model.load_state_dict(lowerCAmelCase__ )
			model.eval()

			# Check outputs on an image, prepared by ConvNextImageProcessor
			__UpperCAmelCase   :      Optional[int]										=     224 if '''224''' in checkpoint_url else 384
			__UpperCAmelCase   :      int										=     ConvNextImageProcessor(size=lowerCAmelCase__ )
			__UpperCAmelCase   :      Union[str, Any]										=     image_processor(images=prepare_img()			,							return_tensors="""pt""" ).pixel_values

			__UpperCAmelCase   :      str										=     model(lowerCAmelCase__ ).logits

			# note: the logits below were obtained without center cropping
			if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
						__UpperCAmelCase   :      Dict										=     torch.tensor([-0.1_210, -0.6_605, 0.1_918] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
						__UpperCAmelCase   :      Union[str, Any]										=     torch.tensor([-0.4_473, -0.1_847, -0.6_365] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
						__UpperCAmelCase   :      List[Any]										=     torch.tensor([0.4_525, 0.7_539, 0.0_308] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
						__UpperCAmelCase   :      Optional[int]										=     torch.tensor([0.3_561, 0.6_350, -0.0_384] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
						__UpperCAmelCase   :      int										=     torch.tensor([0.4_174, -0.0_989, 0.1_489] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
						__UpperCAmelCase   :      List[Any]										=     torch.tensor([0.2_513, -0.1_349, -0.1_613] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
						__UpperCAmelCase   :      List[Any]										=     torch.tensor([1.2_980, 0.3_631, -0.1_198] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
						__UpperCAmelCase   :      str										=     torch.tensor([1.2_963, 0.1_227, 0.1_723] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
						__UpperCAmelCase   :      Union[str, Any]										=     torch.tensor([1.7_956, 0.8_390, 0.2_820] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
						__UpperCAmelCase   :      Tuple										=     torch.tensor([-0.2_822, -0.0_502, -0.0_878] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
						__UpperCAmelCase   :      Tuple										=     torch.tensor([-0.5_672, -0.0_730, -0.4_348] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
						__UpperCAmelCase   :      Union[str, Any]										=     torch.tensor([0.2_681, 0.2_365, 0.6_246] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
						__UpperCAmelCase   :      Optional[int]										=     torch.tensor([-0.2_642, 0.3_931, 0.5_116] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
						__UpperCAmelCase   :      Any										=     torch.tensor([-0.6_677, -0.1_873, -0.8_379] )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
						__UpperCAmelCase   :      str										=     torch.tensor([-0.7_749, -0.2_967, -0.6_444] )
			else:
						raise ValueError(f'Unknown URL: {checkpoint_url}' )

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

			Path(lowerCAmelCase__ ).mkdir(exist_ok=lowerCAmelCase__ )
			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__ )

			print("""Pushing model to the hub...""" )
			__UpperCAmelCase   :      Any										=     '''convnext'''
			if "tiny" in checkpoint_url:
						model_name += "-tiny"
			elif "small" in checkpoint_url:
						model_name += "-small"
			elif "base" in checkpoint_url:
						model_name += "-base"
			elif "xlarge" in checkpoint_url:
						model_name += "-xlarge"
			elif "large" in checkpoint_url:
						model_name += "-large"
			if "224" in checkpoint_url:
						model_name += "-224"
			elif "384" in checkpoint_url:
						model_name += "-384"
			if "22k" in checkpoint_url and "1k" not in checkpoint_url:
						model_name += "-22k"
			if "22k" in checkpoint_url and "1k" in checkpoint_url:
						model_name += "-22k-1k"

			model.push_to_hub(
			    repo_path_or_name=Path(lowerCAmelCase__			,							lowerCAmelCase__ )			,							organization="""nielsr"""			,							commit_message="""Add model"""			,							)


if __name__ == "__main__":
					_UpperCamelCase											=			argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '''--checkpoint_url''',
					    default='''https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth''',
					    type=str,
					    help='''URL of the original ConvNeXT checkpoint you\'d like to convert.''',
					)
					parser.add_argument(
					    '''--pytorch_dump_folder_path''',
					    default=None,
					    type=str,
					    required=True,
					    help='''Path to the output PyTorch model directory.''',
					)

					_UpperCamelCase											=			parser.parse_args()
					convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









'''simple docstring'''
from math import pi


def   lowercase_	(   lowerCAmelCase__	:  float ):

			"""simple docstring"""





			return degree / (180 / pi)


if __name__ == "__main__":
					from doctest import testmod

					testmod()




