"""simple docstring"""






import shutil
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_tf_cross_test,
    require_tf,
    require_torch,
    require_torchvision,
    require_vision,
)
from transformers.utils import is_tf_available, is_torch_available, is_vision_available


if is_vision_available():
						from PIL import Image

						from transformers import AutoProcessor, SamImageProcessor, SamProcessor

if is_torch_available():
						import torch

if is_tf_available():
						import tensorflow as tf
@require_vision
@require_torchvision
class        snake_case (				unittest.TestCase       ):




					"""simple docstring"""

					def 	__lowerCAmelCase		(       self     :			Optional[Any]    ):
									UpperCAmelCase__    =							tempfile.mkdtemp()
									UpperCAmelCase__    =							SamImageProcessor()
									UpperCAmelCase__    =							SamProcessor(lowerCamelCase__    )
									processor.save_pretrained(self.tmpdirname    )







					def 	__lowerCAmelCase		(       self     :			str					,**lowerCamelCase__     :			Dict    ):
									return AutoProcessor.from_pretrained(self.tmpdirname					,**lowerCamelCase__    ).image_processor







					def 	__lowerCAmelCase		(       self     :			Optional[int]    ):
									shutil.rmtree(self.tmpdirname    )







					def 	__lowerCAmelCase		(       self     :			Dict    ):
									UpperCAmelCase__    =							[np.random.randint(255					,size=(3, 30, 400)					,dtype=np.uinta    )]

									UpperCAmelCase__    =							[Image.fromarray(np.moveaxis(lowerCamelCase__					,0					,-1    )    ) for x in image_inputs]

									return image_inputs







					def 	__lowerCAmelCase		(       self     :			Optional[Any]    ):
									UpperCAmelCase__    =							SamProcessor(image_processor=self.get_image_processor()    )
									processor.save_pretrained(self.tmpdirname    )

									UpperCAmelCase__    =							self.get_image_processor(do_normalize=lowerCamelCase__					,padding_value=1.0    )

									UpperCAmelCase__    =							SamProcessor.from_pretrained(self.tmpdirname					,do_normalize=lowerCamelCase__					,padding_value=1.0    )

									self.assertEqual(processor.image_processor.to_json_string()					,image_processor_add_kwargs.to_json_string()    )
									self.assertIsInstance(processor.image_processor					,lowerCamelCase__    )







					def 	__lowerCAmelCase		(       self     :			Optional[int]    ):
									UpperCAmelCase__    =							self.get_image_processor()

									UpperCAmelCase__    =							SamProcessor(image_processor=lowerCamelCase__    )

									UpperCAmelCase__    =							self.prepare_image_inputs()

									UpperCAmelCase__    =							image_processor(lowerCamelCase__					,return_tensors='np'    )
									UpperCAmelCase__    =							processor(images=lowerCamelCase__					,return_tensors='np'    )

									input_feat_extract.pop('original_sizes'    )  # pop original_sizes as it is popped in the processor
									input_feat_extract.pop('reshaped_input_sizes'    )  # pop original_sizes as it is popped in the processor

									for key in input_feat_extract.keys():
													self.assertAlmostEqual(input_feat_extract[key].sum()					,input_processor[key].sum()					,delta=1e-2    )







					@require_torch
					def 	__lowerCAmelCase		(       self     :			Dict    ):
									UpperCAmelCase__    =							self.get_image_processor()

									UpperCAmelCase__    =							SamProcessor(image_processor=lowerCamelCase__    )
									UpperCAmelCase__    =							[torch.ones((1, 3, 5, 5)    )]

									UpperCAmelCase__    =							[[1_764, 2_646]]

									UpperCAmelCase__    =							[[683, 1_024]]
									UpperCAmelCase__    =							processor.post_process_masks(lowerCamelCase__					,lowerCamelCase__					,lowerCamelCase__    )
									self.assertEqual(masks[0].shape					,(1, 3, 1_764, 2_646)    )

									UpperCAmelCase__    =							processor.post_process_masks(
									    lowerCamelCase__					,torch.tensor(lowerCamelCase__    )					,torch.tensor(lowerCamelCase__    )    )
									self.assertEqual(masks[0].shape					,(1, 3, 1_764, 2_646)    )

									# should also work with np
									UpperCAmelCase__    =							[np.ones((1, 3, 5, 5)    )]
									UpperCAmelCase__    =							processor.post_process_masks(lowerCamelCase__					,np.array(lowerCamelCase__    )					,np.array(lowerCamelCase__    )    )

									self.assertEqual(masks[0].shape					,(1, 3, 1_764, 2_646)    )

									UpperCAmelCase__    =							[[1, 0], [0, 1]]
									with self.assertRaises(lowerCamelCase__    ):
													UpperCAmelCase__    =							processor.post_process_masks(lowerCamelCase__					,np.array(lowerCamelCase__    )					,np.array(lowerCamelCase__    )    )






