

'''simple docstring'''

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

def     __magic_name__       ( A							)			->		int:
	snake_case							 =				ConvNextVaConfig()

	if "atto" in checkpoint_url:
		snake_case							 =				[2, 2, 6, 2]
		snake_case							 =				[4_0, 8_0, 1_6_0, 3_2_0]
	if "femto" in checkpoint_url:
		snake_case							 =				[2, 2, 6, 2]
		snake_case							 =				[4_8, 9_6, 1_9_2, 3_8_4]
	if "pico" in checkpoint_url:
		snake_case							 =				[2, 2, 6, 2]
		snake_case							 =				[6_4, 1_2_8, 2_5_6, 5_1_2]
	if "nano" in checkpoint_url:
		snake_case							 =				[2, 2, 8, 2]
		snake_case							 =				[8_0, 1_6_0, 3_2_0, 6_4_0]
	if "tiny" in checkpoint_url:
		snake_case							 =				[3, 3, 9, 3]
		snake_case							 =				[9_6, 1_9_2, 3_8_4, 7_6_8]
	if "base" in checkpoint_url:
		snake_case							 =				[3, 3, 2_7, 3]
		snake_case							 =				[1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
	if "large" in checkpoint_url:
		snake_case							 =				[3, 3, 2_7, 3]
		snake_case							 =				[1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
	if "huge" in checkpoint_url:
		snake_case							 =				[3, 3, 2_7, 3]
		snake_case							 =				[3_5_2, 7_0_4, 1_4_0_8, 2_8_1_6]

	snake_case							 =				1_0_0_0
	snake_case							 =				'imagenet-1k-id2label.json'
	snake_case							 =				(1, 1_0_0_0)

	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     __magic_name__       ( A							)			->		List[str]:
	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     __magic_name__       ( )			->		List[str]:
	snake_case							 =				'http://images.cocodataset.org/val2017/000000039769.jpg'
	snake_case							 =				Image.open(requests.get(A					,   stream=A							).raw							)
	return im

def     __magic_name__       ( A							)			->		List[str]:
	if "224" in checkpoint_url:
		snake_case							 =				2_2_4
		snake_case							 =				2_2_4 / 2_5_6
	elif "384" in checkpoint_url:
		snake_case							 =				3_8_4
		snake_case							 =				None
	else:
		snake_case							 =				5_1_2
		snake_case							 =				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     __magic_name__       ( A					,   A					,   A					,   A							)			->		Optional[int]:
	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.3_930, 0.1_747, -0.5_246, 0.4_177, 0.4_295]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
		snake_case							 =				torch.tensor([-0.1_727, -0.5_341, -0.7_818, -0.4_745, -0.6_566]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
		snake_case							 =				torch.tensor([-0.0_333, 0.1_563, -0.9_137, 0.1_054, 0.0_381]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
		snake_case							 =				torch.tensor([-0.1_744, -0.1_555, -0.0_713, 0.0_950, -0.1_431]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
		snake_case							 =				torch.tensor([0.9_996, 0.1_966, -0.4_386, -0.3_472, 0.6_661]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
		snake_case							 =				torch.tensor([-0.2_553, -0.6_708, -0.1_359, 0.2_518, -0.2_488]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
		snake_case							 =				torch.tensor([-0.0_673, -0.5_627, -0.3_753, -0.2_722, 0.0_178]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
		snake_case							 =				torch.tensor([-0.6_377, -0.7_458, -0.2_150, 0.1_184, -0.0_597]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
		snake_case							 =				torch.tensor([1.0_799, 0.2_322, -0.8_860, 1.0_219, 0.6_231]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
		snake_case							 =				torch.tensor([0.3_766, 0.4_917, -1.1_426, 0.9_942, 0.6_024]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
		snake_case							 =				torch.tensor([0.4_220, -0.6_919, -0.4_317, -0.2_881, -0.6_609]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
		snake_case							 =				torch.tensor([0.1_082, -0.8_286, -0.5_095, 0.4_681, -0.8_085]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
		snake_case							 =				torch.tensor([-0.2_419, -0.6_221, 0.2_176, -0.0_980, -0.7_527]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
		snake_case							 =				torch.tensor([0.0_391, -0.4_371, 0.3_786, 0.1_251, -0.2_784]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
		snake_case							 =				torch.tensor([-0.0_504, 0.5_636, -0.1_729, -0.6_507, -0.3_949]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
		snake_case							 =				torch.tensor([0.3_560, 0.9_486, 0.3_149, -0.2_667, -0.5_138]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
		snake_case							 =				torch.tensor([-0.2_469, -0.4_550, -0.5_853, -0.0_810, 0.0_309]							)
	elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
		snake_case							 =				torch.tensor([-0.3_090, 0.0_802, -0.0_682, -0.1_979, -0.2_826]							)
	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__":
						lowerCAmelCase_						   =			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")

						lowerCAmelCase_						   =			parser.parse_args()
						convert_convnextva_checkpoint(
						    args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
						)








'''simple docstring'''

import copy
import os
from typing import Dict, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import add_start_docstrings, logging
from ..auto import AutoConfig


lowerCAmelCase_						   =			logging.get_logger(__name__)


lowerCAmelCase_						   =			{
    "suno/bark-small": "https://huggingface.co/suno/bark-small/resolve/main/config.json",
    "suno/bark": "https://huggingface.co/suno/bark/resolve/main/config.json",
}

lowerCAmelCase_						   =			"\n    This is the configuration class to store the configuration of a [`{model}`]. It is used to instantiate the model\n    according to the specified arguments, defining the model architecture. Instantiating a configuration with the\n    defaults will yield a similar configuration to that of the Bark [suno/bark](https://huggingface.co/suno/bark)\n    architecture.\n\n    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the\n    documentation from [`PretrainedConfig`] for more information.\n\n    Args:\n        block_size (`int`, *optional*, defaults to 1024):\n            The maximum sequence length that this model might ever be used with. Typically set this to something large\n            just in case (e.g., 512 or 1024 or 2048).\n        input_vocab_size (`int`, *optional*, defaults to 10_048):\n            Vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented by the\n            `inputs_ids` passed when calling [`{model}`]. Defaults to 10_048 but should be carefully thought with\n            regards to the chosen sub-model.\n        output_vocab_size (`int`, *optional*, defaults to 10_048):\n            Output vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented\n            by the: `output_ids` when passing forward a [`{model}`]. Defaults to 10_048 but should be carefully thought\n            with regards to the chosen sub-model.\n        num_layers (`int`, *optional*, defaults to 12):\n            Number of hidden layers in the given sub-model.\n        num_heads (`int`, *optional*, defaults to 12):\n            Number of attention heads for each attention layer in the Transformer architecture.\n        hidden_size (`int`, *optional*, defaults to 768):\n            Dimensionality of the \"intermediate\" (often named feed-forward) layer in the architecture.\n        dropout (`float`, *optional*, defaults to 0.0):\n            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.\n        bias (`bool`, *optional*, defaults to `True`):\n            Whether or not to use bias in the linear layers and layer norm layers.\n        initializer_range (`float`, *optional*, defaults to 0.02):\n            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.\n        use_cache (`bool`, *optional*, defaults to `True`):\n            Whether or not the model should return the last key/values attentions (not used by all models).\n"







class lowerCamelCase      (						__lowerCAmelCase							):
				snake_case_           =   '''bark_module'''
				snake_case_           =   ['''past_key_values''']

				snake_case_           =   {
				    '''num_attention_heads''': '''num_heads''',
				    '''num_hidden_layers''': '''num_layers''',
				    '''vocab_size''': '''input_vocab_size''',
				    '''window_size''': '''block_size''',
				}

				def __init__(					self,	lowercase_=1024,	lowercase_=10048,	lowercase_=10048,	lowercase_=12,	lowercase_=12,	lowercase_=768,	lowercase_=0.0,	lowercase_=True,	lowercase_=0.02,	lowercase_=True,	**lowercase_,	)						->       List[str]:
					snake_case							 =				block_size
					snake_case							 =				input_vocab_size
					snake_case							 =				output_vocab_size
					snake_case							 =				num_layers
					snake_case							 =				num_heads
					snake_case							 =				hidden_size
					snake_case							 =				dropout
					snake_case							 =				bias
					snake_case							 =				use_cache
					snake_case							 =				initializer_range

					super().__init__(**lowercase_   )






				@classmethod
				def 				_lowerCamelCase						(					cls,	lowercase_,	lowercase_ = None,	lowercase_ = False,	lowercase_ = False,	lowercase_ = None,	lowercase_ = "main",	**lowercase_,	)						->       "PretrainedConfig":
					snake_case							 =				cache_dir
					snake_case							 =				force_download
					snake_case							 =				local_files_only
					snake_case							 =				revision

					cls._set_token_in_kwargs(lowercase_,	lowercase_   )

					snake_case , snake_case							 =				cls.get_config_dict(lowercase_,	**lowercase_   )

					# get the config dict if we are loading from Bark
					if config_dict.get('model_type'   ) == "bark":
						snake_case							 =				config_dict[F'''{cls.model_type}_config''']

					if "model_type" in config_dict and hasattr(cls,	'model_type'   ) and config_dict["model_type"] != cls.model_type:
						logger.warning(
						    F'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
						    F'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''   )

					return cls.from_dict(lowercase_,	**lowercase_   )







@add_start_docstrings(
    BARK_SUBMODELCONFIG_START_DOCSTRING.format(config='''BarkSemanticConfig'''							,							model='''BarkSemanticModel'''							)							,							'''
    Example:

    ```python
    >>> from transformers import BarkSemanticConfig, BarkSemanticModel

    >>> # Initializing a Bark sub-module style configuration
    >>> configuration = BarkSemanticConfig()

    >>> # Initializing a model (with random weights) from the suno/bark style configuration
    >>> model = BarkSemanticModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```'''							,							)
class lowerCamelCase      (						__lowerCAmelCase							):
				snake_case_           =   '''semantic'''







@add_start_docstrings(
    BARK_SUBMODELCONFIG_START_DOCSTRING.format(config='''BarkCoarseConfig'''							,							model='''BarkCoarseModel'''							)							,							'''
    Example:

    ```python
    >>> from transformers import BarkCoarseConfig, BarkCoarseModel

    >>> # Initializing a Bark sub-module style configuration
    >>> configuration = BarkCoarseConfig()

    >>> # Initializing a model (with random weights) from the suno/bark style configuration
    >>> model = BarkCoarseModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```'''							,							)
class lowerCamelCase      (						__lowerCAmelCase							):
				snake_case_           =   '''coarse_acoustics'''







@add_start_docstrings(
    BARK_SUBMODELCONFIG_START_DOCSTRING.format(config='''BarkFineConfig'''							,							model='''BarkFineModel'''							)							,							'''
        n_codes_total (`int`, *optional*, defaults to 8):
            The total number of audio codebooks predicted. Used in the fine acoustics sub-model.
        n_codes_given (`int`, *optional*, defaults to 1):
            The number of audio codebooks predicted in the coarse acoustics sub-model. Used in the acoustics
            sub-models.
    Example:

    ```python
    >>> from transformers import BarkFineConfig, BarkFineModel

    >>> # Initializing a Bark sub-module style configuration
    >>> configuration = BarkFineConfig()

    >>> # Initializing a model (with random weights) from the suno/bark style configuration
    >>> model = BarkFineModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```'''							,							)
class lowerCamelCase      (						__lowerCAmelCase							):
				snake_case_           =   '''fine_acoustics'''

				def __init__(					self,	lowercase_=True,	lowercase_=8,	lowercase_=1,	**lowercase_   )						->       Tuple:
					snake_case							 =				n_codes_total
					snake_case							 =				n_codes_given

					super().__init__(tie_word_embeddings=lowercase_,	**lowercase_   )







class lowerCamelCase      (						__lowerCAmelCase							):
				snake_case_           =   '''bark'''
				snake_case_           =   True

				def __init__(					self,	lowercase_ = None,	lowercase_ = None,	lowercase_ = None,	lowercase_ = None,	lowercase_=0.02,	**lowercase_,	)						->       List[str]:
					if semantic_config is None:
						snake_case							 =				{}
						logger.info('semantic_config is None. initializing the semantic model with default values.'   )

					if coarse_acoustics_config is None:
						snake_case							 =				{}
						logger.info('coarse_acoustics_config is None. initializing the coarse model with default values.'   )

					if fine_acoustics_config is None:
						snake_case							 =				{}
						logger.info('fine_acoustics_config is None. initializing the fine model with default values.'   )

					if codec_config is None:
						snake_case							 =				{}
						logger.info('codec_config is None. initializing the codec model with default values.'   )

					snake_case							 =				BarkSemanticConfig(**lowercase_   )
					snake_case							 =				BarkCoarseConfig(**lowercase_   )
					snake_case							 =				BarkFineConfig(**lowercase_   )
					snake_case							 =				AutoConfig.for_model(**lowercase_   )

					snake_case							 =				initializer_range

					super().__init__(**lowercase_   )

				@classmethod
				def 				_lowerCamelCase						(					cls,	lowercase_,	lowercase_,	lowercase_,	lowercase_,	**lowercase_,	)						->       Any:
					return cls(
					    semantic_config=semantic_config.to_dict(),	coarse_acoustics_config=coarse_acoustics_config.to_dict(),	fine_acoustics_config=fine_acoustics_config.to_dict(),	codec_config=codec_config.to_dict(),	**lowercase_,	)






				def 				_lowerCamelCase						(					self   )						->       Dict:
					snake_case							 =				copy.deepcopy(self.__dict__   )

					snake_case							 =				self.semantic_config.to_dict()
					snake_case							 =				self.coarse_acoustics_config.to_dict()
					snake_case							 =				self.fine_acoustics_config.to_dict()
					snake_case							 =				self.codec_config.to_dict()

					snake_case							 =				self.__class__.model_type
					return output





