





from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_outputs import (
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import logging
from .configuration_regnet import RegNetConfig


UpperCAmelCase__						:			List[str]									=logging.get_logger(__name__)

# General docstring
UpperCAmelCase__						:			Union[str, Any]									='''RegNetConfig'''

# Base docstring
UpperCAmelCase__						:			Optional[Any]									='''facebook/regnet-y-040'''
UpperCAmelCase__						:			Tuple									=[1, 10_88, 7, 7]

# Image classification docstring
UpperCAmelCase__						:			Union[str, Any]									='''facebook/regnet-y-040'''
UpperCAmelCase__						:			int									='''tabby, tabby cat'''

UpperCAmelCase__						:			Union[str, Any]									=[
    '''facebook/regnet-y-040''',
    # See all regnet models at https://huggingface.co/models?filter=regnet
]

class       __A     (						nn.Module    ):


				def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_ = 3   , UpperCAmelCase_ = 1   , UpperCAmelCase_ = 1   , UpperCAmelCase_ = "relu"   , ):
								super().__init__()
								lowerCamelCase	 =nn.Convad(
								    _lowerCamelCase   , _lowerCamelCase   , kernel_size=_lowerCamelCase   , stride=_lowerCamelCase   , padding=kernel_size // 2   , groups=_lowerCamelCase   , bias=_lowerCamelCase   , )
								lowerCamelCase	 =nn.BatchNormad(_lowerCamelCase							)
								lowerCamelCase	 =ACTaFN[activation] if activation is not None else nn.Identity()


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =self.convolution(_lowerCamelCase							)
								lowerCamelCase	 =self.normalization(_lowerCamelCase							)
								lowerCamelCase	 =self.activation(_lowerCamelCase							)
								return hidden_state




class       __A     (						nn.Module    ):


				def __init__(		self   , UpperCAmelCase_							):
								super().__init__()
								lowerCamelCase	 =RegNetConvLayer(
								    config.num_channels   , config.embedding_size   , kernel_size=3   , stride=2   , activation=config.hidden_act							)
								lowerCamelCase	 =config.num_channels


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =pixel_values.shape[1]
								if num_channels != self.num_channels:
												raise ValueError(
												    """Make sure that the channel dimension of the pixel values match with the one set in the configuration."""							)
								lowerCamelCase	 =self.embedder(_lowerCamelCase							)
								return hidden_state




class       __A     (						nn.Module    ):


				def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_ = 2							):
								super().__init__()
								lowerCamelCase	 =nn.Convad(_lowerCamelCase   , _lowerCamelCase   , kernel_size=1   , stride=_lowerCamelCase   , bias=_lowerCamelCase							)
								lowerCamelCase	 =nn.BatchNormad(_lowerCamelCase							)


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =self.convolution(_lowerCamelCase							)
								lowerCamelCase	 =self.normalization(_lowerCamelCase							)
								return hidden_state




class       __A     (						nn.Module    ):


				def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_							):
								super().__init__()

								lowerCamelCase	 =nn.AdaptiveAvgPoolad((1, 1)							)
								lowerCamelCase	 =nn.Sequential(
								    nn.Convad(_lowerCamelCase   , _lowerCamelCase   , kernel_size=1							)   , nn.ReLU()   , nn.Convad(_lowerCamelCase   , _lowerCamelCase   , kernel_size=1							)   , nn.Sigmoid()   , )


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =self.pooler(_lowerCamelCase							)
								lowerCamelCase	 =self.attention(_lowerCamelCase							)
								lowerCamelCase	 =hidden_state * attention
								return hidden_state




