

from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


A									=  {
    'configuration_blenderbot_small': [
        'BLENDERBOT_SMALL_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'BlenderbotSmallConfig',
        'BlenderbotSmallOnnxConfig',
    ],
    'tokenization_blenderbot_small': ['BlenderbotSmallTokenizer'],
}

try:
		if not is_tokenizers_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		A									=  ['BlenderbotSmallTokenizerFast']

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		A									=  [
		    'BLENDERBOT_SMALL_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'BlenderbotSmallForCausalLM',
		    'BlenderbotSmallForConditionalGeneration',
		    'BlenderbotSmallModel',
		    'BlenderbotSmallPreTrainedModel',
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		A									=  [
		    'TFBlenderbotSmallForConditionalGeneration',
		    'TFBlenderbotSmallModel',
		    'TFBlenderbotSmallPreTrainedModel',
		]

try:
		if not is_flax_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		A									=  [
		    'FlaxBlenderbotSmallForConditionalGeneration',
		    'FlaxBlenderbotSmallModel',
		    'FlaxBlenderbotSmallPreTrainedModel',
		]

if TYPE_CHECKING:
		from .configuration_blenderbot_small import (
		    BLENDERBOT_SMALL_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    BlenderbotSmallConfig,
		    BlenderbotSmallOnnxConfig,
		)
		from .tokenization_blenderbot_small import BlenderbotSmallTokenizer

		try:
				if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tokenization_blenderbot_small_fast import BlenderbotSmallTokenizerFast

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_blenderbot_small import (
				    BLENDERBOT_SMALL_PRETRAINED_MODEL_ARCHIVE_LIST,
				    BlenderbotSmallForCausalLM,
				    BlenderbotSmallForConditionalGeneration,
				    BlenderbotSmallModel,
				    BlenderbotSmallPreTrainedModel,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_blenderbot_small import (
				    TFBlenderbotSmallForConditionalGeneration,
				    TFBlenderbotSmallModel,
				    TFBlenderbotSmallPreTrainedModel,
				)

		try:
				if not is_flax_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_flax_blenderbot_small import (
				    FlaxBlenderbotSmallForConditionalGeneration,
				    FlaxBlenderbotSmallModel,
				    FlaxBlenderbotSmallPreTrainedModel,
				)

else:
		import sys

		A									=  _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



import gc
import unittest

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    ClapTextConfig,
    ClapTextModelWithProjection,
    RobertaTokenizer,
    SpeechTaHifiGan,
    SpeechTaHifiGanConfig,
)

from diffusers import (
    AudioLDMPipeline,
    AutoencoderKL,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class SCREAMING_SNAKE_CASE	(		__snake_case				,    unittest.TestCase				):






			"""simple docstring"""




			__A    	=	AudioLDMPipeline
			__A    	=	TEXT_TO_AUDIO_PARAMS
			__A    	=	TEXT_TO_AUDIO_BATCH_PARAMS
			__A    	=	frozenset(
			    [
			        """num_inference_steps""",
			        """num_waveforms_per_prompt""",
			        """generator""",
			        """latents""",
			        """output_type""",
			        """return_dict""",
			        """callback""",
			        """callback_steps""",
			    ]				)


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				torch.manual_seed(0				)
				snake_case_   		=			UNetaDConditionModel(
				    block_out_channels=(32, 64)			,					layers_per_block=2			,					sample_size=32			,					in_channels=4			,					out_channels=4			,					down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')			,					up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')			,					cross_attention_dim=(32, 64)			,					class_embed_type='simple_projection'			,					projection_class_embeddings_input_dim=32			,					class_embeddings_concat=__UpperCamelCase			,					)
				snake_case_   		=			DDIMScheduler(
				    beta_start=0.0_0085			,					beta_end=0.012			,					beta_schedule='scaled_linear'			,					clip_sample=__UpperCamelCase			,					set_alpha_to_one=__UpperCamelCase			,					)
				torch.manual_seed(0				)
				snake_case_   		=			AutoencoderKL(
				    block_out_channels=[32, 64]			,					in_channels=1			,					out_channels=1			,					down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']			,					up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']			,					latent_channels=4			,					)
				torch.manual_seed(0				)
				snake_case_   		=			ClapTextConfig(
				    bos_token_id=0			,					eos_token_id=2			,					hidden_size=32			,					intermediate_size=37			,					layer_norm_eps=1E-05			,					num_attention_heads=4			,					num_hidden_layers=5			,					pad_token_id=1			,					vocab_size=10_00			,					projection_dim=32			,					)
				snake_case_   		=			ClapTextModelWithProjection(__UpperCamelCase				)
				snake_case_   		=			RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta'			,					model_max_length=77				)

				snake_case_   		=			SpeechTaHifiGanConfig(
				    model_in_dim=8			,					sampling_rate=1_60_00			,					upsample_initial_channel=16			,					upsample_rates=[2, 2]			,					upsample_kernel_sizes=[4, 4]			,					resblock_kernel_sizes=[3, 7]			,					resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]]			,					normalize_before=__UpperCamelCase			,					)

				snake_case_   		=			SpeechTaHifiGan(__UpperCamelCase				)

				snake_case_   		=			{
				    'unet': unet,
				    'scheduler': scheduler,
				    'vae': vae,
				    'text_encoder': text_encoder,
				    'tokenizer': tokenizer,
				    'vocoder': vocoder,
				}
				return components


			def       __lowerCAmelCase			(     self			,					__UpperCamelCase			,					__UpperCamelCase=0				):

				"""simple docstring"""


				if str(__UpperCamelCase				).startswith('mps'				):
					snake_case_   		=			torch.manual_seed(__UpperCamelCase				)
				else:
					snake_case_   		=			torch.Generator(device=__UpperCamelCase				).manual_seed(__UpperCamelCase				)
				snake_case_   		=			{
				    'prompt': 'A hammer hitting a wooden surface',
				    'generator': generator,
				    'num_inference_steps': 2,
				    'guidance_scale': 6.0,
				}
				return inputs


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			'cpu'  # ensure determinism for the device-dependent torch.Generator

				snake_case_   		=			self.get_dummy_components()
				snake_case_   		=			AudioLDMPipeline(**__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe.to(__UpperCamelCase				)
				audioldm_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

				snake_case_   		=			self.get_dummy_inputs(__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe(**__UpperCamelCase				)
				snake_case_   		=			output.audios[0]

				assert audio.ndim == 1
				assert len(__UpperCamelCase				) == 2_56

				snake_case_   		=			audio[:10]
				snake_case_   		=			np.array(
				    [-0.0050, 0.0050, -0.0060, 0.0033, -0.0026, 0.0033, -0.0027, 0.0033, -0.0028, 0.0033]				)

				assert np.abs(audio_slice - expected_slice				).max() < 1E-2


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self.get_dummy_components()
				snake_case_   		=			AudioLDMPipeline(**__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe.to(__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe.to(__UpperCamelCase				)
				audioldm_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

				snake_case_   		=			self.get_dummy_inputs(__UpperCamelCase				)
				snake_case_   		=			3 * [inputs['prompt']]

				# forward
				snake_case_   		=			audioldm_pipe(**__UpperCamelCase				)
				snake_case_   		=			output.audios[0]

				snake_case_   		=			self.get_dummy_inputs(__UpperCamelCase				)
				snake_case_   		=			3 * [inputs.pop('prompt'				)]

				snake_case_   		=			audioldm_pipe.tokenizer(
				    __UpperCamelCase			,					padding='max_length'			,					max_length=audioldm_pipe.tokenizer.model_max_length			,					truncation=__UpperCamelCase			,					return_tensors='pt'			,					)
				snake_case_   		=			text_inputs['input_ids'].to(__UpperCamelCase				)

				snake_case_   		=			audioldm_pipe.text_encoder(
				    __UpperCamelCase			,					)
				snake_case_   		=			prompt_embeds.text_embeds
				# additional L_2 normalization over each hidden-state
				snake_case_   		=			F.normalize(__UpperCamelCase			,					dim=-1				)

				snake_case_   		=			prompt_embeds

				# forward
				snake_case_   		=			audioldm_pipe(**__UpperCamelCase				)
				snake_case_   		=			output.audios[0]

				assert np.abs(audio_a - audio_a				).max() < 1E-2


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self.get_dummy_components()
				snake_case_   		=			AudioLDMPipeline(**__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe.to(__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe.to(__UpperCamelCase				)
				audioldm_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

				snake_case_   		=			self.get_dummy_inputs(__UpperCamelCase				)
				snake_case_   		=			3 * ['this is a negative prompt']
				snake_case_   		=			negative_prompt
				snake_case_   		=			3 * [inputs['prompt']]

				# forward
				snake_case_   		=			audioldm_pipe(**__UpperCamelCase				)
				snake_case_   		=			output.audios[0]

				snake_case_   		=			self.get_dummy_inputs(__UpperCamelCase				)
				snake_case_   		=			3 * [inputs.pop('prompt'				)]

				snake_case_   		=			[]
				for p in [prompt, negative_prompt]:
					snake_case_   		=			audioldm_pipe.tokenizer(
					    __UpperCamelCase			,					padding='max_length'			,					max_length=audioldm_pipe.tokenizer.model_max_length			,					truncation=__UpperCamelCase			,					return_tensors='pt'			,					)
					snake_case_   		=			text_inputs['input_ids'].to(__UpperCamelCase				)

					snake_case_   		=			audioldm_pipe.text_encoder(
					    __UpperCamelCase			,					)
					snake_case_   		=			text_embeds.text_embeds
					# additional L_2 normalization over each hidden-state
					snake_case_   		=			F.normalize(__UpperCamelCase			,					dim=-1				)

					embeds.append(__UpperCamelCase				)

				snake_case_		,							snake_case_   		=			embeds

				# forward
				snake_case_   		=			audioldm_pipe(**__UpperCamelCase				)
				snake_case_   		=			output.audios[0]

				assert np.abs(audio_a - audio_a				).max() < 1E-2


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			'cpu'  # ensure determinism for the device-dependent torch.Generator
				snake_case_   		=			self.get_dummy_components()
				snake_case_   		=			PNDMScheduler(skip_prk_steps=__UpperCamelCase				)
				snake_case_   		=			AudioLDMPipeline(**__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe.to(__UpperCamelCase				)
				audioldm_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

				snake_case_   		=			self.get_dummy_inputs(__UpperCamelCase				)
				snake_case_   		=			'egg cracking'
				snake_case_   		=			audioldm_pipe(**__UpperCamelCase			,					negative_prompt=__UpperCamelCase				)
				snake_case_   		=			output.audios[0]

				assert audio.ndim == 1
				assert len(__UpperCamelCase				) == 2_56

				snake_case_   		=			audio[:10]
				snake_case_   		=			np.array(
				    [-0.0051, 0.0050, -0.0060, 0.0034, -0.0026, 0.0033, -0.0027, 0.0033, -0.0028, 0.0032]				)

				assert np.abs(audio_slice - expected_slice				).max() < 1E-2


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			'cpu'  # ensure determinism for the device-dependent torch.Generator
				snake_case_   		=			self.get_dummy_components()
				snake_case_   		=			PNDMScheduler(skip_prk_steps=__UpperCamelCase				)
				snake_case_   		=			AudioLDMPipeline(**__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe.to(__UpperCamelCase				)
				audioldm_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

				snake_case_   		=			'A hammer hitting a wooden surface'

				# test num_waveforms_per_prompt=1 (default)
				snake_case_   		=			audioldm_pipe(__UpperCamelCase			,					num_inference_steps=2				).audios

				assert audios.shape == (1, 2_56)

				# test num_waveforms_per_prompt=1 (default) for batch of prompts
				snake_case_   		=			2
				snake_case_   		=			audioldm_pipe([prompt] * batch_size			,					num_inference_steps=2				).audios

				assert audios.shape == (batch_size, 2_56)

				# test num_waveforms_per_prompt for single prompt
				snake_case_   		=			2
				snake_case_   		=			audioldm_pipe(__UpperCamelCase			,					num_inference_steps=2			,					num_waveforms_per_prompt=__UpperCamelCase				).audios

				assert audios.shape == (num_waveforms_per_prompt, 2_56)

				# test num_waveforms_per_prompt for batch of prompts
				snake_case_   		=			2
				snake_case_   		=			audioldm_pipe(
				    [prompt] * batch_size			,					num_inference_steps=2			,					num_waveforms_per_prompt=__UpperCamelCase				).audios

				assert audios.shape == (batch_size * num_waveforms_per_prompt, 2_56)


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			'cpu'  # ensure determinism for the device-dependent torch.Generator
				snake_case_   		=			self.get_dummy_components()
				snake_case_   		=			AudioLDMPipeline(**__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe.to(__UpperCamelCase				)
				audioldm_pipe.set_progress_bar_config(disable=__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe.vocoder.config.sampling_rate

				snake_case_   		=			self.get_dummy_inputs(__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe(audio_length_in_s=0.016			,					**__UpperCamelCase				)
				snake_case_   		=			output.audios[0]

				assert audio.ndim == 1
				assert len(__UpperCamelCase				) / vocoder_sampling_rate == 0.016

				snake_case_   		=			audioldm_pipe(audio_length_in_s=0.032			,					**__UpperCamelCase				)
				snake_case_   		=			output.audios[0]

				assert audio.ndim == 1
				assert len(__UpperCamelCase				) / vocoder_sampling_rate == 0.032


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			self.get_dummy_components()
				snake_case_   		=			AudioLDMPipeline(**__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe.to(__UpperCamelCase				)
				audioldm_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

				snake_case_   		=			['hey']

				snake_case_   		=			audioldm_pipe(__UpperCamelCase			,					num_inference_steps=1				)
				snake_case_   		=			output.audios.shape
				assert audio_shape == (1, 2_56)

				snake_case_   		=			audioldm_pipe.vocoder.config
				config.model_in_dim *= 2
				snake_case_   		=			SpeechTaHifiGan(__UpperCamelCase				).to(__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe(__UpperCamelCase			,					num_inference_steps=1				)
				snake_case_   		=			output.audios.shape
				# waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
				assert audio_shape == (1, 2_56)


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__UpperCamelCase				)


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				self._test_inference_batch_single_identical(test_mean_pixel_difference=__UpperCamelCase				)







			@unittest.skipIf(
			    torch_device != 'cuda' or not is_xformers_available()			,					reason='XFormers attention is only available with CUDA and `xformers` installed'			,					)
			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__UpperCamelCase				)







@slow
class SCREAMING_SNAKE_CASE	(		unittest.TestCase				):






			"""simple docstring"""




			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


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


			def       __lowerCAmelCase			(     self			,					__UpperCamelCase			,					__UpperCamelCase="cpu"			,					__UpperCamelCase=torch.floataa			,					__UpperCamelCase=0				):

				"""simple docstring"""


				snake_case_   		=			torch.Generator(device=__UpperCamelCase				).manual_seed(__UpperCamelCase				)
				snake_case_   		=			np.random.RandomState(__UpperCamelCase				).standard_normal((1, 8, 1_28, 16)				)
				snake_case_   		=			torch.from_numpy(__UpperCamelCase				).to(device=__UpperCamelCase			,					dtype=__UpperCamelCase				)
				snake_case_   		=			{
				    'prompt': 'A hammer hitting a wooden surface',
				    'latents': latents,
				    'generator': generator,
				    'num_inference_steps': 3,
				    'guidance_scale': 2.5,
				}
				return inputs


			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			AudioLDMPipeline.from_pretrained('cvssp/audioldm'				)
				snake_case_   		=			audioldm_pipe.to(__UpperCamelCase				)
				audioldm_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

				snake_case_   		=			self.get_inputs(__UpperCamelCase				)
				snake_case_   		=			25
				snake_case_   		=			audioldm_pipe(**__UpperCamelCase				).audios[0]

				assert audio.ndim == 1
				assert len(__UpperCamelCase				) == 8_19_20

				snake_case_   		=			audio[7_72_30:7_72_40]
				snake_case_   		=			np.array(
				    [-0.4884, -0.4607, 0.0023, 0.5007, 0.5896, 0.5151, 0.3813, -0.0208, -0.3687, -0.4315]				)
				snake_case_   		=			np.abs(expected_slice - audio_slice				).max()
				assert max_diff < 1E-2







			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				snake_case_   		=			AudioLDMPipeline.from_pretrained('cvssp/audioldm'				)
				snake_case_   		=			LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config				)
				snake_case_   		=			audioldm_pipe.to(__UpperCamelCase				)
				audioldm_pipe.set_progress_bar_config(disable=__UpperCamelCase				)

				snake_case_   		=			self.get_inputs(__UpperCamelCase				)
				snake_case_   		=			audioldm_pipe(**__UpperCamelCase				).audios[0]

				assert audio.ndim == 1
				assert len(__UpperCamelCase				) == 8_19_20

				snake_case_   		=			audio[2_77_80:2_77_90]
				snake_case_   		=			np.array([-0.2131, -0.0873, -0.0124, -0.0189, 0.0569, 0.1373, 0.1883, 0.2886, 0.3297, 0.2212]				)
				snake_case_   		=			np.abs(expected_slice - audio_slice				).max()
				assert max_diff < 3E-2
