


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


lowercase							=['model.decoder.embed_positions.weights']





def lowerCamelCase__      (						__lowerCamelCase	:							List[Any]							):







   '''simple docstring'''




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





def lowerCamelCase__      (						__lowerCamelCase	:							OrderedDict	,   __lowerCamelCase	:							int							):







   '''simple docstring'''




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





def lowerCamelCase__      (						__lowerCamelCase	:							str							):







   '''simple docstring'''




   if checkpoint == "small":
      # default config values
      _UpperCAmelCase       :       Tuple							   =1_0_2_4
      _UpperCAmelCase       :       List[Any]							   =2_4
      _UpperCAmelCase       :       List[Any]							   =1_6
   elif checkpoint == "medium":
      _UpperCAmelCase       :       List[Any]							   =1_5_3_6
      _UpperCAmelCase       :       List[str]							   =4_8
      _UpperCAmelCase       :       str							   =2_4
   elif checkpoint == "large":
      _UpperCAmelCase       :       Optional[Any]							   =2_0_4_8
      _UpperCAmelCase       :       int							   =4_8
      _UpperCAmelCase       :       Union[str, Any]							   =3_2
   else:
      raise ValueError(f"Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."							)
   _UpperCAmelCase       :       Any							   =MusicgenDecoderConfig(
       hidden_size=__lowerCamelCase	,   ffn_dim=hidden_size * 4	,   num_hidden_layers=__lowerCamelCase	,   num_attention_heads=__lowerCamelCase	,   )
   return config





@torch.no_grad()
def lowerCamelCase__      (						__lowerCamelCase	:							Optional[Any]	,   __lowerCamelCase	:							Union[str, Any]=None	,   __lowerCamelCase	:							Dict=None	,   __lowerCamelCase	:							Optional[Any]="cpu"							):







   '''simple docstring'''




   _UpperCAmelCase       :       List[Any]							   =MusicGen.get_pretrained(__lowerCamelCase	,   device=__lowerCamelCase							)
   _UpperCAmelCase       :       List[str]							   =decoder_config_from_checkpoint(__lowerCamelCase							)

   _UpperCAmelCase       :       Union[str, Any]							   =fairseq_model.lm.state_dict()
   _UpperCAmelCase   ,	_UpperCAmelCase       :       Optional[int]							   =rename_state_dict(
       __lowerCamelCase	,   hidden_size=decoder_config.hidden_size							)

   _UpperCAmelCase       :       Union[str, Any]							   =TaEncoderModel.from_pretrained('t5-base'							)
   _UpperCAmelCase       :       int							   =EncodecModel.from_pretrained('facebook/encodec_32khz'							)
   _UpperCAmelCase       :       List[Any]							   =MusicgenForCausalLM(__lowerCamelCase							).eval()

   # load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
   _UpperCAmelCase   ,	_UpperCAmelCase       :       List[Any]							   =decoder.load_state_dict(__lowerCamelCase	,   strict=__lowerCamelCase							)

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

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

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

   # init the composite model
   _UpperCAmelCase       :       Tuple							   =MusicgenForConditionalGeneration(text_encoder=__lowerCamelCase	,   audio_encoder=__lowerCamelCase	,   decoder=__lowerCamelCase							)

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

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

   with torch.no_grad():
      _UpperCAmelCase       :       str							   =model(input_ids=__lowerCamelCase	,   decoder_input_ids=__lowerCamelCase							).logits

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

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

   _UpperCAmelCase       :       Tuple							   =MusicgenProcessor(feature_extractor=__lowerCamelCase	,   tokenizer=__lowerCamelCase							)

   # set the appropriate bos/pad token ids
   _UpperCAmelCase       :       Dict							   =2_0_4_8
   _UpperCAmelCase       :       int							   =2_0_4_8

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

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

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


if __name__ == "__main__":
 lowercase							=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.'
 )

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





'''simple docstring'''






from typing import Any





def lowerCamelCase__      (						__lowerCamelCase	:							list	,   __lowerCamelCase	:							list	,   __lowerCamelCase	:							dict	,   __lowerCamelCase	:							dict	,   __lowerCamelCase	:							dict	,   ):







   '''simple docstring'''




   _validation(
       __lowerCamelCase	,   __lowerCamelCase	,   __lowerCamelCase	,   __lowerCamelCase	,   __lowerCamelCase	,   )
   # Creates data structures and fill initial step
   _UpperCAmelCase       :       dict							   ={}
   _UpperCAmelCase       :       dict							   ={}
   for state in states_space:
      _UpperCAmelCase       :       int							   =observations_space[0]
      _UpperCAmelCase       :       int							   =(
          initial_probabilities[state] * emission_probabilities[state][observation]
      )
      _UpperCAmelCase       :       int							   =None

   # Fills the data structure with the probabilities of
   # different transitions and pointers to previous states
   for o in range(1	,   len(__lowerCamelCase							)							):
      _UpperCAmelCase       :       List[Any]							   =observations_space[o]
      _UpperCAmelCase       :       Optional[int]							   =observations_space[o - 1]
      for state in states_space:
         # Calculates the argmax for probability function
         _UpperCAmelCase       :       List[str]							   =''
         _UpperCAmelCase       :       Dict							   =-1
         for k_state in states_space:
            _UpperCAmelCase       :       List[str]							   =(
                probabilities[(k_state, prior_observation)]
                * transition_probabilities[k_state][state]
                * emission_probabilities[state][observation]
            )
            if probability > max_probability:
               _UpperCAmelCase       :       int							   =probability
               _UpperCAmelCase       :       List[Any]							   =k_state

            # Update probabilities and pointers dicts
         _UpperCAmelCase       :       str							   =(
             probabilities[(arg_max, prior_observation)]
             * transition_probabilities[arg_max][state]
             * emission_probabilities[state][observation]
         )

         _UpperCAmelCase       :       List[Any]							   =arg_max

    # The final observation
   _UpperCAmelCase       :       int							   =observations_space[len(__lowerCamelCase							) - 1]

   # argmax for given final observation
   _UpperCAmelCase       :       Any							   =''
   _UpperCAmelCase       :       Union[str, Any]							   =-1
   for k_state in states_space:
      _UpperCAmelCase       :       Optional[int]							   =probabilities[(k_state, final_observation)]
      if probability > max_probability:
         _UpperCAmelCase       :       Union[str, Any]							   =probability
         _UpperCAmelCase       :       int							   =k_state
   _UpperCAmelCase       :       int							   =arg_max

   # Process pointers backwards
   _UpperCAmelCase       :       List[str]							   =last_state
   _UpperCAmelCase       :       Optional[int]							   =[]
   for o in range(len(__lowerCamelCase							) - 1	,   -1	,   -1							):
      result.append(__lowerCamelCase							)
      _UpperCAmelCase       :       Optional[Any]							   =pointers[previous, observations_space[o]]
   result.reverse()

   return result





