

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


def       a__    (  _SCREAMING_SNAKE_CASE					:    Union[str, Any]							)			-> List[str]:


	"""simple docstring"""





	UpperCAmelCase_			:	Tuple      =   ConvNextConfig()

	if "tiny" in checkpoint_url:
		UpperCAmelCase_			:	Tuple      =   [3, 3, 9, 3]
		UpperCAmelCase_			:	Optional[Any]      =   [96, 1_92, 3_84, 7_68]
	if "small" in checkpoint_url:
		UpperCAmelCase_			:	Tuple      =   [3, 3, 27, 3]
		UpperCAmelCase_			:	Dict      =   [96, 1_92, 3_84, 7_68]
	if "base" in checkpoint_url:
		UpperCAmelCase_			:	Optional[int]      =   [3, 3, 27, 3]
		UpperCAmelCase_			:	List[Any]      =   [1_28, 2_56, 5_12, 10_24]
	if "large" in checkpoint_url:
		UpperCAmelCase_			:	List[Any]      =   [3, 3, 27, 3]
		UpperCAmelCase_			:	str      =   [1_92, 3_84, 7_68, 15_36]
	if "xlarge" in checkpoint_url:
		UpperCAmelCase_			:	Any      =   [3, 3, 27, 3]
		UpperCAmelCase_			:	List[str]      =   [2_56, 5_12, 10_24, 20_48]

	if "1k" in checkpoint_url:
		UpperCAmelCase_			:	Tuple      =   10_00
		UpperCAmelCase_			:	str      =   "imagenet-1k-id2label.json"
		UpperCAmelCase_			:	List[Any]      =   (1, 10_00)
	else:
		UpperCAmelCase_			:	str      =   2_18_41
		UpperCAmelCase_			:	Any      =   "imagenet-22k-id2label.json"
		UpperCAmelCase_			:	List[Any]      =   (1, 2_18_41)

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

	return config, expected_shape


def       a__    (  _SCREAMING_SNAKE_CASE					:    int							)			-> Union[str, Any]:


	"""simple docstring"""





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

	return name


def       a__    (  )			-> Tuple:


	"""simple docstring"""





	UpperCAmelCase_			:	Tuple      =   "http://images.cocodataset.org/val2017/000000039769.jpg"
	UpperCAmelCase_			:	List[str]      =   Image.open(requests.get(_SCREAMING_SNAKE_CASE     ,							stream=_SCREAMING_SNAKE_CASE							).raw							)
	return im


