




'''simple docstring'''





import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




class 			_snake_case		(     unittest.TestCase							):




			def        lowerCAmelCase__    ( self							)    ->    List[str]:

										'''simple docstring'''




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




			@property
			def        lowerCAmelCase__    ( self							)    ->    Any:

										'''simple docstring'''




										torch.manual_seed(0							)
										snake_case_      =				UNetaDModel(
										    sample_size=(32, 64)    ,							in_channels=1    ,							out_channels=1    ,							layers_per_block=2    ,							block_out_channels=(128, 128)    ,							down_block_types=("AttnDownBlock2D", "DownBlock2D")    ,							up_block_types=("UpBlock2D", "AttnUpBlock2D")    ,							)
										return model




			@property
			def        lowerCAmelCase__    ( self							)    ->    str:

										'''simple docstring'''




										torch.manual_seed(0							)
										snake_case_      =				UNetaDConditionModel(
										    sample_size=(64, 32)    ,							in_channels=1    ,							out_channels=1    ,							layers_per_block=2    ,							block_out_channels=(128, 128)    ,							down_block_types=("CrossAttnDownBlock2D", "DownBlock2D")    ,							up_block_types=("UpBlock2D", "CrossAttnUpBlock2D")    ,							cross_attention_dim=10    ,							)
										return model




			@property
			def        lowerCAmelCase__    ( self							)    ->    Optional[int]:

										'''simple docstring'''




										torch.manual_seed(0							)
										snake_case_      =				AutoencoderKL(
										    sample_size=(128, 64)    ,							in_channels=1    ,							out_channels=1    ,							latent_channels=1    ,							layers_per_block=2    ,							block_out_channels=(128, 128)    ,							down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D")    ,							up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D")    ,							)
										snake_case_      =				UNetaDModel(
										    sample_size=(64, 32)    ,							in_channels=1    ,							out_channels=1    ,							layers_per_block=2    ,							block_out_channels=(128, 128)    ,							down_block_types=("AttnDownBlock2D", "DownBlock2D")    ,							up_block_types=("UpBlock2D", "AttnUpBlock2D")    ,							)
										return vqvae, unet




			@slow
			def        lowerCAmelCase__    ( self							)    ->    List[str]:

										'''simple docstring'''




										snake_case_      =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										snake_case_      =				Mel(
										    x_res=self.dummy_unet.config.sample_size[1]    ,							y_res=self.dummy_unet.config.sample_size[0]    ,							)

										snake_case_      =				DDPMScheduler()
										snake_case_      =				AudioDiffusionPipeline(vqvae=__lowerCamelCase    ,							unet=self.dummy_unet    ,							mel=__lowerCamelCase    ,							scheduler=__lowerCamelCase							)
										snake_case_      =				pipe.to(__lowerCamelCase							)
										pipe.set_progress_bar_config(disable=__lowerCamelCase							)

										snake_case_      =				torch.Generator(device=__lowerCamelCase							).manual_seed(42							)
										snake_case_      =				pipe(generator=__lowerCamelCase    ,							steps=4							)
										snake_case_      =				output.audios[0]
										snake_case_      =				output.images[0]

										snake_case_      =				torch.Generator(device=__lowerCamelCase							).manual_seed(42							)
										snake_case_      =				pipe(generator=__lowerCamelCase    ,							steps=4    ,							return_dict=__lowerCamelCase							)
										snake_case_      =				output[0][0]

										assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
										assert (
										    image.height == self.dummy_unet.config.sample_size[0]
										    and image.width == self.dummy_unet.config.sample_size[1]
										)
										snake_case_      =				np.frombuffer(image.tobytes()    ,							dtype="uint8"							)[:10]
										snake_case_      =				np.frombuffer(image_from_tuple.tobytes()    ,							dtype="uint8"							)[:10]
										snake_case_      =				np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]							)

										assert np.abs(image_slice.flatten() - expected_slice							).max() == 0
										assert np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() == 0

										snake_case_      =				Mel(
										    x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]    ,							y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]    ,							)

										snake_case_      =				DDIMScheduler()
										snake_case_      =				self.dummy_vqvae_and_unet
										snake_case_      =				AudioDiffusionPipeline(
										    vqvae=self.dummy_vqvae_and_unet[0]    ,							unet=dummy_vqvae_and_unet[1]    ,							mel=__lowerCamelCase    ,							scheduler=__lowerCamelCase							)
										snake_case_      =				pipe.to(__lowerCamelCase							)
										pipe.set_progress_bar_config(disable=__lowerCamelCase							)

										np.random.seed(0							)
										snake_case_      =				np.random.uniform(-1    ,							1    ,							((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)							)
										snake_case_      =				torch.Generator(device=__lowerCamelCase							).manual_seed(42							)
										snake_case_      =				pipe(raw_audio=__lowerCamelCase    ,							generator=__lowerCamelCase    ,							start_step=5    ,							steps=10							)
										snake_case_      =				output.images[0]

										assert (
										    image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
										    and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
										)
										snake_case_      =				np.frombuffer(image.tobytes()    ,							dtype="uint8"							)[:10]
										snake_case_      =				np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]							)

										assert np.abs(image_slice.flatten() - expected_slice							).max() == 0

										snake_case_      =				self.dummy_unet_condition
										snake_case_      =				AudioDiffusionPipeline(
										    vqvae=self.dummy_vqvae_and_unet[0]    ,							unet=__lowerCamelCase    ,							mel=__lowerCamelCase    ,							scheduler=__lowerCamelCase							)
										snake_case_      =				pipe.to(__lowerCamelCase							)
										pipe.set_progress_bar_config(disable=__lowerCamelCase							)

										np.random.seed(0							)
										snake_case_      =				torch.rand((1, 1, 10)							)
										snake_case_      =				pipe(generator=__lowerCamelCase    ,							encoding=__lowerCamelCase							)
										snake_case_      =				output.images[0]
										snake_case_      =				np.frombuffer(image.tobytes()    ,							dtype="uint8"							)[:10]
										snake_case_      =				np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]							)

										assert np.abs(image_slice.flatten() - expected_slice							).max() == 0