class       __A     (						nn.Module    ):


				def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_ = 1							):
								super().__init__()
								lowerCamelCase	 =in_channels != out_channels or stride != 1
								lowerCamelCase	 =max(1   , out_channels // config.groups_width							)
								lowerCamelCase	 =(
								    RegNetShortCut(_lowerCamelCase   , _lowerCamelCase   , stride=_lowerCamelCase							) if should_apply_shortcut else nn.Identity()
								)
								lowerCamelCase	 =nn.Sequential(
								    RegNetConvLayer(_lowerCamelCase   , _lowerCamelCase   , kernel_size=1   , activation=config.hidden_act							)   , RegNetConvLayer(_lowerCamelCase   , _lowerCamelCase   , stride=_lowerCamelCase   , groups=_lowerCamelCase   , activation=config.hidden_act							)   , RegNetConvLayer(_lowerCamelCase   , _lowerCamelCase   , kernel_size=1   , activation=_lowerCamelCase							)   , )
								lowerCamelCase	 =ACTaFN[config.hidden_act]


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =hidden_state
								lowerCamelCase	 =self.layer(_lowerCamelCase							)
								lowerCamelCase	 =self.shortcut(_lowerCamelCase							)
								hidden_state += residual
								lowerCamelCase	 =self.activation(_lowerCamelCase							)
								return hidden_state




class       __A     (						nn.Module    ):


				def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_ = 1							):
								super().__init__()
								lowerCamelCase	 =in_channels != out_channels or stride != 1
								lowerCamelCase	 =max(1   , out_channels // config.groups_width							)
								lowerCamelCase	 =(
								    RegNetShortCut(_lowerCamelCase   , _lowerCamelCase   , stride=_lowerCamelCase							) if should_apply_shortcut else nn.Identity()
								)
								lowerCamelCase	 =nn.Sequential(
								    RegNetConvLayer(_lowerCamelCase   , _lowerCamelCase   , kernel_size=1   , activation=config.hidden_act							)   , RegNetConvLayer(_lowerCamelCase   , _lowerCamelCase   , stride=_lowerCamelCase   , groups=_lowerCamelCase   , activation=config.hidden_act							)   , RegNetSELayer(_lowerCamelCase   , reduced_channels=int(round(in_channels / 4							)							)							)   , RegNetConvLayer(_lowerCamelCase   , _lowerCamelCase   , kernel_size=1   , activation=_lowerCamelCase							)   , )
								lowerCamelCase	 =ACTaFN[config.hidden_act]


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =hidden_state
								lowerCamelCase	 =self.layer(_lowerCamelCase							)
								lowerCamelCase	 =self.shortcut(_lowerCamelCase							)
								hidden_state += residual
								lowerCamelCase	 =self.activation(_lowerCamelCase							)
								return hidden_state




class       __A     (						nn.Module    ):


				def __init__(		self   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_   , UpperCAmelCase_ = 2   , UpperCAmelCase_ = 2   , ):
								super().__init__()

								lowerCamelCase	 =RegNetXLayer if config.layer_type == """x""" else RegNetYLayer

								lowerCamelCase	 =nn.Sequential(
								    # downsampling is done in the first layer with stride of 2
								    layer(
								        _lowerCamelCase   , _lowerCamelCase   , _lowerCamelCase   , stride=_lowerCamelCase   , )   , *[layer(_lowerCamelCase   , _lowerCamelCase   , _lowerCamelCase							) for _ in range(depth - 1							)]   , )


				def _snake_case       (		self   , UpperCAmelCase_							):
								lowerCamelCase	 =self.layers(_lowerCamelCase							)
								return hidden_state




class       __A     (						nn.Module    ):


				def __init__(		self   , UpperCAmelCase_							):
								super().__init__()
								lowerCamelCase	 =nn.ModuleList([]							)
								# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
								self.stages.append(
								    RegNetStage(
								        _lowerCamelCase   , config.embedding_size   , config.hidden_sizes[0]   , stride=2 if config.downsample_in_first_stage else 1   , depth=config.depths[0]   , )							)
								lowerCamelCase	 =zip(config.hidden_sizes   , config.hidden_sizes[1:]							)
								for (in_channels, out_channels), depth in zip(_lowerCamelCase   , config.depths[1:]							):
												self.stages.append(RegNetStage(_lowerCamelCase   , _lowerCamelCase   , _lowerCamelCase   , depth=_lowerCamelCase							)							)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_ = False   , UpperCAmelCase_ = True							):
								lowerCamelCase	 =() if output_hidden_states else None

								for stage_module in self.stages:
												if output_hidden_states:
																lowerCamelCase	 =hidden_states + (hidden_state,)

												lowerCamelCase	 =stage_module(_lowerCamelCase							)

								if output_hidden_states:
												lowerCamelCase	 =hidden_states + (hidden_state,)

								if not return_dict:
												return tuple(v for v in [hidden_state, hidden_states] if v is not None							)

								return BaseModelOutputWithNoAttention(last_hidden_state=_lowerCamelCase   , hidden_states=_lowerCamelCase							)




