

"""simple docstring"""



import inspect
import unittest

from transformers import SegformerConfig, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     MODEL_MAPPING,
     SegformerForImageClassification,
     SegformerForSemanticSegmentation,
     SegformerModel,
 )
 from transformers.models.segformer.modeling_segformer import SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import SegformerImageProcessor



class 					a						(       _lowerCamelCase				):



    """simple docstring"""



    def 				UpperCamelCase	(     self: Dict       ):






         """simple docstring"""


         A__										=  self.config_class(**self.inputs_dict       )
         self.parent.assertTrue(hasattr(UpperCamelCase    ,						"""hidden_sizes"""       )       )
         self.parent.assertTrue(hasattr(UpperCamelCase    ,						"""num_attention_heads"""       )       )
         self.parent.assertTrue(hasattr(UpperCamelCase    ,						"""num_encoder_blocks"""       )       )



class 					a						:



    """simple docstring"""



    def __init__(     self: str    ,						UpperCamelCase: Dict    ,						UpperCamelCase: int=13    ,						UpperCamelCase: Optional[int]=64    ,						UpperCamelCase: List[Any]=3    ,						UpperCamelCase: List[Any]=4    ,						UpperCamelCase: Optional[Any]=[2, 2, 2, 2]    ,						UpperCamelCase: Any=[8, 4, 2, 1]    ,						UpperCamelCase: Optional[int]=[16, 32, 64, 1_28]    ,						UpperCamelCase: str=[1, 4, 8, 16]    ,						UpperCamelCase: Dict=[1, 2, 4, 8]    ,						UpperCamelCase: Optional[Any]=True    ,						UpperCamelCase: Union[str, Any]=True    ,						UpperCamelCase: List[str]="gelu"    ,						UpperCamelCase: Tuple=0.1    ,						UpperCamelCase: Optional[int]=0.1    ,						UpperCamelCase: Tuple=0.02    ,						UpperCamelCase: int=3    ,						UpperCamelCase: str=None    ,						):






         """simple docstring"""


         A__										=  parent
         A__										=  batch_size
         A__										=  image_size
         A__										=  num_channels
         A__										=  num_encoder_blocks
         A__										=  sr_ratios
         A__										=  depths
         A__										=  hidden_sizes
         A__										=  downsampling_rates
         A__										=  num_attention_heads
         A__										=  is_training
         A__										=  use_labels
         A__										=  hidden_act
         A__										=  hidden_dropout_prob
         A__										=  attention_probs_dropout_prob
         A__										=  initializer_range
         A__										=  num_labels
         A__										=  scope

    def 				UpperCamelCase	(     self: Tuple       ):






         """simple docstring"""


         A__										=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

         A__										=  None
         if self.use_labels:
              A__										=  ids_tensor([self.batch_size, self.image_size, self.image_size]    ,						self.num_labels       )

         A__										=  self.get_config()
         return config, pixel_values, labels

    def 				UpperCamelCase	(     self: Optional[int]       ):






         """simple docstring"""


         return SegformerConfig(
             image_size=self.image_size    ,						num_channels=self.num_channels    ,						num_encoder_blocks=self.num_encoder_blocks    ,						depths=self.depths    ,						hidden_sizes=self.hidden_sizes    ,						num_attention_heads=self.num_attention_heads    ,						hidden_act=self.hidden_act    ,						hidden_dropout_prob=self.hidden_dropout_prob    ,						attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,						initializer_range=self.initializer_range    ,						)

    def 				UpperCamelCase	(     self: Tuple    ,						UpperCamelCase: str    ,						UpperCamelCase: Optional[Any]    ,						UpperCamelCase: int       ):






         """simple docstring"""


         A__										=  SegformerModel(config=UpperCamelCase       )
         model.to(UpperCamelCase       )
         model.eval()
         A__										=  model(UpperCamelCase       )
         A__										=  A__										=  self.image_size // (self.downsampling_rates[-1] * 2)
         self.parent.assertEqual(
             result.last_hidden_state.shape    ,						(self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)       )

    def 				UpperCamelCase	(     self: List[str]    ,						UpperCamelCase: Tuple    ,						UpperCamelCase: str    ,						UpperCamelCase: List[str]       ):






         """simple docstring"""


         A__										=  self.num_labels
         A__										=  SegformerForSemanticSegmentation(UpperCamelCase       )
         model.to(UpperCamelCase       )
         model.eval()
         A__										=  model(UpperCamelCase       )
         self.parent.assertEqual(
             result.logits.shape    ,						(self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)       )
         A__										=  model(UpperCamelCase    ,						labels=UpperCamelCase       )
         self.parent.assertEqual(
             result.logits.shape    ,						(self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)       )
         self.parent.assertGreater(result.loss    ,						0.0       )

    def 				UpperCamelCase	(     self: Optional[Any]    ,						UpperCamelCase: List[str]    ,						UpperCamelCase: str    ,						UpperCamelCase: Tuple       ):






         """simple docstring"""


         A__										=  1
         A__										=  SegformerForSemanticSegmentation(config=UpperCamelCase       )
         model.to(UpperCamelCase       )
         model.eval()
         A__										=  torch.randint(0    ,						1    ,						(self.batch_size, self.image_size, self.image_size)       ).to(UpperCamelCase       )
         A__										=  model(UpperCamelCase    ,						labels=UpperCamelCase       )
         self.parent.assertGreater(result.loss    ,						0.0       )





    def 				UpperCamelCase	(     self: Tuple       ):






         """simple docstring"""


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



