





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





def 			a_							(				_lowercase  ):
						_UpperCamelCase     :   List[Any]      =       ConvNextConfig()

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

						if "1k" in checkpoint_url:
												_UpperCamelCase     :   List[Any]      =       1000
												_UpperCamelCase     :   List[Any]      =       '''imagenet-1k-id2label.json'''
												_UpperCamelCase     :   str      =       (1, 1000)
						else:
												_UpperCamelCase     :   Union[str, Any]      =       2_1841
												_UpperCamelCase     :   Dict      =       '''imagenet-22k-id2label.json'''
												_UpperCamelCase     :   Optional[int]      =       (1, 2_1841)

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

						return config, expected_shape





def 			a_							(				_lowercase  ):
						if "downsample_layers.0.0" in name:
												_UpperCamelCase     :   Tuple      =       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     :   Tuple      =       name.replace('''downsample_layers.1.0'''			,		'''stages.1.downsampling_layer.0'''  )
						if "downsample_layers.1.1" in name:
												_UpperCamelCase     :   int      =       name.replace('''downsample_layers.1.1'''			,		'''stages.1.downsampling_layer.1'''  )
						if "downsample_layers.2.0" in name:
												_UpperCamelCase     :   Union[str, Any]      =       name.replace('''downsample_layers.2.0'''			,		'''stages.2.downsampling_layer.0'''  )
						if "downsample_layers.2.1" in name:
												_UpperCamelCase     :   int      =       name.replace('''downsample_layers.2.1'''			,		'''stages.2.downsampling_layer.1'''  )
						if "downsample_layers.3.0" in name:
												_UpperCamelCase     :   int      =       name.replace('''downsample_layers.3.0'''			,		'''stages.3.downsampling_layer.0'''  )
						if "downsample_layers.3.1" in name:
												_UpperCamelCase     :   List[Any]      =       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[Any]      =       name[: len('''stages.0'''  )] + '''.layers''' + name[len('''stages.0'''  ) :]
						if "stages" in name:
												_UpperCamelCase     :   Dict      =       name.replace('''stages'''			,		'''encoder.stages'''  )
						if "norm" in name:
												_UpperCamelCase     :   str      =       name.replace('''norm'''			,		'''layernorm'''  )
						if "gamma" in name:
												_UpperCamelCase     :   List[str]      =       name.replace('''gamma'''			,		'''layer_scale_parameter'''  )
						if "head" in name:
												_UpperCamelCase     :   List[Any]      =       name.replace('''head'''			,		'''classifier'''  )

						return name





