





"""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()


SCREAMING_SNAKE_CASE__ =       '''./tests/fixtures/elise_format0.mid'''

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





						def    _snake_case      (	self			)		->					Optional[Any]:
												torch.manual_seed(0			)
												lowerCAmelCase  						=	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"""				,						)

												lowerCAmelCase  						=	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"""				,						)

												lowerCAmelCase  						=	TaFilmDecoder(
												    input_dims=128				,						targets_length=256				,						max_decoder_noise_time=20_000.0				,						d_model=768				,						num_layers=1				,						num_heads=1				,						d_kv=4				,						d_ff=2_048				,						dropout_rate=0.1				,						)

												lowerCAmelCase  						=	DDPMScheduler()

												lowerCAmelCase  						=	{
												    'notes_encoder': notes_encoder.eval(),
												    'continuous_encoder': continuous_encoder.eval(),
												    'decoder': decoder.eval(),
												    'scheduler': scheduler,
												    'melgan': None,
												}
												return components





						def    _snake_case      (	self				,						lowercase				,						lowercase=0			)		->					Optional[Any]:
												if str(_A			).startswith("""mps"""			):
																		lowerCAmelCase  						=	torch.manual_seed(_A			)
												else:
																		lowerCAmelCase  						=	torch.Generator(device=_A			).manual_seed(_A			)
												lowerCAmelCase  						=	{
												    '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			)		->					Optional[Any]:
												lowerCAmelCase  						=	'cpu'  # ensure determinism for the device-dependent torch.Generator
												lowerCAmelCase  						=	self.get_dummy_components()
												lowerCAmelCase  						=	SpectrogramDiffusionPipeline(**_A			)
												lowerCAmelCase  						=	pipe.to(_A			)
												pipe.set_progress_bar_config(disable=_A			)

												lowerCAmelCase  						=	self.get_dummy_inputs(_A			)
												lowerCAmelCase  						=	pipe(**_A			)
												lowerCAmelCase  						=	output.audios

												lowerCAmelCase  						=	mel[0, -3:, -3:]

												assert mel_slice.shape == (3, 3)
												lowerCAmelCase  						=	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    _snake_case      (	self			)		->					str:
												return super().test_save_load_local()





						@skip_mps
						def    _snake_case      (	self			)		->					List[str]:
												return super().test_dict_tuple_outputs_equivalent()





						@skip_mps
						def    _snake_case      (	self			)		->					str:
												return super().test_save_load_optional_components()





						@skip_mps
						def    _snake_case      (	self			)		->					List[Any]:
												return super().test_attention_slicing_forward_pass()





						def    _snake_case      (	self			)		->					Optional[int]:
												pass





						def    _snake_case      (	self			)		->					Optional[int]:
												pass





						@skip_mps
						def    _snake_case      (	self			)		->					int:
												return super().test_progress_bar()







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





						def    _snake_case      (	self			)		->					Dict:
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()





						def    _snake_case      (	self			)		->					int:
												# TODO - test that pipeline can decode tokens in a callback
												# so that music can be played live
												lowerCAmelCase  						=	torch_device

												lowerCAmelCase  						=	SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""			)
												lowerCAmelCase  						=	pipe.melgan
												lowerCAmelCase  						=	None

												lowerCAmelCase  						=	pipe.to(_A			)
												pipe.set_progress_bar_config(disable=_A			)

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

												lowerCAmelCase  						=	MidiProcessor()
												lowerCAmelCase  						=	processor(_A			)

												lowerCAmelCase  						=	input_tokens[:3]
												lowerCAmelCase  						=	torch.manual_seed(0			)
												pipe(_A				,						num_inference_steps=5				,						generator=_A				,						callback=_A				,						output_type="""mel"""			)





						def    _snake_case      (	self			)		->					List[Any]:
												lowerCAmelCase  						=	torch_device

												lowerCAmelCase  						=	SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""			)
												lowerCAmelCase  						=	pipe.to(_A			)
												pipe.set_progress_bar_config(disable=_A			)
												lowerCAmelCase  						=	MidiProcessor()

												lowerCAmelCase  						=	processor(_A			)
												# just run two denoising loops
												lowerCAmelCase  						=	input_tokens[:2]

												lowerCAmelCase  						=	torch.manual_seed(0			)
												lowerCAmelCase  						=	pipe(_A				,						num_inference_steps=2				,						generator=_A			)

												lowerCAmelCase  						=	output.audios[0]

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





						def    _snake_case      (	self			)		->					Optional[Any]:
												lowerCAmelCase  						=	torch_device

												lowerCAmelCase  						=	SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""			)
												lowerCAmelCase  						=	pipe.to(_A			)
												pipe.set_progress_bar_config(disable=_A			)

												lowerCAmelCase  						=	MidiProcessor()

												lowerCAmelCase  						=	processor(_A			)

												# just run 4 denoising loops
												lowerCAmelCase  						=	input_tokens[:4]

												lowerCAmelCase  						=	torch.manual_seed(0			)
												lowerCAmelCase  						=	pipe(_A				,						num_inference_steps=100				,						generator=_A			)

												lowerCAmelCase  						=	output.audios[0]
												assert abs(np.abs(_A			).sum() - 9_389.1_111			) < 5e-2













"""simple docstring"""





import unittest
from dataclasses import dataclass
from typing import Optional

from transformers.utils import ModelOutput

@dataclass
class 		lowercase ( _UpperCAmelCase ):
      _SCREAMING_SNAKE_CASE         =				42
      _SCREAMING_SNAKE_CASE         =				None
      _SCREAMING_SNAKE_CASE         =				None







class 		lowercase ( unittest.TestCase ):





      def    _snake_case      (	self			)		->					Optional[Any]:
            lowerCAmelCase  						=	ModelOutputTest(a=30			)
            self.assertEqual(x.a				,						30			)
            self.assertIsNone(x.b			)
            self.assertIsNone(x.c			)
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x.d





      def    _snake_case      (	self			)		->					Optional[Any]:
            lowerCAmelCase  						=	ModelOutputTest(a=30				,						b=10			)
            self.assertEqual(x[0]				,						30			)
            self.assertEqual(x[1]				,						10			)
            self.assertEqual(x[:2]				,						(30, 10)			)
            self.assertEqual(x[:]				,						(30, 10)			)

            lowerCAmelCase  						=	ModelOutputTest(a=30				,						c=10			)
            self.assertEqual(x[0]				,						30			)
            self.assertEqual(x[1]				,						10			)
            self.assertEqual(x[:2]				,						(30, 10)			)
            self.assertEqual(x[:]				,						(30, 10)			)





      def    _snake_case      (	self			)		->					Tuple:
            lowerCAmelCase  						=	ModelOutputTest(a=30				,						b=10			)
            self.assertEqual(x["""a"""]				,						30			)
            self.assertEqual(x["""b"""]				,						10			)
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x["""c"""]

            lowerCAmelCase  						=	ModelOutputTest(a=30				,						c=10			)
            self.assertEqual(x["""a"""]				,						30			)
            self.assertEqual(x["""c"""]				,						10			)
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x["""b"""]





      def    _snake_case      (	self			)		->					Union[str, Any]:
            lowerCAmelCase  						=	ModelOutputTest(a=30			)
            self.assertEqual(list(x.keys()			)				,						["""a"""]			)
            self.assertEqual(list(x.values()			)				,						[30]			)
            self.assertEqual(list(x.items()			)				,						[("""a""", 30)]			)
            self.assertEqual(list(lowercase			)				,						["""a"""]			)

            lowerCAmelCase  						=	ModelOutputTest(a=30				,						b=10			)
            self.assertEqual(list(x.keys()			)				,						["""a""", """b"""]			)
            self.assertEqual(list(x.values()			)				,						[30, 10]			)
            self.assertEqual(list(x.items()			)				,						[("""a""", 30), ("""b""", 10)]			)
            self.assertEqual(list(lowercase			)				,						["""a""", """b"""]			)

            lowerCAmelCase  						=	ModelOutputTest(a=30				,						c=10			)
            self.assertEqual(list(x.keys()			)				,						["""a""", """c"""]			)
            self.assertEqual(list(x.values()			)				,						[30, 10]			)
            self.assertEqual(list(x.items()			)				,						[("""a""", 30), ("""c""", 10)]			)
            self.assertEqual(list(lowercase			)				,						["""a""", """c"""]			)

            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x.update({"""d""": 20}			)
            with self.assertRaises(lowercase			):
                  del x["a"]
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x.pop("""a"""			)
            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	x.setdefault("""d"""				,						32			)





      def    _snake_case      (	self			)		->					Dict:
            lowerCAmelCase  						=	ModelOutputTest(a=30			)
            lowerCAmelCase  						=	10
            self.assertEqual(x.a				,						10			)
            self.assertEqual(x["""a"""]				,						10			)





      def    _snake_case      (	self			)		->					List[str]:
            lowerCAmelCase  						=	ModelOutputTest(a=30			)
            lowerCAmelCase  						=	10
            self.assertEqual(x.a				,						10			)
            self.assertEqual(x["""a"""]				,						10			)





      def    _snake_case      (	self			)		->					str:
            lowerCAmelCase  						=	ModelOutputTest({"""a""": 30, """b""": 10}			)
            self.assertEqual(list(x.keys()			)				,						["""a""", """b"""]			)
            self.assertEqual(x.a				,						30			)
            self.assertEqual(x.b				,						10			)





      def    _snake_case      (	self			)		->					Dict:
            lowerCAmelCase  						=	ModelOutputTest([("""a""", 30), ("""b""", 10)]			)
            self.assertEqual(list(x.keys()			)				,						["""a""", """b"""]			)
            self.assertEqual(x.a				,						30			)
            self.assertEqual(x.b				,						10			)

            with self.assertRaises(lowercase			):
                  lowerCAmelCase  						=	ModelOutputTest([("""a""", 30), (10, 10)]			)

            lowerCAmelCase  						=	ModelOutputTest(a=(30, 30)			)
            self.assertEqual(list(x.keys()			)				,						["""a"""]			)
            self.assertEqual(x.a				,						(30, 30)			)






