


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      :	List[Any]  =logging.get_logger(__name__)



def       lowerCAmelCase			(		lowerCAmelCase_			)->			Tuple:
				lowerCAmelCase_			:   Dict     					=					ConvNextConfig()

				if "tiny" in checkpoint_url:
								lowerCAmelCase_			:   List[str]     					=					[3, 3, 9, 3]
								lowerCAmelCase_			:   int     					=					[96, 192, 384, 768]
				if "small" in checkpoint_url:
								lowerCAmelCase_			:   Optional[int]     					=					[3, 3, 27, 3]
								lowerCAmelCase_			:   Any     					=					[96, 192, 384, 768]
				if "base" in checkpoint_url:
								lowerCAmelCase_			:   int     					=					[3, 3, 27, 3]
								lowerCAmelCase_			:   int     					=					[128, 256, 512, 1_024]
				if "large" in checkpoint_url:
								lowerCAmelCase_			:   str     					=					[3, 3, 27, 3]
								lowerCAmelCase_			:   Any     					=					[192, 384, 768, 1_536]
				if "xlarge" in checkpoint_url:
								lowerCAmelCase_			:   Optional[Any]     					=					[3, 3, 27, 3]
								lowerCAmelCase_			:   str     					=					[256, 512, 1_024, 2_048]

				if "1k" in checkpoint_url:
								lowerCAmelCase_			:   Tuple     					=					1_000
								lowerCAmelCase_			:   str     					=					'''imagenet-1k-id2label.json'''
								lowerCAmelCase_			:   List[Any]     					=					(1, 1_000)
				else:
								lowerCAmelCase_			:   Union[str, Any]     					=					21_841
								lowerCAmelCase_			:   List[str]     					=					'''imagenet-22k-id2label.json'''
								lowerCAmelCase_			:   Optional[int]     					=					(1, 21_841)

				lowerCAmelCase_			:   List[Any]     					=					'''huggingface/label-files'''
				lowerCAmelCase_			:   str     					=					num_labels
				lowerCAmelCase_			:   List[Any]     					=					json.load(open(hf_hub_download(lowerCAmelCase_			, lowerCAmelCase_			, repo_type='''dataset'''			)			, '''r'''			)			)
				lowerCAmelCase_			:   Tuple     					=					{int(lowerCAmelCase_			): 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[9_205]
								del idalabel[15_027]
				lowerCAmelCase_			:   Tuple     					=					idalabel
				lowerCAmelCase_			:   Any     					=					{v: k for k, v in idalabel.items()}
				lowerCAmelCase_			:   Union[str, Any]     					=					hidden_sizes
				lowerCAmelCase_			:   Dict     					=					depths

				return config, expected_shape



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

				return name



def       lowerCAmelCase			(		)->			Optional[Any]:
				lowerCAmelCase_			:   List[str]     					=					'''http://images.cocodataset.org/val2017/000000039769.jpg'''
				lowerCAmelCase_			:   Tuple     					=					Image.open(requests.get(lowerCAmelCase_			, stream=lowerCAmelCase_			).raw			)
				return im



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

				# load HuggingFace model
				lowerCAmelCase_			:   Optional[Any]     					=					ConvNextForImageClassification(lowerCAmelCase_			)
				model.load_state_dict(lowerCAmelCase_			)
				model.eval()

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

				lowerCAmelCase_			:   Optional[int]     					=					model(lowerCAmelCase_			).logits

				# note: the logits below were obtained without center cropping
				if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
								lowerCAmelCase_			:   Dict     					=					torch.tensor([-0.1210, -0.6605, 0.1918]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
								lowerCAmelCase_			:   Any     					=					torch.tensor([-0.4473, -0.1847, -0.6365]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
								lowerCAmelCase_			:   Any     					=					torch.tensor([0.4525, 0.7539, 0.0308]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
								lowerCAmelCase_			:   Optional[Any]     					=					torch.tensor([0.3561, 0.6350, -0.0384]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
								lowerCAmelCase_			:   Tuple     					=					torch.tensor([0.4174, -0.0989, 0.1489]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
								lowerCAmelCase_			:   Optional[Any]     					=					torch.tensor([0.2513, -0.1349, -0.1613]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
								lowerCAmelCase_			:   Optional[int]     					=					torch.tensor([1.2980, 0.3631, -0.1198]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
								lowerCAmelCase_			:   str     					=					torch.tensor([1.2963, 0.1227, 0.1723]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
								lowerCAmelCase_			:   Tuple     					=					torch.tensor([1.7956, 0.8390, 0.2820]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
								lowerCAmelCase_			:   Optional[int]     					=					torch.tensor([-0.2822, -0.0502, -0.0878]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
								lowerCAmelCase_			:   List[str]     					=					torch.tensor([-0.5672, -0.0730, -0.4348]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
								lowerCAmelCase_			:   int     					=					torch.tensor([0.2681, 0.2365, 0.6246]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
								lowerCAmelCase_			:   Any     					=					torch.tensor([-0.2642, 0.3931, 0.5116]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
								lowerCAmelCase_			:   str     					=					torch.tensor([-0.6677, -0.1873, -0.8379]			)
				elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
								lowerCAmelCase_			:   Optional[int]     					=					torch.tensor([-0.7749, -0.2967, -0.6444]			)
				else:
								raise ValueError(f"""Unknown URL: {checkpoint_url}"""			)

				assert torch.allclose(logits[0, :3]			, lowerCAmelCase_			, atol=1e-3			)
				assert logits.shape == expected_shape

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

				print('''Pushing model to the hub...'''			)
				lowerCAmelCase_			:   int     					=					'''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(lowerCAmelCase_			, lowerCAmelCase_			)			, organization='''nielsr'''			, commit_message='''Add model'''			, )


if __name__ == "__main__":
						_UpperCAmelCase      :	List[str]  =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      :	List[Any]  =parser.parse_args()
						convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)



import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
						import transformer_engine.pytorch as te



def       lowerCAmelCase			(		lowerCAmelCase_			, lowerCAmelCase_=True			, lowerCAmelCase_=True			, lowerCAmelCase_=True			)->			Union[str, Any]:
				if not is_fpa_available():
								raise ImportError('''Using `convert_model` requires transformer_engine to be installed.'''			)
				for name, module in model.named_children():
								if isinstance(lowerCAmelCase_			, nn.Linear			) and to_transformer_engine and _convert_linear:
												# Return early if the linear layer weights are not multiples of 16
												if any(p % 16 != 0 for p in module.weight.shape			):
																return
												lowerCAmelCase_			:   Optional[int]     					=					module.bias is not None
												lowerCAmelCase_			:   Any     					=					te.Linear(
												    module.in_features			, module.out_features			, bias=lowerCAmelCase_			, params_dtype=module.weight.dtype			)
												lowerCAmelCase_			:   Dict     					=					module.weight.data.clone()
												if has_bias:
																lowerCAmelCase_			:   Any     					=					module.bias.data.clone()

												setattr(lowerCAmelCase_			, lowerCAmelCase_			, lowerCAmelCase_			)
								elif isinstance(lowerCAmelCase_			, nn.LayerNorm			) and to_transformer_engine and _convert_ln:
												lowerCAmelCase_			:   List[str]     					=					te.LayerNorm(module.normalized_shape[0]			, eps=module.eps			, params_dtype=module.weight.dtype			)
												lowerCAmelCase_			:   Union[str, Any]     					=					module.weight.data.clone()
												lowerCAmelCase_			:   List[Any]     					=					module.bias.data.clone()

												setattr(lowerCAmelCase_			, lowerCAmelCase_			, lowerCAmelCase_			)
								elif isinstance(lowerCAmelCase_			, te.Linear			) and not to_transformer_engine and _convert_linear:
												lowerCAmelCase_			:   Dict     					=					module.bias is not None
												lowerCAmelCase_			:   Union[str, Any]     					=					nn.Linear(
												    module.in_features			, module.out_features			, bias=lowerCAmelCase_			, params_dtype=module.weight.dtype			)
												lowerCAmelCase_			:   Optional[Any]     					=					module.weight.data.clone()
												if has_bias:
																lowerCAmelCase_			:   Union[str, Any]     					=					module.bias.data.clone()

												setattr(lowerCAmelCase_			, lowerCAmelCase_			, lowerCAmelCase_			)
								elif isinstance(lowerCAmelCase_			, te.LayerNorm			) and not to_transformer_engine and _convert_ln:
												lowerCAmelCase_			:   Union[str, Any]     					=					nn.LayerNorm(module.normalized_shape[0]			, eps=module.eps			, params_dtype=module.weight.dtype			)
												lowerCAmelCase_			:   List[str]     					=					module.weight.data.clone()
												lowerCAmelCase_			:   List[Any]     					=					module.bias.data.clone()

												setattr(lowerCAmelCase_			, lowerCAmelCase_			, lowerCAmelCase_			)
								else:
												convert_model(
												    lowerCAmelCase_			, to_transformer_engine=lowerCAmelCase_			, _convert_linear=_convert_linear			, _convert_ln=_convert_ln			, )




def       lowerCAmelCase			(		lowerCAmelCase_			)->			str:
				if not is_fpa_available():
								raise ImportError('''Using `has_transformer_engine_layers` requires transformer_engine to be installed.'''			)
				for m in model.modules():
								if isinstance(lowerCAmelCase_			, (te.LayerNorm, te.Linear)			):
												return True
				return False