






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

    """simple docstring"""




    def        lowercase__     (  self		:					Any						):

        '''simple docstring'''
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()




    @property
    def        lowercase__     (  self		:					List[Any]						):

        '''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        lowercase__     (  self		:					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        lowercase__     (  self		:					Optional[Any]						):

        '''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        lowercase__     (  self		:					Optional[int]						):

        '''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=lowerCamelCase,						unet=self.dummy_unet,						mel=lowerCamelCase,						scheduler=lowerCamelCase						)
        lowercase__     				=    pipe.to(lowerCamelCase						)
        pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

        lowercase__     				=    torch.Generator(device=lowerCamelCase						).manual_seed(42						)
        lowercase__     				=    pipe(generator=lowerCamelCase,						steps=4,						return_dict=lowerCamelCase						)
        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=lowerCamelCase,						scheduler=lowerCamelCase						)
        lowercase__     				=    pipe.to(lowerCamelCase						)
        pipe.set_progress_bar_config(disable=lowerCamelCase						)

        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=lowerCamelCase						).manual_seed(42						)
        lowercase__     				=    pipe(raw_audio=lowerCamelCase,						generator=lowerCamelCase,						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=lowerCamelCase,						mel=lowerCamelCase,						scheduler=lowerCamelCase						)
        lowercase__     				=    pipe.to(lowerCamelCase						)
        pipe.set_progress_bar_config(disable=lowerCamelCase						)

        np.random.seed(0						)
        lowercase__     				=    torch.rand((1, 1, 10)						)
        lowercase__     				=    pipe(generator=lowerCamelCase,						encoding=lowerCamelCase						)
        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    _UpperCAmelCase      ( unittest.TestCase			):

    """simple docstring"""




    def        lowercase__     (  self		:					Optional[int]						):

        '''simple docstring'''
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()




    def        lowercase__     (  self		:					Any						):

        '''simple docstring'''
        lowercase__     				=    torch_device

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

        lowercase__     				=    torch.Generator(device=lowerCamelCase						).manual_seed(42						)
        lowercase__     				=    pipe(generator=lowerCamelCase						)
        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












def 						a    (		lowerCamelCase_	,     lowerCamelCase_	,     lowerCamelCase_       ):






				'''simple docstring'''

				if exponent == 1:
								return base
				if exponent % 2 == 0:
								lowercase__     				=    _modexpt(lowerCamelCase_	,     exponent // 2	,     lowerCamelCase_       ) % modulo_value
								return (x * x) % modulo_value
				else:
								return (base * _modexpt(lowerCamelCase_	,     exponent - 1	,     lowerCamelCase_       )) % modulo_value






def 						a    (		lowerCamelCase_ = 1777	,     lowerCamelCase_ = 1855	,     lowerCamelCase_ = 8       ):






				'''simple docstring'''
				lowercase__     				=    base
				for _ in range(1	,     lowerCamelCase_       ):
								lowercase__     				=    _modexpt(lowerCamelCase_	,     lowerCamelCase_	,     10**digits       )

				return result


if __name__ == "__main__":
							print(F"{solution() = }")




