





"""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   __lowerCamelCase     (							a_  :	int							)    ->   Tuple:
      __SCREAMING_SNAKE_CASE     :List[str]	       =						ConvNextVaConfig()

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

      __SCREAMING_SNAKE_CASE     :Dict	       =						10_00
      __SCREAMING_SNAKE_CASE     :List[Any]	       =						'''imagenet-1k-id2label.json'''
      __SCREAMING_SNAKE_CASE     :List[str]	       =						(1, 10_00)

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

      __SCREAMING_SNAKE_CASE     :Any	       =						idalabel
      __SCREAMING_SNAKE_CASE     :List[Any]	       =						{v: k for k, v in idalabel.items()}
      __SCREAMING_SNAKE_CASE     :List[Any]	       =						hidden_sizes
      __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						depths

      return config, expected_shape



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

      return name



def   __lowerCamelCase     (							)    ->   Any:
      __SCREAMING_SNAKE_CASE     :Any	       =						'''http://images.cocodataset.org/val2017/000000039769.jpg'''
      __SCREAMING_SNAKE_CASE     :List[str]	       =						Image.open(requests.get(a_				,     stream=a_							).raw							)
      return im



def   __lowerCamelCase     (							a_  :	str							)    ->   Optional[Any]:
      if "224" in checkpoint_url:
            __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						2_24
            __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						2_24 / 2_56
      elif "384" in checkpoint_url:
            __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						3_84
            __SCREAMING_SNAKE_CASE     :Tuple	       =						None
      else:
            __SCREAMING_SNAKE_CASE     :str	       =						5_12
            __SCREAMING_SNAKE_CASE     :int	       =						None

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



@torch.no_grad()
def   __lowerCamelCase     (							a_  :	Tuple				,     a_  :	Optional[Any]				,     a_  :	Optional[Any]				,     a_  :	Tuple							)    ->   int:
      print('''Downloading original model from checkpoint...'''							)
      # define ConvNeXTV2 configuration based on URL
      __SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE     :Tuple	       =						get_convnextva_config(a_							)
      # load original state_dict from URL
      __SCREAMING_SNAKE_CASE     :Tuple	       =						torch.hub.load_state_dict_from_url(a_							)['''model''']

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

      # load HuggingFace model
      __SCREAMING_SNAKE_CASE     :int	       =						ConvNextVaForImageClassification(a_							)
      model.load_state_dict(a_							)
      model.eval()

      # Check outputs on an image, prepared by ConvNextImageProcessor
      __SCREAMING_SNAKE_CASE     :Dict	       =						convert_preprocessor(a_							)
      __SCREAMING_SNAKE_CASE     :Dict	       =						preprocessor(images=prepare_img()				,     return_tensors='''pt'''							)
      __SCREAMING_SNAKE_CASE     :int	       =						model(**a_							).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":
            __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						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":
            __SCREAMING_SNAKE_CASE     :Optional[int]	       =						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":
            __SCREAMING_SNAKE_CASE     :List[str]	       =						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":
            __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						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":
            __SCREAMING_SNAKE_CASE     :Tuple	       =						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":
            __SCREAMING_SNAKE_CASE     :int	       =						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":
            __SCREAMING_SNAKE_CASE     :Optional[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":
            __SCREAMING_SNAKE_CASE     :Dict	       =						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":
            __SCREAMING_SNAKE_CASE     :List[str]	       =						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":
            __SCREAMING_SNAKE_CASE     :Optional[int]	       =						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":
            __SCREAMING_SNAKE_CASE     :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":
            __SCREAMING_SNAKE_CASE     :Tuple	       =						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":
            __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						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":
            __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						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":
            __SCREAMING_SNAKE_CASE     :Dict	       =						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":
            __SCREAMING_SNAKE_CASE     :Any	       =						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":
            __SCREAMING_SNAKE_CASE     :List[str]	       =						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":
            __SCREAMING_SNAKE_CASE     :Tuple	       =						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]				,     a_				,     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(a_							):
                  os.mkdir(a_							)

            model.save_pretrained(a_							)
            preprocessor.save_pretrained(a_							)

      __SCREAMING_SNAKE_CASE     :str	       =						'''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(a_							)
            preprocessor.push_to_hub(a_							)


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"""
# Lint as: python3

import hashlib
from typing import Union



def   __lowerCamelCase     (							a_  :	Union[str, int, bytes]							)    ->   bytes:
      if isinstance(a_				,     a_							):
            # Data already in bytes, returns as it as
            return hash_data
      elif isinstance(a_				,     a_							):
            # We keep the data as it as for it ot be later encoded to UTF-8
            # However replace `\\` with `/` for Windows compatibility
            __SCREAMING_SNAKE_CASE     :str	       =						hash_data.replace('''\\'''				,     '''/'''							)
      elif isinstance(a_				,     a_							):
            __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						str(a_							)
      else:
            # If data is not of the required type, raise error
            raise InvalidKeyError(a_							)

      return hash_data.encode('''utf-8'''							)






class       _SCREAMING_SNAKE_CASE(     A						):




       def __init__(		self	,SCREAMING_SNAKE_CASE__							)    ->   Optional[int]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE     :Tuple	       =						'''\nFAILURE TO GENERATE DATASET: Invalid key type detected'''
             __SCREAMING_SNAKE_CASE     :int	       =						f'''\nFound Key {hash_data} of type {type(SCREAMING_SNAKE_CASE__							)}'''
             __SCREAMING_SNAKE_CASE     :Optional[int]	       =						'''\nKeys should be either str, int or bytes type'''
             super().__init__(f'''{self.prefix}{self.err_msg}{self.suffix}'''							)






class       _SCREAMING_SNAKE_CASE(     A						):




       def __init__(		self	,SCREAMING_SNAKE_CASE__	,SCREAMING_SNAKE_CASE__	,SCREAMING_SNAKE_CASE__=""							)    ->   List[str]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE     :str	       =						key
             __SCREAMING_SNAKE_CASE     :str	       =						duplicate_key_indices
             __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						fix_msg
             __SCREAMING_SNAKE_CASE     :int	       =						'''Found multiple examples generated with the same key'''
             if len(SCREAMING_SNAKE_CASE__							) <= 20:
                   __SCREAMING_SNAKE_CASE     :Dict	       =						f'''\nThe examples at index {", ".join(SCREAMING_SNAKE_CASE__							)} have the key {key}'''
             else:
                   __SCREAMING_SNAKE_CASE     :List[str]	       =						f'''\nThe examples at index {", ".join(duplicate_key_indices[:20]							)}... ({len(SCREAMING_SNAKE_CASE__							) - 20} more) have the key {key}'''
             __SCREAMING_SNAKE_CASE     :Union[str, Any]	       =						'''\n''' + fix_msg if fix_msg else ''''''
             super().__init__(f'''{self.prefix}{self.err_msg}{self.suffix}'''							)






class       _SCREAMING_SNAKE_CASE:




       def __init__(		self	,SCREAMING_SNAKE_CASE__							)    ->   Union[str, Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE     :Dict	       =						hashlib.mda(_as_bytes(SCREAMING_SNAKE_CASE__							)							)




       def 					_UpperCamelCase  (		self	,SCREAMING_SNAKE_CASE__							)    ->   int:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE     :Tuple	       =						self._split_mda.copy()
             __SCREAMING_SNAKE_CASE     :Optional[Any]	       =						_as_bytes(SCREAMING_SNAKE_CASE__							)
             mda.update(SCREAMING_SNAKE_CASE__							)
             # Convert to integer with hexadecimal conversion
             return int(mda.hexdigest()	,16							)