



'''simple docstring'''





import gc
import unittest

import numpy as np
import torch

from diffusers import DDPMScheduler, MidiProcessor, SpectrogramDiffusionPipeline
from diffusers.pipelines.spectrogram_diffusion import SpectrogramContEncoder, SpectrogramNotesEncoder, TaFilmDecoder
from diffusers.utils import require_torch_gpu, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_note_seq, require_onnxruntime

from ..pipeline_params import TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS, TOKENS_TO_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


A			:   Union[str, Any]			=      """./tests/fixtures/elise_format0.mid"""





@unittest.skip('The note-seq package currently throws an error on import'						)
class 							lowerCAmelCase_		(	a_  ,      unittest.TestCase						):
    __UpperCAmelCase							=	SpectrogramDiffusionPipeline
    __UpperCAmelCase							=	PipelineTesterMixin.required_optional_params - {
        'callback',
        'latents',
        'callback_steps',
        'output_type',
        'num_images_per_prompt',
    }
    __UpperCAmelCase							=	False
    __UpperCAmelCase							=	False
    __UpperCAmelCase							=	TOKENS_TO_AUDIO_GENERATION_PARAMS
    __UpperCAmelCase							=	TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS





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



       '''simple docstring'''







       torch.manual_seed(0     )
       snake_case   :							Tuple						  =SpectrogramNotesEncoder(
           max_length=2_048,    vocab_size=1_536,    d_model=768,    dropout_rate=0.1,    num_layers=1,    num_heads=1,    d_kv=4,    d_ff=2_048,    feed_forward_proj='''gated-gelu''',    )

       snake_case   :							Tuple						  =SpectrogramContEncoder(
           input_dims=128,    targets_context_length=256,    d_model=768,    dropout_rate=0.1,    num_layers=1,    num_heads=1,    d_kv=4,    d_ff=2_048,    feed_forward_proj='''gated-gelu''',    )

       snake_case   :							List[Any]						  =TaFilmDecoder(
           input_dims=128,    targets_length=256,    max_decoder_noise_time=2_0000.0,    d_model=768,    num_layers=1,    num_heads=1,    d_kv=4,    d_ff=2_048,    dropout_rate=0.1,    )

       snake_case   :							Dict						  =DDPMScheduler()

       snake_case   :							Union[str, Any]						  ={
           '''notes_encoder''': notes_encoder.eval(),
           '''continuous_encoder''': continuous_encoder.eval(),
           '''decoder''': decoder.eval(),
           '''scheduler''': scheduler,
           '''melgan''': None,
       }
       return components





    def 					__snake_case       (    self    :    Optional[int],    _snake_case    :    Tuple,    _snake_case    :    List[Any]=0     ):



       '''simple docstring'''







       if str(_snake_case     ).startswith('''mps'''     ):
          snake_case   :							List[Any]						  =torch.manual_seed(_snake_case     )
       else:
          snake_case   :							Any						  =torch.Generator(device=_snake_case     ).manual_seed(_snake_case     )
       snake_case   :							Any						  ={
           '''input_tokens''': [
               [1_134, 90, 1_135, 1_133, 1_080, 112, 1_132, 1_080, 1_133, 1_079, 133, 1_132, 1_079, 1_133, 1] + [0] * 2_033
           ],
           '''generator''': generator,
           '''num_inference_steps''': 4,
           '''output_type''': '''mel''',
       }
       return inputs





    def 					__snake_case       (    self    :    Dict     ):



       '''simple docstring'''







       snake_case   :							List[Any]						  ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
       snake_case   :							str						  =self.get_dummy_components()
       snake_case   :							Any						  =SpectrogramDiffusionPipeline(**_snake_case     )
       snake_case   :							List[Any]						  =pipe.to(_snake_case     )
       pipe.set_progress_bar_config(disable=_snake_case     )

       snake_case   :							Optional[Any]						  =self.get_dummy_inputs(_snake_case     )
       snake_case   :							Union[str, Any]						  =pipe(**_snake_case     )
       snake_case   :							Union[str, Any]						  =output.audios

       snake_case   :							Optional[int]						  =mel[0, -3:, -3:]

       assert mel_slice.shape == (3, 3)
       snake_case   :							Union[str, Any]						  =np.array(
           [-11.51_2925, -4.78_8215, -0.4617_2905, -2.05_1715, -10.53_9147, -10.97_0963, -9.09_1634, 4.0, 4.0]     )
       assert np.abs(mel_slice.flatten() - expected_slice     ).max() < 1E-2





    @skip_mps
    def 					__snake_case       (    self    :    Union[str, Any]     ):



       '''simple docstring'''







       return super().test_save_load_local()





    @skip_mps
    def 					__snake_case       (    self    :    str     ):



       '''simple docstring'''







       return super().test_dict_tuple_outputs_equivalent()





    @skip_mps
    def 					__snake_case       (    self    :    Optional[Any]     ):



       '''simple docstring'''







       return super().test_save_load_optional_components()





    @skip_mps
    def 					__snake_case       (    self    :    Optional[int]     ):



       '''simple docstring'''







       return super().test_attention_slicing_forward_pass()





    def 					__snake_case       (    self    :    Dict     ):



       '''simple docstring'''







       pass





    def 					__snake_case       (    self    :    Optional[int]     ):



       '''simple docstring'''







       pass






    @skip_mps
    def 					__snake_case       (    self    :    Union[str, Any]     ):



       '''simple docstring'''







       return super().test_progress_bar()