def 			a_							(				):
						_UpperCamelCase     :   List[str]      =       '''http://images.cocodataset.org/val2017/000000039769.jpg'''
						_UpperCamelCase     :   List[Any]      =       Image.open(requests.get(_lowercase			,		stream=_lowercase  ).raw  )
						return im





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

						# load HuggingFace model
						_UpperCamelCase     :   Tuple      =       ConvNextForImageClassification(_lowercase  )
						model.load_state_dict(_lowercase  )
						model.eval()

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

						_UpperCamelCase     :   str      =       model(_lowercase  ).logits

						# note: the logits below were obtained without center cropping
						if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
												_UpperCamelCase     :   Any      =       torch.tensor([-0.1210, -0.6605, 0.1918]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
												_UpperCamelCase     :   int      =       torch.tensor([-0.4473, -0.1847, -0.6365]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
												_UpperCamelCase     :   Dict      =       torch.tensor([0.4525, 0.7539, 0.0308]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
												_UpperCamelCase     :   List[Any]      =       torch.tensor([0.3561, 0.6350, -0.0384]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
												_UpperCamelCase     :   Dict      =       torch.tensor([0.4174, -0.0989, 0.1489]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
												_UpperCamelCase     :   Any      =       torch.tensor([0.2513, -0.1349, -0.1613]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
												_UpperCamelCase     :   Dict      =       torch.tensor([1.2980, 0.3631, -0.1198]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
												_UpperCamelCase     :   Union[str, Any]      =       torch.tensor([1.2963, 0.1227, 0.1723]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
												_UpperCamelCase     :   Union[str, Any]      =       torch.tensor([1.7956, 0.8390, 0.2820]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
												_UpperCamelCase     :   Dict      =       torch.tensor([-0.2822, -0.0502, -0.0878]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
												_UpperCamelCase     :   Tuple      =       torch.tensor([-0.5672, -0.0730, -0.4348]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
												_UpperCamelCase     :   List[str]      =       torch.tensor([0.2681, 0.2365, 0.6246]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
												_UpperCamelCase     :   int      =       torch.tensor([-0.2642, 0.3931, 0.5116]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
												_UpperCamelCase     :   str      =       torch.tensor([-0.6677, -0.1873, -0.8379]  )
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
												_UpperCamelCase     :   List[str]      =       torch.tensor([-0.7749, -0.2967, -0.6444]  )
						else:
												raise ValueError(F"""Unknown URL: {checkpoint_url}"""  )

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

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

						print('''Pushing model to the hub...'''  )
						_UpperCamelCase     :   List[str]      =       '''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(_lowercase			,		_lowercase  )			,		organization='''nielsr'''			,		commit_message='''Add model'''			,		)


if __name__ == "__main__":
						UpperCamelCase_	 =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_	 =parser.parse_args()
						convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









'''simple docstring'''


from __future__ import annotations

from collections.abc import Callable
from typing import Generic, TypeVar

UpperCamelCase			     =					TypeVar('''T''')
UpperCamelCase			     =					TypeVar('''U''')






class 						lowerCAmelCase_	(       Generic[T, U]							):







  '''simple docstring'''

  def __init__( self				:							Tuple					,  SCREAMING_SNAKE_CASE_				:							T | None					,  SCREAMING_SNAKE_CASE_				:							U | None	)			->		Any:







       '''simple docstring'''



       A:		Union[str, Any]          =		key
       A:		Union[str, Any]          =		val
       A:		int          =		0
       A:		DoubleLinkedListNode[T, U] | None          =		None
       A:		DoubleLinkedListNode[T, U] | None          =		None




  def __repr__( self				:							Any	)			->		str:







       '''simple docstring'''



       return "Node: key: {}, val: {}, freq: {}, has next: {}, has prev: {}".format(
           self.key					,  self.val					,  self.freq					,  self.next is not None					,  self.prev is not None	)






class 						lowerCAmelCase_	(       Generic[T, U]							):







  '''simple docstring'''

  def __init__( self				:							int	)			->		None:







       '''simple docstring'''



       A:		DoubleLinkedListNode[T, U]          =		DoubleLinkedListNode(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       A:		DoubleLinkedListNode[T, U]          =		DoubleLinkedListNode(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
       A     ,    A:		Any          =		self.rear, self.head



  def __repr__( self				:							Optional[int]	)			->		str:







       '''simple docstring'''



       A:		int          =		['''DoubleLinkedList''']
       A:		Dict          =		self.head
       while node.next is not None:
            rep.append(str(SCREAMING_SNAKE_CASE_	)	)
            A:		int          =		node.next
       rep.append(str(self.rear	)	)
       return ",\n    ".join(SCREAMING_SNAKE_CASE_	)



  def 					_snake_case     ( self				:							Any					,  SCREAMING_SNAKE_CASE_				:							DoubleLinkedListNode[T, U]	)			->		None:







       '''simple docstring'''



       A:		Union[str, Any]          =		self.rear.prev

       # All nodes other than self.head are guaranteed to have non-None previous
       assert previous is not None

       A:		Optional[int]          =		node
       A:		int          =		previous
       A:		Dict          =		node
       A:		int          =		self.rear
       node.freq += 1
       self._position_node(SCREAMING_SNAKE_CASE_	)



  def 					_snake_case     ( self				:							int					,  SCREAMING_SNAKE_CASE_				:							DoubleLinkedListNode[T, U]	)			->		None:







       '''simple docstring'''



       while node.prev is not None and node.prev.freq > node.freq:
            # swap node with previous node
            A:		Any          =		node.prev

            A:		Optional[Any]          =		previous_node.prev
            A:		int          =		node.prev
            A:		int          =		previous_node
            A:		Optional[int]          =		node




  def 					_snake_case     ( self				:							Dict					,  SCREAMING_SNAKE_CASE_				:							DoubleLinkedListNode[T, U]	)			->		DoubleLinkedListNode[T, U] | None:







       '''simple docstring'''



       if node.prev is None or node.next is None:
            return None

       A:		List[Any]          =		node.next
       A:		Dict          =		node.prev
       A:		str          =		None
       A:		Union[str, Any]          =		None
       return node






class 						lowerCAmelCase_	(       Generic[T, U]							):







  '''simple docstring'''

  UpperCamelCase_			:      dict[Callable[[T], U], LFUCache[T, U]]										=				{}



  def __init__( self				:							List[str]					,  SCREAMING_SNAKE_CASE_				:							int	)			->		int:







       '''simple docstring'''



       A:		DoubleLinkedList[T, U]          =		DoubleLinkedList()
       A:		Any          =		capacity
       A:		List[Any]          =		0
       A:		Dict          =		0
       A:		Dict          =		0
       A:		dict[T, DoubleLinkedListNode[T, U]]          =		{}



  def __repr__( self				:							Optional[int]	)			->		str:







       '''simple docstring'''



       return (
           f"""CacheInfo(hits={self.hits}, misses={self.miss}, """
           f"""capacity={self.capacity}, current_size={self.num_keys})"""
       )



  def __contains__( self				:							str					,  SCREAMING_SNAKE_CASE_				:							T	)			->		bool:







       '''simple docstring'''



       return key in self.cache



  def 					_snake_case     ( self				:							Dict					,  SCREAMING_SNAKE_CASE_				:							T	)			->		U | None:







       '''simple docstring'''



       if key in self.cache:
            self.hits += 1
            A:		DoubleLinkedListNode[T, U]          =		self.cache[key]
            A:		List[str]          =		self.list.remove(self.cache[key]	)
            assert node == value_node

            # node is guaranteed not None because it is in self.cache
            assert node is not None
            self.list.add(SCREAMING_SNAKE_CASE_	)
            return node.val
       self.miss += 1
       return None



  def 					_snake_case     ( self				:							List[str]					,  SCREAMING_SNAKE_CASE_				:							T					,  SCREAMING_SNAKE_CASE_				:							U	)			->		None:







       '''simple docstring'''



       if key not in self.cache:
            if self.num_keys >= self.capacity:
                 # delete first node when over capacity
                 A:		List[Any]          =		self.list.head.next

                 # guaranteed to have a non-None first node when num_keys > 0
                 # explain to type checker via assertions
                 assert first_node is not None
                 assert first_node.key is not None
                 assert self.list.remove(SCREAMING_SNAKE_CASE_	) is not None
                 # first_node guaranteed to be in list

                 del self.cache[first_node.key]
                 self.num_keys -= 1
            A:		Tuple          =		DoubleLinkedListNode(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_	)
            self.list.add(self.cache[key]	)
            self.num_keys += 1

       else:
            A:		Optional[Any]          =		self.list.remove(self.cache[key]	)
            assert node is not None  # node guaranteed to be in list
            A:		List[str]          =		value
            self.list.add(SCREAMING_SNAKE_CASE_	)




  @classmethod
  def 					_snake_case     ( cls				:							type[LFUCache[T, U]]					,  SCREAMING_SNAKE_CASE_				:							int = 1_28	)			->		Callable[[Callable[[T], U]], Callable[..., U]]:







       '''simple docstring'''



       def cache_decorator_inner(SCREAMING_SNAKE_CASE_				:							Callable[[T], U]	) -> Callable[..., U]:
            def cache_decorator_wrapper(*SCREAMING_SNAKE_CASE_				:							T	) -> U:
                 if func not in cls.decorator_function_to_instance_map:
                      A:		Tuple          =		LFUCache(SCREAMING_SNAKE_CASE_	)

                 A:		List[Any]          =		cls.decorator_function_to_instance_map[func].get(args[0]	)
                 if result is None:
                      A:		List[str]          =		func(*SCREAMING_SNAKE_CASE_	)
                      cls.decorator_function_to_instance_map[func].put(args[0]					,  SCREAMING_SNAKE_CASE_	)
                 return result

            def cache_info() -> LFUCache[T, U]:
                 return cls.decorator_function_to_instance_map[func]

            setattr(SCREAMING_SNAKE_CASE_					,  '''cache_info'''					,  SCREAMING_SNAKE_CASE_	)  # noqa: B010

            return cache_decorator_wrapper

       return cache_decorator_inner


if __name__ == "__main__":
     import doctest

     doctest.testmod()