@torch.no_grad()
def       a__    (  _SCREAMING_SNAKE_CASE					:    Dict     ,							_SCREAMING_SNAKE_CASE					:    Optional[Any]							)			-> List[Any]:


	"""simple docstring"""





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

	# load HuggingFace model
	UpperCAmelCase_			:	Dict      =   ConvNextForImageClassification(_SCREAMING_SNAKE_CASE							)
	model.load_state_dict(_SCREAMING_SNAKE_CASE							)
	model.eval()

	# Check outputs on an image, prepared by ConvNextImageProcessor
	UpperCAmelCase_			:	Union[str, Any]      =   2_24 if "224" in checkpoint_url else 3_84
	UpperCAmelCase_			:	List[Any]      =   ConvNextImageProcessor(size=_SCREAMING_SNAKE_CASE							)
	UpperCAmelCase_			:	int      =   image_processor(images=prepare_img()     ,							return_tensors="pt"							).pixel_values

	UpperCAmelCase_			:	List[str]      =   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_			:	Optional[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":
		UpperCAmelCase_			:	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":
		UpperCAmelCase_			:	Optional[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":
		UpperCAmelCase_			:	str      =   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":
		UpperCAmelCase_			:	Optional[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":
		UpperCAmelCase_			:	Union[str, Any]      =   torch.tensor([0.2_513, -0.1_349, -0.1_613]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
		UpperCAmelCase_			:	Any      =   torch.tensor([1.2_980, 0.3_631, -0.1_198]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
		UpperCAmelCase_			:	int      =   torch.tensor([1.2_963, 0.1_227, 0.1_723]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
		UpperCAmelCase_			:	List[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":
		UpperCAmelCase_			:	Tuple      =   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":
		UpperCAmelCase_			:	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":
		UpperCAmelCase_			:	Dict      =   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":
		UpperCAmelCase_			:	Tuple      =   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":
		UpperCAmelCase_			:	Optional[int]      =   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":
		UpperCAmelCase_			:	Tuple      =   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]     ,							_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_			:	Dict      =   "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__":
					_lowerCamelCase	 =			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.""",
					)

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







'''simple docstring'''
# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from typing import List


def       a__    (  _SCREAMING_SNAKE_CASE					:    List[str]     ,							_SCREAMING_SNAKE_CASE					:    str							)			-> Optional[Any]:


	"""simple docstring"""





	if name != "training_function" and name != "main":
		raise ValueError(F'''Incorrect function name passed: {name}, choose either \'main\' or \'training_function\''''							)
	UpperCAmelCase_ , UpperCAmelCase_			:	int      =   [], False
	for line in lines:
		if not found_start and F'''def {name}''' in line:
			UpperCAmelCase_			:	List[str]      =   True
			good_lines.append(_SCREAMING_SNAKE_CASE							)
			continue
		if found_start:
			if name == "training_function" and "def main" in line:
				return good_lines
			if name == "main" and "if __name__" in line:
				return good_lines
			good_lines.append(_SCREAMING_SNAKE_CASE							)


def       a__    (  _SCREAMING_SNAKE_CASE					:    List[str]							)			-> Optional[Any]:


	"""simple docstring"""





	return [line for line in lines if not line.lstrip().startswith("#"							) and line != "\n"]





def       a__    (  _SCREAMING_SNAKE_CASE					:    str     ,							_SCREAMING_SNAKE_CASE					:    str     ,							_SCREAMING_SNAKE_CASE					:    bool     ,							_SCREAMING_SNAKE_CASE					:    str = None							)			-> int:


	"""simple docstring"""





	with open(_SCREAMING_SNAKE_CASE     ,							"r"							) as f:
		UpperCAmelCase_			:	Union[str, Any]      =   f.readlines()
	with open(os.path.abspath(os.path.join("examples"     ,							"nlp_example.py"							)							)     ,							"r"							) as f:
		UpperCAmelCase_			:	Dict      =   f.readlines()
	with open(_SCREAMING_SNAKE_CASE     ,							"r"							) as f:
		UpperCAmelCase_			:	Any      =   f.readlines()
	if secondary_filename is not None:
		with open(_SCREAMING_SNAKE_CASE     ,							"r"							) as f:
			UpperCAmelCase_			:	Tuple      =   f.readlines()

    # This is our base, we remove all the code from here in our `full_filename` and `feature_filename` to find the new content
	if parser_only:
		UpperCAmelCase_			:	List[str]      =   clean_lines(get_function_contents_by_name(_SCREAMING_SNAKE_CASE     ,							"main"							)							)
		UpperCAmelCase_			:	Tuple      =   clean_lines(get_function_contents_by_name(_SCREAMING_SNAKE_CASE     ,							"main"							)							)
		UpperCAmelCase_			:	Tuple      =   clean_lines(get_function_contents_by_name(_SCREAMING_SNAKE_CASE     ,							"main"							)							)
		if secondary_filename is not None:
			UpperCAmelCase_			:	List[Any]      =   clean_lines(get_function_contents_by_name(_SCREAMING_SNAKE_CASE     ,							"main"							)							)
	else:
		UpperCAmelCase_			:	Union[str, Any]      =   clean_lines(get_function_contents_by_name(_SCREAMING_SNAKE_CASE     ,							"training_function"							)							)
		UpperCAmelCase_			:	str      =   clean_lines(get_function_contents_by_name(_SCREAMING_SNAKE_CASE     ,							"training_function"							)							)
		UpperCAmelCase_			:	List[Any]      =   clean_lines(get_function_contents_by_name(_SCREAMING_SNAKE_CASE     ,							"training_function"							)							)
		if secondary_filename is not None:
			UpperCAmelCase_			:	Tuple      =   clean_lines(
			    get_function_contents_by_name(_SCREAMING_SNAKE_CASE     ,							"training_function"							)							)

	UpperCAmelCase_			:	List[Any]      =   "train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n"

	# Specific code in our script that differs from the full version, aka what is new
	UpperCAmelCase_			:	Dict      =   []
	UpperCAmelCase_			:	Any      =   []  # We keep track of the idxs just in case it's a repeated statement
	UpperCAmelCase_			:	Optional[int]      =   iter(_SCREAMING_SNAKE_CASE							)
	for i in range(len(_SCREAMING_SNAKE_CASE							) - 1							):
		if i not in passed_idxs:
			UpperCAmelCase_			:	str      =   next(_SCREAMING_SNAKE_CASE							)
			if (line not in full_file_func) and (line.lstrip() != _dl_line):
				if "TESTING_MOCKED_DATALOADERS" not in line:
					new_feature_code.append(_SCREAMING_SNAKE_CASE							)
					passed_idxs.append(_SCREAMING_SNAKE_CASE							)
				else:
					# Skip over the `config['num_epochs'] = 2` statement
					UpperCAmelCase_			:	List[Any]      =   next(_SCREAMING_SNAKE_CASE							)

    # Extract out just the new parts from the full_file_training_func
	UpperCAmelCase_			:	Tuple      =   []
	UpperCAmelCase_			:	Optional[Any]      =   []  # We keep track of the idxs just in case it's a repeated statement
	for i, line in enumerate(_SCREAMING_SNAKE_CASE							):
		if i not in passed_idxs:
			if (line not in full_file_func) and (line.lstrip() != _dl_line):
				if "TESTING_MOCKED_DATALOADERS" not in line:
					new_full_example_parts.append(_SCREAMING_SNAKE_CASE							)
					passed_idxs.append(_SCREAMING_SNAKE_CASE							)

    # Finally, get the overall diff
	UpperCAmelCase_			:	Union[str, Any]      =   [line for line in new_feature_code if line not in new_full_example_parts]
	if secondary_filename is not None:
		UpperCAmelCase_			:	int      =   [line for line in full_file_contents if line not in secondary_file_func]
		UpperCAmelCase_			:	str      =   [line for line in diff_from_example if line not in diff_from_two]

	return diff_from_example




