

'''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				: Tuple 							=						logging.get_logger(__name__)


def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   Union[str, Any]						):
  UpperCAmelCase													=	ConvNextConfig()

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

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

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

  return config, expected_shape


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

  return name


def       lowerCamelCase__		(						):
  UpperCAmelCase													=	'http://images.cocodataset.org/val2017/000000039769.jpg'
  UpperCAmelCase													=	Image.open(requests.get(SCREAMING_SNAKE_CASE       ,   stream=SCREAMING_SNAKE_CASE						).raw						)
  return im


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

  # load HuggingFace model
  UpperCAmelCase													=	ConvNextForImageClassification(SCREAMING_SNAKE_CASE						)
  model.load_state_dict(SCREAMING_SNAKE_CASE						)
  model.eval()

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

  UpperCAmelCase													=	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":
    UpperCAmelCase													=	torch.tensor([-0.1210, -0.6605, 0.1918]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
    UpperCAmelCase													=	torch.tensor([-0.4473, -0.1847, -0.6365]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
    UpperCAmelCase													=	torch.tensor([0.4525, 0.7539, 0.0308]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
    UpperCAmelCase													=	torch.tensor([0.3561, 0.6350, -0.0384]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
    UpperCAmelCase													=	torch.tensor([0.4174, -0.0989, 0.1489]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
    UpperCAmelCase													=	torch.tensor([0.2513, -0.1349, -0.1613]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
    UpperCAmelCase													=	torch.tensor([1.2980, 0.3631, -0.1198]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
    UpperCAmelCase													=	torch.tensor([1.2963, 0.1227, 0.1723]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
    UpperCAmelCase													=	torch.tensor([1.7956, 0.8390, 0.2820]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
    UpperCAmelCase													=	torch.tensor([-0.2822, -0.0502, -0.0878]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
    UpperCAmelCase													=	torch.tensor([-0.5672, -0.0730, -0.4348]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
    UpperCAmelCase													=	torch.tensor([0.2681, 0.2365, 0.6246]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
    UpperCAmelCase													=	torch.tensor([-0.2642, 0.3931, 0.5116]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
    UpperCAmelCase													=	torch.tensor([-0.6677, -0.1873, -0.8379]						)
  elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
    UpperCAmelCase													=	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...'						)
  UpperCAmelCase													=	'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__":
  _a				: str 							=						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'''




import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
  import torch

if is_vision_available():
  from PIL import Image

  from transformers import MobileNetVaImageProcessor



class        lowercase_	(		unittest.TestCase ):

  '''simple docstring'''





  def __init__(     self   , a_   , a_=7   , a_=3   , a_=1_8   , a_=3_0   , a_=4_0_0   , a_=True   , a_=None   , a_=True   , a_=None   , )							->				List[str]:


    """simple docstring"""
    UpperCAmelCase													=	size if size is not None else {'shortest_edge': 2_0}
    UpperCAmelCase													=	crop_size if crop_size is not None else {'height': 1_8, 'width': 1_8}
    UpperCAmelCase													=	parent
    UpperCAmelCase													=	batch_size
    UpperCAmelCase													=	num_channels
    UpperCAmelCase													=	image_size
    UpperCAmelCase													=	min_resolution
    UpperCAmelCase													=	max_resolution
    UpperCAmelCase													=	do_resize
    UpperCAmelCase													=	size
    UpperCAmelCase													=	do_center_crop
    UpperCAmelCase													=	crop_size





  def 						snake_case_			(     self     )							->				Optional[int]:


    """simple docstring"""
    return {
        "do_resize": self.do_resize,
        "size": self.size,
        "do_center_crop": self.do_center_crop,
        "crop_size": self.crop_size,
    }