@require_vision
@require_tf
class        snake_case (				unittest.TestCase       ):




					"""simple docstring"""

					def 	__lowerCAmelCase		(       self     :			List[str]    ):
									UpperCAmelCase__    =							tempfile.mkdtemp()
									UpperCAmelCase__    =							SamImageProcessor()
									UpperCAmelCase__    =							SamProcessor(lowerCamelCase__    )
									processor.save_pretrained(self.tmpdirname    )







					def 	__lowerCAmelCase		(       self     :			str					,**lowerCamelCase__     :			Union[str, Any]    ):
									return AutoProcessor.from_pretrained(self.tmpdirname					,**lowerCamelCase__    ).image_processor







					def 	__lowerCAmelCase		(       self     :			List[Any]    ):
									shutil.rmtree(self.tmpdirname    )







					def 	__lowerCAmelCase		(       self     :			Any    ):
									UpperCAmelCase__    =							[np.random.randint(255					,size=(3, 30, 400)					,dtype=np.uinta    )]

									UpperCAmelCase__    =							[Image.fromarray(np.moveaxis(lowerCamelCase__					,0					,-1    )    ) for x in image_inputs]

									return image_inputs







					def 	__lowerCAmelCase		(       self     :			Optional[int]    ):
									UpperCAmelCase__    =							SamProcessor(image_processor=self.get_image_processor()    )
									processor.save_pretrained(self.tmpdirname    )

									UpperCAmelCase__    =							self.get_image_processor(do_normalize=lowerCamelCase__					,padding_value=1.0    )

									UpperCAmelCase__    =							SamProcessor.from_pretrained(self.tmpdirname					,do_normalize=lowerCamelCase__					,padding_value=1.0    )

									self.assertEqual(processor.image_processor.to_json_string()					,image_processor_add_kwargs.to_json_string()    )
									self.assertIsInstance(processor.image_processor					,lowerCamelCase__    )







					def 	__lowerCAmelCase		(       self     :			Union[str, Any]    ):
									UpperCAmelCase__    =							self.get_image_processor()

									UpperCAmelCase__    =							SamProcessor(image_processor=lowerCamelCase__    )

									UpperCAmelCase__    =							self.prepare_image_inputs()

									UpperCAmelCase__    =							image_processor(lowerCamelCase__					,return_tensors='np'    )
									UpperCAmelCase__    =							processor(images=lowerCamelCase__					,return_tensors='np'    )

									input_feat_extract.pop('original_sizes'    )  # pop original_sizes as it is popped in the processor
									input_feat_extract.pop('reshaped_input_sizes'    )  # pop reshaped_input_sizes as it is popped in the processor

									for key in input_feat_extract.keys():
													self.assertAlmostEqual(input_feat_extract[key].sum()					,input_processor[key].sum()					,delta=1e-2    )







					@require_tf
					def 	__lowerCAmelCase		(       self     :			Tuple    ):
									UpperCAmelCase__    =							self.get_image_processor()

									UpperCAmelCase__    =							SamProcessor(image_processor=lowerCamelCase__    )
									UpperCAmelCase__    =							[tf.ones((1, 3, 5, 5)    )]

									UpperCAmelCase__    =							[[1_764, 2_646]]

									UpperCAmelCase__    =							[[683, 1_024]]
									UpperCAmelCase__    =							processor.post_process_masks(lowerCamelCase__					,lowerCamelCase__					,lowerCamelCase__					,return_tensors='tf'    )
									self.assertEqual(masks[0].shape					,(1, 3, 1_764, 2_646)    )

									UpperCAmelCase__    =							processor.post_process_masks(
									    lowerCamelCase__					,tf.convert_to_tensor(lowerCamelCase__    )					,tf.convert_to_tensor(lowerCamelCase__    )					,return_tensors='tf'					,)
									self.assertEqual(masks[0].shape					,(1, 3, 1_764, 2_646)    )

									# should also work with np
									UpperCAmelCase__    =							[np.ones((1, 3, 5, 5)    )]
									UpperCAmelCase__    =							processor.post_process_masks(
									    lowerCamelCase__					,np.array(lowerCamelCase__    )					,np.array(lowerCamelCase__    )					,return_tensors='tf'    )

									self.assertEqual(masks[0].shape					,(1, 3, 1_764, 2_646)    )

									UpperCAmelCase__    =							[[1, 0], [0, 1]]
									with self.assertRaises(tf.errors.InvalidArgumentError    ):
													UpperCAmelCase__    =							processor.post_process_masks(
													    lowerCamelCase__					,np.array(lowerCamelCase__    )					,np.array(lowerCamelCase__    )					,return_tensors='tf'    )