class       __A     (						lowerCAmelCase__    ):
				__A      =				RegNetConfig
				__A      =				"regnet"
				__A      =				"pixel_values"
				__A      =				True


				def _snake_case       (		self   , UpperCAmelCase_							):





								if isinstance(_lowerCamelCase   , nn.Convad							):
												nn.init.kaiming_normal_(module.weight   , mode="""fan_out"""   , nonlinearity="""relu"""							)
								elif isinstance(_lowerCamelCase   , (nn.BatchNormad, nn.GroupNorm)							):
												nn.init.constant_(module.weight   , 1							)
												nn.init.constant_(module.bias   , 0							)


				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_=False							):





								if isinstance(_lowerCamelCase   , _lowerCamelCase							):
												lowerCamelCase	 =value




UpperCAmelCase__						:			List[Any]									=r'''
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
'''

UpperCAmelCase__						:			Optional[int]									=r'''
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
'''

@add_start_docstrings(
    """The bare RegNet model outputting raw features without any specific head on top."""    ,					lowerCAmelCase__    ,					)
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class       __A     (						lowerCAmelCase__    ):


				def __init__(		self   , UpperCAmelCase_							):
								super().__init__(_lowerCamelCase							)
								lowerCamelCase	 =config
								lowerCamelCase	 =RegNetEmbeddings(_lowerCamelCase							)
								lowerCamelCase	 =RegNetEncoder(_lowerCamelCase							)
								lowerCamelCase	 =nn.AdaptiveAvgPoolad((1, 1)							)
								# Initialize weights and apply final processing
								self.post_init()


				@add_start_docstrings_to_model_forward(_lowerCamelCase							)
				@add_code_sample_docstrings(
				    checkpoint=_CHECKPOINT_FOR_DOC   , output_type=_lowerCamelCase   , config_class=_CONFIG_FOR_DOC   , modality="""vision"""   , expected_output=_EXPECTED_OUTPUT_SHAPE   , )
				def _snake_case       (		self   , UpperCAmelCase_   , UpperCAmelCase_ = None   , UpperCAmelCase_ = None							):
								lowerCamelCase	 =(
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)
								lowerCamelCase	 =return_dict if return_dict is not None else self.config.use_return_dict

								lowerCamelCase	 =self.embedder(_lowerCamelCase							)

								lowerCamelCase	 =self.encoder(
								    _lowerCamelCase   , output_hidden_states=_lowerCamelCase   , return_dict=_lowerCamelCase							)

								lowerCamelCase	 =encoder_outputs[0]

								lowerCamelCase	 =self.pooler(_lowerCamelCase							)

								if not return_dict:
												return (last_hidden_state, pooled_output) + encoder_outputs[1:]

								return BaseModelOutputWithPoolingAndNoAttention(
								    last_hidden_state=_lowerCamelCase   , pooler_output=_lowerCamelCase   , hidden_states=encoder_outputs.hidden_states   , )




@add_start_docstrings(
    """\n    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    """    ,					lowerCAmelCase__    ,					)
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class       __A     (						lowerCAmelCase__    ):


				def __init__(		self   , UpperCAmelCase_							):
								super().__init__(_lowerCamelCase							)
								lowerCamelCase	 =config.num_labels
								lowerCamelCase	 =RegNetModel(_lowerCamelCase							)
								# classification head
								lowerCamelCase	 =nn.Sequential(
								    nn.Flatten()   , nn.Linear(config.hidden_sizes[-1]   , config.num_labels							) if config.num_labels > 0 else nn.Identity()   , )
								# initialize weights and apply final processing
								self.post_init()


				@add_start_docstrings_to_model_forward(_lowerCamelCase							)
				@add_code_sample_docstrings(
				    checkpoint=_IMAGE_CLASS_CHECKPOINT   , output_type=_lowerCamelCase   , config_class=_CONFIG_FOR_DOC   , expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   , )
				def _snake_case       (		self   , UpperCAmelCase_ = None   , UpperCAmelCase_ = None   , UpperCAmelCase_ = None   , UpperCAmelCase_ = None   , ):
								lowerCamelCase	 =return_dict if return_dict is not None else self.config.use_return_dict

								lowerCamelCase	 =self.regnet(_lowerCamelCase   , output_hidden_states=_lowerCamelCase   , return_dict=_lowerCamelCase							)

								lowerCamelCase	 =outputs.pooler_output if return_dict else outputs[1]

								lowerCamelCase	 =self.classifier(_lowerCamelCase							)

								lowerCamelCase	 =None

								if labels is not None:
												if self.config.problem_type is None:
																if self.num_labels == 1:
																				lowerCamelCase	 ="""regression"""
																elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																				lowerCamelCase	 ="""single_label_classification"""
																else:
																				lowerCamelCase	 ="""multi_label_classification"""
												if self.config.problem_type == "regression":
																lowerCamelCase	 =MSELoss()
																if self.num_labels == 1:
																				lowerCamelCase	 =loss_fct(logits.squeeze()   , labels.squeeze()							)
																else:
																				lowerCamelCase	 =loss_fct(_lowerCamelCase   , _lowerCamelCase							)
												elif self.config.problem_type == "single_label_classification":
																lowerCamelCase	 =CrossEntropyLoss()
																lowerCamelCase	 =loss_fct(logits.view(-1   , self.num_labels							)   , labels.view(-1							)							)
												elif self.config.problem_type == "multi_label_classification":
																lowerCamelCase	 =BCEWithLogitsLoss()
																lowerCamelCase	 =loss_fct(_lowerCamelCase   , _lowerCamelCase							)

								if not return_dict:
												lowerCamelCase	 =(logits,) + outputs[2:]
												return (loss,) + output if loss is not None else output

								return ImageClassifierOutputWithNoAttention(loss=_lowerCamelCase   , logits=_lowerCamelCase   , hidden_states=outputs.hidden_states							)













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


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





def      _lowercase					(						_UpperCAmelCase				)						->  List[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      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase				)						->  Tuple[Dict, Dict]:
				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      _lowercase					(						_UpperCAmelCase				)						->  MusicgenDecoderConfig:
				if checkpoint == "small":
								# default config values
								lowerCamelCase	 =10_24
								lowerCamelCase	 =24
								lowerCamelCase	 =16
				elif checkpoint == "medium":
								lowerCamelCase	 =15_36
								lowerCamelCase	 =48
								lowerCamelCase	 =24
				elif checkpoint == "large":
								lowerCamelCase	 =20_48
								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      _lowercase					(						_UpperCAmelCase			,      _UpperCAmelCase=None			,      _UpperCAmelCase=None			,      _UpperCAmelCase="cpu"				)						->  Dict:
				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, 20_48):
								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	 =20_48
				lowerCamelCase	 =20_48

				# 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__":
							UpperCAmelCase__						:			List[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.'''
							)

							UpperCAmelCase__						:			Optional[Any]									=parser.parse_args()
							convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)





