

'''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 lowerCamelCase   (    unittest.TestCase				):

				@property
				def 					_lowerCamelCase    (       self )    ->							Any:
					snake_case							=					1
					snake_case							=					3
					snake_case							=					(32, 32)

					snake_case							=					floats_tensor((batch_size, num_channels) + sizes,						rng=random.Random(0 ) ).to(lowercase_ )
					return image

				@property
				def 					_lowerCamelCase    (       self )    ->							Any:
					torch.manual_seed(0 )
					snake_case							=					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 					_lowerCamelCase    (       self )    ->							Optional[Any]:
					torch.manual_seed(0 )
					snake_case							=					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 					_lowerCamelCase    (       self )    ->							str:
					snake_case							=					'cpu'
					snake_case							=					self.dummy_uncond_unet
					snake_case							=					DDIMScheduler()
					snake_case							=					self.dummy_vq_model

					snake_case							=					LDMSuperResolutionPipeline(unet=lowercase_,						vqvae=lowercase_,						scheduler=lowercase_ )
					ldm.to(lowercase_ )
					ldm.set_progress_bar_config(disable=lowercase_ )

					snake_case							=					self.dummy_image.to(lowercase_ )

					snake_case							=					torch.Generator(device=lowercase_ ).manual_seed(0 )
					snake_case							=					ldm(image=lowercase_,						generator=lowercase_,						num_inference_steps=2,						output_type='numpy' ).images

					snake_case							=					image[0, -3:, -3:, -1]

					assert image.shape == (1, 64, 64, 3)
					snake_case							=					np.array([0.8_678, 0.8_245, 0.6_381, 0.6_830, 0.4_385, 0.5_599, 0.4_641, 0.6_201, 0.5_150] )

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






				@unittest.skipIf(torch_device != 'cuda',						'This test requires a GPU' )
				def 					_lowerCamelCase    (       self )    ->							Union[str, Any]:
					snake_case							=					self.dummy_uncond_unet
					snake_case							=					DDIMScheduler()
					snake_case							=					self.dummy_vq_model

					# put models in fp16
					snake_case							=					unet.half()
					snake_case							=					vqvae.half()

					snake_case							=					LDMSuperResolutionPipeline(unet=lowercase_,						vqvae=lowercase_,						scheduler=lowercase_ )
					ldm.to(lowercase_ )
					ldm.set_progress_bar_config(disable=lowercase_ )

					snake_case							=					self.dummy_image.to(lowercase_ )

					snake_case							=					ldm(lowercase_,						num_inference_steps=2,						output_type='numpy' ).images

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







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

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

					snake_case							=					LDMSuperResolutionPipeline.from_pretrained('duongna/ldm-super-resolution',						device_map='auto' )
					ldm.set_progress_bar_config(disable=lowercase_ )

					snake_case							=					torch.manual_seed(0 )
					snake_case							=					ldm(image=lowercase_,						generator=lowercase_,						num_inference_steps=20,						output_type='numpy' ).images

					snake_case							=					image[0, -3:, -3:, -1]

					assert image.shape == (1, 256, 256, 3)
					snake_case							=					np.array([0.7_644, 0.7_679, 0.7_642, 0.7_633, 0.7_666, 0.7_560, 0.7_425, 0.7_257, 0.6_907] )

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







import numpy as np  # noqa: E402

from ...configuration_utils import ConfigMixin, register_to_config
from ...schedulers.scheduling_utils import SchedulerMixin


try:
					import librosa  # noqa: E402

					_SCREAMING_SNAKE_CASE     :   Tuple							 =      True
					_SCREAMING_SNAKE_CASE     :   List[Any]							 =      ''
except Exception as e:
					_SCREAMING_SNAKE_CASE     :   Optional[Any]							 =      False
					_SCREAMING_SNAKE_CASE     :   Any							 =      (
					    f"""Cannot import librosa because {e}. Make sure to correctly install librosa to be able to install it."""
					)


from PIL import Image  # noqa: E402







class 			UpperCamelCase__		(     SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ ):







	'''simple docstring'''
	_lowerCAmelCase											=					"mel_config.json"





	@register_to_config
	def __init__(      self					,							UpperCamelCase__ = 256					,							UpperCamelCase__ = 256					,							UpperCamelCase__ = 2_2050					,							UpperCamelCase__ = 2048					,							UpperCamelCase__ = 512					,							UpperCamelCase__ = 80					,							UpperCamelCase__ = 32					,							):
					A__	:						Any								=				hop_length
					A__	:						List[Any]								=				sample_rate
					A__	:						Union[str, Any]								=				n_fft
					A__	:						Optional[int]								=				top_db
					A__	:						Optional[Any]								=				n_iter
					self.set_resolution(UpperCamelCase__					,							UpperCamelCase__		)
					A__	:						str								=				None

					if not _librosa_can_be_imported:
									raise ValueError(_import_error		)





	def 	__snake_case					(      self					,							UpperCamelCase__					,							UpperCamelCase__		):
					A__	:						int								=				x_res
					A__	:						List[Any]								=				y_res
					A__	:						List[str]								=				self.y_res
					A__	:						int								=				self.x_res * self.hop_length - 1





	def 	__snake_case					(      self					,							UpperCamelCase__ = None					,							UpperCamelCase__ = None		):
					if audio_file is not None:
									A__							,				A__	:						Union[str, Any]								=				librosa.load(UpperCamelCase__					,							mono=UpperCamelCase__					,							sr=self.sr		)
					else:
									A__	:						Tuple								=				raw_audio

					# Pad with silence if necessary.
					if len(self.audio		) < self.x_res * self.hop_length:
									A__	:						List[Any]								=				np.concatenate([self.audio, np.zeros((self.x_res * self.hop_length - len(self.audio		),)		)]		)





	def 	__snake_case					(      self		):
					return len(self.audio		) // self.slice_size





	def 	__snake_case					(      self					,							UpperCamelCase__ = 0		):
					return self.audio[self.slice_size * slice : self.slice_size * (slice + 1)]





	def 	__snake_case					(      self		):
					return self.sr





	def 	__snake_case					(      self					,							UpperCamelCase__		):
					A__	:						Optional[Any]								=				librosa.feature.melspectrogram(
					    y=self.get_audio_slice(UpperCamelCase__		)					,							sr=self.sr					,							n_fft=self.n_fft					,							hop_length=self.hop_length					,							n_mels=self.n_mels		)
					A__	:						str								=				librosa.power_to_db(UpperCamelCase__					,							ref=np.max					,							top_db=self.top_db		)
					A__	:						str								=				(((log_S + self.top_db) * 255 / self.top_db).clip(0					,							255		) + 0.5).astype(np.uinta		)
					A__	:						Dict								=				Image.fromarray(UpperCamelCase__		)
					return image





	def 	__snake_case					(      self					,							UpperCamelCase__		):
					A__	:						List[Any]								=				np.frombuffer(image.tobytes()					,							dtype='''uint8'''		).reshape((image.height, image.width)		)
					A__	:						Dict								=				bytedata.astype('''float'''		) * self.top_db / 255 - self.top_db
					A__	:						Union[str, Any]								=				librosa.db_to_power(UpperCamelCase__		)
					A__	:						Dict								=				librosa.feature.inverse.mel_to_audio(
					    UpperCamelCase__					,							sr=self.sr					,							n_fft=self.n_fft					,							hop_length=self.hop_length					,							n_iter=self.n_iter		)
					return audio