

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


snake_case_       =       './tests/fixtures/elise_format0.mid'
@unittest.skip("The note-seq package currently throws an error on import"      )
class 	SCREAMING_SNAKE_CASE__			(				_UpperCAmelCase,unittest.TestCase      ):
    _A			    =							SpectrogramDiffusionPipeline
    _A			    =							PipelineTesterMixin.required_optional_params - {
        "callback",
        "latents",
        "callback_steps",
        "output_type",
        "num_images_per_prompt",
    }
    _A			    =							False
    _A			    =							False
    _A			    =							TOKENS_TO_AUDIO_GENERATION_PARAMS
    _A			    =							TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          torch.manual_seed(0			)
          SCREAMING_SNAKE_CASE_      :			int     =       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"				,     )

          SCREAMING_SNAKE_CASE_      :			Tuple     =       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"				,     )

          SCREAMING_SNAKE_CASE_      :			Any     =       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				,     )

          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       DDPMScheduler()

          SCREAMING_SNAKE_CASE_      :			Dict     =       {
              "notes_encoder": notes_encoder.eval(),
              "continuous_encoder": continuous_encoder.eval(),
              "decoder": decoder.eval(),
              "scheduler": scheduler,
              "melgan": None,
          }
          return components




    def   __lowerCamelCase      (					self				,     lowercase__				,     lowercase__=0			):






          """simple docstring"""






          if str(__A			).startswith("mps"			):
                SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       torch.manual_seed(__A			)
          else:
                SCREAMING_SNAKE_CASE_      :			str     =       torch.Generator(device=__A			).manual_seed(__A			)
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       {
              "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   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Tuple     =       "cpu"  # ensure determinism for the device-dependent torch.Generator
          SCREAMING_SNAKE_CASE_      :			Tuple     =       self.get_dummy_components()
          SCREAMING_SNAKE_CASE_      :			str     =       SpectrogramDiffusionPipeline(**__A			)
          SCREAMING_SNAKE_CASE_      :			Any     =       pipe.to(__A			)
          pipe.set_progress_bar_config(disable=__A			)

          SCREAMING_SNAKE_CASE_      :			List[str]     =       self.get_dummy_inputs(__A			)
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       pipe(**__A			)
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       output.audios

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

          assert mel_slice.shape == (3, 3)
          SCREAMING_SNAKE_CASE_      :			List[str]     =       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   __lowerCamelCase      (					self			):






          """simple docstring"""






          return super().test_save_load_local()




    @skip_mps
    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          return super().test_dict_tuple_outputs_equivalent()




    @skip_mps
    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          return super().test_save_load_optional_components()




    @skip_mps
    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          return super().test_attention_slicing_forward_pass()




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          pass




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          pass







    @skip_mps
    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          return super().test_progress_bar()






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




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






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




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Dict     =       torch_device

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

          SCREAMING_SNAKE_CASE_      :			List[Any]     =       pipe.to(__A			)
          pipe.set_progress_bar_config(disable=__A			)

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

          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       MidiProcessor()
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       processor(__A			)

          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       input_tokens[:3]
          SCREAMING_SNAKE_CASE_      :			Any     =       torch.manual_seed(0			)
          pipe(__A				,     num_inference_steps=5				,     generator=__A				,     callback=__A				,     output_type="mel"			)




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       torch_device

          SCREAMING_SNAKE_CASE_      :			Any     =       SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion"			)
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       pipe.to(__A			)
          pipe.set_progress_bar_config(disable=__A			)
          SCREAMING_SNAKE_CASE_      :			Tuple     =       MidiProcessor()

          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       processor(__A			)
          # just run two denoising loops
          SCREAMING_SNAKE_CASE_      :			str     =       input_tokens[:2]

          SCREAMING_SNAKE_CASE_      :			str     =       torch.manual_seed(0			)
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       pipe(__A				,     num_inference_steps=2				,     generator=__A			)

          SCREAMING_SNAKE_CASE_      :			int     =       output.audios[0]

          assert abs(np.abs(__A			).sum() - 3612.841			) < 1e-1







    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Dict     =       torch_device

          SCREAMING_SNAKE_CASE_      :			int     =       SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion"			)
          SCREAMING_SNAKE_CASE_      :			str     =       pipe.to(__A			)
          pipe.set_progress_bar_config(disable=__A			)

          SCREAMING_SNAKE_CASE_      :			Any     =       MidiProcessor()

          SCREAMING_SNAKE_CASE_      :			List[Any]     =       processor(__A			)

          # just run 4 denoising loops
          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       input_tokens[:4]

          SCREAMING_SNAKE_CASE_      :			Dict     =       torch.manual_seed(0			)
          SCREAMING_SNAKE_CASE_      :			int     =       pipe(__A				,     num_inference_steps=100				,     generator=__A			)

          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       output.audios[0]
          assert abs(np.abs(__A			).sum() - 9389.1111			) < 5e-2









import tempfile
import unittest

import numpy as np
import torch

from diffusers import DiffusionPipeline
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor





class   SCREAMING_SNAKE_CASE (       unittest.TestCase    ):





							"""simple docstring"""

							def 					__A       (     self:      Optional[Any]						,   __A:      bool = False      )			-> Union[str, Any]:
												_A									=       10

												if only_cross_attention:
																	_A									=       12
												else:
																	# when only cross attention is not set, the cross attention dim must be the same as the query dim
																	_A									=       query_dim

												return {
												    "query_dim": query_dim,
												    "cross_attention_dim": cross_attention_dim,
												    "heads": 2,
												    "dim_head": 4,
												    "added_kv_proj_dim": 6,
												    "norm_num_groups": 1,
												    "only_cross_attention": only_cross_attention,
												    "processor": AttnAddedKVProcessor(),
												}

							def 					__A       (     self:      List[Any]						,   __A:      str						,   __A:      Union[str, Any]      )			-> Union[str, Any]:
												_A									=       2

												_A									=       torch.rand(__A						,   __A						,   3						,   2      )
												_A									=       torch.rand(__A						,   4						,   __A      )
												_A									=       None

												return {
												    "hidden_states": hidden_states,
												    "encoder_hidden_states": encoder_hidden_states,
												    "attention_mask": attention_mask,
												}



							def 					__A       (     self:      Optional[int]      )			-> Any:
												# self and cross attention

												torch.manual_seed(0      )

												_A									=       self.get_constructor_arguments(only_cross_attention=__A      )
												_A									=       Attention(**__A      )

												self.assertTrue(attn.to_k is not None      )
												self.assertTrue(attn.to_v is not None      )

												_A									=       self.get_forward_arguments(
												    query_dim=constructor_args['''query_dim''']						,   added_kv_proj_dim=constructor_args['''added_kv_proj_dim''']      )

												_A									=       attn(**__A      )

												# only self attention

												torch.manual_seed(0      )

												_A									=       self.get_constructor_arguments(only_cross_attention=__A      )
												_A									=       Attention(**__A      )

												self.assertTrue(attn.to_k is None      )
												self.assertTrue(attn.to_v is None      )

												_A									=       self.get_forward_arguments(
												    query_dim=constructor_args['''query_dim''']						,   added_kv_proj_dim=constructor_args['''added_kv_proj_dim''']      )

												_A									=       attn(**__A      )

												self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all()      )





