



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 				A			(		__UpperCAmelCase	)       ->       Union[str, Any]:



						'''simple docstring'''






						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(__UpperCAmelCase							,					__UpperCAmelCase							,					repo_type='''dataset'''	)							,					'''r'''	)	)
						UpperCAmelCase_     							=     {int(__UpperCAmelCase	): 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 				A			(		__UpperCAmelCase	)       ->       Optional[Any]:



						'''simple docstring'''






						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 				A			(		)       ->       str:



						'''simple docstring'''






						UpperCAmelCase_     							=     '''http://images.cocodataset.org/val2017/000000039769.jpg'''
						UpperCAmelCase_     							=     Image.open(requests.get(__UpperCAmelCase							,					stream=__UpperCAmelCase	).raw	)
						return im







@torch.no_grad()
def 				A			(		__UpperCAmelCase							,					__UpperCAmelCase	)       ->       Union[str, Any]:



						'''simple docstring'''






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

						# load HuggingFace model
						UpperCAmelCase_     							=     ConvNextForImageClassification(__UpperCAmelCase	)
						model.load_state_dict(__UpperCAmelCase	)
						model.eval()

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

						UpperCAmelCase_     							=     model(__UpperCAmelCase	).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.1_210, -0.6_605, 0.1_918]	)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
												UpperCAmelCase_     							=     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_     							=     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_     							=     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_     							=     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_     							=     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_     							=     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_     							=     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_     							=     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_     							=     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_     							=     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_     							=     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_     							=     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_     							=     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_     							=     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]							,					__UpperCAmelCase							,					atol=1e-3	)
						assert logits.shape == expected_shape

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

						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(__UpperCAmelCase							,					__UpperCAmelCase	)							,					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'''






def 					_snake_case   (     A       )						->    str:
							return " ".join(sorted(set(sentence.split()       )       )       )


if __name__ == "__main__":
		import doctest

		doctest.testmod()