




'''simple docstring'''


import argparse
from pathlib import Path
from typing import Dict, OrderedDict, Tuple

import torch
from audiocraft.models import MusicGen

from transformers import (
    AutoFeatureExtractor,
    AutoTokenizer,
    EncodecModel,
    MusicgenDecoderConfig,
    MusicgenForConditionalGeneration,
    MusicgenProcessor,
    TaEncoderModel,
)
from transformers.models.musicgen.modeling_musicgen import MusicgenForCausalLM
from transformers.utils import logging


logging.set_verbosity_info()
_snake_case				:				Dict       							=		logging.get_logger(__name__)


_snake_case				:				int       							=		['model.decoder.embed_positions.weights']
def 		snake_case_   (UpperCamelCase	:  List[Any]				):





							'''simple docstring'''


							if "emb" in name:
														_a     					=    name.replace('''emb'''			,						'''model.decoder.embed_tokens'''				)
							if "transformer" in name:
														_a     					=    name.replace('''transformer'''			,						'''model.decoder'''				)
							if "cross_attention" in name:
														_a     					=    name.replace('''cross_attention'''			,						'''encoder_attn'''				)
							if "linear1" in name:
														_a     					=    name.replace('''linear1'''			,						'''fc1'''				)
							if "linear2" in name:
														_a     					=    name.replace('''linear2'''			,						'''fc2'''				)
							if "norm1" in name:
														_a     					=    name.replace('''norm1'''			,						'''self_attn_layer_norm'''				)
							if "norm_cross" in name:
														_a     					=    name.replace('''norm_cross'''			,						'''encoder_attn_layer_norm'''				)
							if "norm2" in name:
														_a     					=    name.replace('''norm2'''			,						'''final_layer_norm'''				)
							if "out_norm" in name:
														_a     					=    name.replace('''out_norm'''			,						'''model.decoder.layer_norm'''				)
							if "linears" in name:
														_a     					=    name.replace('''linears'''			,						'''lm_heads'''				)
							if "condition_provider.conditioners.description.output_proj" in name:
														_a     					=    name.replace('''condition_provider.conditioners.description.output_proj'''			,						'''enc_to_dec_proj'''				)
							return name
def 		snake_case_   (UpperCamelCase	:  OrderedDict			,						UpperCamelCase	:  int				):





							'''simple docstring'''


							_a     					=    list(state_dict.keys()				)
							_a     					=    {}
							for key in keys:
														_a     					=    state_dict.pop(UpperCamelCase				)
														_a     					=    rename_keys(UpperCamelCase				)
														if "in_proj_weight" in key:
																					# split fused qkv proj
																					_a     					=    val[:hidden_size, :]
																					_a     					=    val[hidden_size : 2 * hidden_size, :]
																					_a     					=    val[-hidden_size:, :]
														elif "enc_to_dec_proj" in key:
																					_a     					=    val
														else:
																					_a     					=    val
							return state_dict, enc_dec_proj_state_dict
def 		snake_case_   (UpperCamelCase	:  str				):





							'''simple docstring'''


							if checkpoint == "small":
														# default config values
														_a     					=    1024
														_a     					=    24
														_a     					=    16
							elif checkpoint == "medium":
														_a     					=    1536
														_a     					=    48
														_a     					=    24
							elif checkpoint == "large":
														_a     					=    2048
														_a     					=    48
														_a     					=    32
							else:
														raise ValueError(f'Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}.'				)
							_a     					=    MusicgenDecoderConfig(
							    hidden_size=UpperCamelCase			,						ffn_dim=hidden_size * 4			,						num_hidden_layers=UpperCamelCase			,						num_attention_heads=UpperCamelCase			,						)
							return config
