

"""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       :  List[Any]        = """./tests/fixtures/elise_format0.mid"""

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

			"""simple docstring"""

			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    _lowerCAmelCase			(							self     :					Any						):
								torch.manual_seed(0						)
								UpperCamelCase__		:    int         =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__		:    int         =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__		:    List[str]         =TaFilmDecoder(
								    input_dims=128						,  targets_length=256						,  max_decoder_noise_time=2_0_0_0_0.0						,  d_model=768						,  num_layers=1						,  num_heads=1						,  d_kv=4						,  d_ff=2048						,  dropout_rate=0.1						,  )

								UpperCamelCase__		:    Union[str, Any]         =DDPMScheduler()

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






			def    _lowerCAmelCase			(							self     :					Optional[int]						,  lowercase_     :					Optional[Any]						,  lowercase_     :					Dict=0						):
								if str(lowercase_						).startswith('''mps'''						):
													UpperCamelCase__		:    Union[str, Any]         =torch.manual_seed(lowercase_						)
								else:
													UpperCamelCase__		:    List[str]         =torch.Generator(device=lowercase_						).manual_seed(lowercase_						)
								UpperCamelCase__		:    Dict         ={
								    '''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    _lowerCAmelCase			(							self     :					Union[str, Any]						):
								UpperCamelCase__		:    Optional[Any]         ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
								UpperCamelCase__		:    List[Any]         =self.get_dummy_components()
								UpperCamelCase__		:    Union[str, Any]         =SpectrogramDiffusionPipeline(**lowercase_						)
								UpperCamelCase__		:    Optional[Any]         =pipe.to(lowercase_						)
								pipe.set_progress_bar_config(disable=lowercase_						)

								UpperCamelCase__		:    Optional[int]         =self.get_dummy_inputs(lowercase_						)
								UpperCamelCase__		:    Optional[Any]         =pipe(**lowercase_						)
								UpperCamelCase__		:    Tuple         =output.audios

								UpperCamelCase__		:    Union[str, Any]         =mel[0, -3:, -3:]

								assert mel_slice.shape == (3, 3)
								UpperCamelCase__		:    Any         =np.array(
								    [-1_1.5_1_2_9_2_5, -4.7_8_8_2_1_5, -0.4_6_1_7_2_9_0_5, -2.0_5_1_7_1_5, -1_0.5_3_9_1_4_7, -1_0.9_7_0_9_6_3, -9.0_9_1_6_3_4, 4.0, 4.0]						)
								assert np.abs(mel_slice.flatten() - expected_slice						).max() < 1e-2






			@skip_mps
			def    _lowerCAmelCase			(							self     :					str						):
								return super().test_save_load_local()






			@skip_mps
			def    _lowerCAmelCase			(							self     :					Tuple						):
								return super().test_dict_tuple_outputs_equivalent()






			@skip_mps
			def    _lowerCAmelCase			(							self     :					Optional[int]						):
								return super().test_save_load_optional_components()






			@skip_mps
			def    _lowerCAmelCase			(							self     :					Dict						):
								return super().test_attention_slicing_forward_pass()






			def    _lowerCAmelCase			(							self     :					Optional[Any]						):
								pass






			def    _lowerCAmelCase			(							self     :					Dict						):
								pass






			@skip_mps
			def    _lowerCAmelCase			(							self     :					Optional[Any]						):
								return super().test_progress_bar()



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

			"""simple docstring"""

			def    _lowerCAmelCase			(							self     :					Optional[Any]						):
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()






			def    _lowerCAmelCase			(							self     :					int						):
								# TODO - test that pipeline can decode tokens in a callback
								# so that music can be played live
								UpperCamelCase__		:    List[Any]         =torch_device

								UpperCamelCase__		:    str         =SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''						)
								UpperCamelCase__		:    Any         =pipe.melgan
								UpperCamelCase__		:    List[str]         =None

								UpperCamelCase__		:    Optional[Any]         =pipe.to(lowercase_						)
								pipe.set_progress_bar_config(disable=lowercase_						)

								def callback(lowercase_     :					Optional[Any]						,  lowercase_     :					Tuple						):
													# decode mel to audio
													UpperCamelCase__		:    Union[str, Any]         =melgan(input_features=mel_output.astype(np.floataa						)						)[0]
													assert len(audio[0]						) == 8_1920 * (step + 1)
													# simulate that audio is played
													return audio

								UpperCamelCase__		:    Tuple         =MidiProcessor()
								UpperCamelCase__		:    int         =processor(lowercase_						)

								UpperCamelCase__		:    List[Any]         =input_tokens[:3]
								UpperCamelCase__		:    Optional[Any]         =torch.manual_seed(0						)
								pipe(lowercase_						,  num_inference_steps=5						,  generator=lowercase_						,  callback=lowercase_						,  output_type='''mel'''						)






			def    _lowerCAmelCase			(							self     :					List[Any]						):
								UpperCamelCase__		:    int         =torch_device

								UpperCamelCase__		:    Dict         =SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''						)
								UpperCamelCase__		:    int         =pipe.to(lowercase_						)
								pipe.set_progress_bar_config(disable=lowercase_						)
								UpperCamelCase__		:    Optional[Any]         =MidiProcessor()

								UpperCamelCase__		:    List[Any]         =processor(lowercase_						)
								# just run two denoising loops
								UpperCamelCase__		:    List[str]         =input_tokens[:2]

								UpperCamelCase__		:    Dict         =torch.manual_seed(0						)
								UpperCamelCase__		:    int         =pipe(lowercase_						,  num_inference_steps=2						,  generator=lowercase_						)

								UpperCamelCase__		:    List[str]         =output.audios[0]

								assert abs(np.abs(lowercase_						).sum() - 3_6_1_2.8_4_1						) < 1e-1






			def    _lowerCAmelCase			(							self     :					Optional[int]						):
								UpperCamelCase__		:    str         =torch_device

								UpperCamelCase__		:    str         =SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''						)
								UpperCamelCase__		:    Optional[int]         =pipe.to(lowercase_						)
								pipe.set_progress_bar_config(disable=lowercase_						)

								UpperCamelCase__		:    Union[str, Any]         =MidiProcessor()

								UpperCamelCase__		:    int         =processor(lowercase_						)

								# just run 4 denoising loops
								UpperCamelCase__		:    int         =input_tokens[:4]

								UpperCamelCase__		:    Optional[int]         =torch.manual_seed(0						)
								UpperCamelCase__		:    str         =pipe(lowercase_						,  num_inference_steps=100						,  generator=lowercase_						)

								UpperCamelCase__		:    Optional[Any]         =output.audios[0]
								assert abs(np.abs(lowercase_						).sum() - 9_3_8_9.1_1_1_1						) < 5e-2





"""simple docstring"""





import math as m



def        _lowerCAmelCase      (						UpperCAmelCase     :							Optional[int]      ,							UpperCAmelCase     :							Optional[Any]      ,							UpperCAmelCase     :							Optional[Any]=0.001				):


					'''simple docstring'''






					return (f(a + h				) - f(a - h				)) / (2 * h)



def        _lowerCAmelCase      (						UpperCAmelCase     :							List[str]      ,							UpperCAmelCase     :							Optional[Any]=0      ,							UpperCAmelCase     :							Optional[int]=100      ,							UpperCAmelCase     :							Union[str, Any]=0.0001      ,							UpperCAmelCase     :							Dict=1E-6      ,							UpperCAmelCase     :							Any=False				):


					'''simple docstring'''






					UpperCamelCase__		:    Optional[Any]         =xa  # set the initial guess
					UpperCamelCase__		:    List[Any]         =[a]
					UpperCamelCase__		:    Optional[Any]         =abs(f(UpperCAmelCase				)				)
					UpperCamelCase__		:    List[str]         =lambda UpperCAmelCase				: calc_derivative(UpperCAmelCase      ,							UpperCAmelCase      ,							h=UpperCAmelCase				)  # noqa: E731  Derivative of f(x)
					for _ in range(UpperCAmelCase				):
										if fa(UpperCAmelCase				) == 0:
															raise ValueError('''No converging solution found'''				)
										UpperCamelCase__		:    Any         =a - f(UpperCAmelCase				) / fa(UpperCAmelCase				)  # Calculate the next estimate
										if logsteps:
															steps.append(UpperCAmelCase				)
										if error < maxerror:
															break
					else:
										raise ValueError('''Iteration limit reached, no converging solution found'''				)
					if logsteps:
										# If logstep is true, then log intermediate steps
										return a, error, steps
					return a, error


if __name__ == "__main__":
				from matplotlib import pyplot as plt

				_SCREAMING_SNAKE_CASE       :  str        = lambda x: m.tanh(x) ** 2 - m.exp(3 * x)  # noqa: E731
				_SCREAMING_SNAKE_CASE       ,       _SCREAMING_SNAKE_CASE       ,       _SCREAMING_SNAKE_CASE       :  Any        = newton_raphson(
				    f, xa=1_0, maxiter=1_0_0_0, step=1e-6, logsteps=True
				)
				plt.plot([abs(f(x)) for x in steps])
				plt.xlabel("""step""")
				plt.ylabel("""error""")
				plt.show()
				print(F'''solution = {{{solution:f}}}, error = {{{error:f}}}''')


