


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


a_						:       str      			=				['''model.decoder.embed_positions.weights''']





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





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





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





@torch.no_grad()
def __snake_case      (   UpperCAmelCase_		:		str      ,       UpperCAmelCase_		:		Optional[Any]=None      ,       UpperCAmelCase_		:		Union[str, Any]=None      ,       UpperCAmelCase_		:		int="cpu"	):
 lowerCamelCase_  				=						MusicGen.get_pretrained(_UpperCAmelCase      ,       device=_UpperCAmelCase	)
 lowerCamelCase_  				=						decoder_config_from_checkpoint(_UpperCAmelCase	)

 lowerCamelCase_  				=						fairseq_model.lm.state_dict()
 lowerCamelCase_				,lowerCamelCase_  				=						rename_state_dict(
     _UpperCAmelCase      ,       hidden_size=decoder_config.hidden_size	)

 lowerCamelCase_  				=						TaEncoderModel.from_pretrained("t5-base"	)
 lowerCamelCase_  				=						EncodecModel.from_pretrained("facebook/encodec_32khz"	)
 lowerCamelCase_  				=						MusicgenForCausalLM(_UpperCAmelCase	).eval()

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

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

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

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

 # init the composite model
 lowerCamelCase_  				=						MusicgenForConditionalGeneration(text_encoder=_UpperCAmelCase      ,       audio_encoder=_UpperCAmelCase      ,       decoder=_UpperCAmelCase	)

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

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

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

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

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

 lowerCamelCase_  				=						MusicgenProcessor(feature_extractor=_UpperCAmelCase      ,       tokenizer=_UpperCAmelCase	)

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

 # set other default generation config params
 lowerCamelCase_  				=						int(30 * audio_encoder.config.frame_rate	)
 lowerCamelCase_  				=						True
 lowerCamelCase_  				=						3.0

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

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


if __name__ == "__main__":
     a_						:       int      			=				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_						:       Optional[Any]      			=				parser.parse_args()
     convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)







'''simple docstring'''

import copy
import inspect
import unittest

from transformers import AutoBackbone
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import require_timm, require_torch, torch_device
from transformers.utils.import_utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor


if is_torch_available():
     import torch

     from transformers import TimmBackbone, TimmBackboneConfig

from ...test_pipeline_mixin import PipelineTesterMixin





