


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


lowerCamelCase__							:						str							     =   ["""model.decoder.embed_positions.weights"""]







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







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







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase				)    ->				MusicgenDecoderConfig:
  if checkpoint == "small":
    # default config values
    snake_case__						= 1024
    snake_case__						= 24
    snake_case__						= 16
  elif checkpoint == "medium":
    snake_case__						= 1536
    snake_case__						= 48
    snake_case__						= 24
  elif checkpoint == "large":
    snake_case__						= 2048
    snake_case__						= 48
    snake_case__						= 32
  else:
    raise ValueError(F"""Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."""				)
  snake_case__						= MusicgenDecoderConfig(
      hidden_size=__lowerCAmelCase							,   ffn_dim=hidden_size * 4							,   num_hidden_layers=__lowerCAmelCase							,   num_attention_heads=__lowerCAmelCase							,   )
  return config







@torch.no_grad()
def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase=None							,   __lowerCAmelCase=None							,   __lowerCAmelCase="cpu"				)    ->				Union[str, Any]:
  snake_case__						= MusicGen.get_pretrained(__lowerCAmelCase							,   device=__lowerCAmelCase				)
  snake_case__						= decoder_config_from_checkpoint(__lowerCAmelCase				)

  snake_case__						= fairseq_model.lm.state_dict()
  snake_case__				,				snake_case__						= rename_state_dict(
      __lowerCAmelCase							,   hidden_size=decoder_config.hidden_size				)

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

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

  with torch.no_grad():
    snake_case__						= model(input_ids=__lowerCAmelCase							,   decoder_input_ids=__lowerCAmelCase				).logits

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

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

  snake_case__						= MusicgenProcessor(feature_extractor=__lowerCAmelCase							,   tokenizer=__lowerCAmelCase				)

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

  # set other default generation config params
  snake_case__						= int(30 * audio_encoder.config.frame_rate				)
  snake_case__						= True
  snake_case__						= 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__":
   lowerCamelCase__							:						Any							     =   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__							:						Optional[int]							     =   parser.parse_args()
   convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)





import importlib
import shutil
import threading
import warnings
from typing import List

import fsspec
import fsspec.asyn

from . import compression
from .hffilesystem import HfFileSystem


SCREAMING_SNAKE_CASE							    =							importlib.util.find_spec('s3fs') is not None

if _has_safs:
  from .safilesystem import SaFileSystem  # noqa: F401

SCREAMING_SNAKE_CASE							    =							[
    compression.BzaFileSystem,
    compression.GzipFileSystem,
    compression.LzaFileSystem,
    compression.XzFileSystem,
    compression.ZstdFileSystem,
]

# Register custom filesystems
for fs_class in COMPRESSION_FILESYSTEMS + [HfFileSystem]:
  if fs_class.protocol in fsspec.registry and fsspec.registry[fs_class.protocol] is not fs_class:
    warnings.warn(f'''A filesystem protocol was already set for {fs_class.protocol} and will be overwritten.''')
  fsspec.register_implementation(fs_class.protocol, fs_class, clobber=True)




def  a   (lowerCAmelCase__   ):
   if "://" in dataset_path:
      __a					=  dataset_path.split("""://"""   )[1]
   return dataset_path




def  a   (lowerCAmelCase__   ):
   if fs is not None and fs.protocol != "file":
      return True
   else:
      return False




def  a   (lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__   ):
   __a					=  not is_remote_filesystem(lowerCAmelCase__   )
   if is_local:
      # LocalFileSystem.mv does copy + rm, it is more efficient to simply move a local directory
      shutil.move(fs._strip_protocol(lowerCAmelCase__   )		,				fs._strip_protocol(lowerCAmelCase__   )   )
   else:
      fs.mv(lowerCAmelCase__		,				lowerCAmelCase__		,				recursive=lowerCAmelCase__   )




def  a   ():
   if hasattr(fsspec.asyn		,				"""reset_lock"""   ):
      # for future fsspec>2022.05.0
      fsspec.asyn.reset_lock()
   else:
      __a					=  None
      __a					=  None
      __a					=  threading.Lock()



