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


lowerCAmelCase__    :    Optional[Any]			=      './tests/fixtures/elise_format0.mid'
@unittest.skip("The note-seq package currently throws an error on import"							)
class 							snake_case			(					__UpperCAmelCase ,       unittest.TestCase							):




					"""simple docstring"""

					snake_case__							       =					SpectrogramDiffusionPipeline
					snake_case__							       =					PipelineTesterMixin.required_optional_params - {
					    "callback",
					    "latents",
					    "callback_steps",
					    "output_type",
					    "num_images_per_prompt",
					}
					snake_case__							       =					False
					snake_case__							       =					False
					snake_case__							       =					TOKENS_TO_AUDIO_GENERATION_PARAMS
					snake_case__							       =					TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS







					def 						__lowerCAmelCase		(							self  :      int     ):
									torch.manual_seed(0     )
									UpperCAmelCase__			 =				SpectrogramNotesEncoder(
									    max_length=2_048       ,vocab_size=1_536       ,d_model=768       ,dropout_rate=0.1       ,num_layers=1       ,num_heads=1       ,d_kv=4       ,d_ff=2_048       ,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=2_048       ,feed_forward_proj='gated-gelu'       ,)

									UpperCAmelCase__			 =				TaFilmDecoder(
									    input_dims=128       ,targets_length=256       ,max_decoder_noise_time=2_0_0_0_0.0       ,d_model=768       ,num_layers=1       ,num_heads=1       ,d_kv=4       ,d_ff=2_048       ,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 						__lowerCAmelCase		(							self  :      List[Any]       ,lowerCamelCase__  :      int       ,lowerCamelCase__  :      int=0     ):
									if str(lowerCamelCase__     ).startswith('mps'     ):
													UpperCAmelCase__			 =				torch.manual_seed(lowerCamelCase__     )
									else:
													UpperCAmelCase__			 =				torch.Generator(device=lowerCamelCase__     ).manual_seed(lowerCamelCase__     )
									UpperCAmelCase__			 =				{
									    'input_tokens': [
									        [1_134, 90, 1_135, 1_133, 1_080, 112, 1_132, 1_080, 1_133, 1_079, 133, 1_132, 1_079, 1_133, 1] + [0] * 2_033
									    ],
									    'generator': generator,
									    'num_inference_steps': 4,
									    'output_type': 'mel',
									}
									return inputs







					def 						__lowerCAmelCase		(							self  :      List[str]     ):
									UpperCAmelCase__			 =				'cpu'  # ensure determinism for the device-dependent torch.Generator
									UpperCAmelCase__			 =				self.get_dummy_components()
									UpperCAmelCase__			 =				SpectrogramDiffusionPipeline(**lowerCamelCase__     )
									UpperCAmelCase__			 =				pipe.to(lowerCamelCase__     )
									pipe.set_progress_bar_config(disable=lowerCamelCase__     )

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

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

									assert mel_slice.shape == (3, 3)
									UpperCAmelCase__			 =				np.array(
									    [-1_1.5_1_2_9_2_5, -4.7_8_8_2_1_5, -0.4_6_1_7_2_9_0_5, -2.0_5_1_7_1_5, -1_0.5_3_9_1_4_7, -1_0.9_7_0_9_6_3, -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 						__lowerCAmelCase		(							self  :      Union[str, Any]     ):
									return super().test_save_load_local()







					@skip_mps
					def 						__lowerCAmelCase		(							self  :      Dict     ):
									return super().test_dict_tuple_outputs_equivalent()







					@skip_mps
					def 						__lowerCAmelCase		(							self  :      Union[str, Any]     ):
									return super().test_save_load_optional_components()







					@skip_mps
					def 						__lowerCAmelCase		(							self  :      Union[str, Any]     ):
									return super().test_attention_slicing_forward_pass()







					def 						__lowerCAmelCase		(							self  :      List[str]     ):
									pass







					def 						__lowerCAmelCase		(							self  :      Union[str, Any]     ):
									pass







					@skip_mps
					def 						__lowerCAmelCase		(							self  :      str     ):
									return super().test_progress_bar()






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




					"""simple docstring"""

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







					def 						__lowerCAmelCase		(							self  :      List[Any]     ):
									# TODO - test that pipeline can decode tokens in a callback
									# so that music can be played live
									UpperCAmelCase__			 =				torch_device

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

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

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

									UpperCAmelCase__			 =				MidiProcessor()
									UpperCAmelCase__			 =				processor(lowerCamelCase__     )

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







					def 						__lowerCAmelCase		(							self  :      Tuple     ):
									UpperCAmelCase__			 =				torch_device

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

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

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

									UpperCAmelCase__			 =				output.audios[0]

									assert abs(np.abs(lowerCamelCase__     ).sum() - 3_6_1_2.8_4_1     ) < 1e-1







					def 						__lowerCAmelCase		(							self  :      Optional[int]     ):
									UpperCAmelCase__			 =				torch_device

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

									UpperCAmelCase__			 =				MidiProcessor()

									UpperCAmelCase__			 =				processor(lowerCamelCase__     )

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

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

									UpperCAmelCase__			 =				output.audios[0]
									assert abs(np.abs(lowerCamelCase__     ).sum() - 9_3_8_9.1_1_1_1     ) < 5e-2




"""simple docstring"""






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_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
						import PIL


lowerCAmelCase__    :    Tuple			=      logging.get_logger(__name__)
class 							snake_case			(					__UpperCAmelCase							):




					"""simple docstring"""

					snake_case__							       =					["pixel_values"]







					def __init__(							self  :      str       ,lowerCamelCase__  :      bool = True       ,lowerCamelCase__  :      Dict[str, int] = None       ,lowerCamelCase__  :      bool = True       ,lowerCamelCase__  :      Dict[str, int] = None       ,lowerCamelCase__  :      PILImageResampling = PILImageResampling.BICUBIC       ,lowerCamelCase__  :      bool = True       ,lowerCamelCase__  :      Union[int, float] = 1 / 255       ,lowerCamelCase__  :      bool = True       ,lowerCamelCase__  :      Optional[Union[float, List[float]]] = None       ,lowerCamelCase__  :      Optional[Union[float, List[float]]] = None       ,**lowerCamelCase__  :      List[Any]       ,):
									super().__init__(**lowerCamelCase__     )
									UpperCAmelCase__			 =				crop_size if crop_size is not None else {'height': 256, 'width': 256}
									UpperCAmelCase__			 =				get_size_dict(lowerCamelCase__       ,param_name='crop_size'     )
									UpperCAmelCase__			 =				size if size is not None else {'height': 224, 'width': 224}
									UpperCAmelCase__			 =				get_size_dict(lowerCamelCase__     )

									UpperCAmelCase__			 =				do_center_crop
									UpperCAmelCase__			 =				crop_size
									UpperCAmelCase__			 =				do_resize
									UpperCAmelCase__			 =				size
									UpperCAmelCase__			 =				resample
									UpperCAmelCase__			 =				do_rescale
									UpperCAmelCase__			 =				rescale_factor
									UpperCAmelCase__			 =				do_normalize
									UpperCAmelCase__			 =				image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
									UpperCAmelCase__			 =				image_std if image_std is not None else IMAGENET_DEFAULT_STD







					def 						__lowerCAmelCase		(							self  :      Optional[int]       ,lowerCamelCase__  :      np.ndarray       ,lowerCamelCase__  :      Dict[str, int]       ,lowerCamelCase__  :      Optional[int] = None       ,lowerCamelCase__  :      Optional[Union[str, ChannelDimension]] = None       ,**lowerCamelCase__  :      List[str]       ,):
									UpperCAmelCase__			 =				self.size if size is None else size
									UpperCAmelCase__			 =				get_size_dict(lowerCamelCase__     )
									UpperCAmelCase__			 =				get_size_dict(lowerCamelCase__       ,param_name='crop_size'     )

									UpperCAmelCase__ ,      UpperCAmelCase__			 =				get_image_size(lowerCamelCase__     )
									UpperCAmelCase__			 =				min(lowerCamelCase__       ,lowerCamelCase__     )
									UpperCAmelCase__			 =				(size['height'] / crop_size['height']) * min_dim
									UpperCAmelCase__			 =				(size['width'] / crop_size['width']) * min_dim
									return center_crop(lowerCamelCase__       ,size=(cropped_height, cropped_width)       ,data_format=lowerCamelCase__       ,**lowerCamelCase__     )







					def 						__lowerCAmelCase		(							self  :      Any       ,lowerCamelCase__  :      np.ndarray       ,lowerCamelCase__  :      Dict[str, int]       ,lowerCamelCase__  :      PILImageResampling = PIL.Image.BICUBIC       ,lowerCamelCase__  :      Optional[Union[str, ChannelDimension]] = None       ,**lowerCamelCase__  :      Any       ,):
									UpperCAmelCase__			 =				get_size_dict(lowerCamelCase__     )
									if "height" not in size or "width" not in size:
													raise ValueError(f'''The size dictionary must contain the keys \'height\' and \'width\'. Got {size.keys()}'''     )
									return resize(
									    lowerCamelCase__       ,size=(size['height'], size['width'])       ,resample=lowerCamelCase__       ,data_format=lowerCamelCase__       ,**lowerCamelCase__     )







					def 						__lowerCAmelCase		(							self  :      Optional[int]       ,lowerCamelCase__  :      np.ndarray       ,lowerCamelCase__  :      Union[int, float]       ,lowerCamelCase__  :      Optional[Union[str, ChannelDimension]] = None       ,**lowerCamelCase__  :      int       ,):
									return rescale(lowerCamelCase__       ,scale=lowerCamelCase__       ,data_format=lowerCamelCase__       ,**lowerCamelCase__     )







					def 						__lowerCAmelCase		(							self  :      int       ,lowerCamelCase__  :      np.ndarray       ,lowerCamelCase__  :      Union[float, List[float]]       ,lowerCamelCase__  :      Union[float, List[float]]       ,lowerCamelCase__  :      Optional[Union[str, ChannelDimension]] = None       ,**lowerCamelCase__  :      List[Any]       ,):
									return normalize(lowerCamelCase__       ,mean=lowerCamelCase__       ,std=lowerCamelCase__       ,data_format=lowerCamelCase__       ,**lowerCamelCase__     )







					def 						__lowerCAmelCase		(							self  :      Optional[Any]       ,lowerCamelCase__  :      ImageInput       ,lowerCamelCase__  :      Optional[bool] = None       ,lowerCamelCase__  :      Optional[Dict[str, int]] = None       ,lowerCamelCase__  :      Optional[bool] = None       ,lowerCamelCase__  :      Optional[Dict[str, int]] = None       ,lowerCamelCase__  :      PILImageResampling = None       ,lowerCamelCase__  :      Optional[bool] = None       ,lowerCamelCase__  :      Optional[float] = None       ,lowerCamelCase__  :      Optional[bool] = None       ,lowerCamelCase__  :      Optional[Union[float, List[float]]] = None       ,lowerCamelCase__  :      Optional[Union[float, List[float]]] = None       ,lowerCamelCase__  :      Optional[Union[str, TensorType]] = None       ,lowerCamelCase__  :      ChannelDimension = ChannelDimension.FIRST       ,**lowerCamelCase__  :      Optional[int]       ,):
									UpperCAmelCase__			 =				do_center_crop if do_center_crop is not None else self.do_center_crop
									UpperCAmelCase__			 =				crop_size if crop_size is not None else self.crop_size
									UpperCAmelCase__			 =				get_size_dict(lowerCamelCase__       ,param_name='crop_size'     )
									UpperCAmelCase__			 =				do_resize if do_resize is not None else self.do_resize
									UpperCAmelCase__			 =				size if size is not None else self.size
									UpperCAmelCase__			 =				get_size_dict(lowerCamelCase__     )
									UpperCAmelCase__			 =				resample if resample is not None else self.resample
									UpperCAmelCase__			 =				do_rescale if do_rescale is not None else self.do_rescale
									UpperCAmelCase__			 =				rescale_factor if rescale_factor is not None else self.rescale_factor
									UpperCAmelCase__			 =				do_normalize if do_normalize is not None else self.do_normalize
									UpperCAmelCase__			 =				image_mean if image_mean is not None else self.image_mean
									UpperCAmelCase__			 =				image_std if image_std is not None else self.image_std

									UpperCAmelCase__			 =				make_list_of_images(lowerCamelCase__     )

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

									if do_center_crop and crop_size is None:
													raise ValueError('If `do_center_crop` is set to `True`, `crop_size` must be provided.'     )

									if do_resize and size is None:
													raise ValueError('Size must be specified if do_resize 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 image standard deviation must be specified if do_normalize is True.'     )

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

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

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

									if do_rescale:
													UpperCAmelCase__			 =				[self.rescale(image=lowerCamelCase__       ,scale=lowerCamelCase__     ) for image in images]

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

									UpperCAmelCase__			 =				[to_channel_dimension_format(lowerCamelCase__       ,lowerCamelCase__     ) for image in images]

									UpperCAmelCase__			 =				{'pixel_values': images}
									return BatchFeature(data=lowerCamelCase__       ,tensor_type=lowerCamelCase__     )



