

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

def 	lowerCamelCase__	(   _a):
	SCREAMING_SNAKE_CASE				: str			       =    ConvNextVaConfig()

	if "atto" in checkpoint_url:
		SCREAMING_SNAKE_CASE				: Tuple			       =    [2, 2, 6, 2]
		SCREAMING_SNAKE_CASE				: List[str]			       =    [40, 80, 160, 320]
	if "femto" in checkpoint_url:
		SCREAMING_SNAKE_CASE				: List[Any]			       =    [2, 2, 6, 2]
		SCREAMING_SNAKE_CASE				: List[str]			       =    [48, 96, 192, 384]
	if "pico" in checkpoint_url:
		SCREAMING_SNAKE_CASE				: List[Any]			       =    [2, 2, 6, 2]
		SCREAMING_SNAKE_CASE				: Dict			       =    [64, 128, 256, 512]
	if "nano" in checkpoint_url:
		SCREAMING_SNAKE_CASE				: Union[str, Any]			       =    [2, 2, 8, 2]
		SCREAMING_SNAKE_CASE				: Union[str, Any]			       =    [80, 160, 320, 640]
	if "tiny" in checkpoint_url:
		SCREAMING_SNAKE_CASE				: str			       =    [3, 3, 9, 3]
		SCREAMING_SNAKE_CASE				: List[str]			       =    [96, 192, 384, 768]
	if "base" in checkpoint_url:
		SCREAMING_SNAKE_CASE				: Optional[Any]			       =    [3, 3, 27, 3]
		SCREAMING_SNAKE_CASE				: Any			       =    [128, 256, 512, 1024]
	if "large" in checkpoint_url:
		SCREAMING_SNAKE_CASE				: Union[str, Any]			       =    [3, 3, 27, 3]
		SCREAMING_SNAKE_CASE				: Union[str, Any]			       =    [192, 384, 768, 1536]
	if "huge" in checkpoint_url:
		SCREAMING_SNAKE_CASE				: Union[str, Any]			       =    [3, 3, 27, 3]
		SCREAMING_SNAKE_CASE				: Union[str, Any]			       =    [352, 704, 1408, 2816]

	SCREAMING_SNAKE_CASE				: Optional[int]			       =    1000
	SCREAMING_SNAKE_CASE				: Any			       =    "imagenet-1k-id2label.json"
	SCREAMING_SNAKE_CASE				: int			       =    (1, 1000)

	SCREAMING_SNAKE_CASE				: Optional[Any]			       =    "huggingface/label-files"
	SCREAMING_SNAKE_CASE				: List[str]			       =    num_labels
	SCREAMING_SNAKE_CASE				: Dict			       =    json.load(open(hf_hub_download(_a							,     _a							,     repo_type="dataset")							,     "r"))
	SCREAMING_SNAKE_CASE				: Optional[Any]			       =    {int(_a): v for k, v in idalabel.items()}

	SCREAMING_SNAKE_CASE				: Dict			       =    idalabel
	SCREAMING_SNAKE_CASE				: Any			       =    {v: k for k, v in idalabel.items()}
	SCREAMING_SNAKE_CASE				: Dict			       =    hidden_sizes
	SCREAMING_SNAKE_CASE				: List[str]			       =    depths

	return config, expected_shape

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

	return name

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

def 	lowerCamelCase__	(   _a):
	if "224" in checkpoint_url:
		SCREAMING_SNAKE_CASE				: Tuple			       =    224
		SCREAMING_SNAKE_CASE				: Tuple			       =    224 / 256
	elif "384" in checkpoint_url:
		SCREAMING_SNAKE_CASE				: int			       =    384
		SCREAMING_SNAKE_CASE				: List[Any]			       =    None
	else:
		SCREAMING_SNAKE_CASE				: List[Any]			       =    512
		SCREAMING_SNAKE_CASE				: int			       =    None

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

