


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						:       Union[str, Any]        =					"./tests/fixtures/elise_format0.mid"





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


				"""simple docstring"""


				UpperCamelCase   :  Optional[Any]    =     SpectrogramDiffusionPipeline
				UpperCamelCase   :  str    =     PipelineTesterMixin.required_optional_params - {
				    '''callback''',
				    '''latents''',
				    '''callback_steps''',
				    '''output_type''',
				    '''num_images_per_prompt''',
				}
				UpperCamelCase   :  Optional[int]    =     False
				UpperCamelCase   :  int    =     False
				UpperCamelCase   :  Optional[Any]    =     TOKENS_TO_AUDIO_GENERATION_PARAMS
				UpperCamelCase   :  str    =     TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS



				def  UpperCAmelCase_		(      self					):
									torch.manual_seed(0					)
									__A     :						Dict          =    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     :						int          =    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     :						Optional[int]          =    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					, )

									__A     :						str          =    DDPMScheduler()

									__A     :						List[str]          =    {
									    'notes_encoder': notes_encoder.eval(),
									    'continuous_encoder': continuous_encoder.eval(),
									    'decoder': decoder.eval(),
									    'scheduler': scheduler,
									    'melgan': None,
									}
									return components



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

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

									__A     :						Any          =    mel[0, -3:, -3:]

									assert mel_slice.shape == (3, 3)
									__A     :						List[Any]          =    np.array(
									    [-11.512925, -4.7_8_8_2_1_5, -0.4_6_1_7_2_9_0_5, -2.0_5_1_7_1_5, -10.539147, -10.970963, -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  UpperCAmelCase_		(      self					):
									return super().test_save_load_local()



				@skip_mps
				def  UpperCAmelCase_		(      self					):
									return super().test_dict_tuple_outputs_equivalent()



				@skip_mps
				def  UpperCAmelCase_		(      self					):
									return super().test_save_load_optional_components()



				@skip_mps
				def  UpperCAmelCase_		(      self					):
									return super().test_attention_slicing_forward_pass()



				def  UpperCAmelCase_		(      self					):
									pass



				def  UpperCAmelCase_		(      self					):
									pass





				@skip_mps
				def  UpperCAmelCase_		(      self					):
									return super().test_progress_bar()





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


				"""simple docstring"""


				def  UpperCAmelCase_		(      self					):
									# clean up the VRAM after each test
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()



				def  UpperCAmelCase_		(      self					):
									# TODO - test that pipeline can decode tokens in a callback
									# so that music can be played live
									__A     :						Union[str, Any]          =    torch_device

									__A     :						str          =    SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'					)
									__A     :						Union[str, Any]          =    pipe.melgan
									__A     :						str          =    None

									__A     :						Optional[Any]          =    pipe.to(_snake_case					)
									pipe.set_progress_bar_config(disable=_snake_case					)

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

									__A     :						int          =    MidiProcessor()
									__A     :						List[Any]          =    processor(_snake_case					)

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



				def  UpperCAmelCase_		(      self					):
									__A     :						str          =    torch_device

									__A     :						List[str]          =    SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'					)
									__A     :						Union[str, Any]          =    pipe.to(_snake_case					)
									pipe.set_progress_bar_config(disable=_snake_case					)
									__A     :						Dict          =    MidiProcessor()

									__A     :						List[Any]          =    processor(_snake_case					)
									# just run two denoising loops
									__A     :						int          =    input_tokens[:2]

									__A     :						str          =    torch.manual_seed(0					)
									__A     :						Optional[int]          =    pipe(_snake_case					, num_inference_steps=2					, generator=_snake_case					)

									__A     :						Any          =    output.audios[0]

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





				def  UpperCAmelCase_		(      self					):
									__A     :						int          =    torch_device

									__A     :						List[Any]          =    SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'					)
									__A     :						List[str]          =    pipe.to(_snake_case					)
									pipe.set_progress_bar_config(disable=_snake_case					)

									__A     :						Tuple          =    MidiProcessor()

									__A     :						str          =    processor(_snake_case					)

									# just run 4 denoising loops
									__A     :						int          =    input_tokens[:4]

									__A     :						Tuple          =    torch.manual_seed(0					)
									__A     :						Any          =    pipe(_snake_case					, num_inference_steps=100					, generator=_snake_case					)

									__A     :						List[str]          =    output.audios[0]
									assert abs(np.abs(_snake_case					).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 			_SCREAMING_SNAKE_CASE    (       a ,						a=10 ,						a=10 ,						a=False ,						a=False  )				->				Union[str, Any]:
					__A     :						str          =    []
					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
										__A     :						Union[str, Any]          =    [Image.fromarray(np.moveaxis(a ,						0 ,						-1  )  ) for frame in video]

					if torchify:
										__A     :						Dict          =    [torch.from_numpy(a  ) for frame in video]

					return video







def 			_SCREAMING_SNAKE_CASE    (       a ,						a=False ,						a=False ,						a=False  )				->				Union[str, Any]:
					assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

					__A     :						List[Any]          =    []
					for i in range(image_processor_tester.batch_size  ):
										if equal_resolution:
															__A     :						int          =    image_processor_tester.max_resolution
										else:
															__A   ,    __A     :						Tuple          =    np.random.choice(
															    np.arange(image_processor_tester.min_resolution ,						image_processor_tester.max_resolution  ) ,						2  )
															__A     :						Optional[int]          =    prepare_video(
															    image_processor_tester=a ,						width=a ,						height=a ,						numpify=a ,						torchify=a ,						)
										video_inputs.append(a  )

					return video_inputs





class 						_A( unittest.TestCase       ):


				"""simple docstring"""


				def __init__(      self					, _A					, _A=7					, _A=3					, _A=4					, _A=18					, _A=30					, _A=400					, _A=True					, _A=None					, _A=True					, _A=[0.5, 0.5, 0.5]					, _A=[0.5, 0.5, 0.5]					, _A=True					, _A=None					, ):
									__A     :						List[str]          =    size if size is not None else {'shortest_edge': 18}
									__A     :						Tuple          =    crop_size if crop_size is not None else {'height': 18, 'width': 18}

									__A     :						Dict          =    parent
									__A     :						str          =    batch_size
									__A     :						int          =    num_channels
									__A     :						Optional[Any]          =    num_frames
									__A     :						Any          =    image_size
									__A     :						List[str]          =    min_resolution
									__A     :						List[Any]          =    max_resolution
									__A     :						Tuple          =    do_resize
									__A     :						List[str]          =    size
									__A     :						str          =    do_normalize
									__A     :						str          =    image_mean
									__A     :						int          =    image_std
									__A     :						Optional[Any]          =    do_center_crop
									__A     :						Union[str, Any]          =    crop_size





				def  UpperCAmelCase_		(      self					):
									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 						_A( snake_case__    ,			unittest.TestCase       ):


				"""simple docstring"""


				UpperCamelCase   :  Optional[Any]    =     TvltImageProcessor if is_vision_available() else None



				def  UpperCAmelCase_		(      self					):
									__A     :						str          =    TvltImageProcessorTester(self					)



				@property
				def  UpperCAmelCase_		(      self					):
									return self.image_processor_tester.prepare_image_processor_dict()



				def  UpperCAmelCase_		(      self					):
									__A     :						Union[str, Any]          =    self.image_processing_class(**self.image_processor_dict					)
									self.assertTrue(hasattr(_A					, 'image_mean'					)					)
									self.assertTrue(hasattr(_A					, 'image_std'					)					)
									self.assertTrue(hasattr(_A					, 'do_normalize'					)					)
									self.assertTrue(hasattr(_A					, 'do_resize'					)					)
									self.assertTrue(hasattr(_A					, 'do_center_crop'					)					)
									self.assertTrue(hasattr(_A					, 'size'					)					)



				def  UpperCAmelCase_		(      self					):
									# Initialize image_processor
									__A     :						List[str]          =    self.image_processing_class(**self.image_processor_dict					)
									# create random PIL videos
									__A     :						List[str]          =    prepare_video_inputs(self.image_processor_tester					, equal_resolution=_A					)
									for video in video_inputs:
														self.assertIsInstance(_A					, _A					)
														self.assertIsInstance(video[0]					, Image.Image					)

									# Test not batched input
									__A     :						str          =    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
									__A     :						List[Any]          =    image_processor(_A					, 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  UpperCAmelCase_		(      self					):
									# Initialize image_processor
									__A     :						Any          =    self.image_processing_class(**self.image_processor_dict					)
									# create random numpy tensors
									__A     :						Dict          =    prepare_video_inputs(self.image_processor_tester					, equal_resolution=_A					, numpify=_A					)
									for video in video_inputs:
														self.assertIsInstance(_A					, _A					)
														self.assertIsInstance(video[0]					, np.ndarray					)

									# Test not batched input
									__A     :						Dict          =    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
									__A     :						Optional[int]          =    image_processor(_A					, 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  UpperCAmelCase_		(      self					):
									# Initialize image_processor
									__A     :						Any          =    self.image_processing_class(**self.image_processor_dict					)
									# create random PyTorch tensors
									__A     :						Dict          =    prepare_video_inputs(self.image_processor_tester					, equal_resolution=_A					, torchify=_A					)
									for video in video_inputs:
														self.assertIsInstance(_A					, _A					)
														self.assertIsInstance(video[0]					, torch.Tensor					)

									# Test not batched input
									__A     :						Optional[Any]          =    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
									__A     :						int          =    image_processor(_A					, 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'],
									    )					, )




