






"""simple docstring"""





import shutil
import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SpeechaTextTokenizer, is_speech_available
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_torch, require_torchaudio
from transformers.utils import FEATURE_EXTRACTOR_NAME

from .test_feature_extraction_speech_to_text import floats_list


if is_speech_available():
  from transformers import SpeechaTextFeatureExtractor, SpeechaTextProcessor


UpperCamelCase					=     get_tests_dir("""fixtures/test_sentencepiece.model""")







@require_torch
@require_torchaudio
@require_sentencepiece
class 						_a  (		unittest.TestCase       ):





   '''simple docstring'''






   def __lowercase							(       self)			->   List[str]:




          '''simple docstring'''




          lowercase__:							Tuple         =     tempfile.mkdtemp()

          lowercase__:							List[str]         =     ["<s>", "<pad>", "</s>", "<unk>", "▁This", "▁is", "▁a", "▁t", "est"]
          lowercase__:							Any         =     dict(zip(UpperCAmelCase_					, range(len(UpperCAmelCase_))))
          lowercase__:							Tuple         =     Path(self.tmpdirname)
          save_json(UpperCAmelCase_					, save_dir / VOCAB_FILES_NAMES["vocab_file"])
          if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
                 copyfile(UpperCAmelCase_					, save_dir / VOCAB_FILES_NAMES["spm_file"])

          lowercase__:							List[Any]         =     SpeechaTextTokenizer.from_pretrained(self.tmpdirname)
          tokenizer.save_pretrained(self.tmpdirname)

          lowercase__:							int         =     {
              "feature_size": 24,
              "num_mel_bins": 24,
              "padding_value": 0.0,
              "sampling_rate": 16_000,
              "return_attention_mask": False,
              "do_normalize": True,
          }
          save_json(UpperCAmelCase_					, save_dir / FEATURE_EXTRACTOR_NAME)





   def __lowercase							(       self					, **UpperCAmelCase_)			->   List[str]:




          '''simple docstring'''




          return SpeechaTextTokenizer.from_pretrained(self.tmpdirname					, **UpperCAmelCase_)





   def __lowercase							(       self					, **UpperCAmelCase_)			->   Dict:




          '''simple docstring'''




          return SpeechaTextFeatureExtractor.from_pretrained(self.tmpdirname					, **UpperCAmelCase_)





   def __lowercase							(       self)			->   Optional[int]:




          '''simple docstring'''




          shutil.rmtree(self.tmpdirname)





   def __lowercase							(       self)			->   str:




          '''simple docstring'''




          lowercase__:							Union[str, Any]         =     self.get_tokenizer()
          lowercase__:							Union[str, Any]         =     self.get_feature_extractor()

          lowercase__:							str         =     SpeechaTextProcessor(tokenizer=UpperCAmelCase_					, feature_extractor=UpperCAmelCase_)

          processor.save_pretrained(self.tmpdirname)
          lowercase__:							List[Any]         =     SpeechaTextProcessor.from_pretrained(self.tmpdirname)

          self.assertEqual(processor.tokenizer.get_vocab()					, tokenizer.get_vocab())
          self.assertIsInstance(processor.tokenizer					, UpperCAmelCase_)

          self.assertEqual(processor.feature_extractor.to_json_string()					, feature_extractor.to_json_string())
          self.assertIsInstance(processor.feature_extractor					, UpperCAmelCase_)





   def __lowercase							(       self)			->   Tuple:




          '''simple docstring'''




          lowercase__:							Dict         =     SpeechaTextProcessor(
              tokenizer=self.get_tokenizer()					, feature_extractor=self.get_feature_extractor())
          processor.save_pretrained(self.tmpdirname)

          lowercase__:							Tuple         =     self.get_tokenizer(bos_token="(BOS)"					, eos_token="(EOS)")
          lowercase__:							List[str]         =     self.get_feature_extractor(do_normalize=UpperCAmelCase_					, padding_value=1.0)

          lowercase__:							str         =     SpeechaTextProcessor.from_pretrained(
              self.tmpdirname					, bos_token="(BOS)"					, eos_token="(EOS)"					, do_normalize=UpperCAmelCase_					, padding_value=1.0)

          self.assertEqual(processor.tokenizer.get_vocab()					, tokenizer_add_kwargs.get_vocab())
          self.assertIsInstance(processor.tokenizer					, UpperCAmelCase_)

          self.assertEqual(processor.feature_extractor.to_json_string()					, feature_extractor_add_kwargs.to_json_string())
          self.assertIsInstance(processor.feature_extractor					, UpperCAmelCase_)





   def __lowercase							(       self)			->   Any:




          '''simple docstring'''




          lowercase__:							Optional[int]         =     self.get_feature_extractor()
          lowercase__:							str         =     self.get_tokenizer()

          lowercase__:							List[Any]         =     SpeechaTextProcessor(tokenizer=UpperCAmelCase_					, feature_extractor=UpperCAmelCase_)

          lowercase__:							Any         =     floats_list((3, 1_000))

          lowercase__:							List[Any]         =     feature_extractor(UpperCAmelCase_					, return_tensors="np")
          lowercase__:							Optional[Any]         =     processor(UpperCAmelCase_					, return_tensors="np")

          for key in input_feat_extract.keys():
                 self.assertAlmostEqual(input_feat_extract[key].sum()					, input_processor[key].sum()					, delta=1E-2)





   def __lowercase							(       self)			->   Tuple:




          '''simple docstring'''




          lowercase__:							Optional[Any]         =     self.get_feature_extractor()
          lowercase__:							Optional[int]         =     self.get_tokenizer()

          lowercase__:							Dict         =     SpeechaTextProcessor(tokenizer=UpperCAmelCase_					, feature_extractor=UpperCAmelCase_)

          lowercase__:							int         =     "This is a test string"

          lowercase__:							List[Any]         =     processor(text=UpperCAmelCase_)

          lowercase__:							Optional[Any]         =     tokenizer(UpperCAmelCase_)

          for key in encoded_tok.keys():
                 self.assertListEqual(encoded_tok[key]					, encoded_processor[key])





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




          '''simple docstring'''




          lowercase__:							Optional[Any]         =     self.get_feature_extractor()
          lowercase__:							Optional[Any]         =     self.get_tokenizer()

          lowercase__:							Tuple         =     SpeechaTextProcessor(tokenizer=UpperCAmelCase_					, feature_extractor=UpperCAmelCase_)

          lowercase__:							Union[str, Any]         =     [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

          lowercase__:							List[Any]         =     processor.batch_decode(UpperCAmelCase_)
          lowercase__:							Any         =     tokenizer.batch_decode(UpperCAmelCase_)

          self.assertListEqual(UpperCAmelCase_					, UpperCAmelCase_)





   def __lowercase							(       self)			->   Tuple:




          '''simple docstring'''




          lowercase__:							str         =     self.get_feature_extractor()
          lowercase__:							Union[str, Any]         =     self.get_tokenizer()

          lowercase__:							List[str]         =     SpeechaTextProcessor(tokenizer=UpperCAmelCase_					, feature_extractor=UpperCAmelCase_)

          self.assertListEqual(
              processor.model_input_names					, feature_extractor.model_input_names					, msg="`processor` and `feature_extractor` model input names do not match"					, )










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


snake_case__  				=			'./tests/fixtures/elise_format0.mid'



@unittest.skip('''The note-seq package currently throws an error on import'''	)
class      snake_case__	(   _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 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




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

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

										__magic_name__					:				Optional[Any]				=   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  ,   )

										__magic_name__					:				Any				=   DDPMScheduler()

										__magic_name__					:				Tuple				=   {
										    'notes_encoder': notes_encoder.eval(),
										    'continuous_encoder': continuous_encoder.eval(),
										    'decoder': decoder.eval(),
										    'scheduler': scheduler,
										    'melgan': None,
										}
										return components
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase  ,   _UpperCAmelCase=0   ):




										"""simple docstring"""




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




										"""simple docstring"""




										__magic_name__					:				Dict				=   'cpu'  # ensure determinism for the device-dependent torch.Generator
										__magic_name__					:				Optional[Any]				=   self.get_dummy_components()
										__magic_name__					:				Dict				=   SpectrogramDiffusionPipeline(**_UpperCAmelCase   )
										__magic_name__					:				Union[str, Any]				=   pipe.to(_UpperCAmelCase   )
										pipe.set_progress_bar_config(disable=_UpperCAmelCase   )

										__magic_name__					:				List[Any]				=   self.get_dummy_inputs(_UpperCAmelCase   )
										__magic_name__					:				Dict				=   pipe(**_UpperCAmelCase   )
										__magic_name__					:				Optional[Any]				=   output.audios

										__magic_name__					:				Dict				=   mel[0, -3:, -3:]

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




										"""simple docstring"""




										return super().test_save_load_local()
				@skip_mps
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										return super().test_dict_tuple_outputs_equivalent()
				@skip_mps
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										return super().test_save_load_optional_components()
				@skip_mps
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										return super().test_attention_slicing_forward_pass()
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										pass
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										pass





				@skip_mps
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										return super().test_progress_bar()






