





"""simple docstring"""





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




def 		lowerCAmelCase_     (  lowercase_       :							Optional[Any]		):

 '''simple docstring'''






 __SCREAMING_SNAKE_CASE			:      int 							=	ConvNextConfig()

 if "tiny" in checkpoint_url:
  __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	[3, 3, 9, 3]
  __SCREAMING_SNAKE_CASE			:      List[str] 							=	[96, 192, 384, 768]
 if "small" in checkpoint_url:
  __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	[3, 3, 27, 3]
  __SCREAMING_SNAKE_CASE			:      List[str] 							=	[96, 192, 384, 768]
 if "base" in checkpoint_url:
  __SCREAMING_SNAKE_CASE			:      int 							=	[3, 3, 27, 3]
  __SCREAMING_SNAKE_CASE			:      Any 							=	[128, 256, 512, 1024]
 if "large" in checkpoint_url:
  __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	[3, 3, 27, 3]
  __SCREAMING_SNAKE_CASE			:      Tuple 							=	[192, 384, 768, 1536]
 if "xlarge" in checkpoint_url:
  __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	[3, 3, 27, 3]
  __SCREAMING_SNAKE_CASE			:      Any 							=	[256, 512, 1024, 2048]

 if "1k" in checkpoint_url:
  __SCREAMING_SNAKE_CASE			:      int 							=	1000
  __SCREAMING_SNAKE_CASE			:      Any 							=	"imagenet-1k-id2label.json"
  __SCREAMING_SNAKE_CASE			:      List[Any] 							=	(1, 1000)
 else:
  __SCREAMING_SNAKE_CASE			:      Tuple 							=	2_1841
  __SCREAMING_SNAKE_CASE			:      str 							=	"imagenet-22k-id2label.json"
  __SCREAMING_SNAKE_CASE			:      str 							=	(1, 2_1841)

 __SCREAMING_SNAKE_CASE			:      Tuple 							=	"huggingface/label-files"
 __SCREAMING_SNAKE_CASE			:      List[Any] 							=	num_labels
 __SCREAMING_SNAKE_CASE			:      str 							=	json.load(open(hf_hub_download(_a			,				_a			,				repo_type='''dataset'''		)			,				'''r'''		)		)
 __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	{int(_a		): 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[9205]
  del idalabel[1_5027]
 __SCREAMING_SNAKE_CASE			:      Dict 							=	idalabel
 __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	{v: k for k, v in idalabel.items()}
 __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	hidden_sizes
 __SCREAMING_SNAKE_CASE			:      Tuple 							=	depths

 return config, expected_shape




def 		lowerCAmelCase_     (  lowercase_       :							Optional[Any]		):

 '''simple docstring'''






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

 return name




def 		lowerCAmelCase_     (  ):

 '''simple docstring'''






 __SCREAMING_SNAKE_CASE			:      Dict 							=	"http://images.cocodataset.org/val2017/000000039769.jpg"
 __SCREAMING_SNAKE_CASE			:      List[str] 							=	Image.open(requests.get(_a			,				stream=_a		).raw		)
 return im




@torch.no_grad()
def 		lowerCAmelCase_     (  lowercase_       :							List[str]			,				lowercase_       :							List[str]		):

 '''simple docstring'''






 __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	get_convnext_config(_a		)
 # load original state_dict from URL
 __SCREAMING_SNAKE_CASE			:      Any 							=	torch.hub.load_state_dict_from_url(_a		)["model"]
 # rename keys
 for key in state_dict.copy().keys():
  __SCREAMING_SNAKE_CASE			:      str 							=	state_dict.pop(_a		)
  __SCREAMING_SNAKE_CASE			:      List[str] 							=	val
 # add prefix to all keys expect classifier head
 for key in state_dict.copy().keys():
  __SCREAMING_SNAKE_CASE			:      List[Any] 							=	state_dict.pop(_a		)
  if not key.startswith('''classifier'''		):
   __SCREAMING_SNAKE_CASE			:      str 							=	"convnext." + key
  __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	val

 # load HuggingFace model
 __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	ConvNextForImageClassification(_a		)
 model.load_state_dict(_a		)
 model.eval()

 # Check outputs on an image, prepared by ConvNextImageProcessor
 __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	224 if "224" in checkpoint_url else 384
 __SCREAMING_SNAKE_CASE			:      List[Any] 							=	ConvNextImageProcessor(size=_a		)
 __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	image_processor(images=prepare_img()			,				return_tensors='''pt'''		).pixel_values

 __SCREAMING_SNAKE_CASE			:      int 							=	model(_a		).logits

 # note: the logits below were obtained without center cropping
 if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
  __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	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":
  __SCREAMING_SNAKE_CASE			:      List[Any] 							=	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":
  __SCREAMING_SNAKE_CASE			:      Dict 							=	torch.tensor([0.4_525, 0.7_539, 0.0_308]		)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
  __SCREAMING_SNAKE_CASE			:      List[str] 							=	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":
  __SCREAMING_SNAKE_CASE			:      str 							=	torch.tensor([0.4_174, -0.0_989, 0.1_489]		)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
  __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	torch.tensor([0.2_513, -0.1_349, -0.1_613]		)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
  __SCREAMING_SNAKE_CASE			:      int 							=	torch.tensor([1.2_980, 0.3_631, -0.1_198]		)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
  __SCREAMING_SNAKE_CASE			:      Tuple 							=	torch.tensor([1.2_963, 0.1_227, 0.1_723]		)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
  __SCREAMING_SNAKE_CASE			:      List[Any] 							=	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":
  __SCREAMING_SNAKE_CASE			:      List[Any] 							=	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":
  __SCREAMING_SNAKE_CASE			:      int 							=	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":
  __SCREAMING_SNAKE_CASE			:      int 							=	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":
  __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	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":
  __SCREAMING_SNAKE_CASE			:      int 							=	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":
  __SCREAMING_SNAKE_CASE			:      Any 							=	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]			,				_a			,				atol=1E-3		)
 assert logits.shape == expected_shape

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

 print('''Pushing model to the hub...'''		)
 __SCREAMING_SNAKE_CASE			:      int 							=	"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(_a			,				_a		)			,				organization='''nielsr'''			,				commit_message='''Add model'''			,				)


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

    _lowerCamelCase  					=     parser.parse_args()
    convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)








