



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






def 	snake_case__  (						__lowerCamelCase							:       Union[str, Any]      ):
			"""simple docstring"""
			lowerCamelCase__						:						List[str]						      =ConvNextConfig()

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

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

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

			return config, expected_shape






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

			return name






def 	snake_case__  (						):
			"""simple docstring"""
			lowerCamelCase__						:						Any						      ='''http://images.cocodataset.org/val2017/000000039769.jpg'''
			lowerCamelCase__						:						int						      =Image.open(requests.get(__lowerCamelCase				,						stream=__lowerCamelCase      ).raw      )
			return im






@torch.no_grad()
def 	snake_case__  (						__lowerCamelCase							:       Union[str, Any]				,						__lowerCamelCase							:       Optional[int]      ):
			"""simple docstring"""
			lowerCamelCase__      ,  lowerCamelCase__						:						Any						      =get_convnext_config(__lowerCamelCase      )
			# load original state_dict from URL
			lowerCamelCase__						:						int						      =torch.hub.load_state_dict_from_url(__lowerCamelCase      )['''model''']
			# rename keys
			for key in state_dict.copy().keys():
						lowerCamelCase__						:						Optional[Any]						      =state_dict.pop(__lowerCamelCase      )
						lowerCamelCase__						:						Optional[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__						:						List[str]						      ='''convnext.''' + key
						lowerCamelCase__						:						Optional[Any]						      =val

			# load HuggingFace model
			lowerCamelCase__						:						List[Any]						      =ConvNextForImageClassification(__lowerCamelCase      )
			model.load_state_dict(__lowerCamelCase      )
			model.eval()

			# Check outputs on an image, prepared by ConvNextImageProcessor
			lowerCamelCase__						:						Dict						      =224 if '''224''' in checkpoint_url else 384
			lowerCamelCase__						:						int						      =ConvNextImageProcessor(size=__lowerCamelCase      )
			lowerCamelCase__						:						Any						      =image_processor(images=prepare_img()				,						return_tensors='''pt'''      ).pixel_values

			lowerCamelCase__						:						Any						      =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__						:						int						      =torch.tensor([-0.12_10, -0.66_05, 0.19_18]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
						lowerCamelCase__						:						List[str]						      =torch.tensor([-0.44_73, -0.18_47, -0.63_65]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
						lowerCamelCase__						:						Optional[int]						      =torch.tensor([0.45_25, 0.75_39, 0.03_08]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
						lowerCamelCase__						:						Optional[Any]						      =torch.tensor([0.35_61, 0.63_50, -0.03_84]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
						lowerCamelCase__						:						Optional[int]						      =torch.tensor([0.41_74, -0.09_89, 0.14_89]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
						lowerCamelCase__						:						Any						      =torch.tensor([0.25_13, -0.13_49, -0.16_13]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
						lowerCamelCase__						:						Optional[Any]						      =torch.tensor([1.29_80, 0.36_31, -0.11_98]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
						lowerCamelCase__						:						Any						      =torch.tensor([1.29_63, 0.12_27, 0.17_23]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
						lowerCamelCase__						:						List[str]						      =torch.tensor([1.79_56, 0.83_90, 0.28_20]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
						lowerCamelCase__						:						List[Any]						      =torch.tensor([-0.28_22, -0.05_02, -0.08_78]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
						lowerCamelCase__						:						List[Any]						      =torch.tensor([-0.56_72, -0.07_30, -0.43_48]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
						lowerCamelCase__						:						List[str]						      =torch.tensor([0.26_81, 0.23_65, 0.62_46]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
						lowerCamelCase__						:						List[str]						      =torch.tensor([-0.26_42, 0.39_31, 0.51_16]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
						lowerCamelCase__						:						Any						      =torch.tensor([-0.66_77, -0.18_73, -0.83_79]      )
			elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
						lowerCamelCase__						:						Union[str, Any]						      =torch.tensor([-0.77_49, -0.29_67, -0.64_44]      )
			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__						:						Optional[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__":
		_lowercase					:					Tuple 		=		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.",
		)

		_lowercase					:					List[Any] 		=		parser.parse_args()
		convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)









"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


_lowercase					:					Any 		=		{
    "configuration_unispeech_sat": ["UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP", "UniSpeechSatConfig"],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase					:					Any 		=		[
		    "UNISPEECH_SAT_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "UniSpeechSatForAudioFrameClassification",
		    "UniSpeechSatForCTC",
		    "UniSpeechSatForPreTraining",
		    "UniSpeechSatForSequenceClassification",
		    "UniSpeechSatForXVector",
		    "UniSpeechSatModel",
		    "UniSpeechSatPreTrainedModel",
		]

if TYPE_CHECKING:
		from .configuration_unispeech_sat import UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP, UniSpeechSatConfig

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_unispeech_sat import (
				    UNISPEECH_SAT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    UniSpeechSatForAudioFrameClassification,
				    UniSpeechSatForCTC,
				    UniSpeechSatForPreTraining,
				    UniSpeechSatForSequenceClassification,
				    UniSpeechSatForXVector,
				    UniSpeechSatModel,
				    UniSpeechSatPreTrainedModel,
				)

else:
		import sys

		_lowercase					:					int 		=		_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




