





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







def 	__lowercase  (     snake_case_						:						Dict		)     ->Optional[int]:



				'''simple docstring'''







				__A					:		List[Any]						=  ConvNextConfig()

				if "tiny" in checkpoint_url:
								__A					:		List[str]						=  [3, 3, 9, 3]
								__A					:		Optional[Any]						=  [96, 192, 384, 768]
				if "small" in checkpoint_url:
								__A					:		str						=  [3, 3, 27, 3]
								__A					:		List[str]						=  [96, 192, 384, 768]
				if "base" in checkpoint_url:
								__A					:		Any						=  [3, 3, 27, 3]
								__A					:		int						=  [128, 256, 512, 1024]
				if "large" in checkpoint_url:
								__A					:		Dict						=  [3, 3, 27, 3]
								__A					:		Dict						=  [192, 384, 768, 1536]
				if "xlarge" in checkpoint_url:
								__A					:		Optional[Any]						=  [3, 3, 27, 3]
								__A					:		List[str]						=  [256, 512, 1024, 2048]

				if "1k" in checkpoint_url:
								__A					:		Union[str, Any]						=  1000
								__A					:		str						=  '''imagenet-1k-id2label.json'''
								__A					:		int						=  (1, 1000)
				else:
								__A					:		Any						=  21841
								__A					:		Any						=  '''imagenet-22k-id2label.json'''
								__A					:		Union[str, Any]						=  (1, 21841)

				__A					:		Optional[Any]						=  '''huggingface/label-files'''
				__A					:		List[str]						=  num_labels
				__A					:		List[str]						=  json.load(open(hf_hub_download(_UpperCAmelCase       ,_UpperCAmelCase       ,repo_type='''dataset'''		)       ,'''r'''		)		)
				__A					:		Any						=  {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[15027]
				__A					:		Optional[Any]						=  idalabel
				__A					:		Any						=  {v: k for k, v in idalabel.items()}
				__A					:		List[Any]						=  hidden_sizes
				__A					:		List[Any]						=  depths

				return config, expected_shape







def 	__lowercase  (     snake_case_						:						List[Any]		)     ->Union[str, Any]:



				'''simple docstring'''







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

				return name







def 	__lowercase  (     )     ->int:



				'''simple docstring'''







				__A					:		Union[str, Any]						=  '''http://images.cocodataset.org/val2017/000000039769.jpg'''
				__A					:		str						=  Image.open(requests.get(_UpperCAmelCase       ,stream=_UpperCAmelCase		).raw		)
				return im







@torch.no_grad()
def 	__lowercase  (     snake_case_						:						Dict       ,snake_case_						:						Tuple		)     ->Optional[int]:



				'''simple docstring'''







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

				# load HuggingFace model
				__A					:		Union[str, Any]						=  ConvNextForImageClassification(_UpperCAmelCase		)
				model.load_state_dict(_UpperCAmelCase		)
				model.eval()

				# Check outputs on an image, prepared by ConvNextImageProcessor
				__A					:		Tuple						=  224 if '''224''' in checkpoint_url else 384
				__A					:		List[Any]						=  ConvNextImageProcessor(size=_UpperCAmelCase		)
				__A					:		List[Any]						=  image_processor(images=prepare_img()       ,return_tensors='''pt'''		).pixel_values

				__A					:		Dict						=  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":
								__A					:		List[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":
								__A					:		Tuple						=  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":
								__A					:		Optional[Any]						=  torch.tensor([0.4_525, 0.7_539, 0.0_308]		)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
								__A					:		Dict						=  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":
								__A					:		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":
								__A					:		Optional[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":
								__A					:		Dict						=  torch.tensor([1.2_980, 0.3_631, -0.1_198]		)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
								__A					:		Tuple						=  torch.tensor([1.2_963, 0.1_227, 0.1_723]		)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
								__A					:		List[str]						=  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":
								__A					:		Any						=  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":
								__A					:		Optional[int]						=  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":
								__A					:		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":
								__A					:		Any						=  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":
								__A					:		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":
								__A					:		Dict						=  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...'''		)
				__A					:		Any						=  '''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__":
							a_				=    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.""",
							)

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











"""simple docstring"""



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 	__lowercase  (     snake_case_						:						Dict=2       ,snake_case_						:						str=True       ,snake_case_						:						Tuple=False		)     ->Any:



    '''simple docstring'''







    __A					:		List[str]						=  Accelerator(step_scheduler_with_optimizer=snake_case_       ,split_batches=snake_case_		)
    __A					:		Dict						=  torch.nn.Linear(2       ,4		)
    __A					:		Tuple						=  torch.optim.AdamW(model.parameters()       ,lr=1.0		)
    __A					:		int						=  torch.optim.lr_scheduler.OneCycleLR(snake_case_       ,max_lr=0.01       ,steps_per_epoch=2       ,epochs=1		)
    __A				, __A				, __A					:		Optional[int]						=  accelerator.prepare(snake_case_       ,snake_case_       ,snake_case_		)

    # 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 	__lowercase  (     snake_case_						:						str=2       ,snake_case_						:						Dict=True       ,snake_case_						:						Tuple=False		)     ->List[Any]:



    '''simple docstring'''







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

    # Optimizer has stepped
    __A					:		List[Any]						=  False
    scheduler.step()
    __A					:		List[Any]						=  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
    __A					:		List[str]						=  True
    scheduler.step()
    if not step_scheduler_with_optimizer:
        __A					:		List[Any]						=  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 	__lowercase  (     snake_case_						:						int = 2		)     ->str:



    '''simple docstring'''







    from transformers import get_linear_schedule_with_warmup

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

        __A				, __A				, __A					:		Union[str, Any]						=  accelerator.prepare(snake_case_       ,snake_case_       ,snake_case_		)

        for i in range(10 * num_steps		):
            with accelerator.accumulate(snake_case_		):
                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 	__snake_case      (	unittest.TestCase  ):





 """simple docstring"""


 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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







 @require_huggingface_suite
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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




