


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

							def  lowerCamelCase_					(				self:      Union[str, Any]				)     ->  Optional[int]:





									"""simple docstring"""







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

							@property
							def  lowerCamelCase_					(				self:      Any				)     ->  str:





									"""simple docstring"""







									torch.manual_seed(0				)
									lowercase__				= 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				)     ->  str:





									"""simple docstring"""







									torch.manual_seed(0				)
									lowercase__				= 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:      str				)     ->  int:





									"""simple docstring"""







									torch.manual_seed(0				)
									lowercase__				= 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''')       ,		)
									lowercase__				= 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:      Any				)     ->  Dict:





									"""simple docstring"""







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

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

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

									lowercase__				= torch.Generator(device=UpperCamelCase_				).manual_seed(42				)
									lowercase__				= pipe(generator=UpperCamelCase_       ,		steps=4       ,		return_dict=UpperCamelCase_				)
									lowercase__				= 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]
									)
									lowercase__				= np.frombuffer(image.tobytes()       ,		dtype='''uint8'''				)[:10]
									lowercase__				= np.frombuffer(image_from_tuple.tobytes()       ,		dtype='''uint8'''				)[:10]
									lowercase__				= 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

									lowercase__				= 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]       ,		)

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

									np.random.seed(0				)
									lowercase__				= np.random.uniform(-1       ,		1       ,		((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)				)
									lowercase__				= torch.Generator(device=UpperCamelCase_				).manual_seed(42				)
									lowercase__				= pipe(raw_audio=UpperCamelCase_       ,		generator=UpperCamelCase_       ,		start_step=5       ,		steps=10				)
									lowercase__				= 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]
									)
									lowercase__				= np.frombuffer(image.tobytes()       ,		dtype='''uint8'''				)[:10]
									lowercase__				= np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]				)

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

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

									np.random.seed(0				)
									lowercase__				= torch.rand((1, 1, 10)				)
									lowercase__				= pipe(generator=UpperCamelCase_       ,		encoding=UpperCamelCase_				)
									lowercase__				= output.images[0]
									lowercase__				= np.frombuffer(image.tobytes()       ,		dtype='''uint8'''				)[:10]
									lowercase__				= 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 			_a				(     unittest.TestCase     ):

							def  lowerCamelCase_					(				self:      Optional[int]				)     ->  int:





									"""simple docstring"""







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


							def  lowerCamelCase_					(				self:      List[Any]				)     ->  Optional[int]:





									"""simple docstring"""







									lowercase__				= torch_device

									lowercase__				= DiffusionPipeline.from_pretrained('''teticio/audio-diffusion-ddim-256'''				)
									lowercase__				= pipe.to(UpperCamelCase_				)
									pipe.set_progress_bar_config(disable=UpperCamelCase_				)

									lowercase__				= torch.Generator(device=UpperCamelCase_				).manual_seed(42				)
									lowercase__				= pipe(generator=UpperCamelCase_				)
									lowercase__				= output.audios[0]
									lowercase__				= 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]
									lowercase__				= np.frombuffer(image.tobytes()       ,		dtype='''uint8'''				)[:10]
									lowercase__				= np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]				)

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










# tests directory-specific settings - this file is run automatically
# by pytest before any tests are run

import sys
import warnings
from os.path import abspath, dirname, join


# allow having multiple repository checkouts and not needing to remember to rerun
# 'pip install -e .[dev]' when switching between checkouts and running tests.
lowerCamelCase      :       Dict    =						abspath(join(dirname(dirname(__file__)), '''src'''))
sys.path.insert(1, git_repo_path)

# silence FutureWarning warnings in tests since often we can't act on them until
# they become normal warnings - i.e. the tests still need to test the current functionality
warnings.simplefilter(action='''ignore''', category=FutureWarning)





def    __lowerCAmelCase   (    __snake_case							):
     from diffusers.utils.testing_utils import pytest_addoption_shared

     pytest_addoption_shared(__snake_case							)





def    __lowerCAmelCase   (    __snake_case							):
     from diffusers.utils.testing_utils import pytest_terminal_summary_main

     __lowerCAmelCase             =       terminalreporter.config.getoption("--make-reports"							)
     if make_reports:
          pytest_terminal_summary_main(__snake_case       ,			id=__snake_case							)