"""simple docstring"""





from ..utils import DummyObject, requires_backends





class       snake_case		(       metaclass=__UpperCAmelCase       ):
      lowerCamelCase__      =			['''torch''', '''transformers''', '''invisible_watermark''']

      def __init__(      self       :Optional[int]						,							*_lowerCamelCase       :int						,							**_lowerCamelCase       :Any						):
       requires_backends(self						,							['''torch''', '''transformers''', '''invisible_watermark''']						)

      @classmethod
      def 						SCREAMING_SNAKE_CASE_      (      cls       :Dict						,							*_lowerCamelCase       :Optional[Any]						,							**_lowerCamelCase       :int						):
       requires_backends(cls						,							['''torch''', '''transformers''', '''invisible_watermark''']						)



      @classmethod
      def 						SCREAMING_SNAKE_CASE_      (      cls       :Optional[int]						,							*_lowerCamelCase       :Optional[Any]						,							**_lowerCamelCase       :Union[str, Any]						):
       requires_backends(cls						,							['''torch''', '''transformers''', '''invisible_watermark''']						)







class       snake_case		(       metaclass=__UpperCAmelCase       ):
      lowerCamelCase__      =			['''torch''', '''transformers''', '''invisible_watermark''']

      def __init__(      self       :Dict						,							*_lowerCamelCase       :Dict						,							**_lowerCamelCase       :Optional[Any]						):
       requires_backends(self						,							['''torch''', '''transformers''', '''invisible_watermark''']						)

      @classmethod
      def 						SCREAMING_SNAKE_CASE_      (      cls       :Optional[int]						,							*_lowerCamelCase       :int						,							**_lowerCamelCase       :Optional[int]						):
       requires_backends(cls						,							['''torch''', '''transformers''', '''invisible_watermark''']						)



      @classmethod
      def 						SCREAMING_SNAKE_CASE_      (      cls       :List[Any]						,							*_lowerCamelCase       :Dict						,							**_lowerCamelCase       :List[str]						):
       requires_backends(cls						,							['''torch''', '''transformers''', '''invisible_watermark''']						)