@slow
@require_torch_gpu
class 			_snake_case		(     unittest.TestCase							):




			def        lowerCAmelCase__    ( self							)    ->    int:

										'''simple docstring'''




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




			def        lowerCAmelCase__    ( self							)    ->    List[str]:

										'''simple docstring'''




										snake_case_      =				torch_device

										snake_case_      =				DiffusionPipeline.from_pretrained("teticio/audio-diffusion-ddim-256"							)
										snake_case_      =				pipe.to(__lowerCamelCase							)
										pipe.set_progress_bar_config(disable=__lowerCamelCase							)

										snake_case_      =				torch.Generator(device=__lowerCamelCase							).manual_seed(42							)
										snake_case_      =				pipe(generator=__lowerCamelCase							)
										snake_case_      =				output.audios[0]
										snake_case_      =				output.images[0]

										assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
										assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
										snake_case_      =				np.frombuffer(image.tobytes()    ,							dtype="uint8"							)[:10]
										snake_case_      =				np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]							)

										assert np.abs(image_slice.flatten() - expected_slice							).max() == 0










'''simple docstring'''





from ..utils import DummyObject, requires_backends




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Dict							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    List[Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    str:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Dict							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Dict:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[str]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Optional[Any]							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    str:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Optional[int]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[str]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Optional[Any]							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Union[str, Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[str]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Dict:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				List[str]							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Optional[int]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				int							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    List[str]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Any:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Optional[Any]							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Any:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    int:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Dict							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Optional[int]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    List[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				List[str]							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Dict:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Union[str, Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				int							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Dict:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Union[str, Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    int:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Any							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Tuple:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    str:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				str							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    List[str]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Union[str, Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Union[str, Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




class 			_snake_case		(     metaclass=lowercase_							):
			lowerCAmelCase_					:				Tuple							=   ["flax"]




			def __init__( self    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(self    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    int:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




			@classmethod
			def        lowerCAmelCase__    ( cls    ,							*a__    ,							**a__							)    ->    Optional[Any]:

										'''simple docstring'''




										requires_backends(cls    ,							["flax"]							)




