




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__    :					str					=						logging.get_logger(__name__)


def 						lowerCamelCase__   (		_A							):







	'''simple docstring'''





	snake_case_		   =					ConvNextVaConfig()

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

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

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

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

	return config, expected_shape


def 						lowerCamelCase__   (		_A							):







	'''simple docstring'''





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

	return name


def 						lowerCamelCase__   (		):







	'''simple docstring'''





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


def 						lowerCamelCase__   (		_A							):







	'''simple docstring'''





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

	return ConvNextImageProcessor(
	    size=_A					,     crop_pct=_A					,     image_mean=[0.4_85, 0.4_56, 0.4_06]					,     image_std=[0.2_29, 0.2_24, 0.2_25]					,     resample=PILImageResampling.BICUBIC					,     )


@torch.no_grad()
def 						lowerCamelCase__   (		_A					,     _A					,     _A					,     _A							):







	'''simple docstring'''





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

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

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

	# Check outputs on an image, prepared by ConvNextImageProcessor
	snake_case_		   =					convert_preprocessor(_A							)
	snake_case_		   =					preprocessor(images=prepare_img()					,     return_tensors="pt"							)
	snake_case_		   =					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":
		snake_case_		   =					torch.tensor([-0.39_30, 0.17_47, -0.52_46, 0.41_77, 0.42_95]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
		snake_case_		   =					torch.tensor([-0.17_27, -0.53_41, -0.78_18, -0.47_45, -0.65_66]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
		snake_case_		   =					torch.tensor([-0.03_33, 0.15_63, -0.91_37, 0.10_54, 0.03_81]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
		snake_case_		   =					torch.tensor([-0.17_44, -0.15_55, -0.07_13, 0.09_50, -0.14_31]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
		snake_case_		   =					torch.tensor([0.99_96, 0.19_66, -0.43_86, -0.34_72, 0.66_61]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
		snake_case_		   =					torch.tensor([-0.25_53, -0.67_08, -0.13_59, 0.25_18, -0.24_88]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
		snake_case_		   =					torch.tensor([-0.06_73, -0.56_27, -0.37_53, -0.27_22, 0.01_78]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
		snake_case_		   =					torch.tensor([-0.63_77, -0.74_58, -0.21_50, 0.11_84, -0.05_97]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
		snake_case_		   =					torch.tensor([1.07_99, 0.23_22, -0.88_60, 1.02_19, 0.62_31]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
		snake_case_		   =					torch.tensor([0.37_66, 0.49_17, -1.14_26, 0.99_42, 0.60_24]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
		snake_case_		   =					torch.tensor([0.42_20, -0.69_19, -0.43_17, -0.28_81, -0.66_09]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
		snake_case_		   =					torch.tensor([0.10_82, -0.82_86, -0.50_95, 0.46_81, -0.80_85]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
		snake_case_		   =					torch.tensor([-0.24_19, -0.62_21, 0.21_76, -0.09_80, -0.75_27]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
		snake_case_		   =					torch.tensor([0.03_91, -0.43_71, 0.37_86, 0.12_51, -0.27_84]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
		snake_case_		   =					torch.tensor([-0.05_04, 0.56_36, -0.17_29, -0.65_07, -0.39_49]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
		snake_case_		   =					torch.tensor([0.35_60, 0.94_86, 0.31_49, -0.26_67, -0.51_38]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
		snake_case_		   =					torch.tensor([-0.24_69, -0.45_50, -0.58_53, -0.08_10, 0.03_09]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
		snake_case_		   =					torch.tensor([-0.30_90, 0.08_02, -0.06_82, -0.19_79, -0.28_26]							)
	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							)

	snake_case_		   =					"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__":
		lowercase__    :					Optional[int]					=						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__    :					List[str]					=						parser.parse_args()
		convert_convnextva_checkpoint(
		    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
		)






import unittest
from functools import partial

import torch

from accelerate import Accelerator, debug_launcher
from accelerate.state import AcceleratorState, GradientState
from accelerate.test_utils import require_cpu, require_huggingface_suite
from accelerate.utils import GradientAccumulationPlugin


def 						lowerCamelCase__   (		_A=2					,     _A=True					,     _A=False							):







	'''simple docstring'''





	snake_case_		   =					Accelerator(step_scheduler_with_optimizer=_A					,     split_batches=_A							)
	snake_case_		   =					torch.nn.Linear(2					,     4							)
	snake_case_		   =					torch.optim.AdamW(model.parameters()					,     lr=1.0							)
	snake_case_		   =					torch.optim.lr_scheduler.OneCycleLR(_A					,     max_lr=0.01					,     steps_per_epoch=2					,     epochs=1							)
	snake_case_   ,				snake_case_   ,				snake_case_		   =					accelerator.prepare(_A					,     _A					,     _A							)

	# Optimizer has stepped
	scheduler.step()
	if step_scheduler_with_optimizer or (num_processes == 1):
		assert (
		    scheduler.scheduler.last_epoch == num_processes
		), f"Last Epoch ({scheduler.scheduler.last_epoch}) != Num Processes ({num_processes})"
	else:
		assert (
		    scheduler.scheduler.last_epoch != num_processes
		), f"Last Epoch ({scheduler.scheduler.last_epoch}) == Num Processes ({num_processes})"


def 						lowerCamelCase__   (		_A=2					,     _A=True					,     _A=False							):







	'''simple docstring'''





	snake_case_		   =					Accelerator(step_scheduler_with_optimizer=_A					,     split_batches=_A							)
	snake_case_		   =					torch.nn.Linear(2					,     4							)
	snake_case_		   =					torch.optim.AdamW(model.parameters()					,     lr=1.0							)
	snake_case_		   =					torch.optim.lr_scheduler.LambdaLR(_A					,     lr_lambda=lambda _A							: 1 - n / 10							)
	snake_case_   ,				snake_case_   ,				snake_case_		   =					accelerator.prepare(_A					,     _A					,     _A							)

	# Optimizer has stepped
	snake_case_		   =					False
	scheduler.step()
	snake_case_		   =					1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 10
	assert (
	    scheduler.get_last_lr()[0] == expected_lr
	), f"Wrong lr found at first step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"

	# Optimizer has not stepped
	snake_case_		   =					True
	scheduler.step()
	if not step_scheduler_with_optimizer:
		snake_case_		   =					1 - 2 / 10
	assert (
	    scheduler.get_last_lr()[0] == expected_lr
	), f"Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"


def 						lowerCamelCase__   (		_A = 2							):







	'''simple docstring'''





	from transformers import get_linear_schedule_with_warmup

	snake_case_		   =					[1, 2, 4]
	for num_steps in steps:
		snake_case_		   =					GradientAccumulationPlugin(num_steps=_A					,     adjust_scheduler=num_steps > 1							)
		snake_case_		   =					Accelerator(gradient_accumulation_plugin=_A							)
		snake_case_		   =					torch.nn.Linear(2					,     4							)
		snake_case_		   =					torch.optim.AdamW(model.parameters()					,     lr=10.0							)
		snake_case_		   =					get_linear_schedule_with_warmup(optimizer=_A					,     num_warmup_steps=0					,     num_training_steps=20							)

		snake_case_   ,				snake_case_   ,				snake_case_		   =					accelerator.prepare(_A					,     _A					,     _A							)

		for i in range(10 * num_steps							):
			with accelerator.accumulate(_A							):
				optimizer.step()
				scheduler.step()

			if i == (10 * num_steps - 2):
				assert (
				    scheduler.get_last_lr()[0] != 0
				), f"Wrong lr found at second-to-last step, expected non-zero, got {scheduler.get_last_lr()[0]}. num_steps: {num_steps}"
		assert (
		    scheduler.get_last_lr()[0] == 0
		), f"Wrong lr found at last step, expected 0, got {scheduler.get_last_lr()[0]}"
		GradientState._reset_state()






@require_cpu
class 		UpperCAmelCase       (    unittest.TestCase      ):







						'''simple docstring'''





						def 	snake_case__						(						self							:					Tuple						):






							"""simple docstring"""






							debug_launcher(partial(__lowercase	,     num_processes=1						)	,     num_processes=1						)
							debug_launcher(partial(__lowercase	,     num_processes=1	,     split_batches=__lowercase						)	,     num_processes=1						)






						def 	snake_case__						(						self							:					int						):






							"""simple docstring"""






							debug_launcher(partial(__lowercase	,     num_processes=1						)	,     num_processes=1						)
							debug_launcher(partial(__lowercase	,     num_processes=1	,     split_batches=__lowercase						)	,     num_processes=1						)






						def 	snake_case__						(						self							:					Dict						):






							"""simple docstring"""






							debug_launcher(partial(__lowercase	,     num_processes=1	,     step_scheduler_with_optimizer=__lowercase						)	,     num_processes=1						)






						def 	snake_case__						(						self							:					Union[str, Any]						):






							"""simple docstring"""






							debug_launcher(partial(__lowercase	,     num_processes=1	,     step_scheduler_with_optimizer=__lowercase						)	,     num_processes=1						)






						def 	snake_case__						(						self							:					List[Any]						):






							"""simple docstring"""






							AcceleratorState._reset_state(__lowercase						)
							debug_launcher(__lowercase						)
							debug_launcher(partial(__lowercase	,     num_processes=1	,     split_batches=__lowercase						)	,     num_processes=1						)






						def 	snake_case__						(						self							:					Dict						):






							"""simple docstring"""






							AcceleratorState._reset_state(__lowercase						)
							debug_launcher(__lowercase						)
							debug_launcher(partial(__lowercase	,     num_processes=1	,     split_batches=__lowercase						)	,     num_processes=1						)






						def 	snake_case__						(						self							:					List[str]						):






							"""simple docstring"""






							AcceleratorState._reset_state(__lowercase						)
							debug_launcher(partial(__lowercase	,     step_scheduler_with_optimizer=__lowercase						)						)






						def 	snake_case__						(						self							:					Optional[Any]						):






							"""simple docstring"""






							AcceleratorState._reset_state(__lowercase						)
							debug_launcher(partial(__lowercase	,     step_scheduler_with_optimizer=__lowercase						)						)






						@require_huggingface_suite
						def 	snake_case__						(						self							:					Optional[Any]						):






							"""simple docstring"""






							AcceleratorState._reset_state(__lowercase						)
							debug_launcher(partial(__lowercase	,     num_processes=1						)						)
							debug_launcher(__lowercase						)
