






"""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						:    Optional[Any]		=						logging.get_logger(__name__)
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Tuple       ):





  '''simple docstring'''





  lowercase				       =			ConvNextConfig()

  if "tiny" in checkpoint_url:
    lowercase				       =			[3, 3, 9, 3]
    lowercase				       =			[96, 1_92, 3_84, 7_68]
  if "small" in checkpoint_url:
    lowercase				       =			[3, 3, 27, 3]
    lowercase				       =			[96, 1_92, 3_84, 7_68]
  if "base" in checkpoint_url:
    lowercase				       =			[3, 3, 27, 3]
    lowercase				       =			[1_28, 2_56, 5_12, 10_24]
  if "large" in checkpoint_url:
    lowercase				       =			[3, 3, 27, 3]
    lowercase				       =			[1_92, 3_84, 7_68, 15_36]
  if "xlarge" in checkpoint_url:
    lowercase				       =			[3, 3, 27, 3]
    lowercase				       =			[2_56, 5_12, 10_24, 20_48]

  if "1k" in checkpoint_url:
    lowercase				       =			10_00
    lowercase				       =			'imagenet-1k-id2label.json'
    lowercase				       =			(1, 10_00)
  else:
    lowercase				       =			2_18_41
    lowercase				       =			'imagenet-22k-id2label.json'
    lowercase				       =			(1, 2_18_41)

  lowercase				       =			'huggingface/label-files'
  lowercase				       =			num_labels
  lowercase				       =			json.load(open(hf_hub_download(__snake_case						,   __snake_case						,   repo_type='dataset'       )						,   'r'       )       )
  lowercase				       =			{int(__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[92_05]
    del idalabel[1_50_27]
  lowercase				       =			idalabel
  lowercase				       =			{v: k for k, v in idalabel.items()}
  lowercase				       =			hidden_sizes
  lowercase				       =			depths

  return config, expected_shape
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    List[str]       ):





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





  '''simple docstring'''





  lowercase				       =			'http://images.cocodataset.org/val2017/000000039769.jpg'
  lowercase				       =			Image.open(requests.get(__snake_case						,   stream=__snake_case       ).raw       )
  return im
@torch.no_grad()
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    Any						,   __snake_case     :    Any       ):





  '''simple docstring'''





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

  # load HuggingFace model
  lowercase				       =			ConvNextForImageClassification(__snake_case       )
  model.load_state_dict(__snake_case       )
  model.eval()

  # Check outputs on an image, prepared by ConvNextImageProcessor
  lowercase				       =			2_24 if '224' in checkpoint_url else 3_84
  lowercase				       =			ConvNextImageProcessor(size=__snake_case       )
  lowercase				       =			image_processor(images=prepare_img()						,   return_tensors='pt'       ).pixel_values

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

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

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

  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(__snake_case						,   __snake_case       )						,   organization='nielsr'						,   commit_message='Add model'						,   )


if __name__ == "__main__":
    _UpperCamelCase						:    Union[str, 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.',
    )

    _UpperCamelCase						:    List[Any]		=						parser.parse_args()
    convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)




import torch

from transformers import PreTrainedModel

from .custom_configuration import CustomConfig, NoSuperInitConfig





class 				UpperCamelCase__		(	__SCREAMING_SNAKE_CASE	):






       """simple docstring"""


       UpperCAmelCase_				       =CustomConfig



       def __init__(							self ,     _A	)				->				int:
          super().__init__(_A	)
          SCREAMING_SNAKE_CASE_	=      torch.nn.Linear(config.hidden_size ,     config.hidden_size	)



       def     _UpperCamelCase	(							self ,     _A	)				->				Optional[int]:
          return self.linear(_A	)




       def     _UpperCamelCase	(							self ,     _A	)				->				int:
          pass







class 				UpperCamelCase__		(	__SCREAMING_SNAKE_CASE	):






       """simple docstring"""


       UpperCAmelCase_				       =NoSuperInitConfig



       def __init__(							self ,     _A	)				->				Optional[int]:
          super().__init__(_A	)
          SCREAMING_SNAKE_CASE_	=      torch.nn.Linear(config.attribute ,     config.attribute	)



       def     _UpperCamelCase	(							self ,     _A	)				->				Any:
          return self.linear(_A	)




       def     _UpperCamelCase	(							self ,     _A	)				->				Tuple:
          pass


