



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 						lowerCamelCase__			(		_a					,       unittest.TestCase ):
						_lowerCAmelCase             =			SpectrogramDiffusionPipeline
						_lowerCAmelCase             =			PipelineTesterMixin.required_optional_params - {
						    '''callback''',
						    '''latents''',
						    '''callback_steps''',
						    '''output_type''',
						    '''num_images_per_prompt''',
						}
						_lowerCAmelCase             =			False
						_lowerCAmelCase             =			False
						_lowerCAmelCase             =			TOKENS_TO_AUDIO_GENERATION_PARAMS
						_lowerCAmelCase             =			TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS


						def 	_lowerCamelCase      (			self    :	Dict      ):
													torch.manual_seed(0      )
													a__:							List[Any]									=SpectrogramNotesEncoder(
													    max_length=2_0_4_8   ,							vocab_size=1_5_3_6   ,							d_model=7_6_8   ,							dropout_rate=0.1   ,							num_layers=1   ,							num_heads=1   ,							d_kv=4   ,							d_ff=2_0_4_8   ,							feed_forward_proj="gated-gelu"   ,							)

													a__:							int									=SpectrogramContEncoder(
													    input_dims=1_2_8   ,							targets_context_length=2_5_6   ,							d_model=7_6_8   ,							dropout_rate=0.1   ,							num_layers=1   ,							num_heads=1   ,							d_kv=4   ,							d_ff=2_0_4_8   ,							feed_forward_proj="gated-gelu"   ,							)

													a__:							Dict									=TaFilmDecoder(
													    input_dims=1_2_8   ,							targets_length=2_5_6   ,							max_decoder_noise_time=2_0_0_0_0.0   ,							d_model=7_6_8   ,							num_layers=1   ,							num_heads=1   ,							d_kv=4   ,							d_ff=2_0_4_8   ,							dropout_rate=0.1   ,							)

													a__:							str									=DDPMScheduler()

													a__:							List[str]									={
													    "notes_encoder": notes_encoder.eval(),
													    "continuous_encoder": continuous_encoder.eval(),
													    "decoder": decoder.eval(),
													    "scheduler": scheduler,
													    "melgan": None,
													}
													return components


						def 	_lowerCamelCase      (			self    :	int   ,							_a    :	Optional[Any]   ,							_a    :	int=0      ):
													if str(_a      ).startswith("mps"      ):
																				a__:							Optional[int]									=torch.manual_seed(_a      )
													else:
																				a__:							Optional[int]									=torch.Generator(device=_a      ).manual_seed(_a      )
													a__:							Union[str, Any]									={
													    "input_tokens": [
													        [1_1_3_4, 9_0, 1_1_3_5, 1_1_3_3, 1_0_8_0, 1_1_2, 1_1_3_2, 1_0_8_0, 1_1_3_3, 1_0_7_9, 1_3_3, 1_1_3_2, 1_0_7_9, 1_1_3_3, 1] + [0] * 2_0_3_3
													    ],
													    "generator": generator,
													    "num_inference_steps": 4,
													    "output_type": "mel",
													}
													return inputs


						def 	_lowerCamelCase      (			self    :	str      ):
													a__:							Optional[int]									="cpu"  # ensure determinism for the device-dependent torch.Generator
													a__:							List[Any]									=self.get_dummy_components()
													a__:							List[str]									=SpectrogramDiffusionPipeline(**_a      )
													a__:							Any									=pipe.to(_a      )
													pipe.set_progress_bar_config(disable=_a      )

													a__:							Optional[Any]									=self.get_dummy_inputs(_a      )
													a__:							Optional[Any]									=pipe(**_a      )
													a__:							Tuple									=output.audios

													a__:							List[str]									=mel[0, -3:, -3:]

													assert mel_slice.shape == (3, 3)
													a__:							Optional[int]									=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    :	int      ):
													return super().test_save_load_local()


						@skip_mps
						def 	_lowerCamelCase      (			self    :	List[str]      ):
													return super().test_dict_tuple_outputs_equivalent()


						@skip_mps
						def 	_lowerCamelCase      (			self    :	List[Any]      ):
													return super().test_save_load_optional_components()


						@skip_mps
						def 	_lowerCamelCase      (			self    :	Tuple      ):
													return super().test_attention_slicing_forward_pass()


						def 	_lowerCamelCase      (			self    :	Optional[int]      ):
													pass


						def 	_lowerCamelCase      (			self    :	str      ):
													pass







						@skip_mps
						def 	_lowerCamelCase      (			self    :	Dict      ):
													return super().test_progress_bar()


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


						def 	_lowerCamelCase      (			self    :	Dict      ):
													# clean up the VRAM after each test
													super().tearDown()
													gc.collect()
													torch.cuda.empty_cache()


						def 	_lowerCamelCase      (			self    :	List[str]      ):
													# TODO - test that pipeline can decode tokens in a callback
													# so that music can be played live
													a__:							Tuple									=torch_device

													a__:							Tuple									=SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion"      )
													a__:							List[Any]									=pipe.melgan
													a__:							int									=None

													a__:							List[str]									=pipe.to(_a      )
													pipe.set_progress_bar_config(disable=_a      )

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

													a__:							Any									=MidiProcessor()
													a__:							Dict									=processor(_a      )

													a__:							Optional[Any]									=input_tokens[:3]
													a__:							Tuple									=torch.manual_seed(0      )
													pipe(_a   ,							num_inference_steps=5   ,							generator=_a   ,							callback=_a   ,							output_type="mel"      )


						def 	_lowerCamelCase      (			self    :	Dict      ):
													a__:							str									=torch_device

													a__:							Dict									=SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion"      )
													a__:							Optional[Any]									=pipe.to(_a      )
													pipe.set_progress_bar_config(disable=_a      )
													a__:							List[str]									=MidiProcessor()

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

													a__:							int									=torch.manual_seed(0      )
													a__:							Tuple									=pipe(_a   ,							num_inference_steps=2   ,							generator=_a      )

													a__:							Dict									=output.audios[0]

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







						def 	_lowerCamelCase      (			self    :	Optional[Any]      ):
													a__:							List[Any]									=torch_device

													a__:							Dict									=SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion"      )
													a__:							Dict									=pipe.to(_a      )
													pipe.set_progress_bar_config(disable=_a      )

													a__:							Optional[Any]									=MidiProcessor()

													a__:							Any									=processor(_a      )

													# just run 4 denoising loops
													a__:							Tuple									=input_tokens[:4]

													a__:							Optional[Any]									=torch.manual_seed(0      )
													a__:							List[str]									=pipe(_a   ,							num_inference_steps=1_0_0   ,							generator=_a      )

													a__:							Dict									=output.audios[0]
													assert abs(np.abs(_a      ).sum() - 9_3_8_9.1_1_1_1      ) < 5e-2










