


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


UpperCAmelCase__   	=							['''model.decoder.embed_positions.weights''']
def 				UpperCAmelCase_					(							__snake_case						)     ->			List[str]:







 """simple docstring"""


 if "emb" in name:
  _lowercase							      =name.replace('''emb'''     ,						'''model.decoder.embed_tokens'''						)
 if "transformer" in name:
  _lowercase							      =name.replace('''transformer'''     ,						'''model.decoder'''						)
 if "cross_attention" in name:
  _lowercase							      =name.replace('''cross_attention'''     ,						'''encoder_attn'''						)
 if "linear1" in name:
  _lowercase							      =name.replace('''linear1'''     ,						'''fc1'''						)
 if "linear2" in name:
  _lowercase							      =name.replace('''linear2'''     ,						'''fc2'''						)
 if "norm1" in name:
  _lowercase							      =name.replace('''norm1'''     ,						'''self_attn_layer_norm'''						)
 if "norm_cross" in name:
  _lowercase							      =name.replace('''norm_cross'''     ,						'''encoder_attn_layer_norm'''						)
 if "norm2" in name:
  _lowercase							      =name.replace('''norm2'''     ,						'''final_layer_norm'''						)
 if "out_norm" in name:
  _lowercase							      =name.replace('''out_norm'''     ,						'''model.decoder.layer_norm'''						)
 if "linears" in name:
  _lowercase							      =name.replace('''linears'''     ,						'''lm_heads'''						)
 if "condition_provider.conditioners.description.output_proj" in name:
  _lowercase							      =name.replace('''condition_provider.conditioners.description.output_proj'''     ,						'''enc_to_dec_proj'''						)
 return name
def 				UpperCAmelCase_					(							__snake_case     ,						__snake_case						)     ->			Tuple[Dict, Dict]:







 """simple docstring"""


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







 """simple docstring"""


 if checkpoint == "small":
  # default config values
  _lowercase							      =1024
  _lowercase							      =24
  _lowercase							      =16
 elif checkpoint == "medium":
  _lowercase							      =1536
  _lowercase							      =48
  _lowercase							      =24
 elif checkpoint == "large":
  _lowercase							      =2048
  _lowercase							      =48
  _lowercase							      =32
 else:
  raise ValueError(F"Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."						)
 _lowercase							      =MusicgenDecoderConfig(
     hidden_size=__snake_case     ,						ffn_dim=hidden_size * 4     ,						num_hidden_layers=__snake_case     ,						num_attention_heads=__snake_case     ,						)
 return config
@torch.no_grad()
def 				UpperCAmelCase_					(							__snake_case     ,						__snake_case=None     ,						__snake_case=None     ,						__snake_case="cpu"						)     ->			Optional[Any]:







 """simple docstring"""


 _lowercase							      =MusicGen.get_pretrained(__snake_case     ,						device=__snake_case						)
 _lowercase							      =decoder_config_from_checkpoint(__snake_case						)

 _lowercase							      =fairseq_model.lm.state_dict()
 _lowercase      ,    _lowercase							      =rename_state_dict(
     __snake_case     ,						hidden_size=decoder_config.hidden_size						)

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

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

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

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

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

 # init the composite model
 _lowercase							      =MusicgenForConditionalGeneration(text_encoder=__snake_case     ,						audio_encoder=__snake_case     ,						decoder=__snake_case						)

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

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

 with torch.no_grad():
  _lowercase							      =model(input_ids=__snake_case     ,						decoder_input_ids=__snake_case						).logits

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

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

 _lowercase							      =MusicgenProcessor(feature_extractor=__snake_case     ,						tokenizer=__snake_case						)

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

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

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

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


if __name__ == "__main__":
  UpperCAmelCase__   	=							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__   	=							parser.parse_args()
  convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)





from typing import TYPE_CHECKING

from ..utils import _LazyModule


UpperCAmelCase__   	=							{
    '''config''': [
        '''EXTERNAL_DATA_FORMAT_SIZE_LIMIT''',
        '''OnnxConfig''',
        '''OnnxConfigWithPast''',
        '''OnnxSeq2SeqConfigWithPast''',
        '''PatchingSpec''',
    ],
    '''convert''': ['''export''', '''validate_model_outputs'''],
    '''features''': ['''FeaturesManager'''],
    '''utils''': ['''ParameterFormat''', '''compute_serialized_parameters_size'''],
}


if TYPE_CHECKING:
  from .config import (
      EXTERNAL_DATA_FORMAT_SIZE_LIMIT,
      OnnxConfig,
      OnnxConfigWithPast,
      OnnxSeqaSeqConfigWithPast,
      PatchingSpec,
  )
  from .convert import export, validate_model_outputs
  from .features import FeaturesManager
  from .utils import ParameterFormat, compute_serialized_parameters_size

else:
  import sys

  UpperCAmelCase__   	=							_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)

