




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    (							UpperCAmelCase_	:      Union[str, Any]      )->			Any:

       """simple docstring"""




       a			   =ConvNextVaConfig()

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

       a			   =1000
       a			   ="""imagenet-1k-id2label.json"""
       a			   =(1, 1000)

       a			   ="""huggingface/label-files"""
       a			   =num_labels
       a			   =json.load(open(hf_hub_download(UpperCAmelCase_   ,       UpperCAmelCase_   ,       repo_type="""dataset"""      )   ,       """r"""      )      )
       a			   ={int(UpperCAmelCase_      ): 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 							lowerCamelCase    (							UpperCAmelCase_	:      List[str]      )->			List[str]:

       """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 							lowerCamelCase    (							)->			Any:

       """simple docstring"""




       a			   ="""http://images.cocodataset.org/val2017/000000039769.jpg"""
       a			   =Image.open(requests.get(UpperCAmelCase_   ,       stream=UpperCAmelCase_      ).raw      )
       return im





def 							lowerCamelCase    (							UpperCAmelCase_	:      Dict      )->			str:

       """simple docstring"""




       if "224" in checkpoint_url:
              a			   =224
              a			   =224 / 256
       elif "384" in checkpoint_url:
              a			   =384
              a			   =None
       else:
              a			   =512
              a			   =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 							lowerCamelCase    (							UpperCAmelCase_	:      Union[str, Any]   ,       UpperCAmelCase_	:      Optional[Any]   ,       UpperCAmelCase_	:      Dict   ,       UpperCAmelCase_	:      Dict      )->			List[Any]:

       """simple docstring"""




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

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

       # load HuggingFace model
       a			   =ConvNextVaForImageClassification(UpperCAmelCase_      )
       model.load_state_dict(UpperCAmelCase_      )
       model.eval()

       # Check outputs on an image, prepared by ConvNextImageProcessor
       a			   =convert_preprocessor(UpperCAmelCase_      )
       a			   =preprocessor(images=prepare_img()   ,       return_tensors="""pt"""      )
       a			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   =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			   ="""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__":
    _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
    )






from __future__ import annotations





def 							lowerCamelCase    (							UpperCAmelCase_	:      list[int | float]      )->			int | float:

       """simple docstring"""




       if len(UpperCAmelCase_      ) == 0:
              raise ValueError("""find_min() arg is an empty sequence"""      )
       a			   =nums[0]
       for num in nums:
              a			   =min(UpperCAmelCase_   ,       UpperCAmelCase_      )
       return min_num


if __name__ == "__main__":
    import doctest

    doctest.testmod(verbose=True)
