

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


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

@unittest.skip("The note-seq package currently throws an error on import"      )
class 					__UpperCamelCase					(       lowerCAmelCase__	,    unittest.TestCase      ):
	_lowercase	: Union[str, Any]												=	SpectrogramDiffusionPipeline
	_lowercase	: Dict												=	PipelineTesterMixin.required_optional_params - {
	    "callback",
	    "latents",
	    "callback_steps",
	    "output_type",
	    "num_images_per_prompt",
	}
	_lowercase	: str												=	False
	_lowercase	: int												=	False
	_lowercase	: str												=	TOKENS_TO_AUDIO_GENERATION_PARAMS
	_lowercase	: str												=	TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS







	def 						_SCREAMING_SNAKE_CASE				( self:			int   ):

						'''simple docstring'''





						torch.manual_seed(0   )
						__magic_name__				=     SpectrogramNotesEncoder(
						    max_length=20_48	,			vocab_size=15_36	,			d_model=7_68	,			dropout_rate=0.1	,			num_layers=1	,			num_heads=1	,			d_kv=4	,			d_ff=20_48	,			feed_forward_proj='gated-gelu'	,			)

						__magic_name__				=     SpectrogramContEncoder(
						    input_dims=1_28	,			targets_context_length=2_56	,			d_model=7_68	,			dropout_rate=0.1	,			num_layers=1	,			num_heads=1	,			d_kv=4	,			d_ff=20_48	,			feed_forward_proj='gated-gelu'	,			)

						__magic_name__				=     TaFilmDecoder(
						    input_dims=1_28	,			targets_length=2_56	,			max_decoder_noise_time=2_00_00.0	,			d_model=7_68	,			num_layers=1	,			num_heads=1	,			d_kv=4	,			d_ff=20_48	,			dropout_rate=0.1	,			)

						__magic_name__				=     DDPMScheduler()

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







	def 						_SCREAMING_SNAKE_CASE				( self:			Tuple	,			__UpperCamelCase:			int	,			__UpperCamelCase:			Union[str, Any]=0   ):

						'''simple docstring'''





						if str(_lowercase   ).startswith('mps'   ):
											__magic_name__				=     torch.manual_seed(_lowercase   )
						else:
											__magic_name__				=     torch.Generator(device=_lowercase   ).manual_seed(_lowercase   )
						__magic_name__				=     {
						    'input_tokens': [
						        [11_34, 90, 11_35, 11_33, 10_80, 1_12, 11_32, 10_80, 11_33, 10_79, 1_33, 11_32, 10_79, 11_33, 1] + [0] * 20_33
						    ],
						    'generator': generator,
						    'num_inference_steps': 4,
						    'output_type': 'mel',
						}
						return inputs







	def 						_SCREAMING_SNAKE_CASE				( self:			int   ):

						'''simple docstring'''





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

						__magic_name__				=     self.get_dummy_inputs(_lowercase   )
						__magic_name__				=     pipe(**_lowercase   )
						__magic_name__				=     output.audios

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

						assert mel_slice.shape == (3, 3)
						__magic_name__				=     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]   ):

						'''simple docstring'''





						return super().test_save_load_local()







	@skip_mps
	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[int]   ):

						'''simple docstring'''





						return super().test_dict_tuple_outputs_equivalent()







	@skip_mps
	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[int]   ):

						'''simple docstring'''





						return super().test_save_load_optional_components()







	@skip_mps
	def 						_SCREAMING_SNAKE_CASE				( self:			str   ):

						'''simple docstring'''





						return super().test_attention_slicing_forward_pass()







	def 						_SCREAMING_SNAKE_CASE				( self:			List[str]   ):

						'''simple docstring'''





						pass







	def 						_SCREAMING_SNAKE_CASE				( self:			int   ):

						'''simple docstring'''





						pass







	@skip_mps
	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[Any]   ):

						'''simple docstring'''





						return super().test_progress_bar()




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







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

						'''simple docstring'''





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







	def 						_SCREAMING_SNAKE_CASE				( self:			List[Any]   ):

						'''simple docstring'''





						__magic_name__				=     torch_device

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

						__magic_name__				=     pipe.to(_lowercase   )
						pipe.set_progress_bar_config(disable=_lowercase   )

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

						__magic_name__				=     MidiProcessor()
						__magic_name__				=     processor(_lowercase   )

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







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

						'''simple docstring'''





						__magic_name__				=     torch_device

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

						__magic_name__				=     processor(_lowercase   )
						# just run two denoising loops
						__magic_name__				=     input_tokens[:2]

						__magic_name__				=     torch.manual_seed(0   )
						__magic_name__				=     pipe(_lowercase	,			num_inference_steps=2	,			generator=_lowercase   )

						__magic_name__				=     output.audios[0]

						assert abs(np.abs(_lowercase   ).sum() - 36_12.8_41   ) < 1E-1







	def 						_SCREAMING_SNAKE_CASE				( self:			str   ):

						'''simple docstring'''





						__magic_name__				=     torch_device

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

						__magic_name__				=     MidiProcessor()

						__magic_name__				=     processor(_lowercase   )

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

						__magic_name__				=     torch.manual_seed(0   )
						__magic_name__				=     pipe(_lowercase	,			num_inference_steps=1_00	,			generator=_lowercase   )

						__magic_name__				=     output.audios[0]
						assert abs(np.abs(_lowercase   ).sum() - 93_89.11_11   ) < 5E-2










