


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


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




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





							'''simple docstring'''






							lowercase__			:   int							=      SpectrogramDiffusionPipeline
							lowercase__			:   Any							=      PipelineTesterMixin.required_optional_params - {
							    "callback",
							    "latents",
							    "callback_steps",
							    "output_type",
							    "num_images_per_prompt",
							}
							lowercase__			:   Optional[int]							=      False
							lowercase__			:   Optional[int]							=      False
							lowercase__			:   Optional[Any]							=      TOKENS_TO_AUDIO_GENERATION_PARAMS
							lowercase__			:   Any							=      TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Union[str, Any]:
														torch.manual_seed(0   )
														lowerCAmelCase__          =	SpectrogramNotesEncoder(
														    max_length=20_48					,    vocab_size=15_36					,    d_model=7_68					,    dropout_rate=0.1					,    num_layers=1					,    num_heads=1					,    d_kv=4					,    d_ff=20_48					,    feed_forward_proj='''gated-gelu'''					,    )

														lowerCAmelCase__          =	SpectrogramContEncoder(
														    input_dims=1_28					,    targets_context_length=2_56					,    d_model=7_68					,    dropout_rate=0.1					,    num_layers=1					,    num_heads=1					,    d_kv=4					,    d_ff=20_48					,    feed_forward_proj='''gated-gelu'''					,    )

														lowerCAmelCase__          =	TaFilmDecoder(
														    input_dims=1_28					,    targets_length=2_56					,    max_decoder_noise_time=20_000.0					,    d_model=7_68					,    num_layers=1					,    num_heads=1					,    d_kv=4					,    d_ff=20_48					,    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   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_=0   )       ->							Tuple:
														if str(lowerCamelCase_   ).startswith('''mps'''   ):
																					lowerCAmelCase__          =	torch.manual_seed(lowerCamelCase_   )
														else:
																					lowerCAmelCase__          =	torch.Generator(device=lowerCamelCase_   ).manual_seed(lowerCamelCase_   )
														lowerCAmelCase__          =	{
														    '''input_tokens''': [
														        [11_34, 90, 11_35, 11_33, 10_80, 1_12, 11_32, 10_80, 11_33, 10_79, 1_33, 11_32, 10_79, 11_33, 1] + [0] * 20_33
														    ],
														    '''generator''': generator,
														    '''num_inference_steps''': 4,
														    '''output_type''': '''mel''',
														}
														return inputs







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Tuple:
														lowerCAmelCase__          =	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
														lowerCAmelCase__          =	self.get_dummy_components()
														lowerCAmelCase__          =	SpectrogramDiffusionPipeline(**lowerCamelCase_   )
														lowerCAmelCase__          =	pipe.to(lowerCamelCase_   )
														pipe.set_progress_bar_config(disable=lowerCamelCase_   )

														lowerCAmelCase__          =	self.get_dummy_inputs(lowerCamelCase_   )
														lowerCAmelCase__          =	pipe(**lowerCamelCase_   )
														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   __SCREAMING_SNAKE_CASE							(		self   )       ->							Any:
														return super().test_save_load_local()







							@skip_mps
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Tuple:
														return super().test_dict_tuple_outputs_equivalent()







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







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







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Tuple:
														pass







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							int:
														pass







							@skip_mps
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Any:
														return super().test_progress_bar()




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





							'''simple docstring'''






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







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Any:
														# 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(lowerCamelCase_   )
														pipe.set_progress_bar_config(disable=lowerCamelCase_   )

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

														lowerCAmelCase__          =	MidiProcessor()
														lowerCAmelCase__          =	processor(lowerCamelCase_   )

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







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Dict:
														lowerCAmelCase__          =	torch_device

														lowerCAmelCase__          =	SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''   )
														lowerCAmelCase__          =	pipe.to(lowerCamelCase_   )
														pipe.set_progress_bar_config(disable=lowerCamelCase_   )
														lowerCAmelCase__          =	MidiProcessor()

														lowerCAmelCase__          =	processor(lowerCamelCase_   )
														# just run two denoising loops
														lowerCAmelCase__          =	input_tokens[:2]

														lowerCAmelCase__          =	torch.manual_seed(0   )
														lowerCAmelCase__          =	pipe(lowerCamelCase_					,    num_inference_steps=2					,    generator=lowerCamelCase_   )

														lowerCAmelCase__          =	output.audios[0]

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







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							int:
														lowerCAmelCase__          =	torch_device

														lowerCAmelCase__          =	SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''   )
														lowerCAmelCase__          =	pipe.to(lowerCamelCase_   )
														pipe.set_progress_bar_config(disable=lowerCamelCase_   )

														lowerCAmelCase__          =	MidiProcessor()

														lowerCAmelCase__          =	processor(lowerCamelCase_   )

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

														lowerCAmelCase__          =	torch.manual_seed(0   )
														lowerCAmelCase__          =	pipe(lowerCamelCase_					,    num_inference_steps=1_00					,    generator=lowerCamelCase_   )

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



'''simple docstring'''






from __future__ import annotations


def 					_snake_case   (     A ,    A       )						->    list[int]:
							lowerCAmelCase__          =	{}
							for index, val in enumerate(A       ):
														lowerCAmelCase__          =	target - val
														if compl in chk_map:
																					return [chk_map[compl], index]
														lowerCAmelCase__          =	index
							return []


if __name__ == "__main__":
		import doctest

		doctest.testmod()
		print(f"""{two_sum([2, 7, 11, 15], 9) = }""")