






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






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												torch.manual_seed(0   )
												lowerCAmelCase__    						=						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     ,						)
												lowerCAmelCase__    						=						DDIMScheduler(
												    beta_start=0.0_00_85     ,						beta_end=0.0_12     ,						beta_schedule='scaled_linear'     ,						clip_sample=_UpperCamelCase     ,						set_alpha_to_one=_UpperCamelCase     ,						)
												torch.manual_seed(0   )
												lowerCAmelCase__    						=						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   )
												lowerCAmelCase__    						=						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     ,						)
												lowerCAmelCase__    						=						ClapTextModelWithProjection(_UpperCamelCase   )
												lowerCAmelCase__    						=						RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta'     ,						model_max_length=77   )

												lowerCAmelCase__    						=						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     ,						)

												lowerCAmelCase__    						=						SpeechTaHifiGan(_UpperCamelCase   )

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






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase=0   ):






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






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						'cpu'  # ensure determinism for the device-dependent torch.Generator

												lowerCAmelCase__    						=						self.get_dummy_components()
												lowerCAmelCase__    						=						AudioLDMPipeline(**_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe.to(_UpperCamelCase   )
												audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

												lowerCAmelCase__    						=						self.get_dummy_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe(**_UpperCamelCase   )
												lowerCAmelCase__    						=						output.audios[0]

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

												lowerCAmelCase__    						=						audio[:10]
												lowerCAmelCase__    						=						np.array(
												    [-0.00_50, 0.00_50, -0.00_60, 0.00_33, -0.00_26, 0.00_33, -0.00_27, 0.00_33, -0.00_28, 0.00_33]   )

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






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.get_dummy_components()
												lowerCAmelCase__    						=						AudioLDMPipeline(**_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe.to(_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe.to(_UpperCamelCase   )
												audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

												lowerCAmelCase__    						=						self.get_dummy_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						3 * [inputs['prompt']]

												# forward
												lowerCAmelCase__    						=						audioldm_pipe(**_UpperCamelCase   )
												lowerCAmelCase__    						=						output.audios[0]

												lowerCAmelCase__    						=						self.get_dummy_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						3 * [inputs.pop('prompt'   )]

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

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

												lowerCAmelCase__    						=						prompt_embeds

												# forward
												lowerCAmelCase__    						=						audioldm_pipe(**_UpperCamelCase   )
												lowerCAmelCase__    						=						output.audios[0]

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






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.get_dummy_components()
												lowerCAmelCase__    						=						AudioLDMPipeline(**_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe.to(_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe.to(_UpperCamelCase   )
												audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

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

												# forward
												lowerCAmelCase__    						=						audioldm_pipe(**_UpperCamelCase   )
												lowerCAmelCase__    						=						output.audios[0]

												lowerCAmelCase__    						=						self.get_dummy_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						3 * [inputs.pop('prompt'   )]

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

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

																		embeds.append(_UpperCamelCase   )

												lowerCAmelCase__						,       lowerCAmelCase__    						=						embeds

												# forward
												lowerCAmelCase__    						=						audioldm_pipe(**_UpperCamelCase   )
												lowerCAmelCase__    						=						output.audios[0]

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






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						'cpu'  # ensure determinism for the device-dependent torch.Generator
												lowerCAmelCase__    						=						self.get_dummy_components()
												lowerCAmelCase__    						=						PNDMScheduler(skip_prk_steps=_UpperCamelCase   )
												lowerCAmelCase__    						=						AudioLDMPipeline(**_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe.to(_UpperCamelCase   )
												audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

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

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

												lowerCAmelCase__    						=						audio[:10]
												lowerCAmelCase__    						=						np.array(
												    [-0.00_51, 0.00_50, -0.00_60, 0.00_34, -0.00_26, 0.00_33, -0.00_27, 0.00_33, -0.00_28, 0.00_32]   )

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






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						'cpu'  # ensure determinism for the device-dependent torch.Generator
												lowerCAmelCase__    						=						self.get_dummy_components()
												lowerCAmelCase__    						=						PNDMScheduler(skip_prk_steps=_UpperCamelCase   )
												lowerCAmelCase__    						=						AudioLDMPipeline(**_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe.to(_UpperCamelCase   )
												audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

												lowerCAmelCase__    						=						'A hammer hitting a wooden surface'

												# test num_waveforms_per_prompt=1 (default)
												lowerCAmelCase__    						=						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
												lowerCAmelCase__    						=						2
												lowerCAmelCase__    						=						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
												lowerCAmelCase__    						=						2
												lowerCAmelCase__    						=						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
												lowerCAmelCase__    						=						2
												lowerCAmelCase__    						=						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       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						'cpu'  # ensure determinism for the device-dependent torch.Generator
												lowerCAmelCase__    						=						self.get_dummy_components()
												lowerCAmelCase__    						=						AudioLDMPipeline(**_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe.to(_UpperCamelCase   )
												audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe.vocoder.config.sampling_rate

												lowerCAmelCase__    						=						self.get_dummy_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe(audio_length_in_s=0.0_16     ,						**_UpperCamelCase   )
												lowerCAmelCase__    						=						output.audios[0]

												assert audio.ndim == 1
												assert len(_UpperCamelCase   ) / vocoder_sampling_rate == 0.0_16

												lowerCAmelCase__    						=						audioldm_pipe(audio_length_in_s=0.0_32     ,						**_UpperCamelCase   )
												lowerCAmelCase__    						=						output.audios[0]

												assert audio.ndim == 1
												assert len(_UpperCamelCase   ) / vocoder_sampling_rate == 0.0_32






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.get_dummy_components()
												lowerCAmelCase__    						=						AudioLDMPipeline(**_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe.to(_UpperCamelCase   )
												audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

												lowerCAmelCase__    						=						['hey']

												lowerCAmelCase__    						=						audioldm_pipe(_UpperCamelCase     ,						num_inference_steps=1   )
												lowerCAmelCase__    						=						output.audios.shape
												assert audio_shape == (1, 2_56)

												lowerCAmelCase__    						=						audioldm_pipe.vocoder.config
												config.model_in_dim *= 2
												lowerCAmelCase__    						=						SpeechTaHifiGan(_UpperCamelCase   ).to(_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe(_UpperCamelCase     ,						num_inference_steps=1   )
												lowerCAmelCase__    						=						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       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												self._test_attention_slicing_forward_pass(test_mean_pixel_difference=_UpperCamelCase   )






						def       UpperCamelCase__      (	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       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=_UpperCamelCase   )




@slow
class     __SCREAMING_SNAKE_CASE					(   unittest.TestCase):






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase="cpu"     ,						_UpperCamelCase=torch.floataa     ,						_UpperCamelCase=0   ):






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






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						AudioLDMPipeline.from_pretrained('cvssp/audioldm'   )
												lowerCAmelCase__    						=						audioldm_pipe.to(_UpperCamelCase   )
												audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

												lowerCAmelCase__    						=						self.get_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						25
												lowerCAmelCase__    						=						audioldm_pipe(**_UpperCamelCase   ).audios[0]

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

												lowerCAmelCase__    						=						audio[7_72_30:7_72_40]
												lowerCAmelCase__    						=						np.array(
												    [-0.48_84, -0.46_07, 0.00_23, 0.50_07, 0.58_96, 0.51_51, 0.38_13, -0.02_08, -0.36_87, -0.43_15]   )
												lowerCAmelCase__    						=						np.abs(expected_slice - audio_slice   ).max()
												assert max_diff < 1E-2







						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						AudioLDMPipeline.from_pretrained('cvssp/audioldm'   )
												lowerCAmelCase__    						=						LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config   )
												lowerCAmelCase__    						=						audioldm_pipe.to(_UpperCamelCase   )
												audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

												lowerCAmelCase__    						=						self.get_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						audioldm_pipe(**_UpperCamelCase   ).audios[0]

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

												lowerCAmelCase__    						=						audio[2_77_80:2_77_90]
												lowerCAmelCase__    						=						np.array([-0.21_31, -0.08_73, -0.01_24, -0.01_89, 0.05_69, 0.13_73, 0.18_83, 0.28_86, 0.32_97, 0.22_12]   )
												lowerCAmelCase__    						=						np.abs(expected_slice - audio_slice   ).max()
												assert max_diff < 3E-2













from collections import deque

from .hash_table import HashTable


class     __SCREAMING_SNAKE_CASE					(   __lowercase):






						def __init__(	self     ,						*_UpperCamelCase     ,						**_UpperCamelCase   ):






												"""simple docstring"""
												super().__init__(*_UpperCamelCase     ,						**_UpperCamelCase   )






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						deque([]   ) if self.values[key] is None else self.values[key]
												self.values[key].appendleft(_UpperCamelCase   )
												lowerCAmelCase__    						=						self.values[key]






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												return (
												    sum(self.charge_factor - len(_UpperCamelCase   ) for slot in self.values   )
												    / self.size_table
												    * self.charge_factor
												)







						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase=None   ):






												"""simple docstring"""
												if not (
												    len(self.values[key]   ) == self.charge_factor and self.values.count(_UpperCamelCase   ) == 0
												):
																		return key
												return super()._collision_resolution(_UpperCamelCase     ,						_UpperCamelCase   )





