


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







       def 						snake_case_     (       self:      List[Any] ):






            '''simple docstring'''




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







       @property
       def 						snake_case_     (       self:      Optional[Any] ):






            '''simple docstring'''




            torch.manual_seed(0 )
            __UpperCamelCase      =       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 						snake_case_     (       self:      List[Any] ):






            '''simple docstring'''




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






            '''simple docstring'''




            torch.manual_seed(0 )
            __UpperCamelCase      =       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'),)
            __UpperCamelCase      =       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 						snake_case_     (       self:      Dict ):






            '''simple docstring'''




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

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

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

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

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

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

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

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

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

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







       def 						snake_case_     (       self:      List[str] ):






            '''simple docstring'''




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







       def 						snake_case_     (       self:      List[str] ):






            '''simple docstring'''




            __UpperCamelCase      =       torch_device

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

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

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




'''simple docstring'''







import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
					from transformers import AutoModelForSeqaSeqLM, AutoTokenizer


@require_torch
@require_sentencepiece
@require_tokenizers
class 						__UpperCAmelCase      (					unittest.TestCase						):


							@slow
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		AutoModelForSeqaSeqLM.from_pretrained('google/mt5-small'    ,   return_dict=lowerCAmelCase_  ).to(lowerCAmelCase_  )
								_snake_case   						=		AutoTokenizer.from_pretrained('google/mt5-small'  )

								_snake_case   						=		tokenizer('Hello there'    ,   return_tensors='pt'  ).input_ids
								_snake_case   						=		tokenizer('Hi I am'    ,   return_tensors='pt'  ).input_ids

								_snake_case   						=		model(input_ids.to(lowerCAmelCase_  )    ,   labels=labels.to(lowerCAmelCase_  )  ).loss
								_snake_case   						=		-(labels.shape[-1] * loss.item())

								_snake_case   						=		-84.9127
								self.assertTrue(abs(mtf_score - EXPECTED_SCORE  ) < 1E-4  )

