





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


SCREAMING_SNAKE_CASE_		  =			"""./tests/fixtures/elise_format0.mid"""

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



     '''simple docstring'''

     __snake_case       :      Union[str, Any]						      =     SpectrogramDiffusionPipeline
     __snake_case       :      str						      =     PipelineTesterMixin.required_optional_params - {
         "callback",
         "latents",
         "callback_steps",
         "output_type",
         "num_images_per_prompt",
     }
     __snake_case       :      str						      =     False
     __snake_case       :      Tuple						      =     False
     __snake_case       :      List[str]						      =     TOKENS_TO_AUDIO_GENERATION_PARAMS
     __snake_case       :      List[str]						      =     TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Tuple			)	->						Optional[Any]:




           '''simple docstring'''


           torch.manual_seed(0			)
           SCREAMING_SNAKE_CASE						=		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						=		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						=		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		,)

           SCREAMING_SNAKE_CASE						=		DDPMScheduler()

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






     def 							SCREAMING_SNAKE_CASE__    (				self		:       str		,lowerCamelCase__		:       List[str]		,lowerCamelCase__		:       Union[str, Any]=0			)	->						Optional[Any]:




           '''simple docstring'''


           if str(lowerCamelCase__			).startswith("""mps"""			):
                 SCREAMING_SNAKE_CASE						=		torch.manual_seed(lowerCamelCase__			)
           else:
                 SCREAMING_SNAKE_CASE						=		torch.Generator(device=lowerCamelCase__			).manual_seed(lowerCamelCase__			)
           SCREAMING_SNAKE_CASE						=		{
               """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[Any]			)	->						str:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE						=		self.get_dummy_components()
           SCREAMING_SNAKE_CASE						=		SpectrogramDiffusionPipeline(**lowerCamelCase__			)
           SCREAMING_SNAKE_CASE						=		pipe.to(lowerCamelCase__			)
           pipe.set_progress_bar_config(disable=lowerCamelCase__			)

           SCREAMING_SNAKE_CASE						=		self.get_dummy_inputs(lowerCamelCase__			)
           SCREAMING_SNAKE_CASE						=		pipe(**lowerCamelCase__			)
           SCREAMING_SNAKE_CASE						=		output.audios

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

           assert mel_slice.shape == (3, 3)
           SCREAMING_SNAKE_CASE						=		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 							SCREAMING_SNAKE_CASE__    (				self		:       Optional[int]			)	->						int:




           '''simple docstring'''


           return super().test_save_load_local()






     @skip_mps
     def 							SCREAMING_SNAKE_CASE__    (				self		:       Any			)	->						Optional[int]:




           '''simple docstring'''


           return super().test_dict_tuple_outputs_equivalent()






     @skip_mps
     def 							SCREAMING_SNAKE_CASE__    (				self		:       int			)	->						str:




           '''simple docstring'''


           return super().test_save_load_optional_components()






     @skip_mps
     def 							SCREAMING_SNAKE_CASE__    (				self		:       Optional[int]			)	->						Any:




           '''simple docstring'''


           return super().test_attention_slicing_forward_pass()






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Tuple			)	->						Tuple:




           '''simple docstring'''


           pass






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Any			)	->						int:




           '''simple docstring'''


           pass






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




           '''simple docstring'''


           return super().test_progress_bar()

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



     '''simple docstring'''

     def 							SCREAMING_SNAKE_CASE__    (				self		:       Optional[int]			)	->						Tuple:




           '''simple docstring'''


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






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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		torch_device

           SCREAMING_SNAKE_CASE						=		SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""			)
           SCREAMING_SNAKE_CASE						=		pipe.melgan
           SCREAMING_SNAKE_CASE						=		None

           SCREAMING_SNAKE_CASE						=		pipe.to(lowerCamelCase__			)
           pipe.set_progress_bar_config(disable=lowerCamelCase__			)

           def callback(lowerCamelCase__		:       Optional[Any]		,lowerCamelCase__		:       str			):
                 # decode mel to audio
                 SCREAMING_SNAKE_CASE						=		melgan(input_features=mel_output.astype(np.floataa			)			)[0]
                 assert len(audio[0]			) == 81920 * (step + 1)
                 # simulate that audio is played
                 return audio

           SCREAMING_SNAKE_CASE						=		MidiProcessor()
           SCREAMING_SNAKE_CASE						=		processor(lowerCamelCase__			)

           SCREAMING_SNAKE_CASE						=		input_tokens[:3]
           SCREAMING_SNAKE_CASE						=		torch.manual_seed(0			)
           pipe(lowerCamelCase__		,num_inference_steps=5		,generator=lowerCamelCase__		,callback=lowerCamelCase__		,output_type="""mel"""			)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       str			)	->						Optional[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		torch_device

           SCREAMING_SNAKE_CASE						=		SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""			)
           SCREAMING_SNAKE_CASE						=		pipe.to(lowerCamelCase__			)
           pipe.set_progress_bar_config(disable=lowerCamelCase__			)
           SCREAMING_SNAKE_CASE						=		MidiProcessor()

           SCREAMING_SNAKE_CASE						=		processor(lowerCamelCase__			)
           # just run two denoising loops
           SCREAMING_SNAKE_CASE						=		input_tokens[:2]

           SCREAMING_SNAKE_CASE						=		torch.manual_seed(0			)
           SCREAMING_SNAKE_CASE						=		pipe(lowerCamelCase__		,num_inference_steps=2		,generator=lowerCamelCase__			)

           SCREAMING_SNAKE_CASE						=		output.audios[0]

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






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Tuple			)	->						Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		torch_device

           SCREAMING_SNAKE_CASE						=		SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""			)
           SCREAMING_SNAKE_CASE						=		pipe.to(lowerCamelCase__			)
           pipe.set_progress_bar_config(disable=lowerCamelCase__			)

           SCREAMING_SNAKE_CASE						=		MidiProcessor()

           SCREAMING_SNAKE_CASE						=		processor(lowerCamelCase__			)

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

           SCREAMING_SNAKE_CASE						=		torch.manual_seed(0			)
           SCREAMING_SNAKE_CASE						=		pipe(lowerCamelCase__		,num_inference_steps=100		,generator=lowerCamelCase__			)

           SCREAMING_SNAKE_CASE						=		output.audios[0]
           assert abs(np.abs(lowerCamelCase__			).sum() - 9389.1111			) < 5e-2









import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin


if is_torch_available():
      import torch

if is_vision_available():
      from PIL import Image

      from transformers import TvltImageProcessor



def   __lowercase  ( _SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE=10					,		_SCREAMING_SNAKE_CASE=10					,		_SCREAMING_SNAKE_CASE=False					,		_SCREAMING_SNAKE_CASE=False      )					->  Optional[Any]:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE						=		[]
      for i in range(image_processor_tester.num_frames      ):
            video.append(np.random.randint(2_55					,		size=(image_processor_tester.num_channels, width, height)					,		dtype=np.uinta      )      )

      if not numpify and not torchify:
            # PIL expects the channel dimension as last dimension
            SCREAMING_SNAKE_CASE						=		[Image.fromarray(np.moveaxis(_SCREAMING_SNAKE_CASE					,		0					,		-1      )      ) for frame in video]

      if torchify:
            SCREAMING_SNAKE_CASE						=		[torch.from_numpy(_SCREAMING_SNAKE_CASE      ) for frame in video]

      return video



def   __lowercase  ( _SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE=False					,		_SCREAMING_SNAKE_CASE=False					,		_SCREAMING_SNAKE_CASE=False      )					->  Tuple:


      '''simple docstring'''


      assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

      SCREAMING_SNAKE_CASE						=		[]
      for i in range(image_processor_tester.batch_size      ):
            if equal_resolution:
                  SCREAMING_SNAKE_CASE						=		SCREAMING_SNAKE_CASE						=		image_processor_tester.max_resolution
            else:
                  SCREAMING_SNAKE_CASE,   SCREAMING_SNAKE_CASE						=		np.random.choice(
                      np.arange(image_processor_tester.min_resolution					,		image_processor_tester.max_resolution      )					,		2      )
                  SCREAMING_SNAKE_CASE						=		prepare_video(
                      image_processor_tester=_SCREAMING_SNAKE_CASE					,		width=_SCREAMING_SNAKE_CASE					,		height=_SCREAMING_SNAKE_CASE					,		numpify=_SCREAMING_SNAKE_CASE					,		torchify=_SCREAMING_SNAKE_CASE					,		)
            video_inputs.append(_SCREAMING_SNAKE_CASE      )

      return video_inputs

class 						UpperCamelCase__					( unittest.TestCase  ):



     '''simple docstring'''

     def __init__(				self		:       Union[str, Any]		,lowerCamelCase__		:       Tuple		,lowerCamelCase__		:       Any=7		,lowerCamelCase__		:       str=3		,lowerCamelCase__		:       Union[str, Any]=4		,lowerCamelCase__		:       Union[str, Any]=18		,lowerCamelCase__		:       Any=30		,lowerCamelCase__		:       Union[str, Any]=400		,lowerCamelCase__		:       Any=True		,lowerCamelCase__		:       Any=None		,lowerCamelCase__		:       Tuple=True		,lowerCamelCase__		:       int=[0.5, 0.5, 0.5]		,lowerCamelCase__		:       Any=[0.5, 0.5, 0.5]		,lowerCamelCase__		:       Dict=True		,lowerCamelCase__		:       Tuple=None		,)	->						Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		size if size is not None else {"""shortest_edge""": 18}
           SCREAMING_SNAKE_CASE						=		crop_size if crop_size is not None else {"""height""": 18, """width""": 18}

           SCREAMING_SNAKE_CASE						=		parent
           SCREAMING_SNAKE_CASE						=		batch_size
           SCREAMING_SNAKE_CASE						=		num_channels
           SCREAMING_SNAKE_CASE						=		num_frames
           SCREAMING_SNAKE_CASE						=		image_size
           SCREAMING_SNAKE_CASE						=		min_resolution
           SCREAMING_SNAKE_CASE						=		max_resolution
           SCREAMING_SNAKE_CASE						=		do_resize
           SCREAMING_SNAKE_CASE						=		size
           SCREAMING_SNAKE_CASE						=		do_normalize
           SCREAMING_SNAKE_CASE						=		image_mean
           SCREAMING_SNAKE_CASE						=		image_std
           SCREAMING_SNAKE_CASE						=		do_center_crop
           SCREAMING_SNAKE_CASE						=		crop_size






     def 							SCREAMING_SNAKE_CASE__    (				self		:       List[Any]			)	->						int:




           '''simple docstring'''


           return {
               "image_mean": self.image_mean,
               "image_std": self.image_std,
               "do_normalize": self.do_normalize,
               "do_resize": self.do_resize,
               "size": self.size,
               "do_center_crop": self.do_center_crop,
               "crop_size": self.crop_size,
           }

@require_torch
@require_vision
class 						UpperCamelCase__					( lowerCAmelCase_				, unittest.TestCase  ):



     '''simple docstring'''

     __snake_case       :      int						      =     TvltImageProcessor if is_vision_available() else None






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Dict			)	->						Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		TvltImageProcessorTester(self			)






     @property
     def 							SCREAMING_SNAKE_CASE__    (				self		:       Optional[int]			)	->						Dict:




           '''simple docstring'''


           return self.image_processor_tester.prepare_image_processor_dict()






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Any			)	->						List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.image_processing_class(**self.image_processor_dict			)
           self.assertTrue(hasattr(lowerCamelCase__		,"""image_mean"""			)			)
           self.assertTrue(hasattr(lowerCamelCase__		,"""image_std"""			)			)
           self.assertTrue(hasattr(lowerCamelCase__		,"""do_normalize"""			)			)
           self.assertTrue(hasattr(lowerCamelCase__		,"""do_resize"""			)			)
           self.assertTrue(hasattr(lowerCamelCase__		,"""do_center_crop"""			)			)
           self.assertTrue(hasattr(lowerCamelCase__		,"""size"""			)			)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       List[Any]			)	->						List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.image_processing_class(**self.image_processor_dict			)
           # create random PIL videos
           SCREAMING_SNAKE_CASE						=		prepare_video_inputs(self.image_processor_tester		,equal_resolution=lowerCamelCase__			)
           for video in video_inputs:
                 self.assertIsInstance(lowerCamelCase__		,lowerCamelCase__			)
                 self.assertIsInstance(video[0]		,Image.Image			)

           # Test not batched input
           SCREAMING_SNAKE_CASE						=		image_processor(video_inputs[0]		,return_tensors="""pt"""			).pixel_values
           self.assertEqual(
               encoded_videos.shape		,(
                   1,
                   self.image_processor_tester.num_frames,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.crop_size["""height"""],
                   self.image_processor_tester.crop_size["""width"""],
               )		,)

           # Test batched
           SCREAMING_SNAKE_CASE						=		image_processor(lowerCamelCase__		,return_tensors="""pt"""			).pixel_values
           self.assertEqual(
               encoded_videos.shape		,(
                   self.image_processor_tester.batch_size,
                   self.image_processor_tester.num_frames,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.crop_size["""height"""],
                   self.image_processor_tester.crop_size["""width"""],
               )		,)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       int			)	->						int:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.image_processing_class(**self.image_processor_dict			)
           # create random numpy tensors
           SCREAMING_SNAKE_CASE						=		prepare_video_inputs(self.image_processor_tester		,equal_resolution=lowerCamelCase__		,numpify=lowerCamelCase__			)
           for video in video_inputs:
                 self.assertIsInstance(lowerCamelCase__		,lowerCamelCase__			)
                 self.assertIsInstance(video[0]		,np.ndarray			)

           # Test not batched input
           SCREAMING_SNAKE_CASE						=		image_processor(video_inputs[0]		,return_tensors="""pt"""			).pixel_values
           self.assertEqual(
               encoded_videos.shape		,(
                   1,
                   self.image_processor_tester.num_frames,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.crop_size["""height"""],
                   self.image_processor_tester.crop_size["""width"""],
               )		,)

           # Test batched
           SCREAMING_SNAKE_CASE						=		image_processor(lowerCamelCase__		,return_tensors="""pt"""			).pixel_values
           self.assertEqual(
               encoded_videos.shape		,(
                   self.image_processor_tester.batch_size,
                   self.image_processor_tester.num_frames,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.crop_size["""height"""],
                   self.image_processor_tester.crop_size["""width"""],
               )		,)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       str			)	->						Optional[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.image_processing_class(**self.image_processor_dict			)
           # create random PyTorch tensors
           SCREAMING_SNAKE_CASE						=		prepare_video_inputs(self.image_processor_tester		,equal_resolution=lowerCamelCase__		,torchify=lowerCamelCase__			)
           for video in video_inputs:
                 self.assertIsInstance(lowerCamelCase__		,lowerCamelCase__			)
                 self.assertIsInstance(video[0]		,torch.Tensor			)

           # Test not batched input
           SCREAMING_SNAKE_CASE						=		image_processor(video_inputs[0]		,return_tensors="""pt"""			).pixel_values
           self.assertEqual(
               encoded_videos.shape		,(
                   1,
                   self.image_processor_tester.num_frames,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.crop_size["""height"""],
                   self.image_processor_tester.crop_size["""width"""],
               )		,)

           # Test batched
           SCREAMING_SNAKE_CASE						=		image_processor(lowerCamelCase__		,return_tensors="""pt"""			).pixel_values
           self.assertEqual(
               encoded_videos.shape		,(
                   self.image_processor_tester.batch_size,
                   self.image_processor_tester.num_frames,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.crop_size["""height"""],
                   self.image_processor_tester.crop_size["""width"""],
               )		,)