@require_vision
@require_torchvision
class        snake_case (				unittest.TestCase       ):




					"""simple docstring"""

					def 	__lowerCAmelCase		(       self     :			Any    ):
									UpperCAmelCase__    =							tempfile.mkdtemp()
									UpperCAmelCase__    =							SamImageProcessor()
									UpperCAmelCase__    =							SamProcessor(lowerCamelCase__    )
									processor.save_pretrained(self.tmpdirname    )







					def 	__lowerCAmelCase		(       self     :			Dict					,**lowerCamelCase__     :			Any    ):
									return AutoProcessor.from_pretrained(self.tmpdirname					,**lowerCamelCase__    ).image_processor







					def 	__lowerCAmelCase		(       self     :			Optional[Any]    ):
									shutil.rmtree(self.tmpdirname    )







					def 	__lowerCAmelCase		(       self     :			List[str]    ):
									UpperCAmelCase__    =							[np.random.randint(255					,size=(3, 30, 400)					,dtype=np.uinta    )]

									UpperCAmelCase__    =							[Image.fromarray(np.moveaxis(lowerCamelCase__					,0					,-1    )    ) for x in image_inputs]

									return image_inputs







					@is_pt_tf_cross_test
					def 	__lowerCAmelCase		(       self     :			List[Any]    ):
									UpperCAmelCase__    =							self.get_image_processor()

									UpperCAmelCase__    =							SamProcessor(image_processor=lowerCamelCase__    )
									UpperCAmelCase__    =							np.random.randint(0					,2					,size=(1, 3, 5, 5)    ).astype(np.floataa    )
									UpperCAmelCase__    =							[tf.convert_to_tensor(lowerCamelCase__    )]
									UpperCAmelCase__    =							[torch.tensor(lowerCamelCase__    )]

									UpperCAmelCase__    =							[[1_764, 2_646]]

									UpperCAmelCase__    =							[[683, 1_024]]
									UpperCAmelCase__    =							processor.post_process_masks(
									    lowerCamelCase__					,lowerCamelCase__					,lowerCamelCase__					,return_tensors='tf'    )
									UpperCAmelCase__    =							processor.post_process_masks(
									    lowerCamelCase__					,lowerCamelCase__					,lowerCamelCase__					,return_tensors='pt'    )

									self.assertTrue(np.all(tf_masks[0].numpy() == pt_masks[0].numpy()    )    )







					@is_pt_tf_cross_test
					def 	__lowerCAmelCase		(       self     :			Optional[Any]    ):
									UpperCAmelCase__    =							self.get_image_processor()

									UpperCAmelCase__    =							SamProcessor(image_processor=lowerCamelCase__    )

									UpperCAmelCase__    =							self.prepare_image_inputs()

									UpperCAmelCase__    =							image_processor(lowerCamelCase__					,return_tensors='pt'    )['pixel_values'].numpy()
									UpperCAmelCase__    =							processor(images=lowerCamelCase__					,return_tensors='pt'    )['pixel_values'].numpy()

									UpperCAmelCase__    =							image_processor(lowerCamelCase__					,return_tensors='tf'    )['pixel_values'].numpy()
									UpperCAmelCase__    =							processor(images=lowerCamelCase__					,return_tensors='tf'    )['pixel_values'].numpy()

									self.assertTrue(np.allclose(lowerCamelCase__					,lowerCamelCase__    )    )
									self.assertTrue(np.allclose(lowerCamelCase__					,lowerCamelCase__    )    )
									self.assertTrue(np.allclose(lowerCamelCase__					,lowerCamelCase__    )    )




