




import argparse
import json
import os

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
lowercase_     :    int      =			logging.get_logger(__name__)


def 							__SCREAMING_SNAKE_CASE  (      snake_case_						):

						'''simple docstring'''


						_UpperCAmelCase													=					ConvNextVaConfig()

						if "atto" in checkpoint_url:
												_UpperCAmelCase													=					[2, 2, 6, 2]
												_UpperCAmelCase													=					[40, 80, 160, 320]
						if "femto" in checkpoint_url:
												_UpperCAmelCase													=					[2, 2, 6, 2]
												_UpperCAmelCase													=					[48, 96, 192, 384]
						if "pico" in checkpoint_url:
												_UpperCAmelCase													=					[2, 2, 6, 2]
												_UpperCAmelCase													=					[64, 128, 256, 512]
						if "nano" in checkpoint_url:
												_UpperCAmelCase													=					[2, 2, 8, 2]
												_UpperCAmelCase													=					[80, 160, 320, 640]
						if "tiny" in checkpoint_url:
												_UpperCAmelCase													=					[3, 3, 9, 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 "huge" in checkpoint_url:
												_UpperCAmelCase													=					[3, 3, 27, 3]
												_UpperCAmelCase													=					[352, 704, 1408, 2816]

						_UpperCAmelCase													=					1000
						_UpperCAmelCase													=					"imagenet-1k-id2label.json"
						_UpperCAmelCase													=					(1, 1000)

						_UpperCAmelCase													=					"huggingface/label-files"
						_UpperCAmelCase													=					num_labels
						_UpperCAmelCase													=					json.load(open(hf_hub_download(snake_case_     ,   snake_case_     ,   repo_type="dataset"						)     ,   "r"						)						)
						_UpperCAmelCase													=					{int(snake_case_						): v for k, v in idalabel.items()}

						_UpperCAmelCase													=					idalabel
						_UpperCAmelCase													=					{v: k for k, v in idalabel.items()}
						_UpperCAmelCase													=					hidden_sizes
						_UpperCAmelCase													=					depths

						return config, expected_shape


def 							__SCREAMING_SNAKE_CASE  (      snake_case_						):

						'''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 "gamma" in name:
												_UpperCAmelCase													=					name.replace("gamma"     ,   "weight"						)
						if "beta" in name:
												_UpperCAmelCase													=					name.replace("beta"     ,   "bias"						)
						if "stages" in name:
												_UpperCAmelCase													=					name.replace("stages"     ,   "encoder.stages"						)
						if "norm" in name:
												_UpperCAmelCase													=					name.replace("norm"     ,   "layernorm"						)
						if "head" in name:
												_UpperCAmelCase													=					name.replace("head"     ,   "classifier"						)

						return name


def 							__SCREAMING_SNAKE_CASE  (      ):

						'''simple docstring'''


						_UpperCAmelCase													=					"http://images.cocodataset.org/val2017/000000039769.jpg"
						_UpperCAmelCase													=					Image.open(requests.get(snake_case_     ,   stream=snake_case_						).raw						)
						return im


def 							__SCREAMING_SNAKE_CASE  (      snake_case_						):

						'''simple docstring'''


						if "224" in checkpoint_url:
												_UpperCAmelCase													=					224
												_UpperCAmelCase													=					224 / 256
						elif "384" in checkpoint_url:
												_UpperCAmelCase													=					384
												_UpperCAmelCase													=					None
						else:
												_UpperCAmelCase													=					512
												_UpperCAmelCase													=					None

						return ConvNextImageProcessor(
						    size=snake_case_     ,   crop_pct=snake_case_     ,   image_mean=[0.485, 0.456, 0.406]     ,   image_std=[0.229, 0.224, 0.225]     ,   resample=PILImageResampling.BICUBIC     ,   )


@torch.no_grad()
def 							__SCREAMING_SNAKE_CASE  (      snake_case_     ,   snake_case_     ,   snake_case_     ,   snake_case_						):

						'''simple docstring'''


						print("Downloading original model from checkpoint..."						)
						# define ConvNeXTV2 configuration based on URL
						_UpperCAmelCase							,			_UpperCAmelCase													=					get_convnextva_config(snake_case_						)
						# load original state_dict from URL
						_UpperCAmelCase													=					torch.hub.load_state_dict_from_url(snake_case_						)["model"]

						print("Converting model parameters..."						)
						# rename keys
						for key in state_dict.copy().keys():
												_UpperCAmelCase													=					state_dict.pop(snake_case_						)
												_UpperCAmelCase													=					val
						# add prefix to all keys expect classifier head
						for key in state_dict.copy().keys():
												_UpperCAmelCase													=					state_dict.pop(snake_case_						)
												if not key.startswith("classifier"						):
																		_UpperCAmelCase													=					"convnextv2." + key
												_UpperCAmelCase													=					val

						# load HuggingFace model
						_UpperCAmelCase													=					ConvNextVaForImageClassification(snake_case_						)
						model.load_state_dict(snake_case_						)
						model.eval()

						# Check outputs on an image, prepared by ConvNextImageProcessor
						_UpperCAmelCase													=					convert_preprocessor(snake_case_						)
						_UpperCAmelCase													=					preprocessor(images=prepare_img()     ,   return_tensors="pt"						)
						_UpperCAmelCase													=					model(**snake_case_						).logits

						# note: the logits below were obtained without center cropping
						if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
												_UpperCAmelCase													=					torch.tensor([-0.3930, 0.1747, -0.5246, 0.4177, 0.4295]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
												_UpperCAmelCase													=					torch.tensor([-0.1727, -0.5341, -0.7818, -0.4745, -0.6566]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
												_UpperCAmelCase													=					torch.tensor([-0.0333, 0.1563, -0.9137, 0.1054, 0.0381]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
												_UpperCAmelCase													=					torch.tensor([-0.1744, -0.1555, -0.0713, 0.0950, -0.1431]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
												_UpperCAmelCase													=					torch.tensor([0.9996, 0.1966, -0.4386, -0.3472, 0.6661]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
												_UpperCAmelCase													=					torch.tensor([-0.2553, -0.6708, -0.1359, 0.2518, -0.2488]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
												_UpperCAmelCase													=					torch.tensor([-0.0673, -0.5627, -0.3753, -0.2722, 0.0178]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
												_UpperCAmelCase													=					torch.tensor([-0.6377, -0.7458, -0.2150, 0.1184, -0.0597]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
												_UpperCAmelCase													=					torch.tensor([1.0799, 0.2322, -0.8860, 1.0219, 0.6231]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
												_UpperCAmelCase													=					torch.tensor([0.3766, 0.4917, -1.1426, 0.9942, 0.6024]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
												_UpperCAmelCase													=					torch.tensor([0.4220, -0.6919, -0.4317, -0.2881, -0.6609]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
												_UpperCAmelCase													=					torch.tensor([0.1082, -0.8286, -0.5095, 0.4681, -0.8085]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
												_UpperCAmelCase													=					torch.tensor([-0.2419, -0.6221, 0.2176, -0.0980, -0.7527]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
												_UpperCAmelCase													=					torch.tensor([0.0391, -0.4371, 0.3786, 0.1251, -0.2784]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
												_UpperCAmelCase													=					torch.tensor([-0.0504, 0.5636, -0.1729, -0.6507, -0.3949]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
												_UpperCAmelCase													=					torch.tensor([0.3560, 0.9486, 0.3149, -0.2667, -0.5138]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
												_UpperCAmelCase													=					torch.tensor([-0.2469, -0.4550, -0.5853, -0.0810, 0.0309]						)
						elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
												_UpperCAmelCase													=					torch.tensor([-0.3090, 0.0802, -0.0682, -0.1979, -0.2826]						)
						else:
												raise ValueError(f"""Unknown URL: {checkpoint_url}"""						)

						assert torch.allclose(logits[0, :5]     ,   snake_case_     ,   atol=1e-3						)
						assert logits.shape == expected_shape
						print("Model outputs match the original results!"						)

						if save_model:
												print("Saving model to local..."						)
												# Create folder to save model
												if not os.path.isdir(snake_case_						):
																		os.mkdir(snake_case_						)

												model.save_pretrained(snake_case_						)
												preprocessor.save_pretrained(snake_case_						)

						_UpperCAmelCase													=					"convnextv2"
						if "atto" in checkpoint_url:
												model_name += "-atto"
						if "femto" in checkpoint_url:
												model_name += "-femto"
						if "pico" in checkpoint_url:
												model_name += "-pico"
						if "nano" in checkpoint_url:
												model_name += "-nano"
						elif "tiny" in checkpoint_url:
												model_name += "-tiny"
						elif "base" in checkpoint_url:
												model_name += "-base"
						elif "large" in checkpoint_url:
												model_name += "-large"
						elif "huge" in checkpoint_url:
												model_name += "-huge"
						if "22k" in checkpoint_url and "1k" not in checkpoint_url:
												model_name += "-22k"
						elif "22k" in checkpoint_url and "1k" in checkpoint_url:
												model_name += "-22k-1k"
						elif "1k" in checkpoint_url:
												model_name += "-1k"
						if "224" in checkpoint_url:
												model_name += "-224"
						elif "384" in checkpoint_url:
												model_name += "-384"
						elif "512" in checkpoint_url:
												model_name += "-512"

						if push_to_hub:
												print(f"""Pushing {model_name} to the hub..."""						)
												model.push_to_hub(snake_case_						)
												preprocessor.push_to_hub(snake_case_						)


if __name__ == "__main__":
				lowercase_     :    Any      =			argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    '--checkpoint_url',
				    default='https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt',
				    type=str,
				    help='URL of the original ConvNeXTV2 checkpoint you\'d like to convert.',
				)
				parser.add_argument(
				    '--pytorch_dump_folder_path',
				    default='model',
				    type=str,
				    help='Path to the output PyTorch model directory.',
				)
				parser.add_argument('--save_model', action='store_true', help='Save model to local')
				parser.add_argument('--push_to_hub', action='store_true', help='Push model and image preprocessor to the hub')

				lowercase_     :    Any      =			parser.parse_args()
				convert_convnextva_checkpoint(
				    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
				)










from __future__ import annotations


def 							__SCREAMING_SNAKE_CASE  (      snake_case_						):

						'''simple docstring'''


						if len(snake_case_						) == 0:
												raise ValueError("find_max() arg is an empty sequence"						)
						_UpperCAmelCase													=					nums[0]
						for x in nums:
												if x > max_num:
																		_UpperCAmelCase													=					x
						return max_num


if __name__ == "__main__":
				import doctest

				doctest.testmod(verbose=True)




