





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



def _a       (			lowerCamelCase:     List[Any]				)			->		Optional[Any]:







						'''simple docstring'''






						__A      =       ConvNextVaConfig()

						if "atto" in checkpoint_url:
												__A      =       [2, 2, 6, 2]
												__A      =       [40, 80, 1_60, 3_20]
						if "femto" in checkpoint_url:
												__A      =       [2, 2, 6, 2]
												__A      =       [48, 96, 1_92, 3_84]
						if "pico" in checkpoint_url:
												__A      =       [2, 2, 6, 2]
												__A      =       [64, 1_28, 2_56, 5_12]
						if "nano" in checkpoint_url:
												__A      =       [2, 2, 8, 2]
												__A      =       [80, 1_60, 3_20, 6_40]
						if "tiny" in checkpoint_url:
												__A      =       [3, 3, 9, 3]
												__A      =       [96, 1_92, 3_84, 7_68]
						if "base" in checkpoint_url:
												__A      =       [3, 3, 27, 3]
												__A      =       [1_28, 2_56, 5_12, 10_24]
						if "large" in checkpoint_url:
												__A      =       [3, 3, 27, 3]
												__A      =       [1_92, 3_84, 7_68, 15_36]
						if "huge" in checkpoint_url:
												__A      =       [3, 3, 27, 3]
												__A      =       [3_52, 7_04, 14_08, 28_16]

						__A      =       10_00
						__A      =       '''imagenet-1k-id2label.json'''
						__A      =       (1, 10_00)

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

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

						return config, expected_shape



def _a       (			lowerCamelCase:     Any				)			->		List[Any]:







						'''simple docstring'''






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

						return name



def _a       (			)			->		Optional[int]:







						'''simple docstring'''






						__A      =       '''http://images.cocodataset.org/val2017/000000039769.jpg'''
						__A      =       Image.open(requests.get(lowerCamelCase      ,       stream=lowerCamelCase				).raw				)
						return im



def _a       (			lowerCamelCase:     List[Any]				)			->		Any:







						'''simple docstring'''






						if "224" in checkpoint_url:
												__A      =       2_24
												__A      =       2_24 / 2_56
						elif "384" in checkpoint_url:
												__A      =       3_84
												__A      =       None
						else:
												__A      =       5_12
												__A      =       None

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



