


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


__UpperCAmelCase      :Optional[Any]			    =       ["model.decoder.embed_positions.weights"]




def      _a     (						_lowercase :     Optional[Any]  ):





     '''simple docstring'''






     if "emb" in name:
          __UpperCAmelCase     :			Optional[Any]  							=      name.replace('''emb'''      , '''model.decoder.embed_tokens'''  )
     if "transformer" in name:
          __UpperCAmelCase     :			Any  							=      name.replace('''transformer'''      , '''model.decoder'''  )
     if "cross_attention" in name:
          __UpperCAmelCase     :			List[str]  							=      name.replace('''cross_attention'''      , '''encoder_attn'''  )
     if "linear1" in name:
          __UpperCAmelCase     :			List[Any]  							=      name.replace('''linear1'''      , '''fc1'''  )
     if "linear2" in name:
          __UpperCAmelCase     :			Optional[Any]  							=      name.replace('''linear2'''      , '''fc2'''  )
     if "norm1" in name:
          __UpperCAmelCase     :			Optional[Any]  							=      name.replace('''norm1'''      , '''self_attn_layer_norm'''  )
     if "norm_cross" in name:
          __UpperCAmelCase     :			str  							=      name.replace('''norm_cross'''      , '''encoder_attn_layer_norm'''  )
     if "norm2" in name:
          __UpperCAmelCase     :			Tuple  							=      name.replace('''norm2'''      , '''final_layer_norm'''  )
     if "out_norm" in name:
          __UpperCAmelCase     :			Union[str, 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     :			str  							=      name.replace('''condition_provider.conditioners.description.output_proj'''      , '''enc_to_dec_proj'''  )
     return name




def      _a     (						_lowercase :     OrderedDict      , _lowercase :     int  ):





     '''simple docstring'''






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




def      _a     (						_lowercase :     str  ):





     '''simple docstring'''






     if checkpoint == "small":
          # default config values
          __UpperCAmelCase     :			Optional[int]  							=      1024
          __UpperCAmelCase     :			Tuple  							=      24
          __UpperCAmelCase     :			Optional[int]  							=      16
     elif checkpoint == "medium":
          __UpperCAmelCase     :			Optional[Any]  							=      1536
          __UpperCAmelCase     :			int  							=      48
          __UpperCAmelCase     :			List[Any]  							=      24
     elif checkpoint == "large":
          __UpperCAmelCase     :			List[str]  							=      2048
          __UpperCAmelCase     :			Union[str, Any]  							=      48
          __UpperCAmelCase     :			Optional[Any]  							=      32
     else:
          raise ValueError(F'Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}.'  )
     __UpperCAmelCase     :			Union[str, Any]  							=      MusicgenDecoderConfig(
         hidden_size=_lowercase      , ffn_dim=hidden_size * 4      , num_hidden_layers=_lowercase      , num_attention_heads=_lowercase      , )
     return config




@torch.no_grad()
def      _a     (						_lowercase :     Any      , _lowercase :     Optional[Any]=None      , _lowercase :     Optional[Any]=None      , _lowercase :     Dict="cpu"  ):





     '''simple docstring'''






     __UpperCAmelCase     :			str  							=      MusicGen.get_pretrained(_lowercase      , device=_lowercase  )
     __UpperCAmelCase     :			List[str]  							=      decoder_config_from_checkpoint(_lowercase  )

     __UpperCAmelCase     :			Tuple  							=      fairseq_model.lm.state_dict()
     __UpperCAmelCase							,       __UpperCAmelCase     :			int  							=      rename_state_dict(
         _lowercase      , hidden_size=decoder_config.hidden_size  )

     __UpperCAmelCase     :			str  							=      TaEncoderModel.from_pretrained('''t5-base'''  )
     __UpperCAmelCase     :			Dict  							=      EncodecModel.from_pretrained('''facebook/encodec_32khz'''  )
     __UpperCAmelCase     :			Tuple  							=      MusicgenForCausalLM(_lowercase  ).eval()

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

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

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

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

     # init the composite model
     __UpperCAmelCase     :			Any  							=      MusicgenForConditionalGeneration(text_encoder=_lowercase      , audio_encoder=_lowercase      , decoder=_lowercase  )

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

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

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

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

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

     __UpperCAmelCase     :			Optional[Any]  							=      MusicgenProcessor(feature_extractor=_lowercase      , tokenizer=_lowercase  )

     # set the appropriate bos/pad token ids
     __UpperCAmelCase     :			List[str]  							=      2048
     __UpperCAmelCase     :			Optional[Any]  							=      2048

     # set other default generation config params
     __UpperCAmelCase     :			int  							=      int(30 * audio_encoder.config.frame_rate  )
     __UpperCAmelCase     :			str  							=      True
     __UpperCAmelCase     :			List[Any]  							=      3.0

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

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


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



'''simple docstring'''



from argparse import ArgumentParser

from .env import EnvironmentCommand




def      _a     (						):





     '''simple docstring'''






     __UpperCAmelCase     :			Optional[Any]  							=      ArgumentParser('''Diffusers CLI tool'''      , usage='''diffusers-cli <command> [<args>]'''  )
     __UpperCAmelCase     :			Union[str, Any]  							=      parser.add_subparsers(help='''diffusers-cli command helpers'''  )

     # Register commands
     EnvironmentCommand.register_subcommand(_lowercase  )

     # Let's go
     __UpperCAmelCase     :			Union[str, Any]  							=      parser.parse_args()

     if not hasattr(_lowercase      , '''func'''  ):
          parser.print_help()
          exit(1  )

     # Run
     __UpperCAmelCase     :			Any  							=      args.func(_lowercase  )
     service.run()


if __name__ == "__main__":
      main()