@torch.no_grad()
def 		snake_case_   (UpperCamelCase	:  List[str]			,						UpperCamelCase	:  List[Any]=None			,						UpperCamelCase	:  List[Any]=None			,						UpperCamelCase	:  Optional[int]="cpu"				):





							'''simple docstring'''


							_a     					=    MusicGen.get_pretrained(UpperCamelCase			,						device=UpperCamelCase				)
							_a     					=    decoder_config_from_checkpoint(UpperCamelCase				)

							_a     					=    fairseq_model.lm.state_dict()
							_a					,							_a     					=    rename_state_dict(
							    UpperCamelCase			,						hidden_size=decoder_config.hidden_size				)

							_a     					=    TaEncoderModel.from_pretrained('''t5-base'''				)
							_a     					=    EncodecModel.from_pretrained('''facebook/encodec_32khz'''				)
							_a     					=    MusicgenForCausalLM(UpperCamelCase				).eval()

							# load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
							_a					,							_a     					=    decoder.load_state_dict(UpperCamelCase			,						strict=UpperCamelCase				)

							for key in missing_keys.copy():
														if key.startswith(('''text_encoder''', '''audio_encoder''')				) or key in EXPECTED_MISSING_KEYS:
																					missing_keys.remove(UpperCamelCase				)

							if len(UpperCamelCase				) > 0:
														raise ValueError(f'Missing key(s) in state_dict: {missing_keys}'				)

							if len(UpperCamelCase				) > 0:
														raise ValueError(f'Unexpected key(s) in state_dict: {unexpected_keys}'				)

							# init the composite model
							_a     					=    MusicgenForConditionalGeneration(text_encoder=UpperCamelCase			,						audio_encoder=UpperCamelCase			,						decoder=UpperCamelCase				)

							# load the pre-trained enc-dec projection (from the decoder state dict)
							model.enc_to_dec_proj.load_state_dict(UpperCamelCase				)

							# check we can do a forward pass
							_a     					=    torch.arange(0			,						8			,						dtype=torch.long				).reshape(2			,						-1				)
							_a     					=    input_ids.reshape(2 * 4			,						-1				)

							with torch.no_grad():
														_a     					=    model(input_ids=UpperCamelCase			,						decoder_input_ids=UpperCamelCase				).logits

							if logits.shape != (8, 1, 2048):
														raise ValueError('''Incorrect shape for logits'''				)

							# now construct the processor
							_a     					=    AutoTokenizer.from_pretrained('''t5-base'''				)
							_a     					=    AutoFeatureExtractor.from_pretrained('''facebook/encodec_32khz'''			,						padding_side='''left'''				)

							_a     					=    MusicgenProcessor(feature_extractor=UpperCamelCase			,						tokenizer=UpperCamelCase				)

							# set the appropriate bos/pad token ids
							_a     					=    2048
							_a     					=    2048

							# set other default generation config params
							_a     					=    int(30 * audio_encoder.config.frame_rate				)
							_a     					=    True
							_a     					=    3.0

							if pytorch_dump_folder is not None:
														Path(UpperCamelCase				).mkdir(exist_ok=UpperCamelCase				)
														logger.info(f'Saving model {checkpoint} to {pytorch_dump_folder}'				)
														model.save_pretrained(UpperCamelCase				)
														processor.save_pretrained(UpperCamelCase				)

							if repo_id:
														logger.info(f'Pushing model {checkpoint} to {repo_id}'				)
														model.push_to_hub(UpperCamelCase				)
														processor.push_to_hub(UpperCamelCase				)


if __name__ == "__main__":
				_snake_case				:				Optional[Any]       							=		argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    '--checkpoint',
				    default='small',
				    type=str,
				    help='Checkpoint size of the MusicGen model you\'d like to convert. Can be one of: `[\'small\', \'medium\', \'large\']`.',
				)
				parser.add_argument(
				    '--pytorch_dump_folder',
				    required=True,
				    default=None,
				    type=str,
				    help='Path to the output PyTorch model directory.',
				)
				parser.add_argument(
				    '--push_to_hub', default=None, type=str, help='Where to upload the converted model on the 🤗 hub.'
				)
				parser.add_argument(
				    '--device', default='cpu', type=str, help='Torch device to run the conversion, either cpu or cuda.'
				)

				_snake_case				:				int       							=		parser.parse_args()
				convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)







