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


__lowerCAmelCase         =    ["""model.decoder.embed_positions.weights"""]





def 					UpperCAmelCase_		(__a      :    List[Any] ):






   """simple docstring"""

   if "emb" in name:
      _a				:	Union[str, Any]			=   name.replace('emb' ,	'model.decoder.embed_tokens' )
   if "transformer" in name:
      _a				:	Tuple			=   name.replace('transformer' ,	'model.decoder' )
   if "cross_attention" in name:
      _a				:	int			=   name.replace('cross_attention' ,	'encoder_attn' )
   if "linear1" in name:
      _a				:	List[Any]			=   name.replace('linear1' ,	'fc1' )
   if "linear2" in name:
      _a				:	Any			=   name.replace('linear2' ,	'fc2' )
   if "norm1" in name:
      _a				:	Optional[int]			=   name.replace('norm1' ,	'self_attn_layer_norm' )
   if "norm_cross" in name:
      _a				:	List[Any]			=   name.replace('norm_cross' ,	'encoder_attn_layer_norm' )
   if "norm2" in name:
      _a				:	List[str]			=   name.replace('norm2' ,	'final_layer_norm' )
   if "out_norm" in name:
      _a				:	Tuple			=   name.replace('out_norm' ,	'model.decoder.layer_norm' )
   if "linears" in name:
      _a				:	Union[str, Any]			=   name.replace('linears' ,	'lm_heads' )
   if "condition_provider.conditioners.description.output_proj" in name:
      _a				:	List[str]			=   name.replace('condition_provider.conditioners.description.output_proj' ,	'enc_to_dec_proj' )
   return name





def 					UpperCAmelCase_		(__a      :    OrderedDict ,	__a      :    int ):






   """simple docstring"""

   _a				:	List[Any]			=   list(state_dict.keys() )
   _a				:	Optional[Any]			=   {}
   for key in keys:
      _a				:	Tuple			=   state_dict.pop(__a )
      _a				:	List[Any]			=   rename_keys(__a )
      if "in_proj_weight" in key:
         # split fused qkv proj
         _a				:	Dict			=   val[:hidden_size, :]
         _a				:	List[Any]			=   val[hidden_size : 2 * hidden_size, :]
         _a				:	List[Any]			=   val[-hidden_size:, :]
      elif "enc_to_dec_proj" in key:
         _a				:	Any			=   val
      else:
         _a				:	Dict			=   val
   return state_dict, enc_dec_proj_state_dict