@require_torch
class 					a						(       _lowerCamelCase,	_lowerCamelCase,	unittest.TestCase				):



    """simple docstring"""



    UpperCAmelCase						 =					(
        (
            SegformerModel,
            SegformerForSemanticSegmentation,
            SegformerForImageClassification,
        )
        if is_torch_available()
        else ()
    )
    UpperCAmelCase						 =					(
        {
            "feature-extraction": SegformerModel,
            "image-classification": SegformerForImageClassification,
            "image-segmentation": SegformerForSemanticSegmentation,
        }
        if is_torch_available()
        else {}
    )

    UpperCAmelCase						 =					True
    UpperCAmelCase						 =					False
    UpperCAmelCase						 =					False
    UpperCAmelCase						 =					False

    def 				UpperCamelCase	(     self: Any       ):






         """simple docstring"""


         A__										=  SegformerModelTester(self       )
         A__										=  SegformerConfigTester(self    ,						config_class=UpperCamelCase       )

    def 				UpperCamelCase	(     self: Any       ):






         """simple docstring"""


         self.config_tester.run_common_tests()

    def 				UpperCamelCase	(     self: Dict       ):






         """simple docstring"""


         A__										=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*UpperCamelCase       )

    def 				UpperCamelCase	(     self: int       ):






         """simple docstring"""


         A__										=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_binary_image_segmentation(*UpperCamelCase       )

    def 				UpperCamelCase	(     self: Dict       ):






         """simple docstring"""


         A__										=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_segmentation(*UpperCamelCase       )

    @unittest.skip("""SegFormer does not use inputs_embeds"""       )
    def 				UpperCamelCase	(     self: List[str]       ):






         """simple docstring"""


         pass

    @unittest.skip("""SegFormer does not have get_input_embeddings method and get_output_embeddings methods"""       )
    def 				UpperCamelCase	(     self: str       ):






         """simple docstring"""


         pass

    def 				UpperCamelCase	(     self: Optional[Any]       ):






         """simple docstring"""


         A__			,				A__										=  self.model_tester.prepare_config_and_inputs_for_common()

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

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

    def 				UpperCamelCase	(     self: List[Any]       ):






         """simple docstring"""


         A__			,				A__										=  self.model_tester.prepare_config_and_inputs_for_common()
         A__										=  True

         for model_class in self.all_model_classes:
              A__										=  True
              A__										=  False
              A__										=  True
              A__										=  model_class(UpperCamelCase       )
              model.to(UpperCamelCase       )
              model.eval()
              with torch.no_grad():
                   A__										=  model(**self._prepare_for_class(UpperCamelCase    ,						UpperCamelCase       )       )
              A__										=  outputs.attentions

              A__										=  sum(self.model_tester.depths       )
              self.assertEqual(len(UpperCamelCase       )    ,						UpperCamelCase       )

              # check that output_attentions also work using config
              del inputs_dict["output_attentions"]
              A__										=  True
              A__										=  model_class(UpperCamelCase       )
              model.to(UpperCamelCase       )
              model.eval()
              with torch.no_grad():
                   A__										=  model(**self._prepare_for_class(UpperCamelCase    ,						UpperCamelCase       )       )
              A__										=  outputs.attentions

              self.assertEqual(len(UpperCamelCase       )    ,						UpperCamelCase       )

              # verify the first attentions (first block, first layer)
              A__										=  (self.model_tester.image_size // 4) ** 2
              A__										=  (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
              self.assertListEqual(
                  list(attentions[0].shape[-3:]       )    ,						[self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]    ,						)

              # verify the last attentions (last block, last layer)
              A__										=  (self.model_tester.image_size // 32) ** 2
              A__										=  (self.model_tester.image_size // (32 * self.model_tester.sr_ratios[-1])) ** 2
              self.assertListEqual(
                  list(attentions[-1].shape[-3:]       )    ,						[self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]    ,						)
              A__										=  len(UpperCamelCase       )

              # Check attention is always last and order is fine
              A__										=  True
              A__										=  True
              A__										=  model_class(UpperCamelCase       )
              model.to(UpperCamelCase       )
              model.eval()
              with torch.no_grad():
                   A__										=  model(**self._prepare_for_class(UpperCamelCase    ,						UpperCamelCase       )       )

              self.assertEqual(out_len + 1    ,						len(UpperCamelCase       )       )

              A__										=  outputs.attentions

              self.assertEqual(len(UpperCamelCase       )    ,						UpperCamelCase       )
              # verify the first attentions (first block, first layer)
              A__										=  (self.model_tester.image_size // 4) ** 2
              A__										=  (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
              self.assertListEqual(
                  list(self_attentions[0].shape[-3:]       )    ,						[self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]    ,						)

    def 				UpperCamelCase	(     self: Union[str, Any]       ):






         """simple docstring"""


         def check_hidden_states_output(UpperCamelCase: Dict    ,						UpperCamelCase: Tuple    ,						UpperCamelCase: Dict       ):
              A__										=  model_class(UpperCamelCase       )
              model.to(UpperCamelCase       )
              model.eval()

              with torch.no_grad():
                   A__										=  model(**self._prepare_for_class(UpperCamelCase    ,						UpperCamelCase       )       )

              A__										=  outputs.hidden_states

              A__										=  self.model_tester.num_encoder_blocks
              self.assertEqual(len(UpperCamelCase       )    ,						UpperCamelCase       )

              # verify the first hidden states (first block)
              self.assertListEqual(
                  list(hidden_states[0].shape[-3:]       )    ,						[
                      self.model_tester.hidden_sizes[0],
                      self.model_tester.image_size // 4,
                      self.model_tester.image_size // 4,
                  ]    ,						)

         A__			,				A__										=  self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
              A__										=  True
              check_hidden_states_output(UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase       )

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              A__										=  True

              check_hidden_states_output(UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase       )

    def 				UpperCamelCase	(     self: Union[str, Any]       ):






         """simple docstring"""


         if not self.model_tester.is_training:
              return

         A__			,				A__										=  self.model_tester.prepare_config_and_inputs_for_common()
         A__										=  True

         for model_class in self.all_model_classes:
              if model_class in get_values(UpperCamelCase       ):
                   continue

              A__										=  model_class(UpperCamelCase       )
              model.to(UpperCamelCase       )
              model.train()
              A__										=  self._prepare_for_class(UpperCamelCase    ,						UpperCamelCase    ,						return_labels=UpperCamelCase       )
              A__										=  model(**UpperCamelCase       ).loss
              loss.backward()

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






         """simple docstring"""


         pass





    @slow
    def 				UpperCamelCase	(     self: Optional[Any]       ):






         """simple docstring"""


         for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              A__										=  SegformerModel.from_pretrained(UpperCamelCase       )
              self.assertIsNotNone(UpperCamelCase       )





def  _snake_case     (    ):
     A__										=  Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""			)
     return image



@require_torch
class 					a						(       unittest.TestCase				):



    """simple docstring"""



    @slow
    def 				UpperCamelCase	(     self: Any       ):






         """simple docstring"""


         A__										=  SegformerImageProcessor(
             image_scale=(5_12, 5_12)    ,						keep_ratio=UpperCamelCase    ,						align=UpperCamelCase    ,						do_random_crop=UpperCamelCase       )
         A__										=  SegformerForSemanticSegmentation.from_pretrained("""nvidia/segformer-b0-finetuned-ade-512-512"""       ).to(
             UpperCamelCase       )

         A__										=  prepare_img()
         A__										=  image_processor(images=UpperCamelCase    ,						return_tensors="""pt"""       )
         A__										=  encoded_inputs.pixel_values.to(UpperCamelCase       )

         with torch.no_grad():
              A__										=  model(UpperCamelCase       )

         A__										=  torch.Size((1, model.config.num_labels, 1_28, 1_28)       )
         self.assertEqual(outputs.logits.shape    ,						UpperCamelCase       )

         A__										=  torch.tensor(
             [
                 [[-4.6_310, -5.5_232, -6.2_356], [-5.1_921, -6.1_444, -6.5_996], [-5.4_424, -6.2_790, -6.7_574]],
                 [[-12.1_391, -13.3_122, -13.9_554], [-12.8_732, -13.9_352, -14.3_563], [-12.9_438, -13.8_226, -14.2_513]],
                 [[-12.5_134, -13.4_686, -14.4_915], [-12.8_669, -14.4_343, -14.7_758], [-13.2_523, -14.5_819, -15.0_694]],
             ]       ).to(UpperCamelCase       )
         self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]    ,						UpperCamelCase    ,						atol=1e-4       )       )

    @slow
    def 				UpperCamelCase	(     self: Optional[Any]       ):






         """simple docstring"""


         A__										=  SegformerImageProcessor(
             image_scale=(5_12, 5_12)    ,						keep_ratio=UpperCamelCase    ,						align=UpperCamelCase    ,						do_random_crop=UpperCamelCase       )
         A__										=  SegformerForSemanticSegmentation.from_pretrained(
             """nvidia/segformer-b1-finetuned-cityscapes-1024-1024"""       ).to(UpperCamelCase       )

         A__										=  prepare_img()
         A__										=  image_processor(images=UpperCamelCase    ,						return_tensors="""pt"""       )
         A__										=  encoded_inputs.pixel_values.to(UpperCamelCase       )

         with torch.no_grad():
              A__										=  model(UpperCamelCase       )

         A__										=  torch.Size((1, model.config.num_labels, 1_28, 1_28)       )
         self.assertEqual(outputs.logits.shape    ,						UpperCamelCase       )

         A__										=  torch.tensor(
             [
                 [[-13.5_748, -13.9_111, -12.6_500], [-14.3_500, -15.3_683, -14.2_328], [-14.7_532, -16.0_424, -15.6_087]],
                 [[-17.1_651, -15.8_725, -12.9_653], [-17.2_580, -17.3_718, -14.8_223], [-16.6_058, -16.8_783, -16.7_452]],
                 [[-3.6_456, -3.0_209, -1.4_203], [-3.0_797, -3.1_959, -2.0_000], [-1.8_757, -1.9_217, -1.6_997]],
             ]       ).to(UpperCamelCase       )
         self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]    ,						UpperCamelCase    ,						atol=1e-1       )       )





    @slow
    def 				UpperCamelCase	(     self: List[Any]       ):






         """simple docstring"""


         A__										=  SegformerImageProcessor(
             image_scale=(5_12, 5_12)    ,						keep_ratio=UpperCamelCase    ,						align=UpperCamelCase    ,						do_random_crop=UpperCamelCase       )
         A__										=  SegformerForSemanticSegmentation.from_pretrained("""nvidia/segformer-b0-finetuned-ade-512-512"""       ).to(
             UpperCamelCase       )

         A__										=  prepare_img()
         A__										=  image_processor(images=UpperCamelCase    ,						return_tensors="""pt"""       )
         A__										=  encoded_inputs.pixel_values.to(UpperCamelCase       )

         with torch.no_grad():
              A__										=  model(UpperCamelCase       )

         A__										=  outputs.logits.detach().cpu()

         A__										=  image_processor.post_process_semantic_segmentation(outputs=UpperCamelCase    ,						target_sizes=[(5_00, 3_00)]       )
         A__										=  torch.Size((5_00, 3_00)       )
         self.assertEqual(segmentation[0].shape    ,						UpperCamelCase       )

         A__										=  image_processor.post_process_semantic_segmentation(outputs=UpperCamelCase       )
         A__										=  torch.Size((1_28, 1_28)       )
         self.assertEqual(segmentation[0].shape    ,						UpperCamelCase       )









