





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 _snake_case       (		self							):
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()


				@property
				def _snake_case       (		self							):
								torch.manual_seed(0							)
								lowerCamelCase	 =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							):
								torch.manual_seed(0							)
								lowerCamelCase	 =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							):
								torch.manual_seed(0							)
								lowerCamelCase	 =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""")   , )
								lowerCamelCase	 =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							):
								lowerCamelCase	 ="""cpu"""  # ensure determinism for the device-dependent torch.Generator
								lowerCamelCase	 =Mel(
								    x_res=self.dummy_unet.config.sample_size[1]   , y_res=self.dummy_unet.config.sample_size[0]   , )

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

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

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

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

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

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

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

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

								np.random.seed(0							)
								lowerCamelCase	 =torch.rand((1, 1, 10)							)
								lowerCamelCase	 =pipe(generator=UpperCAmelCase_   , encoding=UpperCAmelCase_							)
								lowerCamelCase	 =output.images[0]
								lowerCamelCase	 =np.frombuffer(image.tobytes()   , dtype="""uint8"""							)[:10]
								lowerCamelCase	 =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 _snake_case       (		self							):
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()


				def _snake_case       (		self							):
								lowerCamelCase	 =torch_device

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

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

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












from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase__						:			str									={'''configuration_fnet''': ['''FNET_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''FNetConfig''']}

try:
							if not is_sentencepiece_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__						:			List[Any]									=['''FNetTokenizer''']

try:
							if not is_tokenizers_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__						:			Any									=['''FNetTokenizerFast''']

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__						:			Dict									=[
							    '''FNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
							    '''FNetForMaskedLM''',
							    '''FNetForMultipleChoice''',
							    '''FNetForNextSentencePrediction''',
							    '''FNetForPreTraining''',
							    '''FNetForQuestionAnswering''',
							    '''FNetForSequenceClassification''',
							    '''FNetForTokenClassification''',
							    '''FNetLayer''',
							    '''FNetModel''',
							    '''FNetPreTrainedModel''',
							]


if TYPE_CHECKING:
							from .configuration_fnet import FNET_PRETRAINED_CONFIG_ARCHIVE_MAP, FNetConfig

							try:
														if not is_sentencepiece_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tokenization_fnet import FNetTokenizer

							try:
														if not is_tokenizers_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tokenization_fnet_fast import FNetTokenizerFast

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_fnet import (
														    FNET_PRETRAINED_MODEL_ARCHIVE_LIST,
														    FNetForMaskedLM,
														    FNetForMultipleChoice,
														    FNetForNextSentencePrediction,
														    FNetForPreTraining,
														    FNetForQuestionAnswering,
														    FNetForSequenceClassification,
														    FNetForTokenClassification,
														    FNetLayer,
														    FNetModel,
														    FNetPreTrainedModel,
														)


else:
							import sys

							UpperCAmelCase__						:			Optional[int]									=_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





