





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


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







def 		__lowercase     (   snake_case_  :  Optional[int]						)   ->int:



    '''simple docstring'''







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







def 		__lowercase     (   snake_case_  :  OrderedDict				,snake_case_  :  int						)   ->Tuple[Dict, Dict]:



    '''simple docstring'''







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







def 		__lowercase     (   snake_case_  :  str						)   ->MusicgenDecoderConfig:



    '''simple docstring'''







    if checkpoint == "small":
        # default config values
        __A    :						str						 =     1024
        __A    :						Optional[int]						 =     24
        __A    :						int						 =     16
    elif checkpoint == "medium":
        __A    :						List[Any]						 =     1536
        __A    :						int						 =     48
        __A    :						int						 =     24
    elif checkpoint == "large":
        __A    :						Optional[int]						 =     2048
        __A    :						str						 =     48
        __A    :						Tuple						 =     32
    else:
        raise ValueError(F"""Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."""						)
    __A    :						Union[str, Any]						 =     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 		__lowercase     (   snake_case_  :  Tuple				,snake_case_  :  List[Any]=None				,snake_case_  :  Any=None				,snake_case_  :  str="cpu"						)   ->List[str]:



    '''simple docstring'''







    __A    :						Optional[int]						 =     MusicGen.get_pretrained(snake_case_				,device=snake_case_						)
    __A    :						Optional[Any]						 =     decoder_config_from_checkpoint(snake_case_						)

    __A    :						int						 =     fairseq_model.lm.state_dict()
    __A ,       __A    :						str						 =     rename_state_dict(
        snake_case_				,hidden_size=decoder_config.hidden_size						)

    __A    :						Any						 =     TaEncoderModel.from_pretrained('''t5-base'''						)
    __A    :						Dict						 =     EncodecModel.from_pretrained('''facebook/encodec_32khz'''						)
    __A    :						List[Any]						 =     MusicgenForCausalLM(snake_case_						).eval()

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

    with torch.no_grad():
        __A    :						int						 =     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
    __A    :						Tuple						 =     AutoTokenizer.from_pretrained('''t5-base'''						)
    __A    :						Optional[int]						 =     AutoFeatureExtractor.from_pretrained('''facebook/encodec_32khz'''				,padding_side='''left'''						)

    __A    :						Optional[Any]						 =     MusicgenProcessor(feature_extractor=snake_case_				,tokenizer=snake_case_						)

    # set the appropriate bos/pad token ids
    __A    :						int						 =     2048
    __A    :						int						 =     2048

    # set other default generation config params
    __A    :						Union[str, Any]						 =     int(30 * audio_encoder.config.frame_rate						)
    __A    :						Union[str, Any]						 =     True
    __A    :						Union[str, Any]						 =     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__":
       a_ =			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."""
       )

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











"""simple docstring"""



import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
       import os

       import jax.numpy as jnp
       from jax import jit

       from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
       from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

       a_ =			"""0.12"""  # assumed parallelism: 8


if is_torch_available():
       import torch







def 		__lowercase     (   snake_case_  :  Any				,snake_case_  :  int				,snake_case_  :  List[str]=None						)   ->Union[str, Any]:



    '''simple docstring'''







    if rng is None:
        __A    :						Union[str, Any]						 =     random.Random()

    __A    :						List[str]						 =     1
    for dim in shape:
        total_dims *= dim

    __A    :						List[str]						 =     []
    for _ in range(snake_case_						):
        values.append(rng.randint(0				,vocab_size - 1						)						)

    __A    :						Optional[Any]						 =     np.array(snake_case_				,dtype=jnp.intaa						).reshape(snake_case_						)

    return output







def 		__lowercase     (   snake_case_  :  int				,snake_case_  :  Union[str, Any]=None						)   ->int:



    '''simple docstring'''







    __A    :						int						 =     ids_tensor(snake_case_				,vocab_size=2				,rng=snake_case_						)
    # make sure that at least one token is attended to for each batch
    __A    :						str						 =     1
    return attn_mask






@require_flax
class 		__snake_case					:





 """simple docstring"""


 _lowerCamelCase        =     None
 _lowerCamelCase        =     ()

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A    :						Optional[int]						 =     self.model_tester.prepare_config_and_inputs_for_common()

     # cut to half length & take max batch_size 3
     __A    :						Optional[int]						 =     2
     __A    :						Tuple						 =     inputs['''input_ids'''].shape[-1] // 2
     __A    :						List[Any]						 =     inputs['''input_ids'''][:max_batch_size, :sequence_length]

     __A    :						Optional[Any]						 =     jnp.ones_like(__lowerCamelCase  )
     __A    :						List[Any]						 =     attention_mask[:max_batch_size, :sequence_length]

     # generate max 5 tokens
     __A    :						Optional[int]						 =     input_ids.shape[-1] + 5
     if config.eos_token_id is not None and config.pad_token_id is None:
         # hack to allow generate for models such as GPT2 as is done in `generate()`
         __A    :						Optional[int]						 =     config.eos_token_id
     return config, input_ids, attention_mask, max_length

 @is_pt_flax_cross_test
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						List[Any]						 =     self._get_input_ids_and_config()
     __A    :						List[Any]						 =     False
     __A    :						str						 =     max_length
     __A    :						Tuple						 =     0

     for model_class in self.all_generative_model_classes:
         __A    :						Union[str, Any]						 =     model_class(__lowerCamelCase  )

         __A    :						Any						 =     model_class.__name__[4:]  # Skip the "Flax" at the beginning
         __A    :						Dict						 =     getattr(__lowerCamelCase					,						__lowerCamelCase  )
         __A    :						Tuple						 =     pt_model_class(__lowerCamelCase  ).eval()
         __A    :						Optional[Any]						 =     load_flax_weights_in_pytorch_model(__lowerCamelCase					,						flax_model.params  )

         __A    :						str						 =     flax_model.generate(__lowerCamelCase  ).sequences
         __A    :						List[Any]						 =     pt_model.generate(torch.tensor(__lowerCamelCase					,						dtype=torch.long  )  )

         if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
             __A    :						Any						 =     flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

         self.assertListEqual(pt_generation_outputs.numpy().tolist()					,						flax_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Union[str, Any]						 =     self._get_input_ids_and_config()
     __A    :						List[str]						 =     False
     __A    :						Optional[Any]						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						Optional[int]						 =     model_class(__lowerCamelCase  )

         __A    :						Dict						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Union[str, Any]						 =     jit(model.generate  )
         __A    :						str						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						int						 =     self._get_input_ids_and_config()
     __A    :						List[Any]						 =     True
     __A    :						Union[str, Any]						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						Dict						 =     model_class(__lowerCamelCase  )

         __A    :						List[Any]						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						List[Any]						 =     jit(model.generate  )
         __A    :						List[str]						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						List[Any]						 =     self._get_input_ids_and_config()
     __A    :						List[str]						 =     False
     __A    :						Any						 =     max_length
     __A    :						List[Any]						 =     2

     for model_class in self.all_generative_model_classes:
         __A    :						Tuple						 =     model_class(__lowerCamelCase  )

         __A    :						Optional[Any]						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Union[str, Any]						 =     jit(model.generate  )
         __A    :						Optional[int]						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Tuple						 =     self._get_input_ids_and_config()
     __A    :						str						 =     False
     __A    :						Union[str, Any]						 =     max_length
     __A    :						List[Any]						 =     2
     __A    :						Any						 =     2

     for model_class in self.all_generative_model_classes:
         __A    :						List[str]						 =     model_class(__lowerCamelCase  )

         __A    :						Optional[int]						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[0]					,						input_ids.shape[0] * config.num_return_sequences  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Any						 =     self._get_input_ids_and_config()
     __A    :						Optional[Any]						 =     True
     __A    :						Union[str, Any]						 =     max_length
     __A    :						List[str]						 =     0.8
     __A    :						List[str]						 =     10
     __A    :						Union[str, Any]						 =     0.3
     __A    :						Union[str, Any]						 =     1
     __A    :						Optional[Any]						 =     8
     __A    :						Dict						 =     9

     for model_class in self.all_generative_model_classes:
         __A    :						List[Any]						 =     model_class(__lowerCamelCase  )

         __A    :						List[str]						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Any						 =     jit(model.generate  )
         __A    :						str						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Union[str, Any]						 =     self._get_input_ids_and_config()
     __A    :						Union[str, Any]						 =     max_length
     __A    :						List[str]						 =     1
     __A    :						str						 =     8
     __A    :						Any						 =     9

     for model_class in self.all_generative_model_classes:
         __A    :						Union[str, Any]						 =     model_class(__lowerCamelCase  )

         __A    :						Tuple						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Optional[Any]						 =     jit(model.generate  )
         __A    :						Optional[Any]						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						int						 =     self._get_input_ids_and_config()
     __A    :						Optional[int]						 =     max_length
     __A    :						List[str]						 =     2
     __A    :						List[Any]						 =     1
     __A    :						Optional[Any]						 =     8
     __A    :						str						 =     9

     for model_class in self.all_generative_model_classes:
         __A    :						Optional[int]						 =     model_class(__lowerCamelCase  )

         __A    :						int						 =     model.generate(__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						str						 =     jit(model.generate  )
         __A    :						Any						 =     jit_generate(__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						int						 =     self._get_input_ids_and_config()

     # pad attention mask on the left
     __A    :						Dict						 =     attention_mask.at[(0, 0)].set(0  )

     __A    :						Optional[Any]						 =     False
     __A    :						Union[str, Any]						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						int						 =     model_class(__lowerCamelCase  )

         __A    :						Union[str, Any]						 =     model.generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						Optional[int]						 =     jit(model.generate  )
         __A    :						Any						 =     jit_generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Optional[int]						 =     self._get_input_ids_and_config()

     # pad attention mask on the left
     __A    :						str						 =     attention_mask.at[(0, 0)].set(0  )

     __A    :						List[Any]						 =     True
     __A    :						Any						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						str						 =     model_class(__lowerCamelCase  )

         __A    :						List[Any]						 =     model.generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						List[str]						 =     jit(model.generate  )
         __A    :						Optional[int]						 =     jit_generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )







 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Optional[Any]						 =     self._get_input_ids_and_config()

     # pad attention mask on the left
     __A    :						List[str]						 =     attention_mask.at[(0, 0)].set(0  )

     __A    :						Optional[int]						 =     2
     __A    :						Dict						 =     max_length

     for model_class in self.all_generative_model_classes:
         __A    :						Any						 =     model_class(__lowerCamelCase  )

         __A    :						int						 =     model.generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences
         self.assertEqual(generation_outputs.shape[-1]					,						__lowerCamelCase  )

         __A    :						str						 =     jit(model.generate  )
         __A    :						str						 =     jit_generate(__lowerCamelCase					,						attention_mask=__lowerCamelCase  ).sequences

         self.assertListEqual(generation_outputs.tolist()					,						jit_generation_outputs.tolist()  )






@require_flax
class 		__snake_case					(		unittest.TestCase     ):





 """simple docstring"""


 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Optional[int]						 =     AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert'''  )
     __A    :						Optional[Any]						 =     FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''  )

     __A    :						Any						 =     '''Hello world'''
     __A    :						Dict						 =     tokenizer(__lowerCamelCase					,						return_tensors='''np'''  ).input_ids

     # typos are quickly detected (the correct argument is `do_sample`)
     with self.assertRaisesRegex(__lowerCamelCase					,						'''do_samples'''  ):
         model.generate(__lowerCamelCase					,						do_samples=__lowerCamelCase  )

     # arbitrary arguments that will not be used anywhere are also not accepted
     with self.assertRaisesRegex(__lowerCamelCase					,						'''foo'''  ):
         __A    :						Any						 =     {'''foo''': '''bar'''}
         model.generate(__lowerCamelCase					,						**__lowerCamelCase  )