"""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()
lowerCAmelCase__						:						Optional[int]		  =			logging.get_logger(__name__)


lowerCAmelCase__						:						Optional[Any]		  =			['model.decoder.embed_positions.weights']



def 						a_		(      lowerCamelCase ):
				if "emb" in name:
								UpperCAmelCase__    =							name.replace('emb'      ,    'model.decoder.embed_tokens' )
				if "transformer" in name:
								UpperCAmelCase__    =							name.replace('transformer'      ,    'model.decoder' )
				if "cross_attention" in name:
								UpperCAmelCase__    =							name.replace('cross_attention'      ,    'encoder_attn' )
				if "linear1" in name:
								UpperCAmelCase__    =							name.replace('linear1'      ,    'fc1' )
				if "linear2" in name:
								UpperCAmelCase__    =							name.replace('linear2'      ,    'fc2' )
				if "norm1" in name:
								UpperCAmelCase__    =							name.replace('norm1'      ,    'self_attn_layer_norm' )
				if "norm_cross" in name:
								UpperCAmelCase__    =							name.replace('norm_cross'      ,    'encoder_attn_layer_norm' )
				if "norm2" in name:
								UpperCAmelCase__    =							name.replace('norm2'      ,    'final_layer_norm' )
				if "out_norm" in name:
								UpperCAmelCase__    =							name.replace('out_norm'      ,    'model.decoder.layer_norm' )
				if "linears" in name:
								UpperCAmelCase__    =							name.replace('linears'      ,    'lm_heads' )
				if "condition_provider.conditioners.description.output_proj" in name:
								UpperCAmelCase__    =							name.replace('condition_provider.conditioners.description.output_proj'      ,    'enc_to_dec_proj' )
				return name



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



def 						a_		(      lowerCamelCase ):
				if checkpoint == "small":
								# default config values
								UpperCAmelCase__    =							1_0_2_4
								UpperCAmelCase__    =							2_4
								UpperCAmelCase__    =							1_6
				elif checkpoint == "medium":
								UpperCAmelCase__    =							1_5_3_6
								UpperCAmelCase__    =							4_8
								UpperCAmelCase__    =							2_4
				elif checkpoint == "large":
								UpperCAmelCase__    =							2_0_4_8
								UpperCAmelCase__    =							4_8
								UpperCAmelCase__    =							3_2
				else:
								raise ValueError(f'''Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}.''' )
				UpperCAmelCase__    =							MusicgenDecoderConfig(
				    hidden_size=lowerCamelCase      ,    ffn_dim=hidden_size * 4      ,    num_hidden_layers=lowerCamelCase      ,    num_attention_heads=lowerCamelCase      ,    )
				return config



@torch.no_grad()
def 						a_		(      lowerCamelCase      ,    lowerCamelCase=None      ,    lowerCamelCase=None      ,    lowerCamelCase="cpu" ):
				UpperCAmelCase__    =							MusicGen.get_pretrained(lowerCamelCase      ,    device=lowerCamelCase )
				UpperCAmelCase__    =							decoder_config_from_checkpoint(lowerCamelCase )

				UpperCAmelCase__    =							fairseq_model.lm.state_dict()
				UpperCAmelCase__   ,			UpperCAmelCase__    =							rename_state_dict(
				    lowerCamelCase      ,    hidden_size=decoder_config.hidden_size )

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

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

				with torch.no_grad():
								UpperCAmelCase__    =							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__    =							AutoTokenizer.from_pretrained('t5-base' )
				UpperCAmelCase__    =							AutoFeatureExtractor.from_pretrained('facebook/encodec_32khz'      ,    padding_side='left' )

				UpperCAmelCase__    =							MusicgenProcessor(feature_extractor=lowerCamelCase      ,    tokenizer=lowerCamelCase )

				# set the appropriate bos/pad token ids
				UpperCAmelCase__    =							2_0_4_8
				UpperCAmelCase__    =							2_0_4_8

				# set other default generation config params
				UpperCAmelCase__    =							int(3_0 * audio_encoder.config.frame_rate )
				UpperCAmelCase__    =							True
				UpperCAmelCase__    =							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__						:						Tuple		  =			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__						:						List[str]		  =			parser.parse_args()
						convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)



