


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


lowerCamelCase_  :  Tuple     		=    """./tests/fixtures/elise_format0.mid"""






@unittest.skip("The note-seq package currently throws an error on import"     )
class  __A							(			_SCREAMING_SNAKE_CASE,       unittest.TestCase     ):






 """simple docstring"""




 __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  SCREAMING_SNAKE_CASE				(						self		)   ->	Tuple:
     torch.manual_seed(0		)
     a											=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'''					,       )

     a											=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'''					,       )

     a											=TaFilmDecoder(
         input_dims=128					,       targets_length=256					,       max_decoder_noise_time=20_000.0					,       d_model=768					,       num_layers=1					,       num_heads=1					,       d_kv=4					,       d_ff=2048					,       dropout_rate=0.1					,       )

     a											=DDPMScheduler()

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


 def  SCREAMING_SNAKE_CASE				(						self					,       __A					,       __A=0		)   ->	Tuple:
     if str(__A		).startswith('''mps'''		):
         a											=torch.manual_seed(__A		)
     else:
         a											=torch.Generator(device=__A		).manual_seed(__A		)
     a											={
         '''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  SCREAMING_SNAKE_CASE				(						self		)   ->	List[str]:
     a											='''cpu'''  # ensure determinism for the device-dependent torch.Generator
     a											=self.get_dummy_components()
     a											=SpectrogramDiffusionPipeline(**__A		)
     a											=pipe.to(__A		)
     pipe.set_progress_bar_config(disable=__A		)

     a											=self.get_dummy_inputs(__A		)
     a											=pipe(**__A		)
     a											=output.audios

     a											=mel[0, -3:, -3:]

     assert mel_slice.shape == (3, 3)
     a											=np.array(
         [-11.512_925, -4.788_215, -0.46_172_905, -2.051_715, -10.539_147, -10.970_963, -9.091_634, 4.0, 4.0]		)
     assert np.abs(mel_slice.flatten() - expected_slice		).max() < 1E-2


 @skip_mps
 def  SCREAMING_SNAKE_CASE				(						self		)   ->	int:
     return super().test_save_load_local()


 @skip_mps
 def  SCREAMING_SNAKE_CASE				(						self		)   ->	Any:
     return super().test_dict_tuple_outputs_equivalent()


 @skip_mps
 def  SCREAMING_SNAKE_CASE				(						self		)   ->	Union[str, Any]:
     return super().test_save_load_optional_components()


 @skip_mps
 def  SCREAMING_SNAKE_CASE				(						self		)   ->	List[str]:
     return super().test_attention_slicing_forward_pass()


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	Optional[Any]:
     pass


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	int:
     pass





 @skip_mps
 def  SCREAMING_SNAKE_CASE				(						self		)   ->	str:
     return super().test_progress_bar()






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






 """simple docstring"""




 def  SCREAMING_SNAKE_CASE				(						self		)   ->	List[str]:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()


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

     a											=SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''		)
     a											=pipe.melgan
     a											=None

     a											=pipe.to(__A		)
     pipe.set_progress_bar_config(disable=__A		)

     def callback(__A					,       __A		):
         # decode mel to audio
         a											=melgan(input_features=mel_output.astype(np.floataa		)		)[0]
         assert len(audio[0]		) == 8_1920 * (step + 1)
         # simulate that audio is played
         return audio

     a											=MidiProcessor()
     a											=processor(__A		)

     a											=input_tokens[:3]
     a											=torch.manual_seed(0		)
     pipe(__A					,       num_inference_steps=5					,       generator=__A					,       callback=__A					,       output_type='''mel'''		)


 def  SCREAMING_SNAKE_CASE				(						self		)   ->	int:
     a											=torch_device

     a											=SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''		)
     a											=pipe.to(__A		)
     pipe.set_progress_bar_config(disable=__A		)
     a											=MidiProcessor()

     a											=processor(__A		)
     # just run two denoising loops
     a											=input_tokens[:2]

     a											=torch.manual_seed(0		)
     a											=pipe(__A					,       num_inference_steps=2					,       generator=__A		)

     a											=output.audios[0]

     assert abs(np.abs(__A		).sum() - 3_612.841		) < 1E-1





 def  SCREAMING_SNAKE_CASE				(						self		)   ->	Dict:
     a											=torch_device

     a											=SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''		)
     a											=pipe.to(__A		)
     pipe.set_progress_bar_config(disable=__A		)

     a											=MidiProcessor()

     a											=processor(__A		)

     # just run 4 denoising loops
     a											=input_tokens[:4]

     a											=torch.manual_seed(0		)
     a											=pipe(__A					,       num_inference_steps=100					,       generator=__A		)

     a											=output.audios[0]
     assert abs(np.abs(__A		).sum() - 9_389.1_111		) < 5E-2



"""simple docstring"""







import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import (
    TaBlock,
    TaConfig,
    TaLayerNorm,
)

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin






class  __A							(			_SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE     ):






 """simple docstring"""




 @register_to_config
 def __init__(						self					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A = False					,       )   ->	Optional[Any]:
     super().__init__()

     a											=nn.Linear(__A					,       __A					,       bias=__A		)

     a											=nn.Embedding(__A					,       __A		)
     a											=False

     a											=nn.Dropout(p=__A		)

     a											=TaConfig(
         d_model=__A					,       num_heads=__A					,       d_kv=__A					,       d_ff=__A					,       feed_forward_proj=__A					,       dropout_rate=__A					,       is_decoder=__A					,       is_encoder_decoder=__A					,       )
     a											=nn.ModuleList()
     for lyr_num in range(__A		):
         a											=TaBlock(__A		)
         self.encoders.append(__A		)

     a											=TaLayerNorm(__A		)
     a											=nn.Dropout(p=__A		)





 def  SCREAMING_SNAKE_CASE				(						self					,       __A					,       __A		)   ->	str:
     a											=self.input_proj(__A		)

     # terminal relative positional encodings
     a											=encoder_inputs.shape[1]
     a											=torch.arange(__A					,       device=encoder_inputs.device		)

     a											=encoder_inputs_mask.sum(-1		)
     a											=torch.roll(input_positions.unsqueeze(0		)					,       tuple(seq_lens.tolist()		)					,       dims=0		)
     x += self.position_encoding(__A		)

     a											=self.dropout_pre(__A		)

     # inverted the attention mask
     a											=encoder_inputs.size()
     a											=self.get_extended_attention_mask(__A					,       __A		)

     for lyr in self.encoders:
         a											=lyr(__A					,       __A		)[0]
     a											=self.layer_norm(__A		)

     return self.dropout_post(__A		), encoder_inputs_mask