
'''simple docstring'''







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


							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






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


							@property
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								torch.manual_seed(0  )
								_snake_case   						=		UNetaDModel(
								    sample_size=(32, 64)    ,   in_channels=1    ,   out_channels=1    ,   layers_per_block=2    ,   block_out_channels=(1_28, 1_28)    ,   down_block_types=('AttnDownBlock2D', 'DownBlock2D')    ,   up_block_types=('UpBlock2D', 'AttnUpBlock2D')    ,   )
								return model


							@property
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								torch.manual_seed(0  )
								_snake_case   						=		UNetaDConditionModel(
								    sample_size=(64, 32)    ,   in_channels=1    ,   out_channels=1    ,   layers_per_block=2    ,   block_out_channels=(1_28, 1_28)    ,   down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')    ,   up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')    ,   cross_attention_dim=10    ,   )
								return model


							@property
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								torch.manual_seed(0  )
								_snake_case   						=		AutoencoderKL(
								    sample_size=(1_28, 64)    ,   in_channels=1    ,   out_channels=1    ,   latent_channels=1    ,   layers_per_block=2    ,   block_out_channels=(1_28, 1_28)    ,   down_block_types=('DownEncoderBlock2D', 'DownEncoderBlock2D')    ,   up_block_types=('UpDecoderBlock2D', 'UpDecoderBlock2D')    ,   )
								_snake_case   						=		UNetaDModel(
								    sample_size=(64, 32)    ,   in_channels=1    ,   out_channels=1    ,   layers_per_block=2    ,   block_out_channels=(1_28, 1_28)    ,   down_block_types=('AttnDownBlock2D', 'DownBlock2D')    ,   up_block_types=('UpBlock2D', 'AttnUpBlock2D')    ,   )
								return vqvae, unet




							@slow
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






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

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

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

								_snake_case   						=		torch.Generator(device=lowerCAmelCase_  ).manual_seed(42  )
								_snake_case   						=		pipe(generator=lowerCAmelCase_    ,   steps=4    ,   return_dict=lowerCAmelCase_  )
								_snake_case   						=		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]
								)
								_snake_case   						=		np.frombuffer(image.tobytes()    ,   dtype='uint8'  )[:10]
								_snake_case   						=		np.frombuffer(image_from_tuple.tobytes()    ,   dtype='uint8'  )[:10]
								_snake_case   						=		np.array([69, 2_55, 2_55, 2_55, 0, 0, 77, 1_81, 12, 1_27]  )

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

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

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

								np.random.seed(0  )
								_snake_case   						=		np.random.uniform(-1    ,   1    ,   ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)  )
								_snake_case   						=		torch.Generator(device=lowerCAmelCase_  ).manual_seed(42  )
								_snake_case   						=		pipe(raw_audio=lowerCAmelCase_    ,   generator=lowerCAmelCase_    ,   start_step=5    ,   steps=10  )
								_snake_case   						=		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]
								)
								_snake_case   						=		np.frombuffer(image.tobytes()    ,   dtype='uint8'  )[:10]
								_snake_case   						=		np.array([1_20, 1_17, 1_10, 1_09, 1_38, 1_67, 1_38, 1_48, 1_32, 1_21]  )

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

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

								np.random.seed(0  )
								_snake_case   						=		torch.rand((1, 1, 10)  )
								_snake_case   						=		pipe(generator=lowerCAmelCase_    ,   encoding=lowerCAmelCase_  )
								_snake_case   						=		output.images[0]
								_snake_case   						=		np.frombuffer(image.tobytes()    ,   dtype='uint8'  )[:10]
								_snake_case   						=		np.array([1_07, 1_03, 1_20, 1_27, 1_42, 1_22, 1_13, 1_22, 97, 1_11]  )

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


@slow
@require_torch_gpu
class 						__UpperCAmelCase      (					unittest.TestCase						):


							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






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




							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								_snake_case   						=		torch_device

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

								_snake_case   						=		torch.Generator(device=lowerCAmelCase_  ).manual_seed(42  )
								_snake_case   						=		pipe(generator=lowerCAmelCase_  )
								_snake_case   						=		output.audios[0]
								_snake_case   						=		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]
								_snake_case   						=		np.frombuffer(image.tobytes()    ,   dtype='uint8'  )[:10]
								_snake_case   						=		np.array([1_51, 1_67, 1_54, 1_44, 1_22, 1_34, 1_21, 1_05, 70, 26]  )

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



'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


lowercase    :							Tuple      =					{
    "configuration_xlm": ["XLM_PRETRAINED_CONFIG_ARCHIVE_MAP", "XLMConfig", "XLMOnnxConfig"],
    "tokenization_xlm": ["XLMTokenizer"],
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase    :							int      =					[
					    "XLM_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "XLMForMultipleChoice",
					    "XLMForQuestionAnswering",
					    "XLMForQuestionAnsweringSimple",
					    "XLMForSequenceClassification",
					    "XLMForTokenClassification",
					    "XLMModel",
					    "XLMPreTrainedModel",
					    "XLMWithLMHeadModel",
					]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase    :							Tuple      =					[
					    "TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "TFXLMForMultipleChoice",
					    "TFXLMForQuestionAnsweringSimple",
					    "TFXLMForSequenceClassification",
					    "TFXLMForTokenClassification",
					    "TFXLMMainLayer",
					    "TFXLMModel",
					    "TFXLMPreTrainedModel",
					    "TFXLMWithLMHeadModel",
					]


if TYPE_CHECKING:
					from .configuration_xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig, XLMOnnxConfig
					from .tokenization_xlm import XLMTokenizer

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_xlm import (
										    XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
										    XLMForMultipleChoice,
										    XLMForQuestionAnswering,
										    XLMForQuestionAnsweringSimple,
										    XLMForSequenceClassification,
										    XLMForTokenClassification,
										    XLMModel,
										    XLMPreTrainedModel,
										    XLMWithLMHeadModel,
										)

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_xlm import (
										    TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
										    TFXLMForMultipleChoice,
										    TFXLMForQuestionAnsweringSimple,
										    TFXLMForSequenceClassification,
										    TFXLMForTokenClassification,
										    TFXLMMainLayer,
										    TFXLMModel,
										    TFXLMPreTrainedModel,
										    TFXLMWithLMHeadModel,
										)

else:
					import sys

					lowercase    :							Union[str, Any]      =					_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