@slow
@require_torch_gpu
@require_onnxruntime
@require_note_seq
class 							lowerCAmelCase_		(	unittest.TestCase						):





    def 					__snake_case       (    self    :    Dict     ):



       '''simple docstring'''







       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()





    def 					__snake_case       (    self    :    List[str]     ):



       '''simple docstring'''







       snake_case   :							int						  =torch_device

       snake_case   :							Any						  =SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''     )
       snake_case   :							int						  =pipe.melgan
       snake_case   :							Any						  =None

       snake_case   :							str						  =pipe.to(_snake_case     )
       pipe.set_progress_bar_config(disable=_snake_case     )

       def callback(_snake_case    :    int,    _snake_case    :    List[Any]     ):
          # decode mel to audio
          snake_case   :							Optional[Any]						  =melgan(input_features=mel_output.astype(np.floataa     )     )[0]
          assert len(audio[0]     ) == 81_920 * (step + 1)
          # simulate that audio is played
          return audio

       snake_case   :							Tuple						  =MidiProcessor()
       snake_case   :							int						  =processor(_snake_case     )

       snake_case   :							Any						  =input_tokens[:3]
       snake_case   :							Optional[int]						  =torch.manual_seed(0     )
       pipe(_snake_case,    num_inference_steps=5,    generator=_snake_case,    callback=_snake_case,    output_type='''mel'''     )





    def 					__snake_case       (    self    :    Optional[Any]     ):



       '''simple docstring'''







       snake_case   :							Any						  =torch_device

       snake_case   :							Optional[Any]						  =SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''     )
       snake_case   :							Union[str, Any]						  =pipe.to(_snake_case     )
       pipe.set_progress_bar_config(disable=_snake_case     )
       snake_case   :							Optional[int]						  =MidiProcessor()

       snake_case   :							int						  =processor(_snake_case     )
       # just run two denoising loops
       snake_case   :							List[str]						  =input_tokens[:2]

       snake_case   :							Union[str, Any]						  =torch.manual_seed(0     )
       snake_case   :							str						  =pipe(_snake_case,    num_inference_steps=2,    generator=_snake_case     )

       snake_case   :							Optional[int]						  =output.audios[0]

       assert abs(np.abs(_snake_case     ).sum() - 3612.841     ) < 1E-1






    def 					__snake_case       (    self    :    str     ):



       '''simple docstring'''







       snake_case   :							Optional[int]						  =torch_device

       snake_case   :							str						  =SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''     )
       snake_case   :							Dict						  =pipe.to(_snake_case     )
       pipe.set_progress_bar_config(disable=_snake_case     )

       snake_case   :							Union[str, Any]						  =MidiProcessor()

       snake_case   :							Optional[Any]						  =processor(_snake_case     )

       # just run 4 denoising loops
       snake_case   :							int						  =input_tokens[:4]

       snake_case   :							int						  =torch.manual_seed(0     )
       snake_case   :							List[Any]						  =pipe(_snake_case,    num_inference_steps=100,    generator=_snake_case     )

       snake_case   :							Tuple						  =output.audios[0]
       assert abs(np.abs(_snake_case     ).sum() - 9389.1111     ) < 5E-2










