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


_lowercase		:					List[Any]													=		"./tests/fixtures/elise_format0.mid"





@unittest.skip('''The note-seq package currently throws an error on import''')
class        __magic_name__		(       _UpperCAmelCase,	unittest.TestCase):
			UpperCamelCase__							=			SpectrogramDiffusionPipeline
			UpperCamelCase__							=			PipelineTesterMixin.required_optional_params - {
			    '''callback''',
			    '''latents''',
			    '''callback_steps''',
			    '''output_type''',
			    '''num_images_per_prompt''',
			}
			UpperCamelCase__							=			False
			UpperCamelCase__							=			False
			UpperCamelCase__							=			TOKENS_TO_AUDIO_GENERATION_PARAMS
			UpperCamelCase__							=			TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS







			def       SCREAMING_SNAKE_CASE_	(      self							:    Tuple       ):
										torch.manual_seed(0       )
										lowercase_   :     List[str]									=			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"""							,     )

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

										lowercase_   :     Any									=			TaFilmDecoder(
										    input_dims=128							,     targets_length=256							,     max_decoder_noise_time=2_00_00.0							,     d_model=768							,     num_layers=1							,     num_heads=1							,     d_kv=4							,     d_ff=2048							,     dropout_rate=0.1							,     )

										lowercase_   :     List[Any]									=			DDPMScheduler()

										lowercase_   :     Union[str, Any]									=			{
										    """notes_encoder""": notes_encoder.eval(),
										    """continuous_encoder""": continuous_encoder.eval(),
										    """decoder""": decoder.eval(),
										    """scheduler""": scheduler,
										    """melgan""": None,
										}
										return components







			def       SCREAMING_SNAKE_CASE_	(      self							:    Dict							,     lowercase_							:    Any							,     lowercase_							:    List[Any]=0       ):
										if str(lowercase_       ).startswith("""mps"""       ):
																	lowercase_   :     Tuple									=			torch.manual_seed(lowercase_       )
										else:
																	lowercase_   :     Any									=			torch.Generator(device=lowercase_       ).manual_seed(lowercase_       )
										lowercase_   :     Dict									=			{
										    """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							:    Optional[int]       ):
										lowercase_   :     List[Any]									=			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
										lowercase_   :     Any									=			self.get_dummy_components()
										lowercase_   :     Tuple									=			SpectrogramDiffusionPipeline(**lowercase_       )
										lowercase_   :     Any									=			pipe.to(lowercase_       )
										pipe.set_progress_bar_config(disable=lowercase_       )

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

										lowercase_   :     Optional[int]									=			mel[0, -3:, -3:]

										assert mel_slice.shape == (3, 3)
										lowercase_   :     List[str]									=			np.array(
										    [-11.51_29_25, -4.78_82_15, -0.46_17_29_05, -2.05_17_15, -10.53_91_47, -10.97_09_63, -9.09_16_34, 4.0, 4.0]       )
										assert np.abs(mel_slice.flatten() - expected_slice       ).max() < 1E-2







			@skip_mps
			def       SCREAMING_SNAKE_CASE_	(      self							:    str       ):
										return super().test_save_load_local()







			@skip_mps
			def       SCREAMING_SNAKE_CASE_	(      self							:    Optional[Any]       ):
										return super().test_dict_tuple_outputs_equivalent()







			@skip_mps
			def       SCREAMING_SNAKE_CASE_	(      self							:    Any       ):
										return super().test_save_load_optional_components()







			@skip_mps
			def       SCREAMING_SNAKE_CASE_	(      self							:    str       ):
										return super().test_attention_slicing_forward_pass()







			def       SCREAMING_SNAKE_CASE_	(      self							:    Tuple       ):
										pass







			def       SCREAMING_SNAKE_CASE_	(      self							:    Tuple       ):
										pass







			@skip_mps
			def       SCREAMING_SNAKE_CASE_	(      self							:    List[str]       ):
										return super().test_progress_bar()





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







			def       SCREAMING_SNAKE_CASE_	(      self							:    List[str]       ):
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()







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

										lowercase_   :     int									=			SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""       )
										lowercase_   :     Tuple									=			pipe.melgan
										lowercase_   :     Union[str, Any]									=			None

										lowercase_   :     Union[str, Any]									=			pipe.to(lowercase_       )
										pipe.set_progress_bar_config(disable=lowercase_       )

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

										lowercase_   :     Tuple									=			MidiProcessor()
										lowercase_   :     int									=			processor(lowercase_       )

										lowercase_   :     Tuple									=			input_tokens[:3]
										lowercase_   :     Union[str, Any]									=			torch.manual_seed(0       )
										pipe(lowercase_							,     num_inference_steps=5							,     generator=lowercase_							,     callback=lowercase_							,     output_type="""mel"""       )







			def       SCREAMING_SNAKE_CASE_	(      self							:    Optional[Any]       ):
										lowercase_   :     List[Any]									=			torch_device

										lowercase_   :     Optional[Any]									=			SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""       )
										lowercase_   :     List[Any]									=			pipe.to(lowercase_       )
										pipe.set_progress_bar_config(disable=lowercase_       )
										lowercase_   :     Any									=			MidiProcessor()

										lowercase_   :     Tuple									=			processor(lowercase_       )
										# just run two denoising loops
										lowercase_   :     List[Any]									=			input_tokens[:2]

										lowercase_   :     Optional[int]									=			torch.manual_seed(0       )
										lowercase_   :     Any									=			pipe(lowercase_							,     num_inference_steps=2							,     generator=lowercase_       )

										lowercase_   :     List[str]									=			output.audios[0]

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







			def       SCREAMING_SNAKE_CASE_	(      self							:    str       ):
										lowercase_   :     Any									=			torch_device

										lowercase_   :     List[Any]									=			SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""       )
										lowercase_   :     Union[str, Any]									=			pipe.to(lowercase_       )
										pipe.set_progress_bar_config(disable=lowercase_       )

										lowercase_   :     Dict									=			MidiProcessor()

										lowercase_   :     Dict									=			processor(lowercase_       )

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

										lowercase_   :     Dict									=			torch.manual_seed(0       )
										lowercase_   :     str									=			pipe(lowercase_							,     num_inference_steps=100							,     generator=lowercase_       )

										lowercase_   :     Any									=			output.audios[0]
										assert abs(np.abs(lowercase_       ).sum() - 93_89.11_11       ) < 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,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


