


'''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 		_lowercase   (SCREAMING_SNAKE_CASE       ):




     '''simple docstring'''





     __A  :					Tuple						     =       ConvNextConfig()

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

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

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

     return config, expected_shape


def 		_lowercase   (SCREAMING_SNAKE_CASE       ):




     '''simple docstring'''





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

     return name


def 		_lowercase   ():




     '''simple docstring'''





     __A  :					Optional[int]						     =       "http://images.cocodataset.org/val2017/000000039769.jpg"
     __A  :					Union[str, Any]						     =       Image.open(requests.get(SCREAMING_SNAKE_CASE					,				stream=SCREAMING_SNAKE_CASE       ).raw       )
     return im


@torch.no_grad()
def 		_lowercase   (SCREAMING_SNAKE_CASE					,				SCREAMING_SNAKE_CASE       ):




     '''simple docstring'''





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

     # load HuggingFace model
     __A  :					Any						     =       ConvNextForImageClassification(SCREAMING_SNAKE_CASE       )
     model.load_state_dict(SCREAMING_SNAKE_CASE       )
     model.eval()

     # Check outputs on an image, prepared by ConvNextImageProcessor
     __A  :					Tuple						     =       224 if "224" in checkpoint_url else 384
     __A  :					Optional[Any]						     =       ConvNextImageProcessor(size=SCREAMING_SNAKE_CASE       )
     __A  :					Optional[int]						     =       image_processor(images=prepare_img()					,				return_tensors="pt"       ).pixel_values

     __A  :					Optional[Any]						     =       model(SCREAMING_SNAKE_CASE       ).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  :					List[Any]						     =       torch.tensor([-0.1210, -0.6605, 0.1918]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
          __A  :					Dict						     =       torch.tensor([-0.4473, -0.1847, -0.6365]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
          __A  :					List[str]						     =       torch.tensor([0.4525, 0.7539, 0.0308]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
          __A  :					List[Any]						     =       torch.tensor([0.3561, 0.6350, -0.0384]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
          __A  :					int						     =       torch.tensor([0.4174, -0.0989, 0.1489]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
          __A  :					List[Any]						     =       torch.tensor([0.2513, -0.1349, -0.1613]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
          __A  :					Union[str, Any]						     =       torch.tensor([1.2980, 0.3631, -0.1198]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
          __A  :					Optional[Any]						     =       torch.tensor([1.2963, 0.1227, 0.1723]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
          __A  :					Dict						     =       torch.tensor([1.7956, 0.8390, 0.2820]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
          __A  :					Optional[Any]						     =       torch.tensor([-0.2822, -0.0502, -0.0878]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
          __A  :					List[Any]						     =       torch.tensor([-0.5672, -0.0730, -0.4348]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
          __A  :					List[str]						     =       torch.tensor([0.2681, 0.2365, 0.6246]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
          __A  :					List[str]						     =       torch.tensor([-0.2642, 0.3931, 0.5116]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
          __A  :					List[Any]						     =       torch.tensor([-0.6677, -0.1873, -0.8379]       )
     elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
          __A  :					Union[str, Any]						     =       torch.tensor([-0.7749, -0.2967, -0.6444]       )
     else:
          raise ValueError(f"Unknown URL: {checkpoint_url}"       )

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

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

     print("Pushing model to the hub..."       )
     __A  :					Any						     =       "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(SCREAMING_SNAKE_CASE					,				SCREAMING_SNAKE_CASE       )					,				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 collections import defaultdict


def 		_lowercase   (SCREAMING_SNAKE_CASE = 5       ):




     '''simple docstring'''





     __A  :					int						     =       defaultdict(SCREAMING_SNAKE_CASE       )
     __A  :					Optional[Any]						     =       0

     while True:
          __A  :					List[str]						     =       get_digits(SCREAMING_SNAKE_CASE       )
          freqs[digits].append(SCREAMING_SNAKE_CASE       )

          if len(freqs[digits]       ) == max_base:
               __A  :					Optional[int]						     =       freqs[digits][0] ** 3
               return base

          num += 1


def 		_lowercase   (SCREAMING_SNAKE_CASE       ):




     '''simple docstring'''





     return "".join(sorted(str(num**3       )       )       )


if __name__ == "__main__":
      print(F"""{solution() = }""")

