


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()
A__						:						str          =      logging.get_logger(__name__)


def     UpperCamelCase(						__UpperCamelCase							:	List[str]				):
		lowerCAmelCase_    :  List[str]             =				ConvNextConfig()

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

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

		lowerCAmelCase_    :  Optional[int]             =				'''huggingface/label-files'''
		lowerCAmelCase_    :  Optional[Any]             =				num_labels
		lowerCAmelCase_    :  Tuple             =				json.load(open(hf_hub_download(__UpperCamelCase  ,__UpperCamelCase  ,repo_type='''dataset'''				)  ,'''r'''				)				)
		lowerCAmelCase_    :  Dict             =				{int(__UpperCamelCase				): 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]
		lowerCAmelCase_    :  Union[str, Any]             =				idalabel
		lowerCAmelCase_    :  str             =				{v: k for k, v in idalabel.items()}
		lowerCAmelCase_    :  List[Any]             =				hidden_sizes
		lowerCAmelCase_    :  Any             =				depths

		return config, expected_shape


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

		return name


def     UpperCamelCase(						):
		lowerCAmelCase_    :  List[str]             =				'''http://images.cocodataset.org/val2017/000000039769.jpg'''
		lowerCAmelCase_    :  List[str]             =				Image.open(requests.get(__UpperCamelCase  ,stream=__UpperCamelCase				).raw				)
		return im