'''simple docstring'''





import argparse

import torch
from saprl.hub import distilhubert

from transformers import HubertConfig, HubertModel, WavaVecaFeatureExtractor, logging


logging.set_verbosity_info()
A			:   Dict			=      logging.get_logger(__name__)

A			:   Tuple			=      {
    """post_extract_proj""": """feature_projection.projection""",
    """encoder.pos_conv.0""": """encoder.pos_conv_embed.conv""",
    """self_attn.k_proj""": """encoder.layers.*.attention.k_proj""",
    """self_attn.v_proj""": """encoder.layers.*.attention.v_proj""",
    """self_attn.q_proj""": """encoder.layers.*.attention.q_proj""",
    """self_attn.out_proj""": """encoder.layers.*.attention.out_proj""",
    """self_attn_layer_norm""": """encoder.layers.*.layer_norm""",
    """fc1""": """encoder.layers.*.feed_forward.intermediate_dense""",
    """fc2""": """encoder.layers.*.feed_forward.output_dense""",
    """final_layer_norm""": """encoder.layers.*.final_layer_norm""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """mask_emb""": """masked_spec_embed""",
}
def 		_a   (  lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_							):
   for attribute in key.split('''.'''							):
      snake_case   :							Optional[Any]						  =getattr(lowerCamelCase_					,						lowerCamelCase_							)

   if weight_type is not None:
      snake_case   :							int						  =getattr(lowerCamelCase_					,						lowerCamelCase_							).shape
   else:
      snake_case   :							Any						  =hf_pointer.shape

   assert hf_shape == value.shape, (
       F'''Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'''
       F''' {value.shape} for {full_name}'''
   )

   if weight_type == "weight":
      snake_case   :							Dict						  =value
   elif weight_type == "weight_g":
      snake_case   :							Any						  =value
   elif weight_type == "weight_v":
      snake_case   :							int						  =value
   elif weight_type == "bias":
      snake_case   :							int						  =value
   else:
      snake_case   :							str						  =value

   logger.info(F'''{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'''							)
def 		_a   (  lowerCamelCase_					,						lowerCamelCase_							):
   snake_case   :							str						  =[]
   snake_case   :							Tuple						  =fairseq_model.state_dict()

   snake_case   :							Union[str, Any]						  =hf_model.feature_extractor

   for name, value in fairseq_dict.items():
      snake_case   :							Optional[int]						  =False
      if "conv_layers" in name:
         load_conv_layer(
             lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						hf_model.config.feat_extract_norm == '''group'''					,						)
         snake_case   :							int						  =True
      else:
         for key, mapped_key in MAPPING.items():
            snake_case   :							Dict						  =mapped_key

            if key in name:
               snake_case   :							str						  =True
               if "*" in mapped_key:
                  snake_case   :							Tuple						  =name.split(lowerCamelCase_							)[0].split('''.'''							)[-2]
                  snake_case   :							Optional[int]						  =mapped_key.replace('''*'''					,						lowerCamelCase_							)
               if "weight_g" in name:
                  snake_case   :							int						  ='''weight_g'''
               elif "weight_v" in name:
                  snake_case   :							List[str]						  ='''weight_v'''
               elif "weight" in name:
                  snake_case   :							Any						  ='''weight'''
               elif "bias" in name:
                  snake_case   :							List[str]						  ='''bias'''
               else:
                  snake_case   :							Tuple						  =None
               set_recursively(lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_							)
            continue
      if not is_used:
         unused_weights.append(lowerCamelCase_							)

   logger.warning(F'''Unused weights: {unused_weights}'''							)
def 		_a   (  lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_					,						lowerCamelCase_							):
   snake_case   :							List[str]						  =full_name.split('''conv_layers.'''							)[-1]
   snake_case   :							Any						  =name.split('''.'''							)
   snake_case   :							Tuple						  =int(items[0]							)
   snake_case   :							int						  =int(items[1]							)

   if type_id == 0:
      if "bias" in name:
         assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
             F'''{full_name} has size {value.shape}, but'''
             F''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''
         )
         snake_case   :							Optional[int]						  =value
         logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''							)
      elif "weight" in name:
         assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
             F'''{full_name} has size {value.shape}, but'''
             F''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''
         )
         snake_case   :							str						  =value
         logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''							)
   elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
      if "bias" in name:
         assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
             F'''{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'''
             " found."
         )
         snake_case   :							Optional[int]						  =value
         logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''							)
      elif "weight" in name:
         assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
             F'''{full_name} has size {value.shape}, but'''
             F''' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'''
         )
         snake_case   :							Union[str, Any]						  =value
         logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''							)
   else:
      unused_weights.append(lowerCamelCase_							)
def 		_a   (  lowerCamelCase_							):
   snake_case   :							Any						  =HubertConfig()
   snake_case   :							List[str]						  =model.config

   snake_case   :							List[str]						  =fs_config.activation_dropout
   snake_case   :							List[Any]						  =False
   snake_case   :							Optional[Any]						  =fs_config.attention_dropout
   snake_case   :							Optional[int]						  =False
   snake_case   :							Any						  =eval(fs_config.extractor_conv_feature_layers							)
   snake_case   :							Optional[int]						  =[x[0] for x in conv_layers]
   snake_case   :							Any						  =[x[1] for x in conv_layers]
   snake_case   :							Optional[Any]						  =[x[2] for x in conv_layers]
   snake_case   :							Optional[Any]						  ='''gelu'''
   snake_case   :							Optional[Any]						  ='''layer''' if fs_config.extractor_mode == '''layer_norm''' else '''group'''
   snake_case   :							Optional[int]						  =False
   snake_case   :							Optional[Any]						  =0.0
   snake_case   :							Optional[int]						  =0.0
   snake_case   :							Dict						  =fs_config.activation_fn
   snake_case   :							Any						  =fs_config.dropout
   snake_case   :							Optional[int]						  =fs_config.encoder_embed_dim
   snake_case   :							Dict						  =0.02
   snake_case   :							Dict						  =fs_config.encoder_ffn_embed_dim
   snake_case   :							str						  =1e-5
   snake_case   :							List[Any]						  =0.0
   snake_case   :							Dict						  =fs_config.encoder_attention_heads
   snake_case   :							Union[str, Any]						  =fs_config.conv_pos_groups
   snake_case   :							Union[str, Any]						  =fs_config.conv_pos
   snake_case   :							Dict						  =len(lowerCamelCase_							)
   snake_case   :							str						  =fs_config.encoder_layers

   return config
@torch.no_grad()
def 		_a   (  lowerCamelCase_					,						lowerCamelCase_=None							):
   snake_case   :							Union[str, Any]						  =distilhubert().model.model

   if config_path is not None:
      snake_case   :							int						  =HubertConfig.from_pretrained(lowerCamelCase_							)
   else:
      snake_case   :							Optional[Any]						  =convert_config(lowerCamelCase_							)
   snake_case   :							Tuple						  =model.eval()

   snake_case   :							Optional[Any]						  =WavaVecaFeatureExtractor(
       feature_size=1					,						sampling_rate=1_60_00					,						padding_value=0					,						do_normalize=lowerCamelCase_					,						return_attention_mask=lowerCamelCase_					,						)
   snake_case   :							Any						  =HubertModel(lowerCamelCase_							)

   recursively_load_weights(lowerCamelCase_					,						lowerCamelCase_							)

   feature_extractor.save_pretrained(lowerCamelCase_							)
   hf_model.save_pretrained(lowerCamelCase_							)


if __name__ == "__main__":
      A			:   Tuple			=      argparse.ArgumentParser()
      parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
      parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
      A			:   int			=      parser.parse_args()
      convert_hubert_checkpoint(args.pytorch_dump_folder_path, args.config_path)





