






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






def        lowerCamelCase_  (						_lowerCamelCase				):
 lowerCamelCase__     :			Optional[int]     				=						ConvNextConfig()

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

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

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

 return config, expected_shape






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

 return name






def        lowerCamelCase_  (						):
 lowerCamelCase__     :			Optional[int]     				=						'http://images.cocodataset.org/val2017/000000039769.jpg'
 lowerCamelCase__     :			str     				=						Image.open(requests.get(_lowerCamelCase			,   stream=_lowerCamelCase				).raw				)
 return im






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

 # load HuggingFace model
 lowerCamelCase__     :			Optional[int]     				=						ConvNextForImageClassification(_lowerCamelCase				)
 model.load_state_dict(_lowerCamelCase				)
 model.eval()

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

 lowerCamelCase__     :			Any     				=						model(_lowerCamelCase				).logits

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

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

 print('Pushing model to the hub...'				)
 lowerCamelCase__     :			Optional[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(_lowerCamelCase			,   _lowerCamelCase				)			,   organization='nielsr'			,   commit_message='Add model'			,   )


if __name__ == "__main__":
      A_			: List[Any]				=			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_			: List[str]				=			parser.parse_args()
      convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









"""simple docstring"""
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


A_			: Dict				=			{"configuration_qdqbert": ["QDQBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "QDQBertConfig"]}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A_			: Tuple				=			[
          "QDQBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
          "QDQBertForMaskedLM",
          "QDQBertForMultipleChoice",
          "QDQBertForNextSentencePrediction",
          "QDQBertForQuestionAnswering",
          "QDQBertForSequenceClassification",
          "QDQBertForTokenClassification",
          "QDQBertLayer",
          "QDQBertLMHeadModel",
          "QDQBertModel",
          "QDQBertPreTrainedModel",
          "load_tf_weights_in_qdqbert",
      ]


if TYPE_CHECKING:
      from .configuration_qdqbert import QDQBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, QDQBertConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_qdqbert import (
                QDQBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                QDQBertForMaskedLM,
                QDQBertForMultipleChoice,
                QDQBertForNextSentencePrediction,
                QDQBertForQuestionAnswering,
                QDQBertForSequenceClassification,
                QDQBertForTokenClassification,
                QDQBertLayer,
                QDQBertLMHeadModel,
                QDQBertModel,
                QDQBertPreTrainedModel,
                load_tf_weights_in_qdqbert,
            )


else:
      import sys

      A_			: Tuple				=			_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

