



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



def   a_	(			__a      ):
     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(__a		,							__a		,							repo_type='''dataset'''      )		,							'''r'''      )      )
     A__				=   {int(__a      ): 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   a_	(			__a      ):
     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   a_	(			):
     A__				=   '''http://images.cocodataset.org/val2017/000000039769.jpg'''
     A__				=   Image.open(requests.get(__a		,							stream=__a      ).raw      )
     return im



def   a_	(			__a      ):
     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=__a		,							crop_pct=__a		,							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   a_	(			__a		,							__a		,							__a		,							__a      ):
     print('''Downloading original model from checkpoint...'''      )
     # define ConvNeXTV2 configuration based on URL
     A__	,     A__				=   get_convnextva_config(__a      )
     # load original state_dict from URL
     A__				=   torch.hub.load_state_dict_from_url(__a      )['''model''']

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

     # load HuggingFace model
     A__				=   ConvNextVaForImageClassification(__a      )
     model.load_state_dict(__a      )
     model.eval()

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

     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(__a      )
          preprocessor.push_to_hub(__a      )


if __name__ == "__main__":
 __snake_case	:				Dict  					=			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	:				Any  					=			parser.parse_args()
 convert_convnextva_checkpoint(
     args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
 )








"""simple docstring"""







from typing import TYPE_CHECKING

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


__snake_case	:				List[Any]  					=			{
    'configuration_unispeech_sat': ['UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'UniSpeechSatConfig'],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __snake_case	:				str  					=			[
     'UNISPEECH_SAT_PRETRAINED_MODEL_ARCHIVE_LIST',
     'UniSpeechSatForAudioFrameClassification',
     'UniSpeechSatForCTC',
     'UniSpeechSatForPreTraining',
     'UniSpeechSatForSequenceClassification',
     'UniSpeechSatForXVector',
     'UniSpeechSatModel',
     'UniSpeechSatPreTrainedModel',
 ]

if TYPE_CHECKING:
 from .configuration_unispeech_sat import UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP, UniSpeechSatConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_unispeech_sat import (
      UNISPEECH_SAT_PRETRAINED_MODEL_ARCHIVE_LIST,
      UniSpeechSatForAudioFrameClassification,
      UniSpeechSatForCTC,
      UniSpeechSatForPreTraining,
      UniSpeechSatForSequenceClassification,
      UniSpeechSatForXVector,
      UniSpeechSatModel,
      UniSpeechSatPreTrainedModel,
  )

else:
 import sys

 __snake_case	:				Union[str, Any]  					=			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



