




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


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






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







						"""simple docstring"""


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







						def       _a    (     self      :       str    ):





											"""simple docstring"""







											torch.manual_seed(0    )
											A_ :     Optional[int] =							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_ :     Tuple =							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=2_00_00.0							, d_model=768							, num_layers=1							, num_heads=1							, d_kv=4							, d_ff=2048							, dropout_rate=0.1							, )

											A_ :     List[str] =							DDPMScheduler()

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







						def       _a    (     self      :       Union[str, Any]							, _lowerCamelCase      :       Dict							, _lowerCamelCase      :       str=0    ):





											"""simple docstring"""







											if str(_lowerCamelCase    ).startswith('''mps'''    ):
																A_ :     Optional[int] =							torch.manual_seed(_lowerCamelCase    )
											else:
																A_ :     Optional[Any] =							torch.Generator(device=_lowerCamelCase    ).manual_seed(_lowerCamelCase    )
											A_ :     Optional[int] =							{
											    '''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       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







											A_ :     Tuple =							'''cpu'''  # ensure determinism for the device-dependent torch.Generator
											A_ :     List[str] =							self.get_dummy_components()
											A_ :     List[Any] =							SpectrogramDiffusionPipeline(**_lowerCamelCase    )
											A_ :     Any =							pipe.to(_lowerCamelCase    )
											pipe.set_progress_bar_config(disable=_lowerCamelCase    )

											A_ :     Tuple =							self.get_dummy_inputs(_lowerCamelCase    )
											A_ :     Union[str, Any] =							pipe(**_lowerCamelCase    )
											A_ :     Optional[Any] =							output.audios

											A_ :     Optional[Any] =							mel[0, -3:, -3:]

											assert mel_slice.shape == (3, 3)
											A_ :     Optional[Any] =							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       _a    (     self      :       Dict    ):





											"""simple docstring"""







											return super().test_save_load_local()







						@skip_mps
						def       _a    (     self      :       Any    ):





											"""simple docstring"""







											return super().test_dict_tuple_outputs_equivalent()







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





											"""simple docstring"""







											return super().test_save_load_optional_components()







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





											"""simple docstring"""







											return super().test_attention_slicing_forward_pass()







						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											pass







						def       _a    (     self      :       Tuple    ):





											"""simple docstring"""







											pass







						@skip_mps
						def       _a    (     self      :       int    ):





											"""simple docstring"""







											return super().test_progress_bar()






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







						"""simple docstring"""


						def       _a    (     self      :       int    ):





											"""simple docstring"""







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







						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											A_ :     Tuple =							torch_device

											A_ :     List[Any] =							SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''    )
											A_ :     Dict =							pipe.melgan
											A_ :     List[str] =							None

											A_ :     Optional[int] =							pipe.to(_lowerCamelCase    )
											pipe.set_progress_bar_config(disable=_lowerCamelCase    )

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

											A_ :     str =							MidiProcessor()
											A_ :     Optional[int] =							processor(_lowerCamelCase    )

											A_ :     Optional[int] =							input_tokens[:3]
											A_ :     Dict =							torch.manual_seed(0    )
											pipe(_lowerCamelCase							, num_inference_steps=5							, generator=_lowerCamelCase							, callback=_lowerCamelCase							, output_type='''mel'''    )







						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







											A_ :     Union[str, Any] =							torch_device

											A_ :     List[str] =							SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''    )
											A_ :     Any =							pipe.to(_lowerCamelCase    )
											pipe.set_progress_bar_config(disable=_lowerCamelCase    )
											A_ :     Optional[int] =							MidiProcessor()

											A_ :     Union[str, Any] =							processor(_lowerCamelCase    )
											# just run two denoising loops
											A_ :     Optional[Any] =							input_tokens[:2]

											A_ :     Union[str, Any] =							torch.manual_seed(0    )
											A_ :     Tuple =							pipe(_lowerCamelCase							, num_inference_steps=2							, generator=_lowerCamelCase    )

											A_ :     Dict =							output.audios[0]

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







						def       _a    (     self      :       List[str]    ):





											"""simple docstring"""







											A_ :     Optional[int] =							torch_device

											A_ :     Optional[int] =							SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''    )
											A_ :     str =							pipe.to(_lowerCamelCase    )
											pipe.set_progress_bar_config(disable=_lowerCamelCase    )

											A_ :     Any =							MidiProcessor()

											A_ :     str =							processor(_lowerCamelCase    )

											# just run 4 denoising loops
											A_ :     Union[str, Any] =							input_tokens[:4]

											A_ :     List[str] =							torch.manual_seed(0    )
											A_ :     int =							pipe(_lowerCamelCase							, num_inference_steps=100							, generator=_lowerCamelCase    )

											A_ :     Optional[int] =							output.audios[0]
											assert abs(np.abs(_lowerCamelCase    ).sum() - 93_89.11_11    ) < 5E-2









'''simple docstring'''




import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, LDMSuperResolutionPipeline, UNetaDModel, VQModel
from diffusers.utils import PIL_INTERPOLATION, floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch


enable_full_determinism()






class     UpperCamelCase_     (unittest.TestCase     ):







						"""simple docstring"""


						@property
						def       _a    (     self      :       int    ):





											"""simple docstring"""







											A_ :     Dict =							1
											A_ :     Any =							3
											A_ :     Any =							(32, 32)

											A_ :     Dict =							floats_tensor((batch_size, num_channels) + sizes							, rng=random.Random(0    )    ).to(_lowerCamelCase    )
											return image







						@property
						def       _a    (     self      :       Dict    ):





											"""simple docstring"""







											torch.manual_seed(0    )
											A_ :     Optional[int] =							UNetaDModel(
											    block_out_channels=(32, 64)							, layers_per_block=2							, sample_size=32							, in_channels=6							, out_channels=3							, down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')							, up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')							, )
											return model







						@property
						def       _a    (     self      :       Any    ):





											"""simple docstring"""







											torch.manual_seed(0    )
											A_ :     Dict =							VQModel(
											    block_out_channels=[32, 64]							, in_channels=3							, out_channels=3							, down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']							, up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']							, latent_channels=3							, )
											return model







						def       _a    (     self      :       List[Any]    ):





											"""simple docstring"""







											A_ :     Any =							'''cpu'''
											A_ :     Optional[Any] =							self.dummy_uncond_unet
											A_ :     str =							DDIMScheduler()
											A_ :     Optional[int] =							self.dummy_vq_model

											A_ :     Optional[Any] =							LDMSuperResolutionPipeline(unet=_lowerCamelCase							, vqvae=_lowerCamelCase							, scheduler=_lowerCamelCase    )
											ldm.to(_lowerCamelCase    )
											ldm.set_progress_bar_config(disable=_lowerCamelCase    )

											A_ :     Any =							self.dummy_image.to(_lowerCamelCase    )

											A_ :     str =							torch.Generator(device=_lowerCamelCase    ).manual_seed(0    )
											A_ :     List[Any] =							ldm(image=_lowerCamelCase							, generator=_lowerCamelCase							, num_inference_steps=2							, output_type='''numpy'''    ).images

											A_ :     str =							image[0, -3:, -3:, -1]

											assert image.shape == (1, 64, 64, 3)
											A_ :     List[Any] =							np.array([0.86_78, 0.82_45, 0.63_81, 0.68_30, 0.43_85, 0.55_99, 0.46_41, 0.62_01, 0.51_50]    )

											assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2







						@unittest.skipIf(torch_device != '''cuda'''							, '''This test requires a GPU'''    )
						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







											A_ :     Tuple =							self.dummy_uncond_unet
											A_ :     List[str] =							DDIMScheduler()
											A_ :     str =							self.dummy_vq_model

											# put models in fp16
											A_ :     List[Any] =							unet.half()
											A_ :     Optional[int] =							vqvae.half()

											A_ :     Optional[Any] =							LDMSuperResolutionPipeline(unet=_lowerCamelCase							, vqvae=_lowerCamelCase							, scheduler=_lowerCamelCase    )
											ldm.to(_lowerCamelCase    )
											ldm.set_progress_bar_config(disable=_lowerCamelCase    )

											A_ :     Optional[Any] =							self.dummy_image.to(_lowerCamelCase    )

											A_ :     Optional[Any] =							ldm(_lowerCamelCase							, num_inference_steps=2							, output_type='''numpy'''    ).images

											assert image.shape == (1, 64, 64, 3)






@slow
@require_torch
class     UpperCamelCase_     (unittest.TestCase     ):







						"""simple docstring"""


						def       _a    (     self      :       int    ):





											"""simple docstring"""







											A_ :     Optional[int] =							load_image(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
											    '''/vq_diffusion/teddy_bear_pool.png'''    )
											A_ :     List[str] =							init_image.resize((64, 64)							, resample=PIL_INTERPOLATION['''lanczos''']    )

											A_ :     List[str] =							LDMSuperResolutionPipeline.from_pretrained('''duongna/ldm-super-resolution'''							, device_map='''auto'''    )
											ldm.set_progress_bar_config(disable=_lowerCamelCase    )

											A_ :     List[str] =							torch.manual_seed(0    )
											A_ :     str =							ldm(image=_lowerCamelCase							, generator=_lowerCamelCase							, num_inference_steps=20							, output_type='''numpy'''    ).images

											A_ :     str =							image[0, -3:, -3:, -1]

											assert image.shape == (1, 256, 256, 3)
											A_ :     Optional[Any] =							np.array([0.76_44, 0.76_79, 0.76_42, 0.76_33, 0.76_66, 0.75_60, 0.74_25, 0.72_57, 0.69_07]    )

											assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2