@slow
@require_torch_gpu
@require_onnxruntime
@require_note_seq
class      snake_case__	(   unittest.TestCase	):
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				Tuple				=   torch_device

										__magic_name__					:				Optional[int]				=   SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'   )
										__magic_name__					:				Dict				=   pipe.melgan
										__magic_name__					:				int				=   None

										__magic_name__					:				List[str]				=   pipe.to(_UpperCAmelCase   )
										pipe.set_progress_bar_config(disable=_UpperCAmelCase   )

										def callback(_UpperCAmelCase  ,   _UpperCAmelCase   ):
																# decode mel to audio
																__magic_name__					:				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

										__magic_name__					:				Optional[Any]				=   MidiProcessor()
										__magic_name__					:				str				=   processor(_UpperCAmelCase   )

										__magic_name__					:				Optional[int]				=   input_tokens[:3]
										__magic_name__					:				Optional[Any]				=   torch.manual_seed(0   )
										pipe(_UpperCAmelCase  ,   num_inference_steps=5  ,   generator=_UpperCAmelCase  ,   callback=_UpperCAmelCase  ,   output_type='mel'   )
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				Tuple				=   torch_device

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

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

										__magic_name__					:				int				=   torch.manual_seed(0   )
										__magic_name__					:				int				=   pipe(_UpperCAmelCase  ,   num_inference_steps=2  ,   generator=_UpperCAmelCase   )

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

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





				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				Any				=   torch_device

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

										__magic_name__					:				List[str]				=   MidiProcessor()

										__magic_name__					:				Union[str, Any]				=   processor(_UpperCAmelCase   )

										# just run 4 denoising loops
										__magic_name__					:				List[Any]				=   input_tokens[:4]

										__magic_name__					:				Dict				=   torch.manual_seed(0   )
										__magic_name__					:				Union[str, Any]				=   pipe(_UpperCAmelCase  ,   num_inference_steps=100  ,   generator=_UpperCAmelCase   )

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


