


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


_A      =     """./tests/fixtures/elise_format0.mid"""






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




			def 			_snake_case   (      self       )			->						int:


								'''simple docstring'''






								torch.manual_seed(0       )
								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = 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'''					,      )

								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = 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'''					,      )

								__SCREAMING_SNAKE_CASE					:							List[str]			   = 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					,      )

								__SCREAMING_SNAKE_CASE					:							int			   = DDPMScheduler()

								__SCREAMING_SNAKE_CASE					:							List[Any]			   = {
								    '''notes_encoder''': notes_encoder.eval(),
								    '''continuous_encoder''': continuous_encoder.eval(),
								    '''decoder''': decoder.eval(),
								    '''scheduler''': scheduler,
								    '''melgan''': None,
								}
								return components




			def 			_snake_case   (      self					,      lowercase					,      lowercase=0       )			->						int:


								'''simple docstring'''






								if str(lowercase       ).startswith('''mps'''       ):
													__SCREAMING_SNAKE_CASE					:							List[Any]			   = torch.manual_seed(lowercase       )
								else:
													__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = torch.Generator(device=lowercase       ).manual_seed(lowercase       )
								__SCREAMING_SNAKE_CASE					:							Dict			   = {
								    '''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 			_snake_case   (      self       )			->						Union[str, Any]:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Tuple			   = '''cpu'''  # ensure determinism for the device-dependent torch.Generator
								__SCREAMING_SNAKE_CASE					:							Any			   = self.get_dummy_components()
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = SpectrogramDiffusionPipeline(**lowercase       )
								__SCREAMING_SNAKE_CASE					:							int			   = pipe.to(lowercase       )
								pipe.set_progress_bar_config(disable=lowercase       )

								__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = self.get_dummy_inputs(lowercase       )
								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = pipe(**lowercase       )
								__SCREAMING_SNAKE_CASE					:							str			   = output.audios

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

								assert mel_slice.shape == (3, 3)
								__SCREAMING_SNAKE_CASE					:							Union[str, 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 			_snake_case   (      self       )			->						Any:


								'''simple docstring'''






								return super().test_save_load_local()




			@skip_mps
			def 			_snake_case   (      self       )			->						Union[str, Any]:


								'''simple docstring'''






								return super().test_dict_tuple_outputs_equivalent()




			@skip_mps
			def 			_snake_case   (      self       )			->						Any:


								'''simple docstring'''






								return super().test_save_load_optional_components()




			@skip_mps
			def 			_snake_case   (      self       )			->						int:


								'''simple docstring'''






								return super().test_attention_slicing_forward_pass()




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


								'''simple docstring'''






								pass




			def 			_snake_case   (      self       )			->						Dict:


								'''simple docstring'''






								pass




			@skip_mps
			def 			_snake_case   (      self       )			->						Dict:


								'''simple docstring'''






								return super().test_progress_bar()






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




			def 			_snake_case   (      self       )			->						Tuple:


								'''simple docstring'''






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




			def 			_snake_case   (      self       )			->						Tuple:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							List[Any]			   = torch_device

								__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''       )
								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = pipe.melgan
								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = None

								__SCREAMING_SNAKE_CASE					:							List[Any]			   = pipe.to(lowercase       )
								pipe.set_progress_bar_config(disable=lowercase       )

								def callback(lowercase					,      lowercase       ):
													# decode mel to audio
													__SCREAMING_SNAKE_CASE					:							str			   = 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

								__SCREAMING_SNAKE_CASE					:							str			   = MidiProcessor()
								__SCREAMING_SNAKE_CASE					:							Tuple			   = processor(lowercase       )

								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = input_tokens[:3]
								__SCREAMING_SNAKE_CASE					:							Tuple			   = torch.manual_seed(0       )
								pipe(lowercase					,      num_inference_steps=5					,      generator=lowercase					,      callback=lowercase					,      output_type='''mel'''       )




			def 			_snake_case   (      self       )			->						str:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							int			   = torch_device

								__SCREAMING_SNAKE_CASE					:							str			   = SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''       )
								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = pipe.to(lowercase       )
								pipe.set_progress_bar_config(disable=lowercase       )
								__SCREAMING_SNAKE_CASE					:							Tuple			   = MidiProcessor()

								__SCREAMING_SNAKE_CASE					:							List[Any]			   = processor(lowercase       )
								# just run two denoising loops
								__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = input_tokens[:2]

								__SCREAMING_SNAKE_CASE					:							Any			   = torch.manual_seed(0       )
								__SCREAMING_SNAKE_CASE					:							Tuple			   = pipe(lowercase					,      num_inference_steps=2					,      generator=lowercase       )

								__SCREAMING_SNAKE_CASE					:							Any			   = output.audios[0]

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




			def 			_snake_case   (      self       )			->						Union[str, Any]:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							List[str]			   = torch_device

								__SCREAMING_SNAKE_CASE					:							List[str]			   = SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''       )
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = pipe.to(lowercase       )
								pipe.set_progress_bar_config(disable=lowercase       )

								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = MidiProcessor()

								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = processor(lowercase       )

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

								__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = torch.manual_seed(0       )
								__SCREAMING_SNAKE_CASE					:							str			   = pipe(lowercase					,      num_inference_steps=1_0_0					,      generator=lowercase       )

								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = output.audios[0]
								assert abs(np.abs(lowercase       ).sum() - 9_3_8_9.1_1_1_1       ) < 5e-2






'''simple docstring'''






# this script builds a small sample spm file tests/fixtures/test_sentencepiece_no_bos.model, with features needed by pegasus

# 1. pip install sentencepiece
#
# 2. wget https://raw.githubusercontent.com/google/sentencepiece/master/data/botchan.txt

# 3. build
import sentencepiece as spm

# pegasus:
# 1. no bos
# 2. eos_id is 1
# 3. unk_id is 2
# build a sample spm file accordingly
spm.SentencePieceTrainer.train(
    """--input=botchan.txt --model_prefix=test_sentencepiece_no_bos --bos_id=-1 --unk_id=2  --eos_id=1  --vocab_size=1000"""
)

# 4. now update the fixture
# mv test_sentencepiece_no_bos.model ../../tests/fixtures/