def 					UpperCAmelCase_		(__a      :    str ):






   """simple docstring"""

   if checkpoint == "small":
      # default config values
      _a				:	Union[str, Any]			=   1_0_2_4
      _a				:	Any			=   2_4
      _a				:	str			=   1_6
   elif checkpoint == "medium":
      _a				:	Tuple			=   1_5_3_6
      _a				:	Optional[Any]			=   4_8
      _a				:	Union[str, Any]			=   2_4
   elif checkpoint == "large":
      _a				:	Optional[int]			=   2_0_4_8
      _a				:	Optional[int]			=   4_8
      _a				:	str			=   3_2
   else:
      raise ValueError(f"""Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}.""" )
   _a				:	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      :    List[str] ,	__a      :    Dict=None ,	__a      :    Optional[Any]=None ,	__a      :    Dict="cpu" ):






   """simple docstring"""

   _a				:	Optional[int]			=   MusicGen.get_pretrained(__a ,	device=__a )
   _a				:	int			=   decoder_config_from_checkpoint(__a )

   _a				:	Tuple			=   fairseq_model.lm.state_dict()
   _a,  _a				:	int			=   rename_state_dict(
       __a ,	hidden_size=decoder_config.hidden_size )

   _a				:	List[str]			=   TaEncoderModel.from_pretrained('t5-base' )
   _a				:	List[str]			=   EncodecModel.from_pretrained('facebook/encodec_32khz' )
   _a				:	int			=   MusicgenForCausalLM(__a ).eval()

   # load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
   _a,  _a				:	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
   _a				:	List[Any]			=   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
   _a				:	Optional[int]			=   torch.arange(0 ,	8 ,	dtype=torch.long ).reshape(2 ,	-1 )
   _a				:	Dict			=   input_ids.reshape(2 * 4 ,	-1 )

   with torch.no_grad():
      _a				:	Tuple			=   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
   _a				:	str			=   AutoTokenizer.from_pretrained('t5-base' )
   _a				:	Dict			=   AutoFeatureExtractor.from_pretrained('facebook/encodec_32khz' ,	padding_side='left' )

   _a				:	Optional[Any]			=   MusicgenProcessor(feature_extractor=__a ,	tokenizer=__a )

   # set the appropriate bos/pad token ids
   _a				:	str			=   2_0_4_8
   _a				:	Union[str, Any]			=   2_0_4_8

   # set other default generation config params
   _a				:	List[Any]			=   int(3_0 * audio_encoder.config.frame_rate )
   _a				:	Tuple			=   True
   _a				:	Tuple			=   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__":
  __lowerCAmelCase         =    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."""
  )

  __lowerCAmelCase         =    parser.parse_args()
  convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)





'''simple docstring'''







import inspect
import os
import unittest
from pathlib import Path

import torch

import accelerate
from accelerate.test_utils import execute_subprocess_async
from accelerate.test_utils.testing import run_command







class UpperCAmelCase__	( unittest.TestCase	):





  """simple docstring"""

  __UpperCAmelCase				:				Optional[Any]   =					inspect.getfile(accelerate.test_utils	)
  __UpperCAmelCase				:				List[str]   =					os.path.sep.join(mod_file.split(os.path.sep	)[:-1] + ['''scripts''', '''test_cli.py''']	)

  __UpperCAmelCase				:				Dict   =					['''accelerate''', '''launch''']
  __UpperCAmelCase				:				Dict   =					Path.home() / '''.cache/huggingface/accelerate'''
  __UpperCAmelCase				:				Dict   =					'''default_config.yaml'''
  __UpperCAmelCase				:				Optional[Any]   =					config_folder / config_file
  __UpperCAmelCase				:				Dict   =					config_folder / '''_default_config.yaml'''

  __UpperCAmelCase				:				Any   =					Path('''tests/test_configs'''	)


  @classmethod
  def   __lowercase		(		cls      :     int       ):


     '''simple docstring'''






     if cls.config_path.is_file():
        cls.config_path.rename(cls.changed_path       )


  @classmethod
  def   __lowercase		(		cls      :     List[Any]       ):


     '''simple docstring'''






     if cls.changed_path.is_file():
        cls.changed_path.rename(cls.config_path       )


  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     _a				:	Dict			=   self.base_cmd
     if torch.cuda.is_available() and (torch.cuda.device_count() > 1):
        cmd += ["--multi_gpu"]
     execute_subprocess_async(cmd + [self.test_file_path]							,env=os.environ.copy()       )


  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     for config in sorted(self.test_config_path.glob('**/*.yaml'       )       ):
        with self.subTest(config_file=_a       ):
           execute_subprocess_async(
               self.base_cmd + ['--config_file', str(_a       ), self.test_file_path]							,env=os.environ.copy()       )


  def   __lowercase		(		self      :     Optional[int]       ):


     '''simple docstring'''






     execute_subprocess_async(['accelerate', 'test']							,env=os.environ.copy()       )







class UpperCAmelCase__	( unittest.TestCase	):





  """simple docstring"""

  __UpperCAmelCase				:				Optional[Any]   =					'''test-tpu'''
  __UpperCAmelCase				:				Any   =					'''us-central1-a'''
  __UpperCAmelCase				:				List[Any]   =					'''ls'''
  __UpperCAmelCase				:				Any   =					['''accelerate''', '''tpu-config''']
  __UpperCAmelCase				:				Dict   =					'''cd /usr/share'''
  __UpperCAmelCase				:				Any   =					'''tests/test_samples/test_command_file.sh'''
  __UpperCAmelCase				:				List[Any]   =					'''Running gcloud compute tpus tpu-vm ssh'''


  def   __lowercase		(		self      :     Dict       ):


     '''simple docstring'''






     _a				:	Optional[Any]			=   run_command(
         self.cmd
         + ['--command', self.command, '--tpu_zone', self.tpu_zone, '--tpu_name', self.tpu_name, '--debug']							,return_stdout=_a							,)
     self.assertIn(
         F"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""							,_a							,)


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     _a				:	Any			=   run_command(
         self.cmd
         + [
             '--config_file',
             'tests/test_configs/0_12_0.yaml',
             '--command',
             self.command,
             '--tpu_zone',
             self.tpu_zone,
             '--tpu_name',
             self.tpu_name,
             '--debug',
         ]							,return_stdout=_a							,)
     self.assertIn(
         F"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""							,_a							,)


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     _a				:	Optional[int]			=   run_command(
         self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--debug']							,return_stdout=_a       )
     self.assertIn(
         F"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""							,_a							,)


  def   __lowercase		(		self      :     int       ):


     '''simple docstring'''






     _a				:	Optional[Any]			=   run_command(
         self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--command', self.command, '--debug']							,return_stdout=_a							,)
     self.assertIn(
         F"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""							,_a							,)


  def   __lowercase		(		self      :     str       ):


     '''simple docstring'''






     _a				:	List[str]			=   run_command(
         self.cmd
         + [
             '--config_file',
             'tests/test_configs/latest.yaml',
             '--command',
             self.command,
             '--command',
             'echo "Hello World"',
             '--debug',
         ]							,return_stdout=_a							,)
     self.assertIn(
         F"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls; echo \"Hello World\" --worker all"""							,_a							,)


  def   __lowercase		(		self      :     Union[str, Any]       ):


     '''simple docstring'''






     _a				:	Any			=   run_command(
         self.cmd
         + ['--config_file', 'tests/test_configs/latest.yaml', '--command_file', self.command_file, '--debug']							,return_stdout=_a							,)
     self.assertIn(
         F"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""							,_a							,)


  def   __lowercase		(		self      :     Union[str, Any]       ):


     '''simple docstring'''






     _a				:	Union[str, Any]			=   run_command(
         self.cmd
         + [
             '--config_file',
             'tests/test_configs/0_12_0.yaml',
             '--command_file',
             self.command_file,
             '--tpu_zone',
             self.tpu_zone,
             '--tpu_name',
             self.tpu_name,
             '--debug',
         ]							,return_stdout=_a							,)
     self.assertIn(
         F"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""							,_a							,)


  def   __lowercase		(		self      :     Any       ):


     '''simple docstring'''






     _a				:	Optional[int]			=   run_command(
         self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--install_accelerate', '--debug']							,return_stdout=_a							,)
     self.assertIn(
         F"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate -U; echo \"hello world\"; echo \"this is a second command\" --worker all"""							,_a							,)


  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     _a				:	Optional[int]			=   run_command(
         self.cmd
         + [
             '--config_file',
             'tests/test_configs/latest.yaml',
             '--install_accelerate',
             '--accelerate_version',
             '12.0.0',
             '--debug',
         ]							,return_stdout=_a							,)
     self.assertIn(
         F"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate==12.0.0; echo \"hello world\"; echo \"this is a second command\" --worker all"""							,_a							,)



