






'''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()
_UpperCamelCase						:						Dict  =					logging.get_logger(__name__)


_UpperCamelCase						:						Tuple  =					["""model.decoder.embed_positions.weights"""]
def __UpperCAmelCase   (					A					:			Tuple		)    ->   Optional[int]:



			if "emb" in name:
						UpperCAmelCase_      :       Optional[int]          =       name.replace('''emb'''			,       '''model.decoder.embed_tokens'''		)
			if "transformer" in name:
						UpperCAmelCase_      :       Union[str, Any]          =       name.replace('''transformer'''			,       '''model.decoder'''		)
			if "cross_attention" in name:
						UpperCAmelCase_      :       Union[str, Any]          =       name.replace('''cross_attention'''			,       '''encoder_attn'''		)
			if "linear1" in name:
						UpperCAmelCase_      :       str          =       name.replace('''linear1'''			,       '''fc1'''		)
			if "linear2" in name:
						UpperCAmelCase_      :       List[Any]          =       name.replace('''linear2'''			,       '''fc2'''		)
			if "norm1" in name:
						UpperCAmelCase_      :       Dict          =       name.replace('''norm1'''			,       '''self_attn_layer_norm'''		)
			if "norm_cross" in name:
						UpperCAmelCase_      :       Optional[Any]          =       name.replace('''norm_cross'''			,       '''encoder_attn_layer_norm'''		)
			if "norm2" in name:
						UpperCAmelCase_      :       Any          =       name.replace('''norm2'''			,       '''final_layer_norm'''		)
			if "out_norm" in name:
						UpperCAmelCase_      :       Tuple          =       name.replace('''out_norm'''			,       '''model.decoder.layer_norm'''		)
			if "linears" in name:
						UpperCAmelCase_      :       Optional[int]          =       name.replace('''linears'''			,       '''lm_heads'''		)
			if "condition_provider.conditioners.description.output_proj" in name:
						UpperCAmelCase_      :       str          =       name.replace('''condition_provider.conditioners.description.output_proj'''			,       '''enc_to_dec_proj'''		)
			return name
def __UpperCAmelCase   (					A					:			OrderedDict			,       A					:			int		)    ->   Tuple[Dict, Dict]:
			UpperCAmelCase_      :       str          =       list(state_dict.keys()		)
			UpperCAmelCase_      :       Any          =       {}
			for key in keys:
						UpperCAmelCase_      :       Optional[Any]          =       state_dict.pop(A		)
						UpperCAmelCase_      :       List[Any]          =       rename_keys(A		)
						if "in_proj_weight" in key:
									# split fused qkv proj
									UpperCAmelCase_      :       Optional[Any]          =       val[:hidden_size, :]
									UpperCAmelCase_      :       Any          =       val[hidden_size : 2 * hidden_size, :]
									UpperCAmelCase_      :       Tuple          =       val[-hidden_size:, :]
						elif "enc_to_dec_proj" in key:
									UpperCAmelCase_      :       List[Any]          =       val
						else:
									UpperCAmelCase_      :       Optional[Any]          =       val
			return state_dict, enc_dec_proj_state_dict
def __UpperCAmelCase   (					A					:			str		)    ->   MusicgenDecoderConfig:



			if checkpoint == "small":
						# default config values
						UpperCAmelCase_      :       List[str]          =       1_0_2_4
						UpperCAmelCase_      :       Any          =       2_4
						UpperCAmelCase_      :       Dict          =       1_6
			elif checkpoint == "medium":
						UpperCAmelCase_      :       Any          =       1_5_3_6
						UpperCAmelCase_      :       int          =       4_8
						UpperCAmelCase_      :       Any          =       2_4
			elif checkpoint == "large":
						UpperCAmelCase_      :       List[Any]          =       2_0_4_8
						UpperCAmelCase_      :       Optional[int]          =       4_8
						UpperCAmelCase_      :       Any          =       3_2
			else:
						raise ValueError(F"Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}."		)
			UpperCAmelCase_      :       Optional[int]          =       MusicgenDecoderConfig(
			    hidden_size=A			,       ffn_dim=hidden_size * 4			,       num_hidden_layers=A			,       num_attention_heads=A			,       )
			return config
@torch.no_grad()
def __UpperCAmelCase   (					A					:			int			,       A					:			Dict=None			,       A					:			List[Any]=None			,       A					:			Union[str, Any]="cpu"		)    ->   str:
			UpperCAmelCase_      :       Tuple          =       MusicGen.get_pretrained(A			,       device=A		)
			UpperCAmelCase_      :       List[Any]          =       decoder_config_from_checkpoint(A		)

			UpperCAmelCase_      :       str          =       fairseq_model.lm.state_dict()
			UpperCAmelCase_      :       Tuple          =       rename_state_dict(
			    A			,       hidden_size=decoder_config.hidden_size		)

			UpperCAmelCase_      :       Dict          =       TaEncoderModel.from_pretrained('''t5-base'''		)
			UpperCAmelCase_      :       List[str]          =       EncodecModel.from_pretrained('''facebook/encodec_32khz'''		)
			UpperCAmelCase_      :       Tuple          =       MusicgenForCausalLM(A		).eval()

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

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

			if len(A		) > 0:
						raise ValueError(F"Missing key(s) in state_dict: {missing_keys}"		)

			if len(A		) > 0:
						raise ValueError(F"Unexpected key(s) in state_dict: {unexpected_keys}"		)

			# init the composite model
			UpperCAmelCase_      :       Optional[int]          =       MusicgenForConditionalGeneration(text_encoder=A			,       audio_encoder=A			,       decoder=A		)

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

			# check we can do a forward pass
			UpperCAmelCase_      :       Tuple          =       torch.arange(0			,       8			,       dtype=torch.long		).reshape(2			,       -1		)
			UpperCAmelCase_      :       Union[str, Any]          =       input_ids.reshape(2 * 4			,       -1		)

			with torch.no_grad():
						UpperCAmelCase_      :       Union[str, Any]          =       model(input_ids=A			,       decoder_input_ids=A		).logits

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

			# now construct the processor
			UpperCAmelCase_      :       Optional[Any]          =       AutoTokenizer.from_pretrained('''t5-base'''		)
			UpperCAmelCase_      :       str          =       AutoFeatureExtractor.from_pretrained('''facebook/encodec_32khz'''			,       padding_side='''left'''		)

			UpperCAmelCase_      :       str          =       MusicgenProcessor(feature_extractor=A			,       tokenizer=A		)

			# set the appropriate bos/pad token ids
			UpperCAmelCase_      :       Tuple          =       2_0_4_8
			UpperCAmelCase_      :       List[str]          =       2_0_4_8

			# set other default generation config params
			UpperCAmelCase_      :       Optional[int]          =       int(3_0 * audio_encoder.config.frame_rate		)
			UpperCAmelCase_      :       str          =       True
			UpperCAmelCase_      :       str          =       3.0

			if pytorch_dump_folder is not None:
						Path(A		).mkdir(exist_ok=A		)
						logger.info(F"Saving model {checkpoint} to {pytorch_dump_folder}"		)
						model.save_pretrained(A		)
						processor.save_pretrained(A		)

			if repo_id:
						logger.info(F"Pushing model {checkpoint} to {repo_id}"		)
						model.push_to_hub(A		)
						processor.push_to_hub(A		)


if __name__ == "__main__":
						_UpperCamelCase						:						List[str]  =					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.'
						)

						_UpperCamelCase						:						str  =					parser.parse_args()
						convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)










"""simple docstring"""



import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import TaBlock, TaConfig, TaLayerNorm

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin




class 			lowerCAmelCase__			(    __magic_name__							,						__magic_name__							,						__magic_name__   ):
   @register_to_config
   def __init__(		self      :	List[Any]							,			snake_case__      :	int							,			snake_case__      :	int							,			snake_case__      :	int							,			snake_case__      :	float							,			snake_case__      :	int							,			snake_case__      :	int							,			snake_case__      :	int							,			snake_case__      :	int							,			snake_case__      :	str							,			snake_case__      :	bool = False							,			):


        '''simple docstring'''
        super().__init__()

        UpperCAmelCase__						:    Optional[int]							=							nn.Embedding(snake_case__							,			snake_case__							)

        UpperCAmelCase__						:    List[Any]							=							nn.Embedding(snake_case__							,			snake_case__							)
        UpperCAmelCase__						:    Optional[Any]							=							False

        UpperCAmelCase__						:    Optional[int]							=							nn.Dropout(p=snake_case__							)

        UpperCAmelCase__						:    Optional[int]							=							TaConfig(
            vocab_size=snake_case__							,			d_model=snake_case__							,			num_heads=snake_case__							,			d_kv=snake_case__							,			d_ff=snake_case__							,			dropout_rate=snake_case__							,			feed_forward_proj=snake_case__							,			is_decoder=snake_case__							,			is_encoder_decoder=snake_case__							,			)

        UpperCAmelCase__						:    Tuple							=							nn.ModuleList()
        for lyr_num in range(snake_case__							):
             UpperCAmelCase__						:    Tuple							=							TaBlock(snake_case__							)
             self.encoders.append(snake_case__							)

        UpperCAmelCase__						:    str							=							TaLayerNorm(snake_case__							)
        UpperCAmelCase__						:    Tuple							=							nn.Dropout(p=snake_case__							)




   def       __a							(		self      :	int							,			snake_case__      :	List[str]							,			snake_case__      :	List[Any]							):


        '''simple docstring'''
        UpperCAmelCase__						:    List[Any]							=							self.token_embedder(snake_case__							)

        UpperCAmelCase__						:    Optional[int]							=							encoder_input_tokens.shape[1]
        UpperCAmelCase__						:    List[Any]							=							torch.arange(snake_case__							,			device=encoder_input_tokens.device							)
        x += self.position_encoding(snake_case__							)

        UpperCAmelCase__						:    Union[str, Any]							=							self.dropout_pre(snake_case__							)

        # inverted the attention mask
        UpperCAmelCase__						:    List[Any]							=							encoder_input_tokens.size()
        UpperCAmelCase__						:    Tuple							=							self.get_extended_attention_mask(snake_case__							,			snake_case__							)

        for lyr in self.encoders:
             UpperCAmelCase__						:    Any							=							lyr(snake_case__							,			snake_case__							)[0]
        UpperCAmelCase__						:    Any							=							self.layer_norm(snake_case__							)

        return self.dropout_post(snake_case__							), encoder_inputs_mask


