


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







       def 		snake_case_  ( self:   Optional[Any]				):






            '''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_0000.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 		snake_case_  ( self:   str,A_:   Tuple,A_:   Optional[int]=0				):






            '''simple docstring'''




            if str(A_				).startswith('mps'				):
                 __UpperCamelCase   = torch.manual_seed(A_				)
            else:
                 __UpperCamelCase   = torch.Generator(device=A_				).manual_seed(A_				)
            __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 		snake_case_  ( self:   Optional[int]				):






            '''simple docstring'''




            __UpperCamelCase   = 'cpu'  # ensure determinism for the device-dependent torch.Generator
            __UpperCamelCase   = self.get_dummy_components()
            __UpperCamelCase   = SpectrogramDiffusionPipeline(**A_				)
            __UpperCamelCase   = pipe.to(A_				)
            pipe.set_progress_bar_config(disable=A_				)

            __UpperCamelCase   = self.get_dummy_inputs(A_				)
            __UpperCamelCase   = pipe(**A_				)
            __UpperCamelCase   = output.audios

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

            assert mel_slice.shape == (3, 3)
            __UpperCamelCase   = np.array(
                [-11.51_2925, -4.7_8_8_2_1_5, -0.4_6_1_7_2_9_0_5, -2.0_5_1_7_1_5, -10.53_9147, -10.97_0963, -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:   Tuple				):






            '''simple docstring'''




            return super().test_save_load_local()







       @skip_mps
       def 		snake_case_  ( self:   Any				):






            '''simple docstring'''




            return super().test_dict_tuple_outputs_equivalent()







       @skip_mps
       def 		snake_case_  ( self:   str				):






            '''simple docstring'''




            return super().test_save_load_optional_components()







       @skip_mps
       def 		snake_case_  ( self:   Union[str, Any]				):






            '''simple docstring'''




            return super().test_attention_slicing_forward_pass()







       def 		snake_case_  ( self:   Optional[Any]				):






            '''simple docstring'''




            pass







       def 		snake_case_  ( self:   Any				):






            '''simple docstring'''




            pass







       @skip_mps
       def 		snake_case_  ( self:   Tuple				):






            '''simple docstring'''




            return super().test_progress_bar()



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







       def 		snake_case_  ( self:   List[Any]				):






            '''simple docstring'''




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







       def 		snake_case_  ( self:   Optional[Any]				):






            '''simple docstring'''




            __UpperCamelCase   = torch_device

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

            __UpperCamelCase   = pipe.to(A_				)
            pipe.set_progress_bar_config(disable=A_				)

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

            __UpperCamelCase   = MidiProcessor()
            __UpperCamelCase   = processor(A_				)

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







       def 		snake_case_  ( self:   Optional[Any]				):






            '''simple docstring'''




            __UpperCamelCase   = torch_device

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

            __UpperCamelCase   = processor(A_				)
            # just run two denoising loops
            __UpperCamelCase   = input_tokens[:2]

            __UpperCamelCase   = torch.manual_seed(0				)
            __UpperCamelCase   = pipe(A_,num_inference_steps=2,generator=A_				)

            __UpperCamelCase   = output.audios[0]

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







       def 		snake_case_  ( self:   Union[str, Any]				):






            '''simple docstring'''




            __UpperCamelCase   = torch_device

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

            __UpperCamelCase   = MidiProcessor()

            __UpperCamelCase   = processor(A_				)

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

            __UpperCamelCase   = torch.manual_seed(0				)
            __UpperCamelCase   = pipe(A_,num_inference_steps=100,generator=A_				)

            __UpperCamelCase   = output.audios[0]
            assert abs(np.abs(A_				).sum() - 9389.1111				) < 5E-2






from ..utils import DummyObject, requires_backends



class   __lowerCamelCase (metaclass=_a	):
       _lowercase       							=    ["""torch""", """transformers""", """k_diffusion"""]







       def __init__( self:   Tuple,*A_:   Union[str, Any],**A_:   Tuple				):






            '''simple docstring'''




            requires_backends(self,['torch', 'transformers', 'k_diffusion']				)







       @classmethod
       def 		snake_case_  ( cls:   int,*A_:   Optional[Any],**A_:   Tuple				):






            '''simple docstring'''




            requires_backends(cls,['torch', 'transformers', 'k_diffusion']				)







       @classmethod
       def 		snake_case_  ( cls:   int,*A_:   Tuple,**A_:   Tuple				):






            '''simple docstring'''




            requires_backends(cls,['torch', 'transformers', 'k_diffusion']				)


