






"""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_     :					Optional[Any]							=     "./tests/fixtures/elise_format0.mid"



@unittest.skip('The note-seq package currently throws an error on import'					)
class 						a_ (      snake_case_    ,		unittest.TestCase					):







      '''simple docstring'''

      lowerCamelCase__			:							int					=      SpectrogramDiffusionPipeline
      lowerCamelCase__			:							str					=      PipelineTesterMixin.required_optional_params - {
          'callback',
          'latents',
          'callback_steps',
          'output_type',
          'num_images_per_prompt',
      }
      lowerCamelCase__			:							List[Any]					=      False
      lowerCamelCase__			:							Optional[Any]					=      False
      lowerCamelCase__			:							List[str]					=      TOKENS_TO_AUDIO_GENERATION_PARAMS
      lowerCamelCase__			:							Union[str, Any]					=      TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS




      def      a__				(self   ):

       '''simple docstring'''



       torch.manual_seed(0   )
       lowerCamelCase__	:     Union[str, Any]								=		SpectrogramNotesEncoder(
           max_length=2_0_4_8,	vocab_size=1_5_3_6,	d_model=7_6_8,	dropout_rate=0.1,	num_layers=1,	num_heads=1,	d_kv=4,	d_ff=2_0_4_8,	feed_forward_proj='gated-gelu',	)

       lowerCamelCase__	:     Tuple								=		SpectrogramContEncoder(
           input_dims=1_2_8,	targets_context_length=2_5_6,	d_model=7_6_8,	dropout_rate=0.1,	num_layers=1,	num_heads=1,	d_kv=4,	d_ff=2_0_4_8,	feed_forward_proj='gated-gelu',	)

       lowerCamelCase__	:     Dict								=		TaFilmDecoder(
           input_dims=1_2_8,	targets_length=2_5_6,	max_decoder_noise_time=20_000.0,	d_model=7_6_8,	num_layers=1,	num_heads=1,	d_kv=4,	d_ff=2_0_4_8,	dropout_rate=0.1,	)

       lowerCamelCase__	:     str								=		DDPMScheduler()

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




      def      a__				(self,	lowerCamelCase_,	lowerCamelCase_=0   ):

       '''simple docstring'''



       if str(lowerCamelCase_   ).startswith('mps'   ):
        lowerCamelCase__	:     List[Any]								=		torch.manual_seed(lowerCamelCase_   )
       else:
        lowerCamelCase__	:     List[str]								=		torch.Generator(device=lowerCamelCase_   ).manual_seed(lowerCamelCase_   )
       lowerCamelCase__	:     Optional[int]								=		{
           'input_tokens': [
               [1_1_3_4, 9_0, 1_1_3_5, 1_1_3_3, 1_0_8_0, 1_1_2, 1_1_3_2, 1_0_8_0, 1_1_3_3, 1_0_7_9, 1_3_3, 1_1_3_2, 1_0_7_9, 1_1_3_3, 1] + [0] * 2_0_3_3
           ],
           'generator': generator,
           'num_inference_steps': 4,
           'output_type': 'mel',
       }
       return inputs




      def      a__				(self   ):

       '''simple docstring'''



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

       lowerCamelCase__	:     Any								=		self.get_dummy_inputs(lowerCamelCase_   )
       lowerCamelCase__	:     List[str]								=		pipe(**lowerCamelCase_   )
       lowerCamelCase__	:     Optional[int]								=		output.audios

       lowerCamelCase__	:     List[str]								=		mel[0, -3:, -3:]

       assert mel_slice.shape == (3, 3)
       lowerCamelCase__	:     Union[str, Any]								=		np.array(
           [-11.512_925, -4.788_215, -0.46_172_905, -2.051_715, -10.539_147, -10.970_963, -9.091_634, 4.0, 4.0]   )
       assert np.abs(mel_slice.flatten() - expected_slice   ).max() < 1e-2




      @skip_mps
      def      a__				(self   ):

       '''simple docstring'''



       return super().test_save_load_local()




      @skip_mps
      def      a__				(self   ):

       '''simple docstring'''



       return super().test_dict_tuple_outputs_equivalent()




      @skip_mps
      def      a__				(self   ):

       '''simple docstring'''



       return super().test_save_load_optional_components()




      @skip_mps
      def      a__				(self   ):

       '''simple docstring'''



       return super().test_attention_slicing_forward_pass()




      def      a__				(self   ):

       '''simple docstring'''



       pass




      def      a__				(self   ):

       '''simple docstring'''



       pass





      @skip_mps
      def      a__				(self   ):

       '''simple docstring'''



       return super().test_progress_bar()



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







      '''simple docstring'''

      def      a__				(self   ):

       '''simple docstring'''



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




      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     Optional[int]								=		torch_device

       lowerCamelCase__	:     int								=		SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'   )
       lowerCamelCase__	:     Optional[Any]								=		pipe.melgan
       lowerCamelCase__	:     Optional[Any]								=		None

       lowerCamelCase__	:     Union[str, Any]								=		pipe.to(lowerCamelCase_   )
       pipe.set_progress_bar_config(disable=lowerCamelCase_   )

       def callback(lowerCamelCase_,	lowerCamelCase_   ):
        # decode mel to audio
        lowerCamelCase__	:     List[Any]								=		melgan(input_features=mel_output.astype(np.floataa   )   )[0]
        assert len(audio[0]   ) == 8_1_9_2_0 * (step + 1)
        # simulate that audio is played
        return audio

       lowerCamelCase__	:     Optional[int]								=		MidiProcessor()
       lowerCamelCase__	:     int								=		processor(lowerCamelCase_   )

       lowerCamelCase__	:     str								=		input_tokens[:3]
       lowerCamelCase__	:     Dict								=		torch.manual_seed(0   )
       pipe(lowerCamelCase_,	num_inference_steps=5,	generator=lowerCamelCase_,	callback=lowerCamelCase_,	output_type='mel'   )




      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     Tuple								=		torch_device

       lowerCamelCase__	:     int								=		SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'   )
       lowerCamelCase__	:     Any								=		pipe.to(lowerCamelCase_   )
       pipe.set_progress_bar_config(disable=lowerCamelCase_   )
       lowerCamelCase__	:     List[str]								=		MidiProcessor()

       lowerCamelCase__	:     Any								=		processor(lowerCamelCase_   )
       # just run two denoising loops
       lowerCamelCase__	:     List[str]								=		input_tokens[:2]

       lowerCamelCase__	:     Tuple								=		torch.manual_seed(0   )
       lowerCamelCase__	:     Union[str, Any]								=		pipe(lowerCamelCase_,	num_inference_steps=2,	generator=lowerCamelCase_   )

       lowerCamelCase__	:     str								=		output.audios[0]

       assert abs(np.abs(lowerCamelCase_   ).sum() - 3_612.841   ) < 1e-1





      def      a__				(self   ):

       '''simple docstring'''



       lowerCamelCase__	:     List[str]								=		torch_device

       lowerCamelCase__	:     Optional[Any]								=		SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'   )
       lowerCamelCase__	:     Optional[int]								=		pipe.to(lowerCamelCase_   )
       pipe.set_progress_bar_config(disable=lowerCamelCase_   )

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

       lowerCamelCase__	:     Union[str, Any]								=		processor(lowerCamelCase_   )

       # just run 4 denoising loops
       lowerCamelCase__	:     Tuple								=		input_tokens[:4]

       lowerCamelCase__	:     str								=		torch.manual_seed(0   )
       lowerCamelCase__	:     Optional[int]								=		pipe(lowerCamelCase_,	num_inference_steps=1_0_0,	generator=lowerCamelCase_   )

       lowerCamelCase__	:     Dict								=		output.audios[0]
       assert abs(np.abs(lowerCamelCase_   ).sum() - 9_389.1_111   ) < 5e-2









"""simple docstring"""
from transformers import HfArgumentParser, PyTorchBenchmark, PyTorchBenchmarkArguments






def 			lowerCamelCase_    (		):
 lowerCamelCase__	:     Tuple								=		HfArgumentParser(_lowerCamelCase							)
 try:
  lowerCamelCase__	:     Dict								=		parser.parse_args_into_dataclasses()[0]
 except ValueError as e:
  lowerCamelCase__	:     Optional[int]								=		'Arg --no_{0} is no longer used, please use --no-{0} instead.'
  lowerCamelCase__	:     Optional[Any]								=		' '.join(str(_lowerCamelCase							).split(' '							)[:-1]							)
  lowerCamelCase__	:     int								=		''
  lowerCamelCase__	:     Optional[int]								=		eval(str(_lowerCamelCase							).split(' '							)[-1]							)
  lowerCamelCase__	:     Tuple								=		[]
  for arg in depreciated_args:
   # arg[2:] removes '--'
   if arg[2:] in PyTorchBenchmarkArguments.deprecated_args:
    # arg[5:] removes '--no_'
    full_error_msg += arg_error_msg.format(arg[5:]							)
   else:
    wrong_args.append(_lowerCamelCase							)
  if len(_lowerCamelCase							) > 0:
   lowerCamelCase__	:     Optional[int]								=		full_error_msg + begin_error_msg + str(_lowerCamelCase							)
  raise ValueError(_lowerCamelCase							)

 lowerCamelCase__	:     Any								=		PyTorchBenchmark(args=_lowerCamelCase							)
 benchmark.run()


if __name__ == "__main__":
      main()