'''simple docstring'''


import tempfile
import unittest

import numpy as np
from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_flax


if is_flax_available():
				import os

				from flax.core.frozen_dict import unfreeze
				from flax.traverse_util import flatten_dict

				from transformers import FlaxBertModel

				_snake_case				:				Optional[Any]       							=		'0.12'  # assumed parallelism: 8

@require_flax
@is_staging_test
class       A   (		unittest.TestCase   ):







							@classmethod
							def        __lowerCAmelCase      (  cls	:  Tuple )      ->							int:


														"""simple docstring"""


														_a     					=    TOKEN
														HfFolder.save_token(lowerCAmelCase_ )







							@classmethod
							def        __lowerCAmelCase      (  cls	:  Tuple )      ->							Optional[int]:


														"""simple docstring"""


														try:
																					delete_repo(token=cls._token							,      repo_id='''test-model-flax''' )
														except HTTPError:
																					pass

														try:
																					delete_repo(token=cls._token							,      repo_id='''valid_org/test-model-flax-org''' )
														except HTTPError:
																					pass







							def        __lowerCAmelCase      (  self	:  Any )      ->							Optional[Any]:


														"""simple docstring"""


														_a     					=    BertConfig(
														    vocab_size=99							,      hidden_size=32							,      num_hidden_layers=5							,      num_attention_heads=4							,      intermediate_size=37 )
														_a     					=    FlaxBertModel(lowerCAmelCase_ )
														model.push_to_hub('''test-model-flax'''							,      use_auth_token=self._token )

														_a     					=    FlaxBertModel.from_pretrained(F'{USER}/test-model-flax' )

														_a     					=    flatten_dict(unfreeze(model.params ) )
														_a     					=    flatten_dict(unfreeze(new_model.params ) )

														for key in base_params.keys():
																					_a     					=    (base_params[key] - new_params[key]).sum().item()
																					self.assertLessEqual(lowerCAmelCase_							,      1e-3							,      msg=F'{key} not identical' )

														# Reset repo
														delete_repo(token=self._token							,      repo_id='''test-model-flax''' )

														# Push to hub via save_pretrained
														with tempfile.TemporaryDirectory() as tmp_dir:
																					model.save_pretrained(lowerCAmelCase_							,      repo_id='''test-model-flax'''							,      push_to_hub=lowerCAmelCase_							,      use_auth_token=self._token )

														_a     					=    FlaxBertModel.from_pretrained(F'{USER}/test-model-flax' )

														_a     					=    flatten_dict(unfreeze(model.params ) )
														_a     					=    flatten_dict(unfreeze(new_model.params ) )

														for key in base_params.keys():
																					_a     					=    (base_params[key] - new_params[key]).sum().item()
																					self.assertLessEqual(lowerCAmelCase_							,      1e-3							,      msg=F'{key} not identical' )







							def        __lowerCAmelCase      (  self	:  int )      ->							Tuple:


														"""simple docstring"""


														_a     					=    BertConfig(
														    vocab_size=99							,      hidden_size=32							,      num_hidden_layers=5							,      num_attention_heads=4							,      intermediate_size=37 )
														_a     					=    FlaxBertModel(lowerCAmelCase_ )
														model.push_to_hub('''valid_org/test-model-flax-org'''							,      use_auth_token=self._token )

														_a     					=    FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org''' )

														_a     					=    flatten_dict(unfreeze(model.params ) )
														_a     					=    flatten_dict(unfreeze(new_model.params ) )

														for key in base_params.keys():
																					_a     					=    (base_params[key] - new_params[key]).sum().item()
																					self.assertLessEqual(lowerCAmelCase_							,      1e-3							,      msg=F'{key} not identical' )

														# Reset repo
														delete_repo(token=self._token							,      repo_id='''valid_org/test-model-flax-org''' )

														# Push to hub via save_pretrained
														with tempfile.TemporaryDirectory() as tmp_dir:
																					model.save_pretrained(
																					    lowerCAmelCase_							,      repo_id='''valid_org/test-model-flax-org'''							,      push_to_hub=lowerCAmelCase_							,      use_auth_token=self._token )

														_a     					=    FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org''' )

														_a     					=    flatten_dict(unfreeze(model.params ) )
														_a     					=    flatten_dict(unfreeze(new_model.params ) )

														for key in base_params.keys():
																					_a     					=    (base_params[key] - new_params[key]).sum().item()
																					self.assertLessEqual(lowerCAmelCase_							,      1e-3							,      msg=F'{key} not identical' )



