

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       __snake_case  (    lowerCAmelCase__     , unittest.TestCase       ):
		__lowerCAmelCase     :	Optional[Any]						 =  AudioLDMPipeline
		__lowerCAmelCase     :	Optional[int]						 =  TEXT_TO_AUDIO_PARAMS
		__lowerCAmelCase     :	Union[str, Any]						 =  TEXT_TO_AUDIO_BATCH_PARAMS
		__lowerCAmelCase     :	Union[str, Any]						 =  frozenset(
		    [
		        'num_inference_steps',
		        'num_waveforms_per_prompt',
		        'generator',
		        'latents',
		        'output_type',
		        'return_dict',
		        'callback',
		        'callback_steps',
		    ]       )

		def 						lowerCAmelCase__     (					self):
						torch.manual_seed(0)
						SCREAMING_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=_A						,		)
						SCREAMING_SNAKE_CASE_									=						DDIMScheduler(
						    beta_start=0.0_0_0_8_5						,		beta_end=0.0_1_2						,		beta_schedule='scaled_linear'						,		clip_sample=_A						,		set_alpha_to_one=_A						,		)
						torch.manual_seed(0)
						SCREAMING_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)
						SCREAMING_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=1000						,		projection_dim=32						,		)
						SCREAMING_SNAKE_CASE_									=						ClapTextModelWithProjection(_A)
						SCREAMING_SNAKE_CASE_									=						RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta'						,		model_max_length=77)

						SCREAMING_SNAKE_CASE_									=						SpeechTaHifiGanConfig(
						    model_in_dim=8						,		sampling_rate=16000						,		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=_A						,		)

						SCREAMING_SNAKE_CASE_									=						SpeechTaHifiGan(_A)

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

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

		def 						lowerCAmelCase__     (					self):
						SCREAMING_SNAKE_CASE_									=						'cpu'  # ensure determinism for the device-dependent torch.Generator

						SCREAMING_SNAKE_CASE_									=						self.get_dummy_components()
						SCREAMING_SNAKE_CASE_									=						AudioLDMPipeline(**_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.to(_A)
						audioldm_pipe.set_progress_bar_config(disable=_A)

						SCREAMING_SNAKE_CASE_									=						self.get_dummy_inputs(_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(**_A)
						SCREAMING_SNAKE_CASE_									=						output.audios[0]

						assert audio.ndim == 1
						assert len(_A) == 256

						SCREAMING_SNAKE_CASE_									=						audio[:10]
						SCREAMING_SNAKE_CASE_									=						np.array(
						    [-0.0_0_5_0, 0.0_0_5_0, -0.0_0_6_0, 0.0_0_3_3, -0.0_0_2_6, 0.0_0_3_3, -0.0_0_2_7, 0.0_0_3_3, -0.0_0_2_8, 0.0_0_3_3])

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

		def 						lowerCAmelCase__     (					self):
						SCREAMING_SNAKE_CASE_									=						self.get_dummy_components()
						SCREAMING_SNAKE_CASE_									=						AudioLDMPipeline(**_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.to(_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.to(_A)
						audioldm_pipe.set_progress_bar_config(disable=_A)

						SCREAMING_SNAKE_CASE_									=						self.get_dummy_inputs(_A)
						SCREAMING_SNAKE_CASE_									=						3 * [inputs['prompt']]

						# forward
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(**_A)
						SCREAMING_SNAKE_CASE_									=						output.audios[0]

						SCREAMING_SNAKE_CASE_									=						self.get_dummy_inputs(_A)
						SCREAMING_SNAKE_CASE_									=						3 * [inputs.pop('prompt')]

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

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

						SCREAMING_SNAKE_CASE_									=						prompt_embeds

						# forward
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(**_A)
						SCREAMING_SNAKE_CASE_									=						output.audios[0]

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

		def 						lowerCAmelCase__     (					self):
						SCREAMING_SNAKE_CASE_									=						self.get_dummy_components()
						SCREAMING_SNAKE_CASE_									=						AudioLDMPipeline(**_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.to(_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.to(_A)
						audioldm_pipe.set_progress_bar_config(disable=_A)

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

						# forward
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(**_A)
						SCREAMING_SNAKE_CASE_									=						output.audios[0]

						SCREAMING_SNAKE_CASE_									=						self.get_dummy_inputs(_A)
						SCREAMING_SNAKE_CASE_									=						3 * [inputs.pop('prompt')]

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

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

										embeds.append(_A)

						SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						embeds

						# forward
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(**_A)
						SCREAMING_SNAKE_CASE_									=						output.audios[0]

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

		def 						lowerCAmelCase__     (					self):
						SCREAMING_SNAKE_CASE_									=						'cpu'  # ensure determinism for the device-dependent torch.Generator
						SCREAMING_SNAKE_CASE_									=						self.get_dummy_components()
						SCREAMING_SNAKE_CASE_									=						PNDMScheduler(skip_prk_steps=_A)
						SCREAMING_SNAKE_CASE_									=						AudioLDMPipeline(**_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.to(_A)
						audioldm_pipe.set_progress_bar_config(disable=_A)

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

						assert audio.ndim == 1
						assert len(_A) == 256

						SCREAMING_SNAKE_CASE_									=						audio[:10]
						SCREAMING_SNAKE_CASE_									=						np.array(
						    [-0.0_0_5_1, 0.0_0_5_0, -0.0_0_6_0, 0.0_0_3_4, -0.0_0_2_6, 0.0_0_3_3, -0.0_0_2_7, 0.0_0_3_3, -0.0_0_2_8, 0.0_0_3_2])

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

		def 						lowerCAmelCase__     (					self):
						SCREAMING_SNAKE_CASE_									=						'cpu'  # ensure determinism for the device-dependent torch.Generator
						SCREAMING_SNAKE_CASE_									=						self.get_dummy_components()
						SCREAMING_SNAKE_CASE_									=						PNDMScheduler(skip_prk_steps=_A)
						SCREAMING_SNAKE_CASE_									=						AudioLDMPipeline(**_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.to(_A)
						audioldm_pipe.set_progress_bar_config(disable=_A)

						SCREAMING_SNAKE_CASE_									=						'A hammer hitting a wooden surface'

						# test num_waveforms_per_prompt=1 (default)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(_A						,		num_inference_steps=2).audios

						assert audios.shape == (1, 256)

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

						assert audios.shape == (batch_size, 256)

						# test num_waveforms_per_prompt for single prompt
						SCREAMING_SNAKE_CASE_									=						2
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(_A						,		num_inference_steps=2						,		num_waveforms_per_prompt=_A).audios

						assert audios.shape == (num_waveforms_per_prompt, 256)

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

						assert audios.shape == (batch_size * num_waveforms_per_prompt, 256)

		def 						lowerCAmelCase__     (					self):
						SCREAMING_SNAKE_CASE_									=						'cpu'  # ensure determinism for the device-dependent torch.Generator
						SCREAMING_SNAKE_CASE_									=						self.get_dummy_components()
						SCREAMING_SNAKE_CASE_									=						AudioLDMPipeline(**_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.to(_A)
						audioldm_pipe.set_progress_bar_config(disable=_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.vocoder.config.sampling_rate

						SCREAMING_SNAKE_CASE_									=						self.get_dummy_inputs(_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(audio_length_in_s=0.0_1_6						,		**_A)
						SCREAMING_SNAKE_CASE_									=						output.audios[0]

						assert audio.ndim == 1
						assert len(_A) / vocoder_sampling_rate == 0.0_1_6

						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(audio_length_in_s=0.0_3_2						,		**_A)
						SCREAMING_SNAKE_CASE_									=						output.audios[0]

						assert audio.ndim == 1
						assert len(_A) / vocoder_sampling_rate == 0.0_3_2

		def 						lowerCAmelCase__     (					self):
						SCREAMING_SNAKE_CASE_									=						self.get_dummy_components()
						SCREAMING_SNAKE_CASE_									=						AudioLDMPipeline(**_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.to(_A)
						audioldm_pipe.set_progress_bar_config(disable=_A)

						SCREAMING_SNAKE_CASE_									=						['hey']

						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(_A						,		num_inference_steps=1)
						SCREAMING_SNAKE_CASE_									=						output.audios.shape
						assert audio_shape == (1, 256)

						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.vocoder.config
						config.model_in_dim *= 2
						SCREAMING_SNAKE_CASE_									=						SpeechTaHifiGan(_A).to(_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(_A						,		num_inference_steps=1)
						SCREAMING_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, 256)

		def 						lowerCAmelCase__     (					self):
						self._test_attention_slicing_forward_pass(test_mean_pixel_difference=_A)

		def 						lowerCAmelCase__     (					self):
						self._test_inference_batch_single_identical(test_mean_pixel_difference=_A)


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






@slow
class       __snake_case  (    unittest.TestCase       ):

		def 						lowerCAmelCase__     (					self):
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()

		def 						lowerCAmelCase__     (					self						,		_A						,		_A="cpu"						,		_A=torch.floataa						,		_A=0):
						SCREAMING_SNAKE_CASE_									=						torch.Generator(device=_A).manual_seed(_A)
						SCREAMING_SNAKE_CASE_									=						np.random.RandomState(_A).standard_normal((1, 8, 128, 16))
						SCREAMING_SNAKE_CASE_									=						torch.from_numpy(_A).to(device=_A						,		dtype=_A)
						SCREAMING_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):
						SCREAMING_SNAKE_CASE_									=						AudioLDMPipeline.from_pretrained('cvssp/audioldm')
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.to(_A)
						audioldm_pipe.set_progress_bar_config(disable=_A)

						SCREAMING_SNAKE_CASE_									=						self.get_inputs(_A)
						SCREAMING_SNAKE_CASE_									=						25
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(**_A).audios[0]

						assert audio.ndim == 1
						assert len(_A) == 81920

						SCREAMING_SNAKE_CASE_									=						audio[77230:77240]
						SCREAMING_SNAKE_CASE_									=						np.array(
						    [-0.4_8_8_4, -0.4_6_0_7, 0.0_0_2_3, 0.5_0_0_7, 0.5_8_9_6, 0.5_1_5_1, 0.3_8_1_3, -0.0_2_0_8, -0.3_6_8_7, -0.4_3_1_5])
						SCREAMING_SNAKE_CASE_									=						np.abs(expected_slice - audio_slice).max()
						assert max_diff < 1E-2


		def 						lowerCAmelCase__     (					self):
						SCREAMING_SNAKE_CASE_									=						AudioLDMPipeline.from_pretrained('cvssp/audioldm')
						SCREAMING_SNAKE_CASE_									=						LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe.to(_A)
						audioldm_pipe.set_progress_bar_config(disable=_A)

						SCREAMING_SNAKE_CASE_									=						self.get_inputs(_A)
						SCREAMING_SNAKE_CASE_									=						audioldm_pipe(**_A).audios[0]

						assert audio.ndim == 1
						assert len(_A) == 81920

						SCREAMING_SNAKE_CASE_									=						audio[27780:27790]
						SCREAMING_SNAKE_CASE_									=						np.array([-0.2_1_3_1, -0.0_8_7_3, -0.0_1_2_4, -0.0_1_8_9, 0.0_5_6_9, 0.1_3_7_3, 0.1_8_8_3, 0.2_8_8_6, 0.3_2_9_7, 0.2_2_1_2])
						SCREAMING_SNAKE_CASE_									=						np.abs(expected_slice - audio_slice).max()
						assert max_diff < 3E-2





from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__  :      str			=    logging.get_logger(__name__)

UpperCamelCase__  :      Optional[int]			=    {
    "facebook/dpr-ctx_encoder-single-nq-base": (
        "https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/config.json"
    ),
    "facebook/dpr-question_encoder-single-nq-base": (
        "https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/config.json"
    ),
    "facebook/dpr-reader-single-nq-base": (
        "https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/config.json"
    ),
    "facebook/dpr-ctx_encoder-multiset-base": (
        "https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/config.json"
    ),
    "facebook/dpr-question_encoder-multiset-base": (
        "https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/config.json"
    ),
    "facebook/dpr-reader-multiset-base": (
        "https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/config.json"
    ),
}



class       __snake_case  (    lowerCAmelCase__       ):
  __lowerCAmelCase     :	Optional[int]						 =  'dpr'

  def __init__(					self						,		_A=30522						,		_A=768						,		_A=12						,		_A=12						,		_A=3072						,		_A="gelu"						,		_A=0.1						,		_A=0.1						,		_A=512						,		_A=2						,		_A=0.0_2						,		_A=1E-12						,		_A=0						,		_A="absolute"						,		_A = 0						,		**_A						,		):
      super().__init__(pad_token_id=_A						,		**_A)

      SCREAMING_SNAKE_CASE_									=						vocab_size
      SCREAMING_SNAKE_CASE_									=						hidden_size
      SCREAMING_SNAKE_CASE_									=						num_hidden_layers
      SCREAMING_SNAKE_CASE_									=						num_attention_heads
      SCREAMING_SNAKE_CASE_									=						hidden_act
      SCREAMING_SNAKE_CASE_									=						intermediate_size
      SCREAMING_SNAKE_CASE_									=						hidden_dropout_prob
      SCREAMING_SNAKE_CASE_									=						attention_probs_dropout_prob
      SCREAMING_SNAKE_CASE_									=						max_position_embeddings
      SCREAMING_SNAKE_CASE_									=						type_vocab_size
      SCREAMING_SNAKE_CASE_									=						initializer_range
      SCREAMING_SNAKE_CASE_									=						layer_norm_eps
      SCREAMING_SNAKE_CASE_									=						projection_dim
      SCREAMING_SNAKE_CASE_									=						position_embedding_type


