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







def 			_UpperCamelCase			(  _a						:						int      ):




    """simple docstring"""



    __UpperCamelCase						:						Dict		     =	ConvNextVaConfig()

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

    __UpperCamelCase						:						List[Any]		     =	1_0_0_0
    __UpperCamelCase						:						int		     =	'imagenet-1k-id2label.json'
    __UpperCamelCase						:						str		     =	(1, 1_0_0_0)

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

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

    return config, expected_shape







def 			_UpperCamelCase			(  _a						:						List[Any]      ):




    """simple docstring"""



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

    return name







def 			_UpperCamelCase			(  ):




    """simple docstring"""



    __UpperCamelCase						:						int		     =	'http://images.cocodataset.org/val2017/000000039769.jpg'
    __UpperCamelCase						:						List[Any]		     =	Image.open(requests.get(UpperCAmelCase__					,  stream=UpperCAmelCase__      ).raw      )
    return im







def 			_UpperCamelCase			(  _a						:						Dict      ):




    """simple docstring"""



    if "224" in checkpoint_url:
        __UpperCamelCase						:						Optional[Any]		     =	2_2_4
        __UpperCamelCase						:						Optional[Any]		     =	2_2_4 / 2_5_6
    elif "384" in checkpoint_url:
        __UpperCamelCase						:						Tuple		     =	3_8_4
        __UpperCamelCase						:						Any		     =	None
    else:
        __UpperCamelCase						:						Any		     =	5_1_2
        __UpperCamelCase						:						List[str]		     =	None

    return ConvNextImageProcessor(
        size=UpperCAmelCase__					,  crop_pct=UpperCAmelCase__					,  image_mean=[0.485, 0.456, 0.406]					,  image_std=[0.229, 0.224, 0.225]					,  resample=PILImageResampling.BICUBIC					,  )







@torch.no_grad()
def 			_UpperCamelCase			(  _a						:						Optional[int]					,  _a						:						List[str]					,  _a						:						str					,  _a						:						Union[str, Any]      ):




    """simple docstring"""



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

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

    # load HuggingFace model
    __UpperCamelCase						:						Dict		     =	ConvNextVaForImageClassification(UpperCAmelCase__      )
    model.load_state_dict(UpperCAmelCase__      )
    model.eval()

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

    __UpperCamelCase						:						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(UpperCAmelCase__      )
        preprocessor.push_to_hub(UpperCAmelCase__      )


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

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






import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase   					=  logging.get_logger(__name__)

lowerCamelCase   					=  {}





class      snake_case_      (					_a				):


							"""simple docstring"""





							__UpperCAmelCase      ="""fsmt_decoder"""




							def __init__(			self       ,    _A=0       ,    _A=0				):
												super().__init__()
												__lowerCAmelCase       =				vocab_size
												__lowerCAmelCase       =				bos_token_id





class      snake_case_      (					_a				):


							"""simple docstring"""





							__UpperCAmelCase      ="""fsmt"""
							__UpperCAmelCase      ={"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}




							def __init__(			self       ,    _A=["en", "de"]       ,    _A=4_2_0_2_4       ,    _A=4_2_0_2_4       ,    _A="relu"       ,    _A=1_0_2_4       ,    _A=2_0_0       ,    _A=1_0_2_4       ,    _A=4_0_9_6       ,    _A=1_2       ,    _A=1_6       ,    _A=0.0       ,    _A=4_0_9_6       ,    _A=1_2       ,    _A=1_6       ,    _A=0.0       ,    _A=0.0       ,    _A=0.1       ,    _A=0.0       ,    _A=0.02       ,    _A=2       ,    _A=True       ,    _A=True       ,    _A=False       ,    _A=5       ,    _A=1.0       ,    _A=False       ,    _A=True       ,    _A=1       ,    _A=0       ,    _A=2       ,    _A=2       ,    **_A       ,    ):
												__lowerCAmelCase       =				langs
												__lowerCAmelCase       =				src_vocab_size
												__lowerCAmelCase       =				tgt_vocab_size
												__lowerCAmelCase       =				d_model  # encoder_embed_dim and decoder_embed_dim

												__lowerCAmelCase       =				encoder_ffn_dim
												__lowerCAmelCase       =				__lowerCAmelCase       =				encoder_layers
												__lowerCAmelCase       =				encoder_attention_heads
												__lowerCAmelCase       =				encoder_layerdrop
												__lowerCAmelCase       =				decoder_layerdrop
												__lowerCAmelCase       =				decoder_ffn_dim
												__lowerCAmelCase       =				decoder_layers
												__lowerCAmelCase       =				decoder_attention_heads
												__lowerCAmelCase       =				max_position_embeddings
												__lowerCAmelCase       =				init_std  # Normal(0, this parameter)
												__lowerCAmelCase       =				activation_function

												__lowerCAmelCase       =				DecoderConfig(vocab_size=_A       ,    bos_token_id=_A				)
												if "decoder" in common_kwargs:
																	del common_kwargs["decoder"]

												__lowerCAmelCase       =				scale_embedding  # scale factor will be sqrt(d_model) if True

												# 3 Types of Dropout
												__lowerCAmelCase       =				attention_dropout
												__lowerCAmelCase       =				activation_dropout
												__lowerCAmelCase       =				dropout

												__lowerCAmelCase       =				use_cache
												super().__init__(
												    pad_token_id=_A       ,    bos_token_id=_A       ,    eos_token_id=_A       ,    decoder_start_token_id=_A       ,    is_encoder_decoder=_A       ,    tie_word_embeddings=_A       ,    forced_eos_token_id=_A       ,    max_length=_A       ,    num_beams=_A       ,    length_penalty=_A       ,    early_stopping=_A       ,    **_A       ,    )




							def 							A__	(			self				):
												__lowerCAmelCase       =				copy.deepcopy(self.__dict__				)
												__lowerCAmelCase       =				self.decoder.to_dict()
												__lowerCAmelCase       =				self.__class__.model_type
												return output