_lowercase		:					Optional[int]													=		logging.get_logger(__name__)


if is_vision_available():
						import PIL





class        __magic_name__		(       _UpperCAmelCase):
			UpperCamelCase__							=			['''pixel_values''']







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

										lowercase_   :     Tuple									=			do_resize
										lowercase_   :     int									=			size
										lowercase_   :     List[Any]									=			resample
										lowercase_   :     str									=			do_center_crop
										lowercase_   :     Optional[int]									=			crop_size
										lowercase_   :     str									=			do_rescale
										lowercase_   :     List[Any]									=			rescale_factor
										lowercase_   :     Dict									=			do_normalize
										lowercase_   :     List[str]									=			image_mean if image_mean is not None else OPENAI_CLIP_MEAN
										lowercase_   :     Any									=			image_std if image_std is not None else OPENAI_CLIP_STD
										lowercase_   :     List[str]									=			do_convert_rgb







			def       SCREAMING_SNAKE_CASE_	(      self							:    Union[str, Any]							,     lowercase_							:    np.ndarray							,     lowercase_							:    Dict[str, int]							,     lowercase_							:    PILImageResampling = PILImageResampling.BICUBIC							,     lowercase_							:    Optional[Union[str, ChannelDimension]] = None							,     **lowercase_							:    Dict							,     ):
										lowercase_   :     Tuple									=			get_size_dict(lowercase_							,     default_to_square=lowercase_       )
										lowercase_   :     List[Any]									=			get_resize_output_image_size(
										    lowercase_							,     size=(size["""height"""], size["""width"""])							,     default_to_square=lowercase_       )
										return resize(lowercase_							,     size=lowercase_							,     resample=lowercase_							,     data_format=lowercase_							,     **lowercase_       )







			def       SCREAMING_SNAKE_CASE_	(      self							:    List[Any]							,     lowercase_							:    np.ndarray							,     lowercase_							:    Dict[str, int]							,     lowercase_							:    Optional[Union[str, ChannelDimension]] = None							,     **lowercase_							:    Any							,     ):
										lowercase_   :     int									=			get_size_dict(lowercase_       )
										return center_crop(lowercase_							,     size=(size["""height"""], size["""width"""])							,     data_format=lowercase_							,     **lowercase_       )







			def       SCREAMING_SNAKE_CASE_	(      self							:    Tuple							,     lowercase_							:    np.ndarray							,     lowercase_							:    Union[int, float]							,     lowercase_							:    Optional[Union[str, ChannelDimension]] = None							,     **lowercase_							:    List[Any]							,     ):
										return rescale(lowercase_							,     scale=lowercase_							,     data_format=lowercase_							,     **lowercase_       )







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







			def       SCREAMING_SNAKE_CASE_	(      self							:    List[Any]							,     lowercase_							:    ImageInput							,     lowercase_							:    bool = None							,     lowercase_							:    Dict[str, int] = None							,     lowercase_							:    PILImageResampling = None							,     lowercase_							:    bool = None							,     lowercase_							:    int = None							,     lowercase_							:    bool = None							,     lowercase_							:    float = 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_							:    Optional[ChannelDimension] = ChannelDimension.FIRST							,     **lowercase_							:    Optional[int]							,     ):
										lowercase_   :     Optional[Any]									=			do_resize if do_resize is not None else self.do_resize
										lowercase_   :     Union[str, Any]									=			size if size is not None else self.size
										lowercase_   :     Optional[int]									=			get_size_dict(lowercase_							,     default_to_square=lowercase_       )
										lowercase_   :     int									=			resample if resample is not None else self.resample
										lowercase_   :     str									=			do_center_crop if do_center_crop is not None else self.do_center_crop
										lowercase_   :     int									=			crop_size if crop_size is not None else self.crop_size
										lowercase_   :     int									=			get_size_dict(lowercase_       )
										lowercase_   :     Dict									=			do_rescale if do_rescale is not None else self.do_rescale
										lowercase_   :     List[str]									=			rescale_factor if rescale_factor is not None else self.rescale_factor
										lowercase_   :     Optional[Any]									=			do_normalize if do_normalize is not None else self.do_normalize
										lowercase_   :     List[Any]									=			image_mean if image_mean is not None else self.image_mean
										lowercase_   :     str									=			image_std if image_std is not None else self.image_std
										lowercase_   :     Any									=			do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

										lowercase_   :     int									=			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:
																	raise ValueError("""Size 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."""       )

										# PIL RGBA images are converted to RGB
										if do_convert_rgb:
																	lowercase_   :     Dict									=			[convert_to_rgb(lowercase_       ) for image in images]

										# All transformations expect numpy arrays.
										lowercase_   :     Any									=			[to_numpy_array(lowercase_       ) for image in images]

										if do_resize:
																	lowercase_   :     Union[str, Any]									=			[self.resize(image=lowercase_							,     size=lowercase_							,     resample=lowercase_       ) for image in images]

										if do_center_crop:
																	lowercase_   :     List[str]									=			[self.center_crop(image=lowercase_							,     size=lowercase_       ) for image in images]

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

										if do_normalize:
																	lowercase_   :     List[Any]									=			[self.normalize(image=lowercase_							,     mean=lowercase_							,     std=lowercase_       ) for image in images]

										lowercase_   :     List[Any]									=			[to_channel_dimension_format(lowercase_							,     lowercase_       ) for image in images]

										lowercase_   :     Any									=			{"""pixel_values""": images}
										return BatchFeature(data=lowercase_							,     tensor_type=lowercase_       )