def lowerCamelCase__      (						__lowerCamelCase	:							Any	,   __lowerCamelCase	:							Any	,   __lowerCamelCase	:							Any	,   __lowerCamelCase	:							Any	,   __lowerCamelCase	:							Any	,   ):







   '''simple docstring'''




   _validate_not_empty(
       __lowerCamelCase	,   __lowerCamelCase	,   __lowerCamelCase	,   __lowerCamelCase	,   __lowerCamelCase	,   )
   _validate_lists(__lowerCamelCase	,   __lowerCamelCase							)
   _validate_dicts(
       __lowerCamelCase	,   __lowerCamelCase	,   __lowerCamelCase							)





def lowerCamelCase__      (						__lowerCamelCase	:							Any	,   __lowerCamelCase	:							Any	,   __lowerCamelCase	:							Any	,   __lowerCamelCase	:							Any	,   __lowerCamelCase	:							Any	,   ):







   '''simple docstring'''




   if not all(
       [
           observations_space,
           states_space,
           initial_probabilities,
           transition_probabilities,
           emission_probabilities,
       ]							):
      raise ValueError('There\'s an empty parameter'							)





def lowerCamelCase__      (						__lowerCamelCase	:							Any	,   __lowerCamelCase	:							Any							):







   '''simple docstring'''




   _validate_list(__lowerCamelCase	,   'observations_space'							)
   _validate_list(__lowerCamelCase	,   'states_space'							)





def lowerCamelCase__      (						__lowerCamelCase	:							Any	,   __lowerCamelCase	:							str							):







   '''simple docstring'''




   if not isinstance(_object	,   __lowerCamelCase							):
      _UpperCAmelCase       :       Any							   =f"{var_name} must be a list"
      raise ValueError(__lowerCamelCase							)
   else:
      for x in _object:
         if not isinstance(__lowerCamelCase	,   __lowerCamelCase							):
            _UpperCAmelCase       :       Optional[int]							   =f"{var_name} must be a list of strings"
            raise ValueError(__lowerCamelCase							)





def lowerCamelCase__      (						__lowerCamelCase	:							Any	,   __lowerCamelCase	:							Any	,   __lowerCamelCase	:							Any	,   ):







   '''simple docstring'''




   _validate_dict(__lowerCamelCase	,   'initial_probabilities'	,   __lowerCamelCase							)
   _validate_nested_dict(__lowerCamelCase	,   'transition_probabilities'							)
   _validate_nested_dict(__lowerCamelCase	,   'emission_probabilities'							)





def lowerCamelCase__      (						__lowerCamelCase	:							Any	,   __lowerCamelCase	:							str							):







   '''simple docstring'''




   _validate_dict(_object	,   __lowerCamelCase	,   __lowerCamelCase							)
   for x in _object.values():
      _validate_dict(__lowerCamelCase	,   __lowerCamelCase	,   __lowerCamelCase	,   __lowerCamelCase							)





def lowerCamelCase__      (						__lowerCamelCase	:							Any	,   __lowerCamelCase	:							str	,   __lowerCamelCase	:							type	,   __lowerCamelCase	:							bool = False							):







   '''simple docstring'''




   if not isinstance(_object	,   __lowerCamelCase							):
      _UpperCAmelCase       :       List[str]							   =f"{var_name} must be a dict"
      raise ValueError(__lowerCamelCase							)
   if not all(isinstance(__lowerCamelCase	,   __lowerCamelCase							) for x in _object							):
      _UpperCAmelCase       :       str							   =f"{var_name} all keys must be strings"
      raise ValueError(__lowerCamelCase							)
   if not all(isinstance(__lowerCamelCase	,   __lowerCamelCase							) for x in _object.values()							):
      _UpperCAmelCase       :       int							   ='nested dictionary ' if nested else ''
      _UpperCAmelCase       :       Optional[int]							   =f"{var_name} {nested_text}all values must be {value_type.__name__}"
      raise ValueError(__lowerCamelCase							)


if __name__ == "__main__":
 from doctest import testmod

 testmod()

