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







def 						UpperCAmelCase							(		lowerCAmelCase_							)						->						Optional[int]:
     a	:     Dict						      =    ConvNextVaConfig()

     if "atto" in checkpoint_url:
          a	:     str						      =    [2, 2, 6, 2]
          a	:     List[Any]						      =    [4_0, 8_0, 1_6_0, 3_2_0]
     if "femto" in checkpoint_url:
          a	:     int						      =    [2, 2, 6, 2]
          a	:     str						      =    [4_8, 9_6, 1_9_2, 3_8_4]
     if "pico" in checkpoint_url:
          a	:     Any						      =    [2, 2, 6, 2]
          a	:     Union[str, Any]						      =    [6_4, 1_2_8, 2_5_6, 5_1_2]
     if "nano" in checkpoint_url:
          a	:     Optional[Any]						      =    [2, 2, 8, 2]
          a	:     Dict						      =    [8_0, 1_6_0, 3_2_0, 6_4_0]
     if "tiny" in checkpoint_url:
          a	:     int						      =    [3, 3, 9, 3]
          a	:     Optional[Any]						      =    [9_6, 1_9_2, 3_8_4, 7_6_8]
     if "base" in checkpoint_url:
          a	:     Optional[int]						      =    [3, 3, 2_7, 3]
          a	:     Optional[int]						      =    [1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
     if "large" in checkpoint_url:
          a	:     List[str]						      =    [3, 3, 2_7, 3]
          a	:     Dict						      =    [1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
     if "huge" in checkpoint_url:
          a	:     int						      =    [3, 3, 2_7, 3]
          a	:     int						      =    [3_5_2, 7_0_4, 1_4_0_8, 2_8_1_6]

     a	:     int						      =    1_0_0_0
     a	:     Any						      =    '''imagenet-1k-id2label.json'''
     a	:     Any						      =    (1, 1_0_0_0)

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

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

     return config, expected_shape







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

     return name







def 						UpperCAmelCase							(		)						->						List[Any]:
     a	:     Any						      =    '''http://images.cocodataset.org/val2017/000000039769.jpg'''
     a	:     Tuple						      =    Image.open(requests.get(UpperCamelCase__    ,   stream=UpperCamelCase__							).raw							)
     return im







def 						UpperCAmelCase							(		lowerCAmelCase_							)						->						Any:
     if "224" in checkpoint_url:
          a	:     str						      =    2_2_4
          a	:     str						      =    2_2_4 / 2_5_6
     elif "384" in checkpoint_url:
          a	:     int						      =    3_8_4
          a	:     List[Any]						      =    None
     else:
          a	:     Optional[int]						      =    5_1_2
          a	:     str						      =    None

     return ConvNextImageProcessor(
         size=UpperCamelCase__    ,   crop_pct=UpperCamelCase__    ,   image_mean=[0.4_85, 0.4_56, 0.4_06]    ,   image_std=[0.2_29, 0.2_24, 0.2_25]    ,   resample=PILImageResampling.BICUBIC    ,   )







@torch.no_grad()
def 						UpperCAmelCase							(		lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_							)						->						Union[str, Any]:
     print('''Downloading original model from checkpoint...'''							)
     # define ConvNeXTV2 configuration based on URL
     a	:     Dict						      =    get_convnextva_config(UpperCamelCase__							)
     # load original state_dict from URL
     a	:     Union[str, Any]						      =    torch.hub.load_state_dict_from_url(UpperCamelCase__							)['''model''']

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

     # load HuggingFace model
     a	:     List[Any]						      =    ConvNextVaForImageClassification(UpperCamelCase__							)
     model.load_state_dict(UpperCamelCase__							)
     model.eval()

     # Check outputs on an image, prepared by ConvNextImageProcessor
     a	:     str						      =    convert_preprocessor(UpperCamelCase__							)
     a	:     Union[str, Any]						      =    preprocessor(images=prepare_img()    ,   return_tensors='''pt'''							)
     a	:     int						      =    model(**UpperCamelCase__							).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	:     Union[str, Any]						      =    torch.tensor([-0.39_30, 0.17_47, -0.52_46, 0.41_77, 0.42_95]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
          a	:     Union[str, Any]						      =    torch.tensor([-0.17_27, -0.53_41, -0.78_18, -0.47_45, -0.65_66]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
          a	:     Optional[int]						      =    torch.tensor([-0.03_33, 0.15_63, -0.91_37, 0.10_54, 0.03_81]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
          a	:     List[str]						      =    torch.tensor([-0.17_44, -0.15_55, -0.07_13, 0.09_50, -0.14_31]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
          a	:     Any						      =    torch.tensor([0.99_96, 0.19_66, -0.43_86, -0.34_72, 0.66_61]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
          a	:     str						      =    torch.tensor([-0.25_53, -0.67_08, -0.13_59, 0.25_18, -0.24_88]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
          a	:     List[Any]						      =    torch.tensor([-0.06_73, -0.56_27, -0.37_53, -0.27_22, 0.01_78]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
          a	:     Dict						      =    torch.tensor([-0.63_77, -0.74_58, -0.21_50, 0.11_84, -0.05_97]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
          a	:     Optional[Any]						      =    torch.tensor([1.07_99, 0.23_22, -0.88_60, 1.02_19, 0.62_31]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
          a	:     Optional[int]						      =    torch.tensor([0.37_66, 0.49_17, -1.14_26, 0.99_42, 0.60_24]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
          a	:     Optional[Any]						      =    torch.tensor([0.42_20, -0.69_19, -0.43_17, -0.28_81, -0.66_09]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
          a	:     Tuple						      =    torch.tensor([0.10_82, -0.82_86, -0.50_95, 0.46_81, -0.80_85]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
          a	:     int						      =    torch.tensor([-0.24_19, -0.62_21, 0.21_76, -0.09_80, -0.75_27]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
          a	:     Any						      =    torch.tensor([0.03_91, -0.43_71, 0.37_86, 0.12_51, -0.27_84]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
          a	:     int						      =    torch.tensor([-0.05_04, 0.56_36, -0.17_29, -0.65_07, -0.39_49]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
          a	:     Any						      =    torch.tensor([0.35_60, 0.94_86, 0.31_49, -0.26_67, -0.51_38]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
          a	:     str						      =    torch.tensor([-0.24_69, -0.45_50, -0.58_53, -0.08_10, 0.03_09]							)
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
          a	:     Union[str, Any]						      =    torch.tensor([-0.30_90, 0.08_02, -0.06_82, -0.19_79, -0.28_26]							)
     else:
          raise ValueError(f"""Unknown URL: {checkpoint_url}"""							)

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

          model.save_pretrained(UpperCamelCase__							)
          preprocessor.save_pretrained(UpperCamelCase__							)

     a	:     int						      =    '''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(UpperCamelCase__							)
          preprocessor.push_to_hub(UpperCamelCase__							)


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

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






'''simple docstring'''






import platform
from argparse import ArgumentParser

import huggingface_hub
import pandas
import pyarrow

from datasets import __version__ as version
from datasets.commands import BaseDatasetsCLICommand
def 				_lowerCamelCase(	UpperCamelCase__     :			Any		)		->	List[str]:
    return EnvironmentCommand()





class 	_lowercase    (				a					):


   @staticmethod
   def        snake_case					(			_UpperCAmelCase       ):
       A							:							List[str]				=    parser.add_parser('''env'''     , help='''Print relevant system environment info.'''       )
       download_parser.set_defaults(func=_UpperCAmelCase       )


   def        snake_case					(			self       ):
       A							:							Optional[int]				=    {
           '''`datasets` version''': version,
           '''Platform''': platform.platform(),
           '''Python version''': platform.python_version(),
           '''Huggingface_hub version''': huggingface_hub.__version__,
           '''PyArrow version''': pyarrow.__version__,
           '''Pandas version''': pandas.__version__,
       }

       print('''\nCopy-and-paste the text below in your GitHub issue.\n'''       )
       print(self.format_dict(_UpperCAmelCase       )       )

       return info


   @staticmethod
   def        snake_case					(			_UpperCAmelCase       ):
       return "\n".join([f'''- {prop}: {val}''' for prop, val in d.items()]       ) + "\n"
