


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 						A__							(		UpperCamelCase__   ,     unittest.TestCase  ):
						UpperCAmelCase             =						AudioLDMPipeline
						UpperCAmelCase             =						TEXT_TO_AUDIO_PARAMS
						UpperCAmelCase             =						TEXT_TO_AUDIO_BATCH_PARAMS
						UpperCAmelCase             =						frozenset(
						    [
						        "num_inference_steps",
						        "num_waveforms_per_prompt",
						        "generator",
						        "latents",
						        "output_type",
						        "return_dict",
						        "callback",
						        "callback_steps",
						    ]  )




						def 							__UpperCamelCase				(       self					:      List[Any]  )							->     int:


										"""simple docstring"""





										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_00_85				,    beta_end=0.0_12				,    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=1_6000				,    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 							__UpperCamelCase				(       self					:      Tuple				,    _a					:      List[Any]				,    _a					:      str=0  )							->     Tuple:


										"""simple docstring"""





										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 							__UpperCamelCase				(       self					:      int  )							->     Dict:


										"""simple docstring"""





										_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.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					:      Dict  )							->     List[str]:


										"""simple docstring"""





										_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 							__UpperCamelCase				(       self					:      str  )							->     Any:


										"""simple docstring"""





										_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 							__UpperCamelCase				(       self					:      str  )							->     Optional[Any]:


										"""simple docstring"""





										_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.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					:      Optional[int]  )							->     List[str]:


										"""simple docstring"""





										_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 							__UpperCamelCase				(       self					:      Union[str, Any]  )							->     Optional[int]:


										"""simple docstring"""





										_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_16				,    **_a  )
										_SCREAMING_SNAKE_CASE			  =output.audios[0]

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

										_SCREAMING_SNAKE_CASE			  =audioldm_pipe(audio_length_in_s=0.0_32				,    **_a  )
										_SCREAMING_SNAKE_CASE			  =output.audios[0]

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




						def 							__UpperCamelCase				(       self					:      Optional[int]  )							->     str:


										"""simple docstring"""





										_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 							__UpperCamelCase				(       self					:      Any  )							->     Optional[int]:


										"""simple docstring"""





										self._test_attention_slicing_forward_pass(test_mean_pixel_difference=_a  )




						def 							__UpperCamelCase				(       self					:      Optional[int]  )							->     int:


										"""simple docstring"""





										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 							__UpperCamelCase				(       self					:      List[Any]  )							->     Optional[int]:


										"""simple docstring"""





										self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=_a  )

@slow
class 						A__							(		unittest.TestCase  ):




						def 							__UpperCamelCase				(       self					:      List[Any]  )							->     int:


										"""simple docstring"""





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




						def 							__UpperCamelCase				(       self					:      Union[str, Any]				,    _a					:      Tuple				,    _a					:      Optional[int]="cpu"				,    _a					:      Optional[int]=torch.floataa				,    _a					:      Tuple=0  )							->     str:


										"""simple docstring"""





										_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 							__UpperCamelCase				(       self					:      Any  )							->     Union[str, Any]:


										"""simple docstring"""





										_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  ) == 8_1920

										_SCREAMING_SNAKE_CASE			  =audio[7_7230:7_7240]
										_SCREAMING_SNAKE_CASE			  =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]  )
										_SCREAMING_SNAKE_CASE			  =np.abs(expected_slice - audio_slice  ).max()
										assert max_diff < 1E-2




						def 							__UpperCamelCase				(       self					:      List[str]  )							->     Any:


										"""simple docstring"""





										_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  ) == 8_1920

										_SCREAMING_SNAKE_CASE			  =audio[2_7780:2_7790]
										_SCREAMING_SNAKE_CASE			  =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]  )
										_SCREAMING_SNAKE_CASE			  =np.abs(expected_slice - audio_slice  ).max()
										assert max_diff < 3E-2



from datetime import datetime

import matplotlib.pyplot as plt
import torch







def 						lowerCamelCase(   a__):
				for param in module.parameters():
								_SCREAMING_SNAKE_CASE			  =False







def 						lowerCamelCase(   ):
				_SCREAMING_SNAKE_CASE			  ='''cuda''' if torch.cuda.is_available() else '''cpu'''
				if torch.backends.mps.is_available() and torch.backends.mps.is_built():
								_SCREAMING_SNAKE_CASE			  ='''mps'''
				if device == "mps":
								print(
								    '''WARNING: MPS currently doesn\'t seem to work, and messes up backpropagation without any visible torch'''
								    ''' errors. I recommend using CUDA on a colab notebook or CPU instead if you\'re facing inexplicable issues'''
								    ''' with generations.''')
				return device







def 						lowerCamelCase(   a__):
				_SCREAMING_SNAKE_CASE			  =plt.imshow(a__)
				fig.axes.get_xaxis().set_visible(a__)
				fig.axes.get_yaxis().set_visible(a__)
				plt.show()







def 						lowerCamelCase(   ):
				_SCREAMING_SNAKE_CASE			  =datetime.now()
				_SCREAMING_SNAKE_CASE			  =current_time.strftime('''%H:%M:%S''')
				return timestamp