




"""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()
__lowerCamelCase								=						logging.get_logger(__name__)



def 				a (						__snake_case     :				List[str]  ):



				'''simple docstring'''





				UpperCAmelCase_					:Optional[Any]      			=      ConvNextVaConfig()

				if "atto" in checkpoint_url:
								UpperCAmelCase_					:int      			=      [2, 2, 6, 2]
								UpperCAmelCase_					:int      			=      [40, 80, 160, 320]
				if "femto" in checkpoint_url:
								UpperCAmelCase_					:Dict      			=      [2, 2, 6, 2]
								UpperCAmelCase_					:Optional[Any]      			=      [48, 96, 192, 384]
				if "pico" in checkpoint_url:
								UpperCAmelCase_					:Any      			=      [2, 2, 6, 2]
								UpperCAmelCase_					:List[Any]      			=      [64, 128, 256, 512]
				if "nano" in checkpoint_url:
								UpperCAmelCase_					:Tuple      			=      [2, 2, 8, 2]
								UpperCAmelCase_					:List[Any]      			=      [80, 160, 320, 640]
				if "tiny" in checkpoint_url:
								UpperCAmelCase_					:Tuple      			=      [3, 3, 9, 3]
								UpperCAmelCase_					:Union[str, Any]      			=      [96, 192, 384, 768]
				if "base" in checkpoint_url:
								UpperCAmelCase_					:Dict      			=      [3, 3, 27, 3]
								UpperCAmelCase_					:int      			=      [128, 256, 512, 1024]
				if "large" in checkpoint_url:
								UpperCAmelCase_					:Optional[Any]      			=      [3, 3, 27, 3]
								UpperCAmelCase_					:Optional[int]      			=      [192, 384, 768, 1536]
				if "huge" in checkpoint_url:
								UpperCAmelCase_					:int      			=      [3, 3, 27, 3]
								UpperCAmelCase_					:Tuple      			=      [352, 704, 1408, 2816]

				UpperCAmelCase_					:List[str]      			=      1000
				UpperCAmelCase_					:List[Any]      			=      "imagenet-1k-id2label.json"
				UpperCAmelCase_					:Dict      			=      (1, 1000)

				UpperCAmelCase_					:List[Any]      			=      "huggingface/label-files"
				UpperCAmelCase_					:str      			=      num_labels
				UpperCAmelCase_					:List[str]      			=      json.load(open(hf_hub_download(__snake_case,    __snake_case,    repo_type='''dataset'''  ),    '''r'''  )  )
				UpperCAmelCase_					:Tuple      			=      {int(__snake_case  ): v for k, v in idalabel.items()}

				UpperCAmelCase_					:Optional[Any]      			=      idalabel
				UpperCAmelCase_					:int      			=      {v: k for k, v in idalabel.items()}
				UpperCAmelCase_					:Dict      			=      hidden_sizes
				UpperCAmelCase_					:Optional[Any]      			=      depths

				return config, expected_shape



def 				a (						__snake_case     :				int  ):



				'''simple docstring'''





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

				return name



def 				a (						):



				'''simple docstring'''





				UpperCAmelCase_					:List[str]      			=      "http://images.cocodataset.org/val2017/000000039769.jpg"
				UpperCAmelCase_					:Dict      			=      Image.open(requests.get(__snake_case,    stream=__snake_case  ).raw  )
				return im



def 				a (						__snake_case     :				Any  ):



				'''simple docstring'''





				if "224" in checkpoint_url:
								UpperCAmelCase_					:Optional[int]      			=      224
								UpperCAmelCase_					:str      			=      224 / 256
				elif "384" in checkpoint_url:
								UpperCAmelCase_					:Tuple      			=      384
								UpperCAmelCase_					:Tuple      			=      None
				else:
								UpperCAmelCase_					:Optional[Any]      			=      512
								UpperCAmelCase_					:int      			=      None

				return ConvNextImageProcessor(
				    size=__snake_case,    crop_pct=__snake_case,    image_mean=[0.485, 0.456, 0.406],    image_std=[0.229, 0.224, 0.225],    resample=PILImageResampling.BICUBIC,    )



@torch.no_grad()
def 				a (						__snake_case     :				Optional[Any],    __snake_case     :				Tuple,    __snake_case     :				List[Any],    __snake_case     :				List[Any]  ):



				'''simple docstring'''





				print('''Downloading original model from checkpoint...'''  )
				# define ConvNeXTV2 configuration based on URL
				UpperCAmelCase_					:List[Any]      			=      get_convnextva_config(__snake_case  )
				# load original state_dict from URL
				UpperCAmelCase_					:Optional[int]      			=      torch.hub.load_state_dict_from_url(__snake_case  )["model"]

				print('''Converting model parameters...'''  )
				# rename keys
				for key in state_dict.copy().keys():
								UpperCAmelCase_					:List[Any]      			=      state_dict.pop(__snake_case  )
								UpperCAmelCase_					:Optional[int]      			=      val
				# add prefix to all keys expect classifier head
				for key in state_dict.copy().keys():
								UpperCAmelCase_					:Optional[Any]      			=      state_dict.pop(__snake_case  )
								if not key.startswith('''classifier'''  ):
												UpperCAmelCase_					:str      			=      "convnextv2." + key
								UpperCAmelCase_					:List[str]      			=      val

				# load HuggingFace model
				UpperCAmelCase_					:List[str]      			=      ConvNextVaForImageClassification(__snake_case  )
				model.load_state_dict(__snake_case  )
				model.eval()

				# Check outputs on an image, prepared by ConvNextImageProcessor
				UpperCAmelCase_					:Optional[Any]      			=      convert_preprocessor(__snake_case  )
				UpperCAmelCase_					:Any      			=      preprocessor(images=prepare_img(),    return_tensors='''pt'''  )
				UpperCAmelCase_					:Dict      			=      model(**__snake_case  ).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":
								UpperCAmelCase_					:Dict      			=      torch.tensor([-0.3930, 0.1747, -0.5246, 0.4177, 0.4295]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
								UpperCAmelCase_					:List[Any]      			=      torch.tensor([-0.1727, -0.5341, -0.7818, -0.4745, -0.6566]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
								UpperCAmelCase_					:int      			=      torch.tensor([-0.0333, 0.1563, -0.9137, 0.1054, 0.0381]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
								UpperCAmelCase_					:str      			=      torch.tensor([-0.1744, -0.1555, -0.0713, 0.0950, -0.1431]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
								UpperCAmelCase_					:Optional[int]      			=      torch.tensor([0.9996, 0.1966, -0.4386, -0.3472, 0.6661]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
								UpperCAmelCase_					:Dict      			=      torch.tensor([-0.2553, -0.6708, -0.1359, 0.2518, -0.2488]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
								UpperCAmelCase_					:Union[str, Any]      			=      torch.tensor([-0.0673, -0.5627, -0.3753, -0.2722, 0.0178]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
								UpperCAmelCase_					:List[Any]      			=      torch.tensor([-0.6377, -0.7458, -0.2150, 0.1184, -0.0597]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
								UpperCAmelCase_					:Union[str, Any]      			=      torch.tensor([1.0799, 0.2322, -0.8860, 1.0219, 0.6231]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
								UpperCAmelCase_					:str      			=      torch.tensor([0.3766, 0.4917, -1.1426, 0.9942, 0.6024]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
								UpperCAmelCase_					:Optional[Any]      			=      torch.tensor([0.4220, -0.6919, -0.4317, -0.2881, -0.6609]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
								UpperCAmelCase_					:List[Any]      			=      torch.tensor([0.1082, -0.8286, -0.5095, 0.4681, -0.8085]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
								UpperCAmelCase_					:List[str]      			=      torch.tensor([-0.2419, -0.6221, 0.2176, -0.0980, -0.7527]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
								UpperCAmelCase_					:Tuple      			=      torch.tensor([0.0391, -0.4371, 0.3786, 0.1251, -0.2784]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
								UpperCAmelCase_					:List[Any]      			=      torch.tensor([-0.0504, 0.5636, -0.1729, -0.6507, -0.3949]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
								UpperCAmelCase_					:str      			=      torch.tensor([0.3560, 0.9486, 0.3149, -0.2667, -0.5138]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
								UpperCAmelCase_					:int      			=      torch.tensor([-0.2469, -0.4550, -0.5853, -0.0810, 0.0309]  )
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
								UpperCAmelCase_					:Union[str, Any]      			=      torch.tensor([-0.3090, 0.0802, -0.0682, -0.1979, -0.2826]  )
				else:
								raise ValueError(f'Unknown URL: {checkpoint_url}'  )

				assert torch.allclose(logits[0, :5],    __snake_case,    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(__snake_case  ):
												os.mkdir(__snake_case  )

								model.save_pretrained(__snake_case  )
								preprocessor.save_pretrained(__snake_case  )

				UpperCAmelCase_					:Optional[Any]      			=      "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(__snake_case  )
								preprocessor.push_to_hub(__snake_case  )


if __name__ == "__main__":
						__lowerCamelCase								=						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")

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












"""simple docstring"""

from collections import OrderedDict
from typing import Mapping

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


__lowerCamelCase								=						logging.get_logger(__name__)

__lowerCamelCase								=						{
    "mnaylor/mega-base-wikitext": "https://huggingface.co/mnaylor/mega-base-wikitext/resolve/main/config.json",
}



class  _snake_case				(							A__       ):

				'''simple docstring'''




				UpperCamelCase__		     ="""mega"""





				def __init__(							self      :  Optional[int]					,							snake_case      :  Union[str, Any]=30_522					,							snake_case      :  Optional[Any]=128					,							snake_case      :  str=4					,							snake_case      :  Dict=256					,							snake_case      :  Optional[Any]=16					,							snake_case      :  Any=True					,							snake_case      :  str=64					,							snake_case      :  Optional[Any]=False					,							snake_case      :  Union[str, Any]=-1					,							snake_case      :  List[str]=None					,							snake_case      :  Any=True					,							snake_case      :  Optional[int]="scalenorm"					,							snake_case      :  Tuple=True					,							snake_case      :  str="silu"					,							snake_case      :  Dict="softmax"					,							snake_case      :  Optional[int]=0.1					,							snake_case      :  Tuple=0.1					,							snake_case      :  Dict=0.1					,							snake_case      :  List[str]=False					,							snake_case      :  Any=True					,							snake_case      :  List[Any]=256					,							snake_case      :  Optional[Any]=True					,							snake_case      :  Optional[int]=0.1					,							snake_case      :  str=2_048					,							snake_case      :  Union[str, Any]=False					,							snake_case      :  List[Any]=2					,							snake_case      :  str=0.02					,							snake_case      :  str=0.2					,							snake_case      :  str=0.02					,							snake_case      :  Union[str, Any]=1.0					,							snake_case      :  Optional[Any]=1					,							snake_case      :  List[Any]=0					,							snake_case      :  List[Any]=2					,							snake_case      :  List[Any]="rotary"					,							snake_case      :  Tuple=None					,							snake_case      :  int=True					,							snake_case      :  Dict=True					,							**snake_case      :  List[str]					,							):
								super().__init__(pad_token_id=snake_case					,							bos_token_id=snake_case					,							eos_token_id=snake_case					,							**snake_case     )

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



class  _snake_case				(							A__       ):

				'''simple docstring'''




				@property
				def 				snake_case_ (							self      :  int     ):
								if self.task == "multiple-choice":
												UpperCAmelCase_					:List[str]      			=      {0: '''batch''', 1: '''choice''', 2: '''sequence'''}
								else:
												UpperCAmelCase_					:Tuple      			=      {0: '''batch''', 1: '''sequence'''}
								return OrderedDict(
								    [
								        ('''input_ids''', dynamic_axis),
								        ('''attention_mask''', dynamic_axis),
								    ]     )