def 		snake_case_   (UpperCamelCase	:  List[str]			,						UpperCamelCase	:  Optional[Any]				):





							'''simple docstring'''


							_a     					=    True
							_a     					=    flatten_dict(modela.params				)
							_a     					=    flatten_dict(modela.params				)
							for key in flat_params_a.keys():
														if np.sum(np.abs(flat_params_a[key] - flat_params_a[key]				)				) > 1e-4:
																					_a     					=    False

							return models_are_equal

@require_flax
class       A   (		unittest.TestCase   ):







							def        __lowerCAmelCase      (  self	:  Dict )      ->							int:


														"""simple docstring"""


														_a     					=    BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only''' )
														_a     					=    FlaxBertModel(lowerCAmelCase_ )

														_a     					=    '''bert'''
														with tempfile.TemporaryDirectory() as tmp_dir:
																					model.save_pretrained(os.path.join(lowerCAmelCase_							,      lowerCAmelCase_ ) )

																					with self.assertRaises(lowerCAmelCase_ ):
																												_a     					=    FlaxBertModel.from_pretrained(lowerCAmelCase_ )

																					_a     					=    FlaxBertModel.from_pretrained(lowerCAmelCase_							,      subfolder=lowerCAmelCase_ )

														self.assertTrue(check_models_equal(lowerCAmelCase_							,      lowerCAmelCase_ ) )







							def        __lowerCAmelCase      (  self	:  Union[str, Any] )      ->							Optional[int]:


														"""simple docstring"""


														_a     					=    BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only''' )
														_a     					=    FlaxBertModel(lowerCAmelCase_ )

														_a     					=    '''bert'''
														with tempfile.TemporaryDirectory() as tmp_dir:
																					model.save_pretrained(os.path.join(lowerCAmelCase_							,      lowerCAmelCase_ )							,      max_shard_size='''10KB''' )

																					with self.assertRaises(lowerCAmelCase_ ):
																												_a     					=    FlaxBertModel.from_pretrained(lowerCAmelCase_ )

																					_a     					=    FlaxBertModel.from_pretrained(lowerCAmelCase_							,      subfolder=lowerCAmelCase_ )

														self.assertTrue(check_models_equal(lowerCAmelCase_							,      lowerCAmelCase_ ) )







							def        __lowerCAmelCase      (  self	:  Optional[Any] )      ->							List[str]:


														"""simple docstring"""


														_a     					=    '''bert'''
														_a     					=    '''hf-internal-testing/tiny-random-bert-subfolder'''

														with self.assertRaises(lowerCAmelCase_ ):
																					_a     					=    FlaxBertModel.from_pretrained(lowerCAmelCase_ )

														_a     					=    FlaxBertModel.from_pretrained(lowerCAmelCase_							,      subfolder=lowerCAmelCase_ )

														self.assertIsNotNone(lowerCAmelCase_ )







							def        __lowerCAmelCase      (  self	:  Tuple )      ->							Tuple:


														"""simple docstring"""


														_a     					=    '''bert'''
														_a     					=    '''hf-internal-testing/tiny-random-bert-sharded-subfolder'''
														with self.assertRaises(lowerCAmelCase_ ):
																					_a     					=    FlaxBertModel.from_pretrained(lowerCAmelCase_ )

														_a     					=    FlaxBertModel.from_pretrained(lowerCAmelCase_							,      subfolder=lowerCAmelCase_ )

														self.assertIsNotNone(lowerCAmelCase_ )

