


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[str]    							=					logging.get_logger(__name__)

def       _a	(					SCREAMING_SNAKE_CASE						:					Union[str, Any]  ):







 """simple docstring"""



 UpperCamelCase__       :     Union[str, Any] 							= ConvNextVaConfig()

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

 UpperCamelCase__       :     List[Any] 							= 1000
 UpperCamelCase__       :     str 							= '''imagenet-1k-id2label.json'''
 UpperCamelCase__       :     List[Any] 							= (1, 1000)

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

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

 return config, expected_shape

def       _a	(					SCREAMING_SNAKE_CASE						:					Optional[Any]  ):







 """simple docstring"""



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

 return name

def       _a	(					):







 """simple docstring"""



 UpperCamelCase__       :     str 							= '''http://images.cocodataset.org/val2017/000000039769.jpg'''
 UpperCamelCase__       :     Optional[Any] 							= Image.open(requests.get(SCREAMING_SNAKE_CASE     ,			stream=SCREAMING_SNAKE_CASE  ).raw  )
 return im

def       _a	(					SCREAMING_SNAKE_CASE						:					Optional[int]  ):







 """simple docstring"""



 if "224" in checkpoint_url:
  UpperCamelCase__       :     str 							= 224
  UpperCamelCase__       :     Dict 							= 224 / 256
 elif "384" in checkpoint_url:
  UpperCamelCase__       :     int 							= 384
  UpperCamelCase__       :     Dict 							= None
 else:
  UpperCamelCase__       :     Tuple 							= 512
  UpperCamelCase__       :     Union[str, Any] 							= None

 return ConvNextImageProcessor(
     size=SCREAMING_SNAKE_CASE     ,			crop_pct=SCREAMING_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	(					SCREAMING_SNAKE_CASE						:					List[Any]     ,			SCREAMING_SNAKE_CASE						:					Any     ,			SCREAMING_SNAKE_CASE						:					str     ,			SCREAMING_SNAKE_CASE						:					List[Any]  ):







 """simple docstring"""



 print('''Downloading original model from checkpoint...'''  )
 # define ConvNeXTV2 configuration based on URL
 UpperCamelCase__							,					UpperCamelCase__       :     List[Any] 							= get_convnextva_config(SCREAMING_SNAKE_CASE  )
 # load original state_dict from URL
 UpperCamelCase__       :     Dict 							= torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE  )['''model''']

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

 # load HuggingFace model
 UpperCamelCase__       :     Tuple 							= ConvNextVaForImageClassification(SCREAMING_SNAKE_CASE  )
 model.load_state_dict(SCREAMING_SNAKE_CASE  )
 model.eval()

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

  model.save_pretrained(SCREAMING_SNAKE_CASE  )
  preprocessor.save_pretrained(SCREAMING_SNAKE_CASE  )

 UpperCamelCase__       :     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(SCREAMING_SNAKE_CASE  )
  preprocessor.push_to_hub(SCREAMING_SNAKE_CASE  )


if __name__ == "__main__":
       __UpperCamelCase							:							List[str]    							=					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[int]    							=					parser.parse_args()
       convert_convnextva_checkpoint(
           args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
       )






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCamelCase							:							Dict    							=					{
    "configuration_dpr": ["DPR_PRETRAINED_CONFIG_ARCHIVE_MAP", "DPRConfig"],
    "tokenization_dpr": [
        "DPRContextEncoderTokenizer",
        "DPRQuestionEncoderTokenizer",
        "DPRReaderOutput",
        "DPRReaderTokenizer",
    ],
}


try:
       if not is_tokenizers_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase							:							Any    							=					[
           "DPRContextEncoderTokenizerFast",
           "DPRQuestionEncoderTokenizerFast",
           "DPRReaderTokenizerFast",
       ]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase							:							str    							=					[
           "DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST",
           "DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST",
           "DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST",
           "DPRContextEncoder",
           "DPRPretrainedContextEncoder",
           "DPRPreTrainedModel",
           "DPRPretrainedQuestionEncoder",
           "DPRPretrainedReader",
           "DPRQuestionEncoder",
           "DPRReader",
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase							:							int    							=					[
           "TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST",
           "TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST",
           "TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST",
           "TFDPRContextEncoder",
           "TFDPRPretrainedContextEncoder",
           "TFDPRPretrainedQuestionEncoder",
           "TFDPRPretrainedReader",
           "TFDPRQuestionEncoder",
           "TFDPRReader",
       ]


if TYPE_CHECKING:
       from .configuration_dpr import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP, DPRConfig
       from .tokenization_dpr import (
           DPRContextEncoderTokenizer,
           DPRQuestionEncoderTokenizer,
           DPRReaderOutput,
           DPRReaderTokenizer,
       )

       try:
              if not is_tokenizers_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_dpr_fast import (
                  DPRContextEncoderTokenizerFast,
                  DPRQuestionEncoderTokenizerFast,
                  DPRReaderTokenizerFast,
              )

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_dpr import (
                  DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
                  DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
                  DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
                  DPRContextEncoder,
                  DPRPretrainedContextEncoder,
                  DPRPreTrainedModel,
                  DPRPretrainedQuestionEncoder,
                  DPRPretrainedReader,
                  DPRQuestionEncoder,
                  DPRReader,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_dpr import (
                  TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
                  TFDPRContextEncoder,
                  TFDPRPretrainedContextEncoder,
                  TFDPRPretrainedQuestionEncoder,
                  TFDPRPretrainedReader,
                  TFDPRQuestionEncoder,
                  TFDPRReader,
              )

else:
       import sys

       __UpperCamelCase							:							Dict    							=					_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