@require_torch
@require_vision
class        lowercase_	(		a ,       unittest.TestCase ):

  '''simple docstring'''





  __lowerCAmelCase					: Union[str, Any]      =   MobileNetVaImageProcessor if is_vision_available() else None
  def 						snake_case_			(     self     )							->				Dict:


    """simple docstring"""
    UpperCAmelCase													=	MobileNetVaImageProcessingTester(self     )
  @property
  def 						snake_case_			(     self     )							->				List[str]:


    """simple docstring"""
    return self.image_processor_tester.prepare_image_processor_dict()
  def 						snake_case_			(     self     )							->				Optional[int]:


    """simple docstring"""
    UpperCAmelCase													=	self.image_processing_class(**self.image_processor_dict     )
    self.assertTrue(hasattr(a_   , 'do_resize'     )     )
    self.assertTrue(hasattr(a_   , 'size'     )     )
    self.assertTrue(hasattr(a_   , 'do_center_crop'     )     )
    self.assertTrue(hasattr(a_   , 'center_crop'     )     )
  def 						snake_case_			(     self     )							->				str:


    """simple docstring"""
    UpperCAmelCase													=	self.image_processing_class.from_dict(self.image_processor_dict     )
    self.assertEqual(image_processor.size   , {'shortest_edge': 2_0}     )
    self.assertEqual(image_processor.crop_size   , {'height': 1_8, 'width': 1_8}     )

    UpperCAmelCase													=	self.image_processing_class.from_dict(self.image_processor_dict   , size=4_2   , crop_size=8_4     )
    self.assertEqual(image_processor.size   , {'shortest_edge': 4_2}     )
    self.assertEqual(image_processor.crop_size   , {'height': 8_4, 'width': 8_4}     )
  def 						snake_case_			(     self     )							->				List[str]:


    """simple docstring"""
    pass
  def 						snake_case_			(     self     )							->				List[Any]:


    """simple docstring"""
    # Initialize image_processing
    UpperCAmelCase													=	self.image_processing_class(**self.image_processor_dict     )
    # create random PIL images
    UpperCAmelCase													=	prepare_image_inputs(self.image_processor_tester   , equal_resolution=a_     )
    for image in image_inputs:
      self.assertIsInstance(a_   , Image.Image     )

    # Test not batched input
    UpperCAmelCase													=	image_processing(image_inputs[0]   , return_tensors='pt'     ).pixel_values
    self.assertEqual(
        encoded_images.shape   , (
            1,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['height'],
            self.image_processor_tester.crop_size['width'],
        )   , )

    # Test batched
    UpperCAmelCase													=	image_processing(a_   , return_tensors='pt'     ).pixel_values
    self.assertEqual(
        encoded_images.shape   , (
            self.image_processor_tester.batch_size,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['height'],
            self.image_processor_tester.crop_size['width'],
        )   , )
  def 						snake_case_			(     self     )							->				str:


    """simple docstring"""
    # Initialize image_processing
    UpperCAmelCase													=	self.image_processing_class(**self.image_processor_dict     )
    # create random numpy tensors
    UpperCAmelCase													=	prepare_image_inputs(self.image_processor_tester   , equal_resolution=a_   , numpify=a_     )
    for image in image_inputs:
      self.assertIsInstance(a_   , np.ndarray     )

    # Test not batched input
    UpperCAmelCase													=	image_processing(image_inputs[0]   , return_tensors='pt'     ).pixel_values
    self.assertEqual(
        encoded_images.shape   , (
            1,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['height'],
            self.image_processor_tester.crop_size['width'],
        )   , )

    # Test batched
    UpperCAmelCase													=	image_processing(a_   , return_tensors='pt'     ).pixel_values
    self.assertEqual(
        encoded_images.shape   , (
            self.image_processor_tester.batch_size,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['height'],
            self.image_processor_tester.crop_size['width'],
        )   , )





  def 						snake_case_			(     self     )							->				List[Any]:


    """simple docstring"""
    # Initialize image_processing
    UpperCAmelCase													=	self.image_processing_class(**self.image_processor_dict     )
    # create random PyTorch tensors
    UpperCAmelCase													=	prepare_image_inputs(self.image_processor_tester   , equal_resolution=a_   , torchify=a_     )
    for image in image_inputs:
      self.assertIsInstance(a_   , torch.Tensor     )

    # Test not batched input
    UpperCAmelCase													=	image_processing(image_inputs[0]   , return_tensors='pt'     ).pixel_values
    self.assertEqual(
        encoded_images.shape   , (
            1,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['height'],
            self.image_processor_tester.crop_size['width'],
        )   , )

    # Test batched
    UpperCAmelCase													=	image_processing(a_   , return_tensors='pt'     ).pixel_values
    self.assertEqual(
        encoded_images.shape   , (
            self.image_processor_tester.batch_size,
            self.image_processor_tester.num_channels,
            self.image_processor_tester.crop_size['height'],
            self.image_processor_tester.crop_size['width'],
        )   , )


