




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						:		Dict          =						"./tests/fixtures/elise_format0.mid"

@unittest.skip("The note-seq package currently throws an error on import"				)
class 		lowercase_       (	lowerCAmelCase__						,							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     _lowercase						(   self:		int):


									'''simple docstring'''


									torch.manual_seed(0)
									__lowerCAmelCase					=  SpectrogramNotesEncoder(
									    max_length=2048,     vocab_size=1536,     d_model=768,     dropout_rate=0.1,     num_layers=1,     num_heads=1,     d_kv=4,     d_ff=2048,     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=2048,     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=2048,     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     _lowercase						(   self:		Tuple,     _lowercase:		int,     _lowercase:		Union[str, Any]=0):


									'''simple docstring'''


									if str(_lowercase).startswith("""mps"""):
													__lowerCAmelCase					=  torch.manual_seed(_lowercase)
									else:
													__lowerCAmelCase					=  torch.Generator(device=_lowercase).manual_seed(_lowercase)
									__lowerCAmelCase					=  {
									    """input_tokens""": [
									        [1134, 90, 1135, 1133, 1080, 112, 1132, 1080, 1133, 1079, 133, 1132, 1079, 1133, 1] + [0] * 2033
									    ],
									    """generator""": generator,
									    """num_inference_steps""": 4,
									    """output_type""": """mel""",
									}
									return inputs




					def     _lowercase						(   self:		int):


									'''simple docstring'''


									__lowerCAmelCase					=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
									__lowerCAmelCase					=  self.get_dummy_components()
									__lowerCAmelCase					=  SpectrogramDiffusionPipeline(**_lowercase)
									__lowerCAmelCase					=  pipe.to(_lowercase)
									pipe.set_progress_bar_config(disable=_lowercase)

									__lowerCAmelCase					=  self.get_dummy_inputs(_lowercase)
									__lowerCAmelCase					=  pipe(**_lowercase)
									__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     _lowercase						(   self:		Optional[int]):


									'''simple docstring'''


									return super().test_save_load_local()




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


									'''simple docstring'''


									return super().test_dict_tuple_outputs_equivalent()




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


									'''simple docstring'''


									return super().test_save_load_optional_components()




					@skip_mps
					def     _lowercase						(   self:		str):


									'''simple docstring'''


									return super().test_attention_slicing_forward_pass()




					def     _lowercase						(   self:		List[str]):


									'''simple docstring'''


									pass




					def     _lowercase						(   self:		int):


									'''simple docstring'''


									pass







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


									'''simple docstring'''


									return super().test_progress_bar()





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




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


									'''simple docstring'''


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




					def     _lowercase						(   self:		List[Any]):


									'''simple docstring'''


									__lowerCAmelCase					=  torch_device

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

									__lowerCAmelCase					=  pipe.to(_lowercase)
									pipe.set_progress_bar_config(disable=_lowercase)

									def callback(_lowercase:		Dict,     _lowercase:		List[str]):
													# decode mel to audio
													__lowerCAmelCase					=  melgan(input_features=mel_output.astype(np.floataa))[0]
													assert len(audio[0]) == 81920 * (step + 1)
													# simulate that audio is played
													return audio

									__lowerCAmelCase					=  MidiProcessor()
									__lowerCAmelCase					=  processor(_lowercase)

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




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


									'''simple docstring'''


									__lowerCAmelCase					=  torch_device

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

									__lowerCAmelCase					=  processor(_lowercase)
									# just run two denoising loops
									__lowerCAmelCase					=  input_tokens[:2]

									__lowerCAmelCase					=  torch.manual_seed(0)
									__lowerCAmelCase					=  pipe(_lowercase,     num_inference_steps=2,     generator=_lowercase)

									__lowerCAmelCase					=  output.audios[0]

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







					def     _lowercase						(   self:		str):


									'''simple docstring'''


									__lowerCAmelCase					=  torch_device

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

									__lowerCAmelCase					=  MidiProcessor()

									__lowerCAmelCase					=  processor(_lowercase)

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

									__lowerCAmelCase					=  torch.manual_seed(0)
									__lowerCAmelCase					=  pipe(_lowercase,     num_inference_steps=100,     generator=_lowercase)

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









"""simple docstring"""






def __lowercase  (   lowerCamelCase_	:   str	):
				SCREAMING_SNAKE_CASE__  				=					str(lowerCamelCase_	).strip()
				if not oct_string:
								raise ValueError("Empty string was passed to the function"	)
				SCREAMING_SNAKE_CASE__  				=					oct_string[0] == "-"
				if is_negative:
								SCREAMING_SNAKE_CASE__  				=					oct_string[1:]
				if not oct_string.isdigit() or not all(0 <= int(lowerCamelCase_	) <= 7 for char in oct_string	):
								raise ValueError("Non-octal value was passed to the function"	)
				SCREAMING_SNAKE_CASE__  				=					0
				for char in oct_string:
								SCREAMING_SNAKE_CASE__  				=					8 * decimal_number + int(lowerCamelCase_	)
				if is_negative:
								SCREAMING_SNAKE_CASE__  				=					-decimal_number
				return decimal_number


if __name__ == "__main__":
	from doctest import testmod

	testmod()