"""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()
SCREAMING_SNAKE_CASE_					:	Union[str, Any]							  =						logging.get_logger(__name__)


SCREAMING_SNAKE_CASE_					:	Any							  =						['model.decoder.embed_positions.weights']





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





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





def  _snake_case     (    UpperCAmelCase_  :					str			):
     if checkpoint == "small":
          # default config values
          A__										=  1024
          A__										=  24
          A__										=  16
     elif checkpoint == "medium":
          A__										=  1536
          A__										=  48
          A__										=  24
     elif checkpoint == "large":
          A__										=  2048
          A__										=  48
          A__										=  32
     else:
          raise ValueError(F"""Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."""			)
     A__										=  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_  :					int					,   UpperCAmelCase_  :					List[Any]=None					,   UpperCAmelCase_  :					Tuple=None					,   UpperCAmelCase_  :					Any="cpu"			):
     A__										=  MusicGen.get_pretrained(UpperCAmelCase_					,   device=UpperCAmelCase_			)
     A__										=  decoder_config_from_checkpoint(UpperCAmelCase_			)

     A__										=  fairseq_model.lm.state_dict()
     A__			,				A__										=  rename_state_dict(
         UpperCAmelCase_					,   hidden_size=decoder_config.hidden_size			)

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

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

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

     A__										=  MusicgenProcessor(feature_extractor=UpperCAmelCase_					,   tokenizer=UpperCAmelCase_			)

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

     # set other default generation config params
     A__										=  int(30 * audio_encoder.config.frame_rate			)
     A__										=  True
     A__										=  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__":
 SCREAMING_SNAKE_CASE_					:	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.'
 )

 SCREAMING_SNAKE_CASE_					:	Tuple							  =						parser.parse_args()
 convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)