class   SCREAMING_SNAKE_CASE (       unittest.TestCase    ):





							"""simple docstring"""

							def 					__A       (     self:      Any      )			-> Dict:
												_A									=       DiffusionPipeline.from_pretrained('''hf-internal-testing/tiny-stable-diffusion-pipe'''						,   safety_checker=__A      )

												_A									=       pipe(
												    '''foo'''						,   num_inference_steps=2						,   generator=torch.Generator('''cpu'''      ).manual_seed(0      )						,   output_type='''np'''						,   ).images

												# the initial conversion succeeds
												_A									=       DiffusionPipeline.from_pretrained(
												    '''hf-internal-testing/tiny-stable-diffusion-pipe'''						,   device_map='''sequential'''						,   safety_checker=__A      )

												_A									=       pipe(
												    '''foo'''						,   num_inference_steps=2						,   generator=torch.Generator('''cpu'''      ).manual_seed(0      )						,   output_type='''np'''						,   ).images

												with tempfile.TemporaryDirectory() as tmpdir:
																	# save the converted model
																	pipe.save_pretrained(__A      )

																	# can also load the converted weights
																	_A									=       DiffusionPipeline.from_pretrained(__A						,   device_map='''sequential'''						,   safety_checker=__A      )

												_A									=       pipe(
												    '''foo'''						,   num_inference_steps=2						,   generator=torch.Generator('''cpu'''      ).manual_seed(0      )						,   output_type='''np'''						,   ).images

												self.assertTrue(np.allclose(__A						,   __A						,   atol=1e-5      )      )
												self.assertTrue(np.allclose(__A						,   __A						,   atol=1e-5      )      )



