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


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


def A				(__lowerCamelCase				: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 A				(__lowerCamelCase				:OrderedDict			,     __lowerCamelCase				:int		):
						_lowerCAmelCase  						=			list(state_dict.keys()		)
						_lowerCAmelCase  						=			{}
						for key in keys:
												_lowerCAmelCase  						=			state_dict.pop(__lowerCamelCase		)
												_lowerCAmelCase  						=			rename_keys(__lowerCamelCase		)
												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 A				(__lowerCamelCase				: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=__lowerCamelCase			,     ffn_dim=hidden_size * 4			,     num_hidden_layers=__lowerCamelCase			,     num_attention_heads=__lowerCamelCase			,     )
						return config


@torch.no_grad()
def A				(__lowerCamelCase				:str			,     __lowerCamelCase				:List[Any]=None			,     __lowerCamelCase				:Dict=None			,     __lowerCamelCase				:Optional[Any]="cpu"		):
						_lowerCAmelCase  						=			MusicGen.get_pretrained(__lowerCamelCase			,     device=__lowerCamelCase		)
						_lowerCAmelCase  						=			decoder_config_from_checkpoint(__lowerCamelCase		)

						_lowerCAmelCase  						=			fairseq_model.lm.state_dict()
						_lowerCAmelCase						,   _lowerCAmelCase  						=			rename_state_dict(
						    __lowerCamelCase			,     hidden_size=decoder_config.hidden_size		)

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

						# load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
						_lowerCAmelCase						,   _lowerCAmelCase  						=			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
						_lowerCAmelCase  						=			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
						_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=__lowerCamelCase			,     decoder_input_ids=__lowerCamelCase		).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=__lowerCamelCase			,     tokenizer=__lowerCamelCase		)

						# 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(__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__":
						_lowercase    =		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."""
						)

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


'''simple docstring'''



from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)




@flax.struct.dataclass
class   UpperCAmelCase_					(		_SCREAMING_SNAKE_CASE			):






						'''simple docstring'''






						_lowercase		:				jnp.ndarray






@flax_register_to_config
class   UpperCAmelCase_					(		nn.Module	,							_SCREAMING_SNAKE_CASE	,							_SCREAMING_SNAKE_CASE			):






						'''simple docstring'''






						_lowercase		:				int									=       3_2
						_lowercase		:				int									=       4
						_lowercase		:				int									=       4
						_lowercase		:				Tuple[str]									=       (
						    "CrossAttnDownBlock2D",
						    "CrossAttnDownBlock2D",
						    "CrossAttnDownBlock2D",
						    "DownBlock2D",
						)
						_lowercase		:				Tuple[str]									=       ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
						_lowercase		:				Union[bool, Tuple[bool]]									=       False
						_lowercase		:				Tuple[int]									=       (3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
						_lowercase		:				int									=       2
						_lowercase		:				Union[int, Tuple[int]]									=       8
						_lowercase		:				Optional[Union[int, Tuple[int]]]									=       None
						_lowercase		:				int									=       1_2_8_0
						_lowercase		:				float									=       0.0
						_lowercase		:				bool									=       False
						_lowercase		:				jnp.dtype									=       jnp.floataa
						_lowercase		:				bool									=       True
						_lowercase		:				int									=       0
						_lowercase		:				bool									=       False




						def  _lowercase     ( self   ,       _lowercase				):

												"""simple docstring"""



												_lowerCAmelCase  						=			(1, self.in_channels, self.sample_size, self.sample_size)
												_lowerCAmelCase  						=			jnp.zeros(_lowercase   ,       dtype=jnp.floataa				)
												_lowerCAmelCase  						=			jnp.ones((1,)   ,       dtype=jnp.intaa				)
												_lowerCAmelCase  						=			jnp.zeros((1, 1, self.cross_attention_dim)   ,       dtype=jnp.floataa				)

												_lowerCAmelCase						,   _lowerCAmelCase  						=			jax.random.split(_lowercase				)
												_lowerCAmelCase  						=			{"""params""": params_rng, """dropout""": dropout_rng}

												return self.init(_lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase				)["params"]




						def  _lowercase     ( self				):

												"""simple docstring"""



												_lowerCAmelCase  						=			self.block_out_channels
												_lowerCAmelCase  						=			block_out_channels[0] * 4

												if self.num_attention_heads is not None:
																		raise ValueError(
																		    """At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."""				)

												# If `num_attention_heads` is not defined (which is the case for most models)
												# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
												# The reason for this behavior is to correct for incorrectly named variables that were introduced
												# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
												# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
												# which is why we correct for the naming here.
												_lowerCAmelCase  						=			self.num_attention_heads or self.attention_head_dim

												# input
												_lowerCAmelCase  						=			nn.Conv(
												    block_out_channels[0]   ,       kernel_size=(3, 3)   ,       strides=(1, 1)   ,       padding=((1, 1), (1, 1))   ,       dtype=self.dtype   ,       )

												# time
												_lowerCAmelCase  						=			FlaxTimesteps(
												    block_out_channels[0]   ,       flip_sin_to_cos=self.flip_sin_to_cos   ,       freq_shift=self.config.freq_shift				)
												_lowerCAmelCase  						=			FlaxTimestepEmbedding(_lowercase   ,       dtype=self.dtype				)

												_lowerCAmelCase  						=			self.only_cross_attention
												if isinstance(_lowercase   ,       _lowercase				):
																		_lowerCAmelCase  						=			(only_cross_attention,) * len(self.down_block_types				)

												if isinstance(_lowercase   ,       _lowercase				):
																		_lowerCAmelCase  						=			(num_attention_heads,) * len(self.down_block_types				)

												# down
												_lowerCAmelCase  						=			[]
												_lowerCAmelCase  						=			block_out_channels[0]
												for i, down_block_type in enumerate(self.down_block_types				):
																		_lowerCAmelCase  						=			output_channel
																		_lowerCAmelCase  						=			block_out_channels[i]
																		_lowerCAmelCase  						=			i == len(_lowercase				) - 1

																		if down_block_type == "CrossAttnDownBlock2D":
																								_lowerCAmelCase  						=			FlaxCrossAttnDownBlockaD(
																								    in_channels=_lowercase   ,       out_channels=_lowercase   ,       dropout=self.dropout   ,       num_layers=self.layers_per_block   ,       num_attention_heads=num_attention_heads[i]   ,       add_downsample=not is_final_block   ,       use_linear_projection=self.use_linear_projection   ,       only_cross_attention=only_cross_attention[i]   ,       use_memory_efficient_attention=self.use_memory_efficient_attention   ,       dtype=self.dtype   ,       )
																		else:
																								_lowerCAmelCase  						=			FlaxDownBlockaD(
																								    in_channels=_lowercase   ,       out_channels=_lowercase   ,       dropout=self.dropout   ,       num_layers=self.layers_per_block   ,       add_downsample=not is_final_block   ,       dtype=self.dtype   ,       )

																		down_blocks.append(_lowercase				)
												_lowerCAmelCase  						=			down_blocks

												# mid
												_lowerCAmelCase  						=			FlaxUNetMidBlockaDCrossAttn(
												    in_channels=block_out_channels[-1]   ,       dropout=self.dropout   ,       num_attention_heads=num_attention_heads[-1]   ,       use_linear_projection=self.use_linear_projection   ,       use_memory_efficient_attention=self.use_memory_efficient_attention   ,       dtype=self.dtype   ,       )

												# up
												_lowerCAmelCase  						=			[]
												_lowerCAmelCase  						=			list(reversed(_lowercase				)				)
												_lowerCAmelCase  						=			list(reversed(_lowercase				)				)
												_lowerCAmelCase  						=			list(reversed(_lowercase				)				)
												_lowerCAmelCase  						=			reversed_block_out_channels[0]
												for i, up_block_type in enumerate(self.up_block_types				):
																		_lowerCAmelCase  						=			output_channel
																		_lowerCAmelCase  						=			reversed_block_out_channels[i]
																		_lowerCAmelCase  						=			reversed_block_out_channels[min(i + 1   ,       len(_lowercase				) - 1				)]

																		_lowerCAmelCase  						=			i == len(_lowercase				) - 1

																		if up_block_type == "CrossAttnUpBlock2D":
																								_lowerCAmelCase  						=			FlaxCrossAttnUpBlockaD(
																								    in_channels=_lowercase   ,       out_channels=_lowercase   ,       prev_output_channel=_lowercase   ,       num_layers=self.layers_per_block + 1   ,       num_attention_heads=reversed_num_attention_heads[i]   ,       add_upsample=not is_final_block   ,       dropout=self.dropout   ,       use_linear_projection=self.use_linear_projection   ,       only_cross_attention=only_cross_attention[i]   ,       use_memory_efficient_attention=self.use_memory_efficient_attention   ,       dtype=self.dtype   ,       )
																		else:
																								_lowerCAmelCase  						=			FlaxUpBlockaD(
																								    in_channels=_lowercase   ,       out_channels=_lowercase   ,       prev_output_channel=_lowercase   ,       num_layers=self.layers_per_block + 1   ,       add_upsample=not is_final_block   ,       dropout=self.dropout   ,       dtype=self.dtype   ,       )

																		up_blocks.append(_lowercase				)
																		_lowerCAmelCase  						=			output_channel
												_lowerCAmelCase  						=			up_blocks

												# out
												_lowerCAmelCase  						=			nn.GroupNorm(num_groups=32   ,       epsilon=1e-5				)
												_lowerCAmelCase  						=			nn.Conv(
												    self.out_channels   ,       kernel_size=(3, 3)   ,       strides=(1, 1)   ,       padding=((1, 1), (1, 1))   ,       dtype=self.dtype   ,       )




						def __call__( self   ,       _lowercase   ,       _lowercase   ,       _lowercase   ,       _lowercase=None   ,       _lowercase=None   ,       _lowercase = True   ,       _lowercase = False   ,       ):

												"""simple docstring"""



												if not isinstance(_lowercase   ,       jnp.ndarray				):
																		_lowerCAmelCase  						=			jnp.array([timesteps]   ,       dtype=jnp.intaa				)
												elif isinstance(_lowercase   ,       jnp.ndarray				) and len(timesteps.shape				) == 0:
																		_lowerCAmelCase  						=			timesteps.astype(dtype=jnp.floataa				)
																		_lowerCAmelCase  						=			jnp.expand_dims(_lowercase   ,       0				)

												_lowerCAmelCase  						=			self.time_proj(_lowercase				)
												_lowerCAmelCase  						=			self.time_embedding(_lowercase				)

												# 2. pre-process
												_lowerCAmelCase  						=			jnp.transpose(_lowercase   ,       (0, 2, 3, 1)				)
												_lowerCAmelCase  						=			self.conv_in(_lowercase				)

												# 3. down
												_lowerCAmelCase  						=			(sample,)
												for down_block in self.down_blocks:
																		if isinstance(_lowercase   ,       _lowercase				):
																								_lowerCAmelCase						,   _lowerCAmelCase  						=			down_block(_lowercase   ,       _lowercase   ,       _lowercase   ,       deterministic=not train				)
																		else:
																								_lowerCAmelCase						,   _lowerCAmelCase  						=			down_block(_lowercase   ,       _lowercase   ,       deterministic=not train				)
																		down_block_res_samples += res_samples

												if down_block_additional_residuals is not None:
																		_lowerCAmelCase  						=			()

																		for down_block_res_sample, down_block_additional_residual in zip(
																		    _lowercase   ,       _lowercase				):
																								down_block_res_sample += down_block_additional_residual
																								new_down_block_res_samples += (down_block_res_sample,)

																		_lowerCAmelCase  						=			new_down_block_res_samples

												# 4. mid
												_lowerCAmelCase  						=			self.mid_block(_lowercase   ,       _lowercase   ,       _lowercase   ,       deterministic=not train				)

												if mid_block_additional_residual is not None:
																		sample += mid_block_additional_residual

												# 5. up
												for up_block in self.up_blocks:
																		_lowerCAmelCase  						=			down_block_res_samples[-(self.layers_per_block + 1) :]
																		_lowerCAmelCase  						=			down_block_res_samples[: -(self.layers_per_block + 1)]
																		if isinstance(_lowercase   ,       _lowercase				):
																								_lowerCAmelCase  						=			up_block(
																								    _lowercase   ,       temb=_lowercase   ,       encoder_hidden_states=_lowercase   ,       res_hidden_states_tuple=_lowercase   ,       deterministic=not train   ,       )
																		else:
																								_lowerCAmelCase  						=			up_block(_lowercase   ,       temb=_lowercase   ,       res_hidden_states_tuple=_lowercase   ,       deterministic=not train				)

        # 6. post-process
												_lowerCAmelCase  						=			self.conv_norm_out(_lowercase				)
												_lowerCAmelCase  						=			nn.silu(_lowercase				)
												_lowerCAmelCase  						=			self.conv_out(_lowercase				)
												_lowerCAmelCase  						=			jnp.transpose(_lowercase   ,       (0, 3, 1, 2)				)

												if not return_dict:
																		return (sample,)

												return FlaxUNetaDConditionOutput(sample=_lowercase				)