@torch.no_grad()
def _a       (			lowerCamelCase:     Any      ,       lowerCamelCase:     int      ,       lowerCamelCase:     Tuple      ,       lowerCamelCase:     Optional[int]				)			->		Optional[Any]:







						'''simple docstring'''






						print('''Downloading original model from checkpoint...'''				)
						# define ConvNeXTV2 configuration based on URL
						__A     ,		__A      =       get_convnextva_config(lowerCamelCase				)
						# load original state_dict from URL
						__A      =       torch.hub.load_state_dict_from_url(lowerCamelCase				)['''model''']

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

						# load HuggingFace model
						__A      =       ConvNextVaForImageClassification(lowerCamelCase				)
						model.load_state_dict(lowerCamelCase				)
						model.eval()

						# Check outputs on an image, prepared by ConvNextImageProcessor
						__A      =       convert_preprocessor(lowerCamelCase				)
						__A      =       preprocessor(images=prepare_img()      ,       return_tensors='''pt'''				)
						__A      =       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":
												__A      =       torch.tensor([-0.3_930, 0.1_747, -0.5_246, 0.4_177, 0.4_295]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
												__A      =       torch.tensor([-0.1_727, -0.5_341, -0.7_818, -0.4_745, -0.6_566]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
												__A      =       torch.tensor([-0.0_333, 0.1_563, -0.9_137, 0.1_054, 0.0_381]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
												__A      =       torch.tensor([-0.1_744, -0.1_555, -0.0_713, 0.0_950, -0.1_431]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
												__A      =       torch.tensor([0.9_996, 0.1_966, -0.4_386, -0.3_472, 0.6_661]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
												__A      =       torch.tensor([-0.2_553, -0.6_708, -0.1_359, 0.2_518, -0.2_488]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
												__A      =       torch.tensor([-0.0_673, -0.5_627, -0.3_753, -0.2_722, 0.0_178]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
												__A      =       torch.tensor([-0.6_377, -0.7_458, -0.2_150, 0.1_184, -0.0_597]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
												__A      =       torch.tensor([1.0_799, 0.2_322, -0.8_860, 1.0_219, 0.6_231]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
												__A      =       torch.tensor([0.3_766, 0.4_917, -1.1_426, 0.9_942, 0.6_024]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
												__A      =       torch.tensor([0.4_220, -0.6_919, -0.4_317, -0.2_881, -0.6_609]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
												__A      =       torch.tensor([0.1_082, -0.8_286, -0.5_095, 0.4_681, -0.8_085]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
												__A      =       torch.tensor([-0.2_419, -0.6_221, 0.2_176, -0.0_980, -0.7_527]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
												__A      =       torch.tensor([0.0_391, -0.4_371, 0.3_786, 0.1_251, -0.2_784]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
												__A      =       torch.tensor([-0.0_504, 0.5_636, -0.1_729, -0.6_507, -0.3_949]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
												__A      =       torch.tensor([0.3_560, 0.9_486, 0.3_149, -0.2_667, -0.5_138]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
												__A      =       torch.tensor([-0.2_469, -0.4_550, -0.5_853, -0.0_810, 0.0_309]				)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
												__A      =       torch.tensor([-0.3_090, 0.0_802, -0.0_682, -0.1_979, -0.2_826]				)
						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				)

						__A      =       '''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__":
							snake_case__      :							Any									=  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')

							snake_case__      :							int									=  parser.parse_args()
							convert_convnextva_checkpoint(
							    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
							)












class A_      :




   def __init__(self			:Union[str, Any]				)->      None:
         # numbermap keys are the number and its values are lists of indexes sorted
         # in ascending order
         __A      =       {}
         # indexmap keys are an index and it's values are the number at that index
         __A      =       {}




   def _lowerCAmelCase   (self			:Any					,     _UpperCamelCase			:list | str | range					,     _UpperCamelCase			:int				)->      list[int]:
         if isinstance(_UpperCamelCase					,     (range, str)				):
               __A      =       list(_UpperCamelCase				)
         elif not isinstance(_UpperCamelCase					,     _UpperCamelCase				):
               raise TypeError(
                   '''binary_search_delete() only accepts either a list, range or str'''				)

         __A      =       0
         __A      =       len(_UpperCamelCase				) - 1

         while low <= high:
               __A      =       (low + high) // 2
               if array[mid] == item:
                     array.pop(_UpperCamelCase				)
                     return array
               elif array[mid] < item:
                     __A      =       mid + 1
               else:
                     __A      =       mid - 1
         raise ValueError(
             '''Either the item is not in the array or the array was unsorted'''				)




   def _lowerCAmelCase   (self			:Optional[Any]					,     _UpperCamelCase			:list | str | range					,     _UpperCamelCase			:int				)->      list[int]:
         if isinstance(_UpperCamelCase					,     (range, str)				):
               __A      =       list(_UpperCamelCase				)
         elif not isinstance(_UpperCamelCase					,     _UpperCamelCase				):
               raise TypeError(
                   '''binary_search_insert() only accepts either a list, range or str'''				)

         __A      =       0
         __A      =       len(_UpperCamelCase				) - 1

         while low <= high:
               __A      =       (low + high) // 2
               if array[mid] == index:
                     # If the item already exists in the array,
                     # insert it after the existing item
                     array.insert(mid + 1					,     _UpperCamelCase				)
                     return array
               elif array[mid] < index:
                     __A      =       mid + 1
               else:
                     __A      =       mid - 1

        # If the item doesn't exist in the array, insert it at the appropriate position
         array.insert(_UpperCamelCase					,     _UpperCamelCase				)
         return array




   def _lowerCAmelCase   (self			:Union[str, Any]					,     _UpperCamelCase			:int					,     _UpperCamelCase			:int				)->      None:
         # Remove previous index
         if index in self.indexmap:
               __A      =       self.indexmap[index]
               if len(self.numbermap[n]				) == 1:
                     del self.numbermap[n]
               else:
                     __A      =       self.binary_search_delete(self.numbermap[n]					,     _UpperCamelCase				)

        # Set new index
         __A      =       number

         # Number not seen before or empty so insert number value
         if number not in self.numbermap:
               __A      =       [index]

         # Here we need to perform a binary search insertion in order to insert
         # The item in the correct place
         else:
               __A      =       self.binary_search_insert(
                   self.numbermap[number]					,     _UpperCamelCase				)






   def _lowerCAmelCase   (self			:List[Any]					,     _UpperCamelCase			:int				)->      int:
         return self.numbermap.get(_UpperCamelCase					,     [-1]				)[0]


if __name__ == "__main__":
       import doctest

       doctest.testmod()