@torch.no_grad()
def     UpperCamelCase(						__UpperCamelCase							:	int  ,__UpperCamelCase							:	Optional[int]				):
		lowerCAmelCase_    :  Optional[int]             =				get_convnext_config(__UpperCamelCase				)
		# load original state_dict from URL
		lowerCAmelCase_    :  List[str]             =				torch.hub.load_state_dict_from_url(__UpperCamelCase				)['''model''']
		# rename keys
		for key in state_dict.copy().keys():
				lowerCAmelCase_    :  List[str]             =				state_dict.pop(__UpperCamelCase				)
				lowerCAmelCase_    :  List[Any]             =				val
		# add prefix to all keys expect classifier head
		for key in state_dict.copy().keys():
				lowerCAmelCase_    :  Tuple             =				state_dict.pop(__UpperCamelCase				)
				if not key.startswith('''classifier'''				):
						lowerCAmelCase_    :  Union[str, Any]             =				'''convnext.''' + key
				lowerCAmelCase_    :  int             =				val

		# load HuggingFace model
		lowerCAmelCase_    :  List[Any]             =				ConvNextForImageClassification(__UpperCamelCase				)
		model.load_state_dict(__UpperCamelCase				)
		model.eval()

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

		lowerCAmelCase_    :  List[Any]             =				model(__UpperCamelCase				).logits

		# note: the logits below were obtained without center cropping
		if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
				lowerCAmelCase_    :  Dict             =				torch.tensor([-0.1_2_1_0, -0.6_6_0_5, 0.1_9_1_8]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
				lowerCAmelCase_    :  Any             =				torch.tensor([-0.4_4_7_3, -0.1_8_4_7, -0.6_3_6_5]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
				lowerCAmelCase_    :  List[str]             =				torch.tensor([0.4_5_2_5, 0.7_5_3_9, 0.0_3_0_8]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
				lowerCAmelCase_    :  Optional[int]             =				torch.tensor([0.3_5_6_1, 0.6_3_5_0, -0.0_3_8_4]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
				lowerCAmelCase_    :  Tuple             =				torch.tensor([0.4_1_7_4, -0.0_9_8_9, 0.1_4_8_9]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
				lowerCAmelCase_    :  Any             =				torch.tensor([0.2_5_1_3, -0.1_3_4_9, -0.1_6_1_3]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
				lowerCAmelCase_    :  str             =				torch.tensor([1.2_9_8_0, 0.3_6_3_1, -0.1_1_9_8]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
				lowerCAmelCase_    :  List[Any]             =				torch.tensor([1.2_9_6_3, 0.1_2_2_7, 0.1_7_2_3]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
				lowerCAmelCase_    :  Union[str, Any]             =				torch.tensor([1.7_9_5_6, 0.8_3_9_0, 0.2_8_2_0]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
				lowerCAmelCase_    :  List[Any]             =				torch.tensor([-0.2_8_2_2, -0.0_5_0_2, -0.0_8_7_8]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
				lowerCAmelCase_    :  Union[str, Any]             =				torch.tensor([-0.5_6_7_2, -0.0_7_3_0, -0.4_3_4_8]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
				lowerCAmelCase_    :  Union[str, Any]             =				torch.tensor([0.2_6_8_1, 0.2_3_6_5, 0.6_2_4_6]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
				lowerCAmelCase_    :  Optional[Any]             =				torch.tensor([-0.2_6_4_2, 0.3_9_3_1, 0.5_1_1_6]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
				lowerCAmelCase_    :  int             =				torch.tensor([-0.6_6_7_7, -0.1_8_7_3, -0.8_3_7_9]				)
		elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
				lowerCAmelCase_    :  int             =				torch.tensor([-0.7_7_4_9, -0.2_9_6_7, -0.6_4_4_4]				)
		else:
				raise ValueError(f"""Unknown URL: {checkpoint_url}"""				)

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

		Path(__UpperCamelCase				).mkdir(exist_ok=__UpperCamelCase				)
		print(f"""Saving model to {pytorch_dump_folder_path}"""				)
		model.save_pretrained(__UpperCamelCase				)
		print(f"""Saving image processor to {pytorch_dump_folder_path}"""				)
		image_processor.save_pretrained(__UpperCamelCase				)

		print('''Pushing model to the hub...'''				)
		lowerCAmelCase_    :  Union[str, 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(__UpperCamelCase  ,__UpperCamelCase				)  ,organization='''nielsr'''  ,commit_message='''Add model'''  ,)


if __name__ == "__main__":
						A__						:						Tuple          =      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.''',
						)

						A__						:						Union[str, Any]          =      parser.parse_args()
						convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)










from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


A__						:						int          =      logging.get_logger(__name__)

A__						:						Dict          =      {
    '''mnaylor/mega-base-wikitext''': '''https://huggingface.co/mnaylor/mega-base-wikitext/resolve/main/config.json''',
}



class        __snake_case		(  UpperCamelCase_		):
							_a       =				'''mega'''




							def __init__( self  : Any       ,       A_  : Dict=3_0_5_2_2       ,       A_  : int=1_2_8       ,       A_  : str=4       ,       A_  : int=2_5_6       ,       A_  : int=1_6       ,       A_  : str=True       ,       A_  : List[str]=6_4       ,       A_  : Optional[Any]=False       ,       A_  : str=-1       ,       A_  : Optional[int]=None       ,       A_  : Tuple=True       ,       A_  : Tuple="scalenorm"       ,       A_  : Optional[Any]=True       ,       A_  : int="silu"       ,       A_  : Optional[int]="softmax"       ,       A_  : Optional[int]=0.1       ,       A_  : Union[str, Any]=0.1       ,       A_  : Any=0.1       ,       A_  : Dict=False       ,       A_  : int=True       ,       A_  : Tuple=2_5_6       ,       A_  : int=True       ,       A_  : List[Any]=0.1       ,       A_  : Optional[Any]=2_0_4_8       ,       A_  : Dict=False       ,       A_  : Dict=2       ,       A_  : Tuple=0.02       ,       A_  : List[Any]=0.2       ,       A_  : str=0.02       ,       A_  : List[str]=1.0       ,       A_  : Union[str, Any]=1       ,       A_  : List[Any]=0       ,       A_  : int=2       ,       A_  : Tuple="rotary"       ,       A_  : Any=None       ,       A_  : Tuple=True       ,       A_  : Any=True       ,       **A_  : Optional[int]       ,       ):
									super().__init__(pad_token_id=A_       ,       bos_token_id=A_       ,       eos_token_id=A_       ,       **A_)

									lowerCAmelCase_    :  Tuple             =				vocab_size
									lowerCAmelCase_    :  str             =				hidden_size
									lowerCAmelCase_    :  Any             =				num_hidden_layers
									lowerCAmelCase_    :  Optional[int]             =				activation
									lowerCAmelCase_    :  Any             =				attention_activation
									lowerCAmelCase_    :  List[Any]             =				intermediate_size
									lowerCAmelCase_    :  Any             =				ema_projection_size
									lowerCAmelCase_    :  Optional[int]             =				bidirectional
									lowerCAmelCase_    :  Tuple             =				shared_representation_size
									lowerCAmelCase_    :  Any             =				use_chunking
									lowerCAmelCase_    :  List[str]             =				chunk_size
									lowerCAmelCase_    :  int             =				truncation
									lowerCAmelCase_    :  Tuple             =				normalize_before_mega
									lowerCAmelCase_    :  Tuple             =				normalization_type
									lowerCAmelCase_    :  List[str]             =				norm_affine
									lowerCAmelCase_    :  List[str]             =				dropout_prob
									lowerCAmelCase_    :  List[str]             =				hidden_dropout_prob
									lowerCAmelCase_    :  Tuple             =				attention_probs_dropout_prob
									lowerCAmelCase_    :  str             =				use_feature_dropout
									lowerCAmelCase_    :  Union[str, Any]             =				use_normalized_ffn
									lowerCAmelCase_    :  Union[str, Any]             =				nffn_hidden_size
									lowerCAmelCase_    :  int             =				normalize_before_ffn
									lowerCAmelCase_    :  Tuple             =				nffn_activation_dropout_prob
									lowerCAmelCase_    :  List[str]             =				max_positions
									lowerCAmelCase_    :  Union[str, Any]             =				add_token_type_embeddings
									lowerCAmelCase_    :  str             =				type_vocab_size
									lowerCAmelCase_    :  int             =				initializer_range
									lowerCAmelCase_    :  str             =				ema_delta_alpha_range
									lowerCAmelCase_    :  Tuple             =				ema_beta_range
									lowerCAmelCase_    :  Any             =				ema_gamma_omega_range
									lowerCAmelCase_    :  Any             =				relative_positional_bias
									lowerCAmelCase_    :  Optional[int]             =				use_cache
									lowerCAmelCase_    :  int             =				classifier_dropout
									lowerCAmelCase_    :  List[Any]             =				add_lm_hidden_dense_layer
									lowerCAmelCase_    :  Optional[Any]             =				1  # not used but required by Hugging Face





class        __snake_case		(  UpperCamelCase_		):




							@property
							def      UpperCAmelCase__    ( self  : Any):
									if self.task == "multiple-choice":
											lowerCAmelCase_    :  Dict             =				{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
									else:
											lowerCAmelCase_    :  int             =				{0: '''batch''', 1: '''sequence'''}
									return OrderedDict(
									    [
									        ('''input_ids''', dynamic_axis),
									        ('''attention_mask''', dynamic_axis),
									    ])