class   snake_case					:





  """simple docstring"""




  def __init__(					self    ,						UpperCamelCase    ,						UpperCamelCase=None    ,						UpperCamelCase=None    ,						UpperCamelCase=None    ,						UpperCamelCase="resnet50"    ,						UpperCamelCase=3    ,						UpperCamelCase=32    ,						UpperCamelCase=3    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						):







   """simple docstring"""
   lowerCamelCase_  				=						parent
   lowerCamelCase_  				=						out_indices if out_indices is not None else [4]
   lowerCamelCase_  				=						stage_names
   lowerCamelCase_  				=						out_features
   lowerCamelCase_  				=						backbone
   lowerCamelCase_  				=						batch_size
   lowerCamelCase_  				=						image_size
   lowerCamelCase_  				=						num_channels
   lowerCamelCase_  				=						use_pretrained_backbone
   lowerCamelCase_  				=						is_training






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )
   lowerCamelCase_  				=						self.get_config()

   return config, pixel_values






  def 						snake_case	(					self  ):







   """simple docstring"""
   return TimmBackboneConfig(
       image_size=self.image_size    ,						num_channels=self.num_channels    ,						out_features=self.out_features    ,						out_indices=self.out_indices    ,						stage_names=self.stage_names    ,						use_pretrained_backbone=self.use_pretrained_backbone    ,						backbone=self.backbone    ,						)






  def 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						TimmBackbone(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   with torch.no_grad():
    lowerCamelCase_  				=						model(UpperCamelCase  )
   self.parent.assertEqual(
       result.feature_map[-1].shape    ,						(self.batch_size, model.channels[-1], 14, 14)    ,						)







  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.prepare_config_and_inputs()
   lowerCamelCase_				,lowerCamelCase_  				=						config_and_inputs
   lowerCamelCase_  				=						{"pixel_values": pixel_values}
   return config, inputs_dict





@require_torch
@require_timm
class   snake_case					(					lowercase    ,		lowercase    ,		lowercase    ,		unittest.TestCase    ):





  """simple docstring"""




  _lowerCamelCase						=						(TimmBackbone,) if is_torch_available() else ()
  _lowerCamelCase						=						{"feature-extraction": TimmBackbone} if is_torch_available() else {}
  _lowerCamelCase						=						False
  _lowerCamelCase						=						False
  _lowerCamelCase						=						False
  _lowerCamelCase						=						False






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						TimmBackboneModelTester(self  )
   lowerCamelCase_  				=						ConfigTester(self    ,						config_class=UpperCamelCase    ,						has_text_modality=UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   self.config_tester.create_and_test_config_to_json_string()
   self.config_tester.create_and_test_config_to_json_file()
   self.config_tester.create_and_test_config_from_and_save_pretrained()
   self.config_tester.create_and_test_config_with_num_labels()
   self.config_tester.check_config_can_be_init_without_params()
   self.config_tester.check_config_arguments_init()






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						"resnet18"
   lowerCamelCase_  				=						"microsoft/resnet-18"

   lowerCamelCase_  				=						AutoBackbone.from_pretrained(UpperCamelCase    ,						use_timm_backbone=UpperCamelCase  )
   lowerCamelCase_  				=						AutoBackbone.from_pretrained(UpperCamelCase  )

   self.assertEqual(len(timm_model.out_features  )    ,						len(transformers_model.out_features  )  )
   self.assertEqual(len(timm_model.stage_names  )    ,						len(transformers_model.stage_names  )  )
   self.assertEqual(timm_model.channels    ,						transformers_model.channels  )
   # Out indices are set to the last layer by default. For timm models, we don't know
   # the number of layers in advance, so we set it to (-1,), whereas for transformers
   # models, we set it to [len(stage_names) - 1] (kept for backward compatibility).
   self.assertEqual(timm_model.out_indices    ,						(-1,)  )
   self.assertEqual(transformers_model.out_indices    ,						[len(timm_model.stage_names  ) - 1]  )

   lowerCamelCase_  				=						AutoBackbone.from_pretrained(UpperCamelCase    ,						use_timm_backbone=UpperCamelCase    ,						out_indices=[1, 2, 3]  )
   lowerCamelCase_  				=						AutoBackbone.from_pretrained(UpperCamelCase    ,						out_indices=[1, 2, 3]  )

   self.assertEqual(timm_model.out_indices    ,						transformers_model.out_indices  )
   self.assertEqual(len(timm_model.out_features  )    ,						len(transformers_model.out_features  )  )
   self.assertEqual(timm_model.channels    ,						transformers_model.channels  )






  @unittest.skip("TimmBackbone doesn't support feed forward chunking"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone doesn't have num_hidden_layers attribute"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone initialization is managed on the timm side"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone models doesn't have inputs_embeds"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone models doesn't have inputs_embeds"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone model cannot be created without specifying a backbone checkpoint"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("Only checkpoints on timm can be loaded into TimmBackbone"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("model weights aren't tied in TimmBackbone."  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("model weights aren't tied in TimmBackbone."  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("Only checkpoints on timm can be loaded into TimmBackbone"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("Only checkpoints on timm can be loaded into TimmBackbone"  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone doesn't have hidden size info in its configuration."  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("TimmBackbone doesn't support output_attentions."  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("Safetensors is not supported by timm."  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."  )
  def 						snake_case	(					self  ):







   """simple docstring"""
   pass






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_				,lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    lowerCamelCase_  				=						model_class(UpperCamelCase  )
    lowerCamelCase_  				=						inspect.signature(model.forward  )
    # signature.parameters is an OrderedDict => so arg_names order is deterministic
    lowerCamelCase_  				=						[*signature.parameters.keys()]

    lowerCamelCase_  				=						["pixel_values"]
    self.assertListEqual(arg_names[:1]    ,						UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_				,lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs_for_common()
   lowerCamelCase_  				=						True
   lowerCamelCase_  				=						self.has_attentions

   # no need to test all models as different heads yield the same functionality
   lowerCamelCase_  				=						self.all_model_classes[0]
   lowerCamelCase_  				=						model_class(UpperCamelCase  )
   model.to(UpperCamelCase  )

   lowerCamelCase_  				=						self._prepare_for_class(UpperCamelCase    ,						UpperCamelCase  )
   lowerCamelCase_  				=						model(**UpperCamelCase  )
   lowerCamelCase_  				=						outputs[0][-1]

   # Encoder-/Decoder-only models
   lowerCamelCase_  				=						outputs.hidden_states[0]
   hidden_states.retain_grad()

   if self.has_attentions:
    lowerCamelCase_  				=						outputs.attentions[0]
    attentions.retain_grad()

   output.flatten()[0].backward(retain_graph=UpperCamelCase  )

   self.assertIsNotNone(hidden_states.grad  )

   if self.has_attentions:
    self.assertIsNotNone(attentions.grad  )







  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_				,lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    lowerCamelCase_  				=						model_class(UpperCamelCase  )
    model.to(UpperCamelCase  )
    model.eval()
    lowerCamelCase_  				=						model(**UpperCamelCase  )

    self.assertEqual(len(result.feature_maps  )    ,						len(config.out_indices  )  )
    self.assertEqual(len(model.channels  )    ,						len(config.out_indices  )  )

    # Check output of last stage is taken if out_features=None, out_indices=None
    lowerCamelCase_  				=						copy.deepcopy(UpperCamelCase  )
    lowerCamelCase_  				=						None
    lowerCamelCase_  				=						model_class(UpperCamelCase  )
    model.to(UpperCamelCase  )
    model.eval()
    lowerCamelCase_  				=						model(**UpperCamelCase  )

    self.assertEqual(len(result.feature_maps  )    ,						1  )
    self.assertEqual(len(model.channels  )    ,						1  )

    # Check backbone can be initialized with fresh weights
    lowerCamelCase_  				=						copy.deepcopy(UpperCamelCase  )
    lowerCamelCase_  				=						False
    lowerCamelCase_  				=						model_class(UpperCamelCase  )
    model.to(UpperCamelCase  )
    model.eval()
    lowerCamelCase_  				=						model(**UpperCamelCase  )