@torch.no_grad()
def 	lowerCamelCase__	(   _a							,     _a							,     _a							,     _a):
	print("Downloading original model from checkpoint...")
	# define ConvNeXTV2 configuration based on URL
	SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE				: Union[str, Any]			       =    get_convnextva_config(_a)
	# load original state_dict from URL
	SCREAMING_SNAKE_CASE				: Tuple			       =    torch.hub.load_state_dict_from_url(_a)["model"]

	print("Converting model parameters...")
	# rename keys
	for key in state_dict.copy().keys():
		SCREAMING_SNAKE_CASE				: Optional[Any]			       =    state_dict.pop(_a)
		SCREAMING_SNAKE_CASE				: Optional[Any]			       =    val
	# add prefix to all keys expect classifier head
	for key in state_dict.copy().keys():
		SCREAMING_SNAKE_CASE				: Tuple			       =    state_dict.pop(_a)
		if not key.startswith("classifier"):
			SCREAMING_SNAKE_CASE				: Union[str, Any]			       =    "convnextv2." + key
		SCREAMING_SNAKE_CASE				: Tuple			       =    val

	# load HuggingFace model
	SCREAMING_SNAKE_CASE				: Optional[int]			       =    ConvNextVaForImageClassification(_a)
	model.load_state_dict(_a)
	model.eval()

	# Check outputs on an image, prepared by ConvNextImageProcessor
	SCREAMING_SNAKE_CASE				: Optional[int]			       =    convert_preprocessor(_a)
	SCREAMING_SNAKE_CASE				: str			       =    preprocessor(images=prepare_img()							,     return_tensors="pt")
	SCREAMING_SNAKE_CASE				: Tuple			       =    model(**_a).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":
		SCREAMING_SNAKE_CASE				: List[str]			       =    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":
		SCREAMING_SNAKE_CASE				: Any			       =    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":
		SCREAMING_SNAKE_CASE				: Union[str, Any]			       =    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":
		SCREAMING_SNAKE_CASE				: List[str]			       =    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":
		SCREAMING_SNAKE_CASE				: Any			       =    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":
		SCREAMING_SNAKE_CASE				: List[str]			       =    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":
		SCREAMING_SNAKE_CASE				: Union[str, Any]			       =    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":
		SCREAMING_SNAKE_CASE				: Tuple			       =    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":
		SCREAMING_SNAKE_CASE				: Dict			       =    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":
		SCREAMING_SNAKE_CASE				: Tuple			       =    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":
		SCREAMING_SNAKE_CASE				: Tuple			       =    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":
		SCREAMING_SNAKE_CASE				: Any			       =    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":
		SCREAMING_SNAKE_CASE				: Optional[Any]			       =    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":
		SCREAMING_SNAKE_CASE				: int			       =    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":
		SCREAMING_SNAKE_CASE				: Tuple			       =    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":
		SCREAMING_SNAKE_CASE				: Optional[Any]			       =    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":
		SCREAMING_SNAKE_CASE				: Tuple			       =    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":
		SCREAMING_SNAKE_CASE				: Optional[Any]			       =    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]							,     _a							,     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(_a):
			os.mkdir(_a)

		model.save_pretrained(_a)
		preprocessor.save_pretrained(_a)

	SCREAMING_SNAKE_CASE				: Any			       =    "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(_a)
		preprocessor.push_to_hub(_a)


if __name__ == "__main__":
			a_         =      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')

			a_         =      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

from typing import Any






class _UpperCamelCase							:
			'''simple docstring'''







			def __init__(     self							:	Dict , a							:	Any , a							:	Any		)							->    None:  # noqa: A002

				"""simple docstring"""

				SCREAMING_SNAKE_CASE				: str			       =    item
				SCREAMING_SNAKE_CASE				: List[Any]			       =    next






class _UpperCamelCase							:
			'''simple docstring'''







			def __init__(     self							:	Optional[int]		)							->    None:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE				: Node | None			       =    None
				SCREAMING_SNAKE_CASE				: List[Any]			       =    0
			def 							__UpperCamelCase	(     self							:	Union[str, Any] , a							:	Any		)							->    None:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE				: Any			       =    Node(a , self.head		)
				self.size += 1
			def 							__UpperCamelCase	(     self							:	Dict		)							->    Any:

				"""simple docstring"""

				if self.head is None:
					return None
				else:
					SCREAMING_SNAKE_CASE				: Optional[int]			       =    self.head.item
					SCREAMING_SNAKE_CASE				: Dict			       =    self.head.next
					self.size -= 1
					return item
			def 							__UpperCamelCase	(     self							:	Any		)							->    bool:

				"""simple docstring"""

				return self.head is None
			def __str__(     self							:	Union[str, Any]		)							->    str:

				"""simple docstring"""

				if self.is_empty():
					return ""
				else:
					SCREAMING_SNAKE_CASE				: Union[str, Any]			       =    self.head
					SCREAMING_SNAKE_CASE				: str			       =    ""
					SCREAMING_SNAKE_CASE				: list[str]			       =    []
					while iterate:
						item_list.append(str(iterate.item		)		)
						SCREAMING_SNAKE_CASE				: str			       =    iterate.next

					SCREAMING_SNAKE_CASE				: List[str]			       =    " --> ".join(a		)

					return item_str






			def __len__(     self							:	List[Any]		)							->    int:

				"""simple docstring"""

				return self.size