





import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
_UpperCamelCase       :				Union[str, Any]           =			logging.get_logger(__name__)
def  __UpperCamelCase    (						snake_case      )    ->		int:


   '''simple docstring'''







   __A    =  ConvNextConfig()

   if "tiny" in checkpoint_url:
      __A    =  [3, 3, 9, 3]
      __A    =  [9_6, 1_9_2, 3_8_4, 7_6_8]
   if "small" in checkpoint_url:
      __A    =  [3, 3, 2_7, 3]
      __A    =  [9_6, 1_9_2, 3_8_4, 7_6_8]
   if "base" in checkpoint_url:
      __A    =  [3, 3, 2_7, 3]
      __A    =  [1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
   if "large" in checkpoint_url:
      __A    =  [3, 3, 2_7, 3]
      __A    =  [1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
   if "xlarge" in checkpoint_url:
      __A    =  [3, 3, 2_7, 3]
      __A    =  [2_5_6, 5_1_2, 1_0_2_4, 2_0_4_8]

   if "1k" in checkpoint_url:
      __A    =  1_0_0_0
      __A    =  '''imagenet-1k-id2label.json'''
      __A    =  (1, 1_0_0_0)
   else:
      __A    =  2_1_8_4_1
      __A    =  '''imagenet-22k-id2label.json'''
      __A    =  (1, 2_1_8_4_1)

   __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()}
   if "1k" not in checkpoint_url:
      # this dataset contains 21843 labels but the model only has 21841
      # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
      del idalabel[9_2_0_5]
      del idalabel[1_5_0_2_7]
   __A    =  idalabel
   __A    =  {v: k for k, v in idalabel.items()}
   __A    =  hidden_sizes
   __A    =  depths

   return config, expected_shape
def  __UpperCamelCase    (						snake_case      )    ->		Tuple:


   '''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 "stages" in name:
      __A    =  name.replace('''stages'''							,			'''encoder.stages'''      )
   if "norm" in name:
      __A    =  name.replace('''norm'''							,			'''layernorm'''      )
   if "gamma" in name:
      __A    =  name.replace('''gamma'''							,			'''layer_scale_parameter'''      )
   if "head" in name:
      __A    =  name.replace('''head'''							,			'''classifier'''      )

   return name
def  __UpperCamelCase    (						)    ->		Dict:


   '''simple docstring'''







   __A    =  '''http://images.cocodataset.org/val2017/000000039769.jpg'''
   __A    =  Image.open(requests.get(__UpperCAmelCase							,			stream=__UpperCAmelCase      ).raw      )
   return im
@torch.no_grad()
def  __UpperCamelCase    (						snake_case							,			snake_case      )    ->		Union[str, Any]:


   '''simple docstring'''







   __A   ,		__A    =  get_convnext_config(__UpperCAmelCase      )
   # load original state_dict from URL
   __A    =  torch.hub.load_state_dict_from_url(__UpperCAmelCase      )['''model''']
   # 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    =  '''convnext.''' + key
      __A    =  val

   # load HuggingFace model
   __A    =  ConvNextForImageClassification(__UpperCAmelCase      )
   model.load_state_dict(__UpperCAmelCase      )
   model.eval()

   # Check outputs on an image, prepared by ConvNextImageProcessor
   __A    =  2_2_4 if '''224''' in checkpoint_url else 3_8_4
   __A    =  ConvNextImageProcessor(size=__UpperCAmelCase      )
   __A    =  image_processor(images=prepare_img()							,			return_tensors='''pt'''      ).pixel_values

   __A    =  model(__UpperCAmelCase      ).logits

   # note: the logits below were obtained without center cropping
   if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
      __A    =  torch.tensor([-0.1_210, -0.6_605, 0.1_918]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
      __A    =  torch.tensor([-0.4_473, -0.1_847, -0.6_365]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
      __A    =  torch.tensor([0.4_525, 0.7_539, 0.0_308]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
      __A    =  torch.tensor([0.3_561, 0.6_350, -0.0_384]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
      __A    =  torch.tensor([0.4_174, -0.0_989, 0.1_489]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
      __A    =  torch.tensor([0.2_513, -0.1_349, -0.1_613]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
      __A    =  torch.tensor([1.2_980, 0.3_631, -0.1_198]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
      __A    =  torch.tensor([1.2_963, 0.1_227, 0.1_723]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
      __A    =  torch.tensor([1.7_956, 0.8_390, 0.2_820]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
      __A    =  torch.tensor([-0.2_822, -0.0_502, -0.0_878]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
      __A    =  torch.tensor([-0.5_672, -0.0_730, -0.4_348]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
      __A    =  torch.tensor([0.2_681, 0.2_365, 0.6_246]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
      __A    =  torch.tensor([-0.2_642, 0.3_931, 0.5_116]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
      __A    =  torch.tensor([-0.6_677, -0.1_873, -0.8_379]      )
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
      __A    =  torch.tensor([-0.7_749, -0.2_967, -0.6_444]      )
   else:
      raise ValueError(F"Unknown URL: {checkpoint_url}"      )

   assert torch.allclose(logits[0, :3]							,			__UpperCAmelCase							,			atol=1E-3      )
   assert logits.shape == expected_shape

   Path(__UpperCAmelCase      ).mkdir(exist_ok=__UpperCAmelCase      )
   print(F"Saving model to {pytorch_dump_folder_path}"      )
   model.save_pretrained(__UpperCAmelCase      )
   print(F"Saving image processor to {pytorch_dump_folder_path}"      )
   image_processor.save_pretrained(__UpperCAmelCase      )

   print('''Pushing model to the hub...'''      )
   __A    =  '''convnext'''
   if "tiny" in checkpoint_url:
      model_name += "-tiny"
   elif "small" in checkpoint_url:
      model_name += "-small"
   elif "base" in checkpoint_url:
      model_name += "-base"
   elif "xlarge" in checkpoint_url:
      model_name += "-xlarge"
   elif "large" in checkpoint_url:
      model_name += "-large"
   if "224" in checkpoint_url:
      model_name += "-224"
   elif "384" in checkpoint_url:
      model_name += "-384"
   if "22k" in checkpoint_url and "1k" not in checkpoint_url:
      model_name += "-22k"
   if "22k" in checkpoint_url and "1k" in checkpoint_url:
      model_name += "-22k-1k"

   model.push_to_hub(
       repo_path_or_name=Path(__UpperCAmelCase							,			__UpperCAmelCase      )							,			organization='''nielsr'''							,			commit_message='''Add model'''							,			)


if __name__ == "__main__":
    _UpperCamelCase       :				List[str]           =			argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        """--checkpoint_url""",
        default="""https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth""",
        type=str,
        help="""URL of the original ConvNeXT checkpoint you\'d like to convert.""",
    )
    parser.add_argument(
        """--pytorch_dump_folder_path""",
        default=None,
        type=str,
        required=True,
        help="""Path to the output PyTorch model directory.""",
    )

    _UpperCamelCase       :				Union[str, Any]           =			parser.parse_args()
    convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)


from ..utils import deprecate
from .attention_processor import (  # noqa: F401
    Attention,
    AttentionProcessor,
    AttnAddedKVProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRALinearLayer,
    LoRAXFormersAttnProcessor,
    SlicedAttnAddedKVProcessor,
    SlicedAttnProcessor,
    XFormersAttnProcessor,
)
from .attention_processor import AttnProcessor as AttnProcessorRename  # noqa: F401


deprecate(
    'cross_attention',
    '0.20.0',
    'Importing from cross_attention is deprecated. Please import from diffusers.models.attention_processor instead.',
    standard_warn=False,
)


snake_case__						      =			AttentionProcessor



class 			UpperCamelCase					(			__lowercase			):

     '''simple docstring'''

     def __init__(       self					,					*A_					,					**A_    )			->					Union[str, Any]:






           """simple docstring"""
           _lowerCamelCase           =					F'{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"    )    )} instead.'
           deprecate('''cross_attention'''					,					'''0.20.0'''					,					A_					,					standard_warn=A_    )
           super().__init__(*A_					,					**A_    )





class 			UpperCamelCase					(			__lowercase			):

     '''simple docstring'''

     def __init__(       self					,					*A_					,					**A_    )			->					List[Any]:






           """simple docstring"""
           _lowerCamelCase           =					F'{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"    )    )} instead.'
           deprecate('''cross_attention'''					,					'''0.20.0'''					,					A_					,					standard_warn=A_    )
           super().__init__(*A_					,					**A_    )





class 			UpperCamelCase					(			__lowercase			):

     '''simple docstring'''

     def __init__(       self					,					*A_					,					**A_    )			->					Optional[int]:






           """simple docstring"""
           _lowerCamelCase           =					F'{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"    )    )} instead.'
           deprecate('''cross_attention'''					,					'''0.20.0'''					,					A_					,					standard_warn=A_    )
           super().__init__(*A_					,					**A_    )





class 			UpperCamelCase					(			__lowercase			):

     '''simple docstring'''

     def __init__(       self					,					*A_					,					**A_    )			->					Optional[int]:






           """simple docstring"""
           _lowerCamelCase           =					F'{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"    )    )} instead.'
           deprecate('''cross_attention'''					,					'''0.20.0'''					,					A_					,					standard_warn=A_    )
           super().__init__(*A_					,					**A_    )





class 			UpperCamelCase					(			__lowercase			):

     '''simple docstring'''

     def __init__(       self					,					*A_					,					**A_    )			->					Tuple:






           """simple docstring"""
           _lowerCamelCase           =					F'{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"    )    )} instead.'
           deprecate('''cross_attention'''					,					'''0.20.0'''					,					A_					,					standard_warn=A_    )
           super().__init__(*A_					,					**A_    )





class 			UpperCamelCase					(			__lowercase			):

     '''simple docstring'''

     def __init__(       self					,					*A_					,					**A_    )			->					int:






           """simple docstring"""
           _lowerCamelCase           =					F'{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"    )    )} instead.'
           deprecate('''cross_attention'''					,					'''0.20.0'''					,					A_					,					standard_warn=A_    )
           super().__init__(*A_					,					**A_    )





class 			UpperCamelCase					(			__lowercase			):

     '''simple docstring'''

     def __init__(       self					,					*A_					,					**A_    )			->					Optional[int]:






           """simple docstring"""
           _lowerCamelCase           =					F'{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"    )    )} instead.'
           deprecate('''cross_attention'''					,					'''0.20.0'''					,					A_					,					standard_warn=A_    )
           super().__init__(*A_					,					**A_    )





class 			UpperCamelCase					(			__lowercase			):

     '''simple docstring'''

     def __init__(       self					,					*A_					,					**A_    )			->					Dict:






           """simple docstring"""
           _lowerCamelCase           =					F'{self.__class__.__name__} is deprecated and will be removed in `0.20.0`. Please use `from diffusers.models.attention_processor import {"".join(self.__class__.__name__.split("Cross"    )    )} instead.'
           deprecate('''cross_attention'''					,					'''0.20.0'''					,					A_					,					standard_warn=A_    )
           super().__init__(*A_					,					**A_    )