






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):





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

	@property
	def        UpperCamelCase    (							self    ):
							"""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        UpperCamelCase    (							self    ):
							"""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        UpperCamelCase    (							self    ):
							"""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        UpperCamelCase    (							self    ):
							"""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=UpperCAmelCase       ,    unet=self.dummy_unet       ,    mel=UpperCAmelCase       ,    scheduler=UpperCAmelCase    )
							_UpperCAmelCase									=			pipe.to(UpperCAmelCase    )
							pipe.set_progress_bar_config(disable=UpperCAmelCase    )

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

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

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

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





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			torch_device

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

							_UpperCAmelCase									=			torch.Generator(device=UpperCAmelCase    ).manual_seed(42    )
							_UpperCAmelCase									=			pipe(generator=UpperCAmelCase    )
							_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 re

from filelock import FileLock


try:
  import nltk

  UpperCamelCase__    					=      True
except (ImportError, ModuleNotFoundError):
  UpperCamelCase__    					=      False

if NLTK_AVAILABLE:
  with FileLock('''.lock''') as lock:
    nltk.download('''punkt''', quiet=True)





def 					a__       (					lowerCAmelCase__						)    ->	str:
    re.sub('''<n>'''  ,					''''''  ,					lowerCAmelCase__						)  # remove pegasus newline char
    assert NLTK_AVAILABLE, "nltk must be installed to separate newlines between sentences. (pip install nltk)"
    return "\n".join(nltk.sent_tokenize(lowerCAmelCase__						)						)
