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






@unittest.skip('The note-seq package currently throws an error on import'    )
class    A     (	SCREAMING_SNAKE_CASE__  ,      unittest.TestCase    ):
	snake_case__      :Union[str, Any]        =     SpectrogramDiffusionPipeline
	snake_case__      :Any        =     PipelineTesterMixin.required_optional_params - {
	    'callback',
	    'latents',
	    'callback_steps',
	    'output_type',
	    'num_images_per_prompt',
	}
	snake_case__      :Optional[int]        =     False
	snake_case__      :Optional[Any]        =     False
	snake_case__      :List[Any]        =     TOKENS_TO_AUDIO_GENERATION_PARAMS
	snake_case__      :Union[str, Any]        =     TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS





	def 							__SCREAMING_SNAKE_CASE  (					self			:       Any							):


				"""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=2_0000.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 							__SCREAMING_SNAKE_CASE  (					self			:       Dict      ,							__magic_name__			:       Any      ,							__magic_name__			:       List[str]=0							):


				"""simple docstring"""

				if str(__magic_name__							).startswith("mps"							):
							lowerCAmelCase__								=							torch.manual_seed(__magic_name__							)
				else:
							lowerCAmelCase__								=							torch.Generator(device=__magic_name__							).manual_seed(__magic_name__							)
				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 							__SCREAMING_SNAKE_CASE  (					self			:       Optional[Any]							):


				"""simple docstring"""

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

				lowerCAmelCase__								=							self.get_dummy_inputs(__magic_name__							)
				lowerCAmelCase__								=							pipe(**__magic_name__							)
				lowerCAmelCase__								=							output.audios

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

				assert mel_slice.shape == (3, 3)
				lowerCAmelCase__								=							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 							__SCREAMING_SNAKE_CASE  (					self			:       List[str]							):


				"""simple docstring"""

				return super().test_save_load_local()





	@skip_mps
	def 							__SCREAMING_SNAKE_CASE  (					self			:       str							):


				"""simple docstring"""

				return super().test_dict_tuple_outputs_equivalent()





	@skip_mps
	def 							__SCREAMING_SNAKE_CASE  (					self			:       int							):


				"""simple docstring"""

				return super().test_save_load_optional_components()





	@skip_mps
	def 							__SCREAMING_SNAKE_CASE  (					self			:       Any							):


				"""simple docstring"""

				return super().test_attention_slicing_forward_pass()





	def 							__SCREAMING_SNAKE_CASE  (					self			:       Dict							):


				"""simple docstring"""

				pass





	def 							__SCREAMING_SNAKE_CASE  (					self			:       Optional[int]							):


				"""simple docstring"""

				pass







	@skip_mps
	def 							__SCREAMING_SNAKE_CASE  (					self			:       Any							):


				"""simple docstring"""

				return super().test_progress_bar()






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





	def 							__SCREAMING_SNAKE_CASE  (					self			:       List[str]							):


				"""simple docstring"""

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





	def 							__SCREAMING_SNAKE_CASE  (					self			:       Optional[int]							):


				"""simple docstring"""

				lowerCAmelCase__								=							torch_device

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

				lowerCAmelCase__								=							pipe.to(__magic_name__							)
				pipe.set_progress_bar_config(disable=__magic_name__							)

				def callback(__magic_name__			:       Tuple      ,							__magic_name__			:       Optional[int]							):
							# 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(__magic_name__							)

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





	def 							__SCREAMING_SNAKE_CASE  (					self			:       int							):


				"""simple docstring"""

				lowerCAmelCase__								=							torch_device

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

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

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

				lowerCAmelCase__								=							output.audios[0]

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







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


				"""simple docstring"""

				lowerCAmelCase__								=							torch_device

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

				lowerCAmelCase__								=							MidiProcessor()

				lowerCAmelCase__								=							processor(__magic_name__							)

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

				lowerCAmelCase__								=							torch.manual_seed(0							)
				lowerCAmelCase__								=							pipe(__magic_name__      ,							num_inference_steps=100      ,							generator=__magic_name__							)

				lowerCAmelCase__								=							output.audios[0]
				assert abs(np.abs(__magic_name__							).sum() - 9389.1111							) < 5E-2







'''simple docstring'''

from __future__ import annotations







def   A     (      UpperCamelCase_					:     int      )					->       list[int]:







			'''simple docstring'''



			assert (
			    isinstance(UpperCamelCase_						,	UpperCamelCase_      ) and mask > 0
			), F"""mask needs to be positive integer, your input {mask}"""
			lowerCAmelCase__								=							[]
			lowerCAmelCase__								=							mask

			while submask:
						all_submasks.append(UpperCamelCase_      )
						lowerCAmelCase__								=							(submask - 1) & mask

			return all_submasks


if __name__ == "__main__":
		import doctest

		doctest.testmod()





