

'''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()
_snake_case     						=						logging.get_logger(__name__)







def       _A	(			snake_case  )    ->			Dict:
   _lowercase    :   str      =     ConvNextVaConfig()

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

   _lowercase    :   Any      =     10_00
   _lowercase    :   List[Any]      =     "imagenet-1k-id2label.json"
   _lowercase    :   Tuple      =     (1, 10_00)

   _lowercase    :   int      =     "huggingface/label-files"
   _lowercase    :   Any      =     num_labels
   _lowercase    :   str      =     json.load(open(hf_hub_download(snake_case	,							snake_case	,							repo_type="dataset"  )	,							"r"  )  )
   _lowercase    :   int      =     {int(snake_case  ): v for k, v in idalabel.items()}

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

   return config, expected_shape







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

   return name







def       _A	(			)    ->			Dict:
   _lowercase    :   List[str]      =     "http://images.cocodataset.org/val2017/000000039769.jpg"
   _lowercase    :   str      =     Image.open(requests.get(snake_case	,							stream=snake_case  ).raw  )
   return im







def       _A	(			snake_case  )    ->			List[str]:
   if "224" in checkpoint_url:
      _lowercase    :   List[str]      =     2_24
      _lowercase    :   Dict      =     2_24 / 2_56
   elif "384" in checkpoint_url:
      _lowercase    :   Optional[int]      =     3_84
      _lowercase    :   Optional[Any]      =     None
   else:
      _lowercase    :   Optional[Any]      =     5_12
      _lowercase    :   str      =     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	,							snake_case	,							snake_case	,							snake_case  )    ->			List[str]:
   print("Downloading original model from checkpoint..."  )
   # define ConvNeXTV2 configuration based on URL
   _lowercase  ,   _lowercase    :   List[Any]      =     get_convnextva_config(snake_case  )
   # load original state_dict from URL
   _lowercase    :   int      =     torch.hub.load_state_dict_from_url(snake_case  )["model"]

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

   # load HuggingFace model
   _lowercase    :   Union[str, Any]      =     ConvNextVaForImageClassification(snake_case  )
   model.load_state_dict(snake_case  )
   model.eval()

   # Check outputs on an image, prepared by ConvNextImageProcessor
   _lowercase    :   int      =     convert_preprocessor(snake_case  )
   _lowercase    :   List[Any]      =     preprocessor(images=prepare_img()	,							return_tensors="pt"  )
   _lowercase    :   List[str]      =     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":
      _lowercase    :   Tuple      =     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":
      _lowercase    :   Union[str, 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":
      _lowercase    :   Tuple      =     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":
      _lowercase    :   int      =     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":
      _lowercase    :   List[Any]      =     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":
      _lowercase    :   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":
      _lowercase    :   List[str]      =     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":
      _lowercase    :   int      =     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":
      _lowercase    :   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":
      _lowercase    :   Optional[Any]      =     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":
      _lowercase    :   Tuple      =     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":
      _lowercase    :   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":
      _lowercase    :   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":
      _lowercase    :   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":
      _lowercase    :   Union[str, 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":
      _lowercase    :   Optional[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":
      _lowercase    :   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":
      _lowercase    :   Dict      =     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  )

   _lowercase    :   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__":
    _snake_case     						=						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     						=						parser.parse_args()
    convert_convnextva_checkpoint(
        args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
    )









'''simple docstring'''
import warnings
from dataclasses import dataclass, field
from typing import Optional, Tuple

from .training_args import TrainingArguments
from .utils import cached_property, is_tf_available, logging, requires_backends


_snake_case     						=						logging.get_logger(__name__)

if is_tf_available():
    import tensorflow as tf


@dataclass
class      a__						(							lowerCamelCase_						):
   _SCREAMING_SNAKE_CASE					:  int			 =							'tf'
   _SCREAMING_SNAKE_CASE					:  Optional[str]			 =							field(
       default=lowerCamelCase_						,   metadata={'help': 'Name of TPU'}						,   )

   _SCREAMING_SNAKE_CASE					:  Optional[str]			 =							field(
       default=lowerCamelCase_						,   metadata={'help': 'Zone of TPU'}						,   )

   _SCREAMING_SNAKE_CASE					:  Optional[str]			 =							field(
       default=lowerCamelCase_						,   metadata={'help': 'Name of Cloud TPU-enabled project'}						,   )

   _SCREAMING_SNAKE_CASE					:  float			 =							field(
       default=1.0						,   metadata={'help': 'Power for the Polynomial decay LR scheduler.'}						,   )

   _SCREAMING_SNAKE_CASE					:  bool			 =							field(default=lowerCamelCase_						,   metadata={'help': 'Whether to activate the XLA compilation or not'}						)


   @cached_property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      requires_backends(self    ,       ["tf"]						)
      logger.info("Tensorflow: setting up strategy"						)

      _lowercase    :   Optional[int]      =     tf.config.list_physical_devices("GPU"						)

      # Set to float16 at first
      if self.fpaa:
         tf.keras.mixed_precision.set_global_policy("mixed_float16"						)

      if self.no_cuda:
         _lowercase    :   Optional[int]      =     tf.distribute.OneDeviceStrategy(device="/cpu:0"						)
      else:
         try:
            if self.tpu_name:
               _lowercase    :   int      =     tf.distribute.cluster_resolver.TPUClusterResolver(
                   self.tpu_name    ,       zone=self.tpu_zone    ,       project=self.gcp_project						)
            else:
               _lowercase    :   Any      =     tf.distribute.cluster_resolver.TPUClusterResolver()
         except ValueError:
            if self.tpu_name:
               raise RuntimeError(f'''Couldn\'t connect to TPU {self.tpu_name}!'''						)
            else:
               _lowercase    :   Tuple      =     None

         if tpu:
            # Set to bfloat16 in case of TPU
            if self.fpaa:
               tf.keras.mixed_precision.set_global_policy("mixed_bfloat16"						)

            tf.config.experimental_connect_to_cluster(_UpperCamelCase						)
            tf.tpu.experimental.initialize_tpu_system(_UpperCamelCase						)

            _lowercase    :   List[str]      =     tf.distribute.TPUStrategy(_UpperCamelCase						)

         elif len(_UpperCamelCase						) == 0:
            _lowercase    :   Optional[Any]      =     tf.distribute.OneDeviceStrategy(device="/cpu:0"						)
         elif len(_UpperCamelCase						) == 1:
            _lowercase    :   Optional[int]      =     tf.distribute.OneDeviceStrategy(device="/gpu:0"						)
         elif len(_UpperCamelCase						) > 1:
            # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
            _lowercase    :   Any      =     tf.distribute.MirroredStrategy()
         else:
            raise ValueError("Cannot find the proper strategy, please check your environment properties."						)

      return strategy


   @property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      requires_backends(self    ,       ["tf"]						)
      return self._setup_strategy


   @property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      requires_backends(self    ,       ["tf"]						)
      return self._setup_strategy.num_replicas_in_sync


   @property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      return False  # TF Logging is handled by Keras not the Trainer


   @property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      if self.per_gpu_train_batch_size:
         logger.warning(
             "Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future "
             "version. Using `--per_device_train_batch_size` is preferred."						)
      _lowercase    :   str      =     self.per_gpu_train_batch_size or self.per_device_train_batch_size
      return per_device_batch_size * self.n_replicas


   @property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      if self.per_gpu_eval_batch_size:
         logger.warning(
             "Using deprecated `--per_gpu_eval_batch_size` argument which will be removed in a future "
             "version. Using `--per_device_eval_batch_size` is preferred."						)
      _lowercase    :   List[Any]      =     self.per_gpu_eval_batch_size or self.per_device_eval_batch_size
      return per_device_batch_size * self.n_replicas






   @property
   def 		_lowerCamelCase     (				self						):





      """simple docstring"""





      requires_backends(self    ,       ["tf"]						)
      warnings.warn(
          "The n_gpu argument is deprecated and will be removed in a future version, use n_replicas instead."    ,       _UpperCamelCase    ,       )
      return self._setup_strategy.num_replicas_in_sync