from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_torch_available


__UpperCAmelCase    	=  {
    '''configuration_ernie_m''': ['''ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ErnieMConfig'''],
}

try:
		if not is_sentencepiece_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__UpperCAmelCase    	=  ['''ErnieMTokenizer''']

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__UpperCAmelCase    	=  [
		    '''ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''ErnieMForMultipleChoice''',
		    '''ErnieMForQuestionAnswering''',
		    '''ErnieMForSequenceClassification''',
		    '''ErnieMForTokenClassification''',
		    '''ErnieMModel''',
		    '''ErnieMPreTrainedModel''',
		    '''ErnieMForInformationExtraction''',
		]


if TYPE_CHECKING:
		from .configuration_ernie_m import ERNIE_M_PRETRAINED_CONFIG_ARCHIVE_MAP, ErnieMConfig

		try:
				if not is_sentencepiece_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tokenization_ernie_m import ErnieMTokenizer

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_ernie_m import (
				    ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST,
				    ErnieMForInformationExtraction,
				    ErnieMForMultipleChoice,
				    ErnieMForQuestionAnswering,
				    ErnieMForSequenceClassification,
				    ErnieMForTokenClassification,
				    ErnieMModel,
				    ErnieMPreTrainedModel,
				)


else:
		import sys

		__UpperCAmelCase    	=  _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