from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import center_crop, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
	import PIL


__A						:		Dict          =						logging.get_logger(__name__)

class 		lowercase_       (	lowerCAmelCase__				):
					__UpperCamelCase     	= ["pixel_values"]




					def __init__(   self:		str,     _lowercase:		bool = True,     _lowercase:		Dict[str, int] = None,     _lowercase:		PILImageResampling = PIL.Image.NEAREST,     _lowercase:		bool = False,     _lowercase:		Dict[str, int] = None,     _lowercase:		Union[int, float] = 1 / 255,     _lowercase:		bool = False,     _lowercase:		bool = True,     _lowercase:		bool = True,     _lowercase:		Optional[Union[float, List[float]]] = None,     _lowercase:		Optional[Union[float, List[float]]] = None,     _lowercase:		bool = True,     **_lowercase:		Union[str, Any],     ):


									'''simple docstring'''


									super().__init__(**_lowercase)
									__lowerCAmelCase					=  size if size is not None else {"""height""": 346, """width""": 346}
									__lowerCAmelCase					=  get_size_dict(_lowercase)
									__lowerCAmelCase					=  crop_size if crop_size is not None else {"""height""": 289, """width""": 289}
									__lowerCAmelCase					=  get_size_dict(_lowercase,     param_name="""crop_size""")

									__lowerCAmelCase					=  do_resize
									__lowerCAmelCase					=  size
									__lowerCAmelCase					=  resample
									__lowerCAmelCase					=  do_center_crop
									__lowerCAmelCase					=  crop_size
									__lowerCAmelCase					=  do_rescale
									__lowerCAmelCase					=  rescale_factor
									__lowerCAmelCase					=  rescale_offset
									__lowerCAmelCase					=  do_normalize
									__lowerCAmelCase					=  image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
									__lowerCAmelCase					=  image_std if image_std is not None else IMAGENET_STANDARD_STD
									__lowerCAmelCase					=  include_top




					def     _lowercase						(   self:		Optional[Any],     _lowercase:		np.ndarray,     _lowercase:		Dict[str, int],     _lowercase:		PILImageResampling = PIL.Image.NEAREST,     _lowercase:		Optional[Union[str, ChannelDimension]] = None,     **_lowercase:		Dict,     ):


									'''simple docstring'''


									__lowerCAmelCase					=  get_size_dict(_lowercase)
									if "height" not in size or "width" not in size:
													raise ValueError(f'''The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}''')
									return resize(
									    _lowercase,     size=(size["""height"""], size["""width"""]),     resample=_lowercase,     data_format=_lowercase,     **_lowercase)




					def     _lowercase						(   self:		List[str],     _lowercase:		np.ndarray,     _lowercase:		Dict[str, int],     _lowercase:		Optional[Union[str, ChannelDimension]] = None,     **_lowercase:		Tuple,     ):


									'''simple docstring'''


									__lowerCAmelCase					=  get_size_dict(_lowercase)
									if "height" not in size or "width" not in size:
													raise ValueError(f'''The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}''')
									return center_crop(_lowercase,     size=(size["""height"""], size["""width"""]),     data_format=_lowercase,     **_lowercase)




					def     _lowercase						(   self:		Union[str, Any],     _lowercase:		np.ndarray,     _lowercase:		Union[int, float],     _lowercase:		bool = True,     _lowercase:		Optional[Union[str, ChannelDimension]] = None,     **_lowercase:		Dict,     ):


									'''simple docstring'''


									if offset:
													__lowerCAmelCase					=  (image - 127.5) * scale
													if data_format is not None:
																	__lowerCAmelCase					=  to_channel_dimension_format(_lowercase,     _lowercase)
													__lowerCAmelCase					=  rescaled_image.astype(np.floataa)
									else:
													__lowerCAmelCase					=  rescale(_lowercase,     scale=_lowercase,     data_format=_lowercase,     **_lowercase)
									return rescaled_image




					def     _lowercase						(   self:		Dict,     _lowercase:		np.ndarray,     _lowercase:		Union[float, List[float]],     _lowercase:		Union[float, List[float]],     _lowercase:		Optional[Union[str, ChannelDimension]] = None,     **_lowercase:		Any,     ):


									'''simple docstring'''


									return normalize(_lowercase,     mean=_lowercase,     std=_lowercase,     data_format=_lowercase,     **_lowercase)







					def     _lowercase						(   self:		Dict,     _lowercase:		ImageInput,     _lowercase:		bool = None,     _lowercase:		Dict[str, int] = None,     _lowercase:		str=None,     _lowercase:		bool = None,     _lowercase:		Dict[str, int] = None,     _lowercase:		bool = None,     _lowercase:		float = None,     _lowercase:		bool = None,     _lowercase:		bool = None,     _lowercase:		Optional[Union[float, List[float]]] = None,     _lowercase:		Optional[Union[float, List[float]]] = None,     _lowercase:		bool = None,     _lowercase:		Optional[Union[str, TensorType]] = None,     _lowercase:		ChannelDimension = ChannelDimension.FIRST,     **_lowercase:		Tuple,     ):


									'''simple docstring'''


									__lowerCAmelCase					=  do_resize if do_resize is not None else self.do_resize
									__lowerCAmelCase					=  resample if resample is not None else self.resample
									__lowerCAmelCase					=  do_center_crop if do_center_crop is not None else self.do_center_crop
									__lowerCAmelCase					=  do_rescale if do_rescale is not None else self.do_rescale
									__lowerCAmelCase					=  rescale_factor if rescale_factor is not None else self.rescale_factor
									__lowerCAmelCase					=  rescale_offset if rescale_offset is not None else self.rescale_offset
									__lowerCAmelCase					=  do_normalize if do_normalize is not None else self.do_normalize
									__lowerCAmelCase					=  image_mean if image_mean is not None else self.image_mean
									__lowerCAmelCase					=  image_std if image_std is not None else self.image_std
									__lowerCAmelCase					=  include_top if include_top is not None else self.include_top

									__lowerCAmelCase					=  size if size is not None else self.size
									__lowerCAmelCase					=  get_size_dict(_lowercase)
									__lowerCAmelCase					=  crop_size if crop_size is not None else self.crop_size
									__lowerCAmelCase					=  get_size_dict(_lowercase,     param_name="""crop_size""")

									__lowerCAmelCase					=  make_list_of_images(_lowercase)

									if not valid_images(_lowercase):
													raise ValueError(
													    """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
													    """torch.Tensor, tf.Tensor or jax.ndarray.""")

									if do_resize and size is None or resample is None:
													raise ValueError("""Size and resample must be specified if do_resize is True.""")

									if do_center_crop and crop_size is None:
													raise ValueError("""Crop size must be specified if do_center_crop is True.""")

									if do_rescale and rescale_factor is None:
													raise ValueError("""Rescale factor must be specified if do_rescale is True.""")

									if do_normalize and (image_mean is None or image_std is None):
													raise ValueError("""Image mean and std must be specified if do_normalize is True.""")

									# All transformations expect numpy arrays.
									__lowerCAmelCase					=  [to_numpy_array(_lowercase) for image in images]

									if do_resize:
													__lowerCAmelCase					=  [self.resize(image=_lowercase,     size=_lowercase,     resample=_lowercase) for image in images]

									if do_center_crop:
													__lowerCAmelCase					=  [self.center_crop(image=_lowercase,     size=_lowercase) for image in images]

									if do_rescale:
													__lowerCAmelCase					=  [self.rescale(image=_lowercase,     scale=_lowercase,     offset=_lowercase) for image in images]

									if do_normalize:
													__lowerCAmelCase					=  [self.normalize(image=_lowercase,     mean=_lowercase,     std=_lowercase) for image in images]

									if include_top:
													__lowerCAmelCase					=  [self.normalize(image=_lowercase,     mean=[0, 0, 0],     std=_lowercase) for image in images]

									__lowerCAmelCase					=  [to_channel_dimension_format(_lowercase,     _lowercase) for image in images]

									__lowerCAmelCase					=  {"""pixel_values""": images}
									return BatchFeature(data=_lowercase,     tensor_type=_lowercase)





