





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 				_UpperCamelCase       (						lowerCAmelCase_	,    unittest.TestCase  ):
       _UpperCamelCase					:       Dict       = SpectrogramDiffusionPipeline
       _UpperCamelCase					:       Dict       = PipelineTesterMixin.required_optional_params - {
           '''callback''',
           '''latents''',
           '''callback_steps''',
           '''output_type''',
           '''num_images_per_prompt''',
       }
       _UpperCamelCase					:       str       = False
       _UpperCamelCase					:       Optional[Any]       = False
       _UpperCamelCase					:       str       = TOKENS_TO_AUDIO_GENERATION_PARAMS
       _UpperCamelCase					:       Any       = TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS

       def lowercase			(		self:			Optional[Any]		)							->			int:




              """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=2_00_00.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 lowercase			(		self:			List[str]       , _SCREAMING_SNAKE_CASE:			Tuple       , _SCREAMING_SNAKE_CASE:			List[str]=0		)							->			List[Any]:




              """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 lowercase			(		self:			Union[str, Any]		)							->			Dict:




              """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.51_29_25, -4.78_82_15, -0.46_17_29_05, -2.05_17_15, -10.53_91_47, -10.97_09_63, -9.09_16_34, 4.0, 4.0]		)
              assert np.abs(mel_slice.flatten() - expected_slice		).max() < 1e-2

       @skip_mps
       def lowercase			(		self:			int		)							->			Any:




              """simple docstring"""


              return super().test_save_load_local()

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




              """simple docstring"""


              return super().test_dict_tuple_outputs_equivalent()

       @skip_mps
       def lowercase			(		self:			List[str]		)							->			Optional[int]:




              """simple docstring"""


              return super().test_save_load_optional_components()

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




              """simple docstring"""


              return super().test_attention_slicing_forward_pass()

       def lowercase			(		self:			Dict		)							->			List[str]:




              """simple docstring"""


              pass

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




              """simple docstring"""


              pass



       @skip_mps
       def lowercase			(		self:			str		)							->			Dict:




              """simple docstring"""


              return super().test_progress_bar()



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

       def lowercase			(		self:			Optional[int]		)							->			Union[str, Any]:




              """simple docstring"""


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

       def lowercase			(		self:			Optional[Any]		)							->			int:




              """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:			Union[str, Any]       , _SCREAMING_SNAKE_CASE:			int		):
                     # 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 lowercase			(		self:			Any		)							->			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() - 36_12.8_41		) < 1e-1



       def lowercase			(		self:			Tuple		)							->			Optional[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() - 93_89.11_11		) < 5e-2








from ...configuration_utils import PretrainedConfig
from ...utils import logging


_UpperCAmelCase										=					logging.get_logger(__name__)

_UpperCAmelCase										=					{
    'bert-base-uncased': 'https://huggingface.co/bert-base-uncased/resolve/main/config.json',
    # QDQBERT models can be loaded from any BERT checkpoint, available at https://huggingface.co/models?filter=bert
}



class 				_UpperCamelCase       (						lowerCAmelCase_  ):
       _UpperCamelCase					:       Union[str, Any]       = '''qdqbert'''

       def __init__(		self:			Union[str, Any]       , _SCREAMING_SNAKE_CASE:			Any=30522       , _SCREAMING_SNAKE_CASE:			List[str]=768       , _SCREAMING_SNAKE_CASE:			int=12       , _SCREAMING_SNAKE_CASE:			List[Any]=12       , _SCREAMING_SNAKE_CASE:			Union[str, Any]=3072       , _SCREAMING_SNAKE_CASE:			List[str]="gelu"       , _SCREAMING_SNAKE_CASE:			List[Any]=0.1       , _SCREAMING_SNAKE_CASE:			List[Any]=0.1       , _SCREAMING_SNAKE_CASE:			Dict=512       , _SCREAMING_SNAKE_CASE:			Union[str, Any]=2       , _SCREAMING_SNAKE_CASE:			int=0.02       , _SCREAMING_SNAKE_CASE:			Union[str, Any]=1e-12       , _SCREAMING_SNAKE_CASE:			str=True       , _SCREAMING_SNAKE_CASE:			Any=1       , _SCREAMING_SNAKE_CASE:			int=0       , _SCREAMING_SNAKE_CASE:			Tuple=2       , **_SCREAMING_SNAKE_CASE:			List[str]       , )							->			Tuple:




              """simple docstring"""


              super().__init__(pad_token_id=_SCREAMING_SNAKE_CASE       , bos_token_id=_SCREAMING_SNAKE_CASE       , eos_token_id=_SCREAMING_SNAKE_CASE       , **_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_ 							=	vocab_size
              UpperCamelCase_ 							=	max_position_embeddings
              UpperCamelCase_ 							=	hidden_size
              UpperCamelCase_ 							=	num_hidden_layers
              UpperCamelCase_ 							=	num_attention_heads
              UpperCamelCase_ 							=	intermediate_size
              UpperCamelCase_ 							=	hidden_act
              UpperCamelCase_ 							=	hidden_dropout_prob
              UpperCamelCase_ 							=	attention_probs_dropout_prob
              UpperCamelCase_ 							=	initializer_range
              UpperCamelCase_ 							=	type_vocab_size
              UpperCamelCase_ 							=	layer_norm_eps
              UpperCamelCase_ 							=	use_cache

