


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


lowerCAmelCase__						:Tuple         =							'''./tests/fixtures/elise_format0.mid'''

@unittest.skip('The note-seq package currently throws an error on import'    )
class 	__a   (       UpperCAmelCase					,   unittest.TestCase    ):
		_a			:						Tuple    			=			SpectrogramDiffusionPipeline
		_a			:						str    			=			PipelineTesterMixin.required_optional_params - {
		    'callback',
		    'latents',
		    'callback_steps',
		    'output_type',
		    'num_images_per_prompt',
		}
		_a			:						Tuple    			=			False
		_a			:						Tuple    			=			False
		_a			:						Union[str, Any]    			=			TOKENS_TO_AUDIO_GENERATION_PARAMS
		_a			:						Union[str, Any]    			=			TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS
		def 			UpperCAmelCase__       (    self		)				-> List[str]:




									"""simple docstring"""





									torch.manual_seed(0		)
									_UpperCAmelCase					=    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'    ,					)

									_UpperCAmelCase					=    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'    ,					)

									_UpperCAmelCase					=    TaFilmDecoder(
									    input_dims=128    ,					targets_length=256    ,					max_decoder_noise_time=20000.0    ,					d_model=768    ,					num_layers=1    ,					num_heads=1    ,					d_kv=4    ,					d_ff=2048    ,					dropout_rate=0.1    ,					)

									_UpperCAmelCase					=    DDPMScheduler()

									_UpperCAmelCase					=    {
									    'notes_encoder': notes_encoder.eval(),
									    'continuous_encoder': continuous_encoder.eval(),
									    'decoder': decoder.eval(),
									    'scheduler': scheduler,
									    'melgan': None,
									}
									return components
		def 			UpperCAmelCase__       (    self    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE=0		)				-> List[str]:




									"""simple docstring"""





									if str(_SCREAMING_SNAKE_CASE		).startswith('mps'		):
																_UpperCAmelCase					=    torch.manual_seed(_SCREAMING_SNAKE_CASE		)
									else:
																_UpperCAmelCase					=    torch.Generator(device=_SCREAMING_SNAKE_CASE		).manual_seed(_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase					=    {
									    '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 			UpperCAmelCase__       (    self		)				-> List[str]:




									"""simple docstring"""





									_UpperCAmelCase					=    'cpu'  # ensure determinism for the device-dependent torch.Generator
									_UpperCAmelCase					=    self.get_dummy_components()
									_UpperCAmelCase					=    SpectrogramDiffusionPipeline(**_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase					=    pipe.to(_SCREAMING_SNAKE_CASE		)
									pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase					=    self.get_dummy_inputs(_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase					=    pipe(**_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase					=    output.audios

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

									assert mel_slice.shape == (3, 3)
									_UpperCAmelCase					=    np.array(
									    [-11.512925, -4.788215, -0.46172905, -2.051715, -10.539147, -10.970963, -9.091634, 4.0, 4.0]		)
									assert np.abs(mel_slice.flatten() - expected_slice		).max() < 1e-2
		@skip_mps
		def 			UpperCAmelCase__       (    self		)				-> Any:




									"""simple docstring"""





									return super().test_save_load_local()
		@skip_mps
		def 			UpperCAmelCase__       (    self		)				-> str:




									"""simple docstring"""





									return super().test_dict_tuple_outputs_equivalent()
		@skip_mps
		def 			UpperCAmelCase__       (    self		)				-> Dict:




									"""simple docstring"""





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




									"""simple docstring"""





									return super().test_attention_slicing_forward_pass()
		def 			UpperCAmelCase__       (    self		)				-> Dict:




									"""simple docstring"""





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




									"""simple docstring"""





									pass






		@skip_mps
		def 			UpperCAmelCase__       (    self		)				-> List[Any]:




									"""simple docstring"""





									return super().test_progress_bar()



@slow
@require_torch_gpu
@require_onnxruntime
@require_note_seq
class 	__a   (       unittest.TestCase    ):
		def 			UpperCAmelCase__       (    self		)				-> Dict:




									"""simple docstring"""





									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()
		def 			UpperCAmelCase__       (    self		)				-> List[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=    torch_device

									_UpperCAmelCase					=    SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'		)
									_UpperCAmelCase					=    pipe.melgan
									_UpperCAmelCase					=    None

									_UpperCAmelCase					=    pipe.to(_SCREAMING_SNAKE_CASE		)
									pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

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

									_UpperCAmelCase					=    MidiProcessor()
									_UpperCAmelCase					=    processor(_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase					=    input_tokens[:3]
									_UpperCAmelCase					=    torch.manual_seed(0		)
									pipe(_SCREAMING_SNAKE_CASE    ,					num_inference_steps=5    ,					generator=_SCREAMING_SNAKE_CASE    ,					callback=_SCREAMING_SNAKE_CASE    ,					output_type='mel'		)
		def 			UpperCAmelCase__       (    self		)				-> int:




									"""simple docstring"""





									_UpperCAmelCase					=    torch_device

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

									_UpperCAmelCase					=    processor(_SCREAMING_SNAKE_CASE		)
									# just run two denoising loops
									_UpperCAmelCase					=    input_tokens[:2]

									_UpperCAmelCase					=    torch.manual_seed(0		)
									_UpperCAmelCase					=    pipe(_SCREAMING_SNAKE_CASE    ,					num_inference_steps=2    ,					generator=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase					=    output.audios[0]

									assert abs(np.abs(_SCREAMING_SNAKE_CASE		).sum() - 3612.841		) < 1e-1






		def 			UpperCAmelCase__       (    self		)				-> List[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=    torch_device

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

									_UpperCAmelCase					=    MidiProcessor()

									_UpperCAmelCase					=    processor(_SCREAMING_SNAKE_CASE		)

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

									_UpperCAmelCase					=    torch.manual_seed(0		)
									_UpperCAmelCase					=    pipe(_SCREAMING_SNAKE_CASE    ,					num_inference_steps=100    ,					generator=_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase					=    output.audios[0]
									assert abs(np.abs(_SCREAMING_SNAKE_CASE		).sum() - 9389.1111		) < 5e-2




import tempfile
import unittest

import numpy as np
import torch

from diffusers import ScoreSdeVeScheduler

class 	__a   (       unittest.TestCase    ):
		# TODO adapt with class SchedulerCommonTest (scheduler needs Numpy Integration)
		_a			:						List[Any]    			=			(ScoreSdeVeScheduler,)
		_a			:						Any    			=			()
		@property
		def 			UpperCAmelCase__       (    self		)				-> List[str]:




									"""simple docstring"""





									_UpperCAmelCase					=    4
									_UpperCAmelCase					=    3
									_UpperCAmelCase					=    8
									_UpperCAmelCase					=    8

									_UpperCAmelCase					=    torch.rand((batch_size, num_channels, height, width)		)

									return sample
		@property
		def 			UpperCAmelCase__       (    self		)				-> Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=    4
									_UpperCAmelCase					=    3
									_UpperCAmelCase					=    8
									_UpperCAmelCase					=    8

									_UpperCAmelCase					=    batch_size * num_channels * height * width
									_UpperCAmelCase					=    torch.arange(_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase					=    sample.reshape(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase					=    sample / num_elems
									_UpperCAmelCase					=    sample.permute(3    ,					0    ,					1    ,					2		)

									return sample
		def 			UpperCAmelCase__       (    self		)				-> str:




									"""simple docstring"""





									def model(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					*_SCREAMING_SNAKE_CASE		):
																return sample * t / (t + 1)

									return model
		def 			UpperCAmelCase__       (    self    ,					**_SCREAMING_SNAKE_CASE		)				-> Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase					=    {
									    'num_train_timesteps': 2000,
									    'snr': 0.15,
									    'sigma_min': 0.01,
									    'sigma_max': 1348,
									    'sampling_eps': 1e-5,
									}

									config.update(**_SCREAMING_SNAKE_CASE		)
									return config
		def 			UpperCAmelCase__       (    self    ,					_SCREAMING_SNAKE_CASE=0    ,					**_SCREAMING_SNAKE_CASE		)				-> List[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=    dict(self.forward_default_kwargs		)

									for scheduler_class in self.scheduler_classes:
																_UpperCAmelCase					=    self.dummy_sample
																_UpperCAmelCase					=    0.1 * sample

																_UpperCAmelCase					=    self.get_scheduler_config(**_SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase					=    scheduler_class(**_SCREAMING_SNAKE_CASE		)

																with tempfile.TemporaryDirectory() as tmpdirname:
																							scheduler.save_config(_SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase					=    scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE		)

																_UpperCAmelCase					=    scheduler.step_pred(
																    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample
																_UpperCAmelCase					=    new_scheduler.step_pred(
																    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample

																assert torch.sum(torch.abs(output - new_output		)		) < 1e-5, "Scheduler outputs are not identical"

																_UpperCAmelCase					=    scheduler.step_correct(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample
																_UpperCAmelCase					=    new_scheduler.step_correct(
																    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample

																assert torch.sum(torch.abs(output - new_output		)		) < 1e-5, "Scheduler correction are not identical"
		def 			UpperCAmelCase__       (    self    ,					_SCREAMING_SNAKE_CASE=0    ,					**_SCREAMING_SNAKE_CASE		)				-> List[str]:




									"""simple docstring"""





									_UpperCAmelCase					=    dict(self.forward_default_kwargs		)
									kwargs.update(_SCREAMING_SNAKE_CASE		)

									for scheduler_class in self.scheduler_classes:
																_UpperCAmelCase					=    self.dummy_sample
																_UpperCAmelCase					=    0.1 * sample

																_UpperCAmelCase					=    self.get_scheduler_config()
																_UpperCAmelCase					=    scheduler_class(**_SCREAMING_SNAKE_CASE		)

																with tempfile.TemporaryDirectory() as tmpdirname:
																							scheduler.save_config(_SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase					=    scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE		)

																_UpperCAmelCase					=    scheduler.step_pred(
																    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample
																_UpperCAmelCase					=    new_scheduler.step_pred(
																    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample

																assert torch.sum(torch.abs(output - new_output		)		) < 1e-5, "Scheduler outputs are not identical"

																_UpperCAmelCase					=    scheduler.step_correct(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample
																_UpperCAmelCase					=    new_scheduler.step_correct(
																    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample

																assert torch.sum(torch.abs(output - new_output		)		) < 1e-5, "Scheduler correction are not identical"
		def 			UpperCAmelCase__       (    self		)				-> Any:




									"""simple docstring"""





									for timesteps in [10, 100, 1000]:
																self.check_over_configs(num_train_timesteps=_SCREAMING_SNAKE_CASE		)
		def 			UpperCAmelCase__       (    self		)				-> List[Any]:




									"""simple docstring"""





									for sigma_min, sigma_max in zip([0.0001, 0.001, 0.01]    ,					[1, 100, 1000]		):
																self.check_over_configs(sigma_min=_SCREAMING_SNAKE_CASE    ,					sigma_max=_SCREAMING_SNAKE_CASE		)
		def 			UpperCAmelCase__       (    self		)				-> str:




									"""simple docstring"""





									for t in [0.1, 0.5, 0.75]:
																self.check_over_forward(time_step=_SCREAMING_SNAKE_CASE		)
		def 			UpperCAmelCase__       (    self		)				-> Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase					=    dict(self.forward_default_kwargs		)

									_UpperCAmelCase					=    self.scheduler_classes[0]
									_UpperCAmelCase					=    self.get_scheduler_config()
									_UpperCAmelCase					=    scheduler_class(**_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase					=    3

									_UpperCAmelCase					=    self.dummy_model()
									_UpperCAmelCase					=    self.dummy_sample_deter

									scheduler.set_sigmas(_SCREAMING_SNAKE_CASE		)
									scheduler.set_timesteps(_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase					=    torch.manual_seed(0		)

									for i, t in enumerate(scheduler.timesteps		):
																_UpperCAmelCase					=    scheduler.sigmas[i]

																for _ in range(scheduler.config.correct_steps		):
																							with torch.no_grad():
																														_UpperCAmelCase					=    model(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase					=    scheduler.step_correct(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=_SCREAMING_SNAKE_CASE    ,					**_SCREAMING_SNAKE_CASE		).prev_sample

																with torch.no_grad():
																							_UpperCAmelCase					=    model(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE		)

																_UpperCAmelCase					=    scheduler.step_pred(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					generator=_SCREAMING_SNAKE_CASE    ,					**_SCREAMING_SNAKE_CASE		)
																_UpperCAmelCase	, _UpperCAmelCase					=    output.prev_sample, output.prev_sample_mean

									_UpperCAmelCase					=    torch.sum(torch.abs(_SCREAMING_SNAKE_CASE		)		)
									_UpperCAmelCase					=    torch.mean(torch.abs(_SCREAMING_SNAKE_CASE		)		)

									assert np.isclose(result_sum.item()    ,					14372758528.0		)
									assert np.isclose(result_mean.item()    ,					18714530.0		)






		def 			UpperCAmelCase__       (    self		)				-> Tuple:




									"""simple docstring"""





									_UpperCAmelCase					=    dict(self.forward_default_kwargs		)

									_UpperCAmelCase					=    kwargs.pop('num_inference_steps'    ,					_SCREAMING_SNAKE_CASE		)

									for scheduler_class in self.scheduler_classes:
																_UpperCAmelCase					=    self.get_scheduler_config()
																_UpperCAmelCase					=    scheduler_class(**_SCREAMING_SNAKE_CASE		)

																_UpperCAmelCase					=    self.dummy_sample
																_UpperCAmelCase					=    0.1 * sample

																if num_inference_steps is not None and hasattr(_SCREAMING_SNAKE_CASE    ,					'set_timesteps'		):
																							scheduler.set_timesteps(_SCREAMING_SNAKE_CASE		)
																elif num_inference_steps is not None and not hasattr(_SCREAMING_SNAKE_CASE    ,					'set_timesteps'		):
																							_UpperCAmelCase					=    num_inference_steps

																_UpperCAmelCase					=    scheduler.step_pred(_SCREAMING_SNAKE_CASE    ,					0    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample
																_UpperCAmelCase					=    scheduler.step_pred(_SCREAMING_SNAKE_CASE    ,					1    ,					_SCREAMING_SNAKE_CASE    ,					generator=torch.manual_seed(0		)    ,					**_SCREAMING_SNAKE_CASE		).prev_sample

																self.assertEqual(output_a.shape    ,					sample.shape		)
																self.assertEqual(output_a.shape    ,					output_a.shape		)
