

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__			:  List[Any]					     =     logging.get_logger(__name__)






def        __lowercase   (  _A      )  ->    int:
      SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = ConvNextVaConfig()

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

      SCREAMING_SNAKE_CASE			:   Optional[int]	  = 1000
      SCREAMING_SNAKE_CASE			:   Dict	  = """imagenet-1k-id2label.json"""
      SCREAMING_SNAKE_CASE			:   List[Any]	  = (1, 1000)

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

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

      return config, expected_shape






def        __lowercase   (  _A      )  ->    Union[str, Any]:



      if "downsample_layers.0.0" in name:
            SCREAMING_SNAKE_CASE			:   Optional[int]	  = name.replace("""downsample_layers.0.0"""    ,     """embeddings.patch_embeddings"""      )
      if "downsample_layers.0.1" in name:
            SCREAMING_SNAKE_CASE			:   Any	  = 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			:   List[Any]	  = name.replace("""downsample_layers.1.1"""    ,     """stages.1.downsampling_layer.1"""      )
      if "downsample_layers.2.0" in name:
            SCREAMING_SNAKE_CASE			:   Dict	  = name.replace("""downsample_layers.2.0"""    ,     """stages.2.downsampling_layer.0"""      )
      if "downsample_layers.2.1" in name:
            SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = name.replace("""downsample_layers.2.1"""    ,     """stages.2.downsampling_layer.1"""      )
      if "downsample_layers.3.0" in name:
            SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = name.replace("""downsample_layers.3.0"""    ,     """stages.3.downsampling_layer.0"""      )
      if "downsample_layers.3.1" in name:
            SCREAMING_SNAKE_CASE			:   Dict	  = 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			:   Optional[int]	  = name.replace("""gamma"""    ,     """weight"""      )
      if "beta" in name:
            SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = name.replace("""beta"""    ,     """bias"""      )
      if "stages" in name:
            SCREAMING_SNAKE_CASE			:   str	  = name.replace("""stages"""    ,     """encoder.stages"""      )
      if "norm" in name:
            SCREAMING_SNAKE_CASE			:   Tuple	  = name.replace("""norm"""    ,     """layernorm"""      )
      if "head" in name:
            SCREAMING_SNAKE_CASE			:   Tuple	  = name.replace("""head"""    ,     """classifier"""      )

      return name






def        __lowercase   (  )  ->    Tuple:
      SCREAMING_SNAKE_CASE			:   int	  = """http://images.cocodataset.org/val2017/000000039769.jpg"""
      SCREAMING_SNAKE_CASE			:   int	  = Image.open(requests.get(_A    ,     stream=_A      ).raw      )
      return im






def        __lowercase   (  _A      )  ->    int:



      if "224" in checkpoint_url:
            SCREAMING_SNAKE_CASE			:   str	  = 224
            SCREAMING_SNAKE_CASE			:   List[str]	  = 224 / 256
      elif "384" in checkpoint_url:
            SCREAMING_SNAKE_CASE			:   Optional[int]	  = 384
            SCREAMING_SNAKE_CASE			:   List[str]	  = None
      else:
            SCREAMING_SNAKE_CASE			:   Any	  = 512
            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        __lowercase   (  _A    ,     _A    ,     _A    ,     _A      )  ->    str:
      print("""Downloading original model from checkpoint..."""      )
      # define ConvNeXTV2 configuration based on URL
      SCREAMING_SNAKE_CASE  ,						SCREAMING_SNAKE_CASE			:   List[str]	  = 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			:   List[Any]	  = state_dict.pop(_A      )
            SCREAMING_SNAKE_CASE			:   Dict	  = val
      # add prefix to all keys expect classifier head
      for key in state_dict.copy().keys():
            SCREAMING_SNAKE_CASE			:   Optional[Any]	  = state_dict.pop(_A      )
            if not key.startswith("""classifier"""      ):
                  SCREAMING_SNAKE_CASE			:   List[Any]	  = """convnextv2.""" + key
            SCREAMING_SNAKE_CASE			:   str	  = val

      # load HuggingFace model
      SCREAMING_SNAKE_CASE			:   Optional[int]	  = ConvNextVaForImageClassification(_A      )
      model.load_state_dict(_A      )
      model.eval()

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


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

  UpperCAmelCase__			:  Optional[Any]					     =     parser.parse_args()
  convert_convnextva_checkpoint(
      args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
  )





"""simple docstring"""


import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import AutoConfig


lowerCAmelCase_:       List[str]			     =       logging.get_logger(__name__)






class  a__     ( _a						):
   snake_case_				=   "speech-encoder-decoder"
   snake_case_				=   True


   def __init__(  self,					**_UpperCAmelCase	):
        '''simple docstring'''



        super().__init__(**_UpperCAmelCase	)
        if "encoder" not in kwargs or "decoder" not in kwargs:
             raise ValueError(
                 F'''A configuraton of type {self.model_type} cannot be instantiated because not both `encoder` and'''
                 F''' `decoder` sub-configurations are passed, but only {kwargs}'''	)

        lowercase__								=				kwargs.pop("encoder"	)
        lowercase__								=				encoder_config.pop("model_type"	)
        lowercase__								=				kwargs.pop("decoder"	)
        lowercase__								=				decoder_config.pop("model_type"	)

        lowercase__								=				AutoConfig.for_model(_UpperCAmelCase,					**_UpperCAmelCase	)
        lowercase__								=				AutoConfig.for_model(_UpperCAmelCase,					**_UpperCAmelCase	)
        lowercase__								=				True


   @classmethod
   def 						snake_case__						(  cls,					_UpperCAmelCase,					_UpperCAmelCase,					**_UpperCAmelCase	):
        '''simple docstring'''



        logger.info("Setting `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config"	)
        lowercase__								=				True
        lowercase__								=				True

        return cls(encoder=encoder_config.to_dict(),					decoder=decoder_config.to_dict(),					**_UpperCAmelCase	)


   def 						snake_case__						(  self	):
        '''simple docstring'''



        lowercase__								=				copy.deepcopy(self.__dict__	)
        lowercase__								=				self.encoder.to_dict()
        lowercase__								=				self.decoder.to_dict()
        lowercase__								=				self.__class__.model_type
        return output




