

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




def     _A    (    UpperCamelCase_	:			Optional[Any])      ->		Dict:
   '''simple docstring'''


   __lowercase								=		ConvNextConfig()

   if "tiny" in checkpoint_url:
      __lowercase								=		[3, 3, 9, 3]
      __lowercase								=		[96, 192, 384, 768]
   if "small" in checkpoint_url:
      __lowercase								=		[3, 3, 27, 3]
      __lowercase								=		[96, 192, 384, 768]
   if "base" in checkpoint_url:
      __lowercase								=		[3, 3, 27, 3]
      __lowercase								=		[128, 256, 512, 1024]
   if "large" in checkpoint_url:
      __lowercase								=		[3, 3, 27, 3]
      __lowercase								=		[192, 384, 768, 1536]
   if "xlarge" in checkpoint_url:
      __lowercase								=		[3, 3, 27, 3]
      __lowercase								=		[256, 512, 1024, 2048]

   if "1k" in checkpoint_url:
      __lowercase								=		1000
      __lowercase								=		'''imagenet-1k-id2label.json'''
      __lowercase								=		(1, 1000)
   else:
      __lowercase								=		21841
      __lowercase								=		'''imagenet-22k-id2label.json'''
      __lowercase								=		(1, 21841)

   __lowercase								=		'''huggingface/label-files'''
   __lowercase								=		num_labels
   __lowercase								=		json.load(open(hf_hub_download(_lowercase,     _lowercase,     repo_type="dataset"),     "r"))
   __lowercase								=		{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[15027]
   __lowercase								=		idalabel
   __lowercase								=		{v: k for k, v in idalabel.items()}
   __lowercase								=		hidden_sizes
   __lowercase								=		depths

   return config, expected_shape




def     _A    (    UpperCamelCase_	:			List[Any])      ->		Optional[Any]:
   '''simple docstring'''


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

   return name




def     _A    (    )      ->		List[Any]:
   '''simple docstring'''


   __lowercase								=		'''http://images.cocodataset.org/val2017/000000039769.jpg'''
   __lowercase								=		Image.open(requests.get(_lowercase,     stream=_lowercase).raw)
   return im




@torch.no_grad()
def     _A    (    UpperCamelCase_	:			Union[str, Any],     UpperCamelCase_	:			Tuple)      ->		Tuple:
   '''simple docstring'''


   __lowercase								=		get_convnext_config(_lowercase)
   # load original state_dict from URL
   __lowercase								=		torch.hub.load_state_dict_from_url(_lowercase)['''model''']
   # rename keys
   for key in state_dict.copy().keys():
      __lowercase								=		state_dict.pop(_lowercase)
      __lowercase								=		val
   # add prefix to all keys expect classifier head
   for key in state_dict.copy().keys():
      __lowercase								=		state_dict.pop(_lowercase)
      if not key.startswith("classifier"):
         __lowercase								=		'''convnext.''' + key
      __lowercase								=		val

   # load HuggingFace model
   __lowercase								=		ConvNextForImageClassification(_lowercase)
   model.load_state_dict(_lowercase)
   model.eval()

   # Check outputs on an image, prepared by ConvNextImageProcessor
   __lowercase								=		224 if '''224''' in checkpoint_url else 384
   __lowercase								=		ConvNextImageProcessor(size=_lowercase)
   __lowercase								=		image_processor(images=prepare_img(),     return_tensors="pt").pixel_values

   __lowercase								=		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":
      __lowercase								=		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":
      __lowercase								=		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":
      __lowercase								=		torch.tensor([0.4_525, 0.7_539, 0.0_308])
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
      __lowercase								=		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":
      __lowercase								=		torch.tensor([0.4_174, -0.0_989, 0.1_489])
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
      __lowercase								=		torch.tensor([0.2_513, -0.1_349, -0.1_613])
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
      __lowercase								=		torch.tensor([1.2_980, 0.3_631, -0.1_198])
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
      __lowercase								=		torch.tensor([1.2_963, 0.1_227, 0.1_723])
   elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
      __lowercase								=		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":
      __lowercase								=		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":
      __lowercase								=		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":
      __lowercase								=		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":
      __lowercase								=		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":
      __lowercase								=		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":
      __lowercase								=		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],     _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...")
   __lowercase								=		'''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__":
   _a					 =				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.',
   )

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











"""simple docstring"""




from .hash_table import HashTable
from .number_theory.prime_numbers import is_prime, next_prime




class 					_a		( _lowerCAmelCase							):






				def __init__(					self						:       Tuple,					*lowerCAmelCase__						:       Tuple,					**lowerCAmelCase__						:       Optional[int]				)      -> Union[str, Any]:




										'''simple docstring'''



										super().__init__(*lowerCAmelCase__,					**lowerCAmelCase__				)






				def 							snake_case      (					self						:       Any,					lowerCAmelCase__						:       Dict,					lowerCAmelCase__						:       Tuple				)      -> Optional[Any]:




										'''simple docstring'''



										_UpperCamelCase     :   Tuple      =       (
										    next_prime(value % self.size_table				)
										    if not is_prime(value % self.size_table				)
										    else value % self.size_table
										)  # gt = bigger than
										return next_prime_gt - (data % next_prime_gt)






				def 							snake_case      (					self						:       int,					lowerCAmelCase__						:       Any,					lowerCAmelCase__						:       List[Any],					lowerCAmelCase__						:       Tuple				)      -> Any:




										'''simple docstring'''



										return (increment * self.__hash_function_a(lowerCAmelCase__,					lowerCAmelCase__				)) % self.size_table






				def 							snake_case      (					self						:       Dict,					lowerCAmelCase__						:       str,					lowerCAmelCase__						:       Optional[Any]=None				)      -> Union[str, Any]:




										'''simple docstring'''



										_UpperCamelCase     :   Tuple      =       1
										_UpperCamelCase     :   Dict      =       self.hash_function(lowerCAmelCase__				)

										while self.values[new_key] is not None and self.values[new_key] != key:
																_UpperCamelCase     :   Optional[int]      =       (
																    self.__hash_double_function(lowerCAmelCase__,					lowerCAmelCase__,					lowerCAmelCase__				)
																    if self.balanced_factor() >= self.lim_charge
																    else None
																)
																if new_key is None:
																						break
																else:
																						i += 1

										return new_key



