

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_   (		__A	,       unittest.TestCase					):







			'''simple docstring'''

			SCREAMING_SNAKE_CASE     :   Any									=				AudioLDMPipeline
			SCREAMING_SNAKE_CASE     :   str									=				TEXT_TO_AUDIO_PARAMS
			SCREAMING_SNAKE_CASE     :   Optional[int]									=				TEXT_TO_AUDIO_BATCH_PARAMS
			SCREAMING_SNAKE_CASE     :   Dict									=				frozenset(
			    [
			        "num_inference_steps",
			        "num_waveforms_per_prompt",
			        "generator",
			        "latents",
			        "output_type",
			        "return_dict",
			        "callback",
			        "callback_steps",
			    ]					)







			def 				snake_case__(     self :				Dict							)      ->str:
									torch.manual_seed(0							)
									snake_case_   =	UNetaDConditionModel(
									    block_out_channels=(3_2, 6_4)							,   layers_per_block=2							,   sample_size=3_2							,   in_channels=4							,   out_channels=4							,   down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')							,   up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')							,   cross_attention_dim=(3_2, 6_4)							,   class_embed_type='''simple_projection'''							,   projection_class_embeddings_input_dim=3_2							,   class_embeddings_concat=_UpperCamelCase							,   )
									snake_case_   =	DDIMScheduler(
									    beta_start=0.00085							,   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=[3_2, 6_4]							,   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=3_2							,   intermediate_size=3_7							,   layer_norm_eps=1e-05							,   num_attention_heads=4							,   num_hidden_layers=5							,   pad_token_id=1							,   vocab_size=1_0_0_0							,   projection_dim=3_2							,   )
									snake_case_   =	ClapTextModelWithProjection(_UpperCamelCase							)
									snake_case_   =	RobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-random-roberta'''							,   model_max_length=7_7							)

									snake_case_   =	SpeechTaHifiGanConfig(
									    model_in_dim=8							,   sampling_rate=1_6_0_0_0							,   upsample_initial_channel=1_6							,   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 				snake_case__(     self :				str							,   _UpperCamelCase :				int							,   _UpperCamelCase :				Union[str, Any]=0							)      ->str:
									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 				snake_case__(     self :				int							)      ->Optional[int]:
									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_5_6

									snake_case_   =	audio[:1_0]
									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 				snake_case__(     self :				List[str]							)      ->Any:
									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 				snake_case__(     self :				Any							)      ->Union[str, Any]:
									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 				snake_case__(     self :				List[Any]							)      ->Optional[int]:
									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_5_6

									snake_case_   =	audio[:1_0]
									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 				snake_case__(     self :				List[Any]							)      ->Union[str, Any]:
									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_5_6)

									# 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_5_6)

									# 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_5_6)

									# 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_5_6)







			def 				snake_case__(     self :				Optional[int]							)      ->List[str]:
									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 				snake_case__(     self :				str							)      ->List[Any]:
									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_5_6)

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







			def 				snake_case__(     self :				Optional[Any]							)      ->int:
									self._test_attention_slicing_forward_pass(test_mean_pixel_difference=_UpperCamelCase							)







			def 				snake_case__(     self :				int							)      ->Optional[int]:
									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 				snake_case__(     self :				Optional[Any]							)      ->int:
									self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=_UpperCamelCase							)






@slow
class 	snake_case_   (		unittest.TestCase					):







			'''simple docstring'''

			def 				snake_case__(     self :				Any							)      ->str:
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()







			def 				snake_case__(     self :				Tuple							,   _UpperCamelCase :				Union[str, Any]							,   _UpperCamelCase :				str="cpu"							,   _UpperCamelCase :				List[Any]=torch.floataa							,   _UpperCamelCase :				Any=0							)      ->List[str]:
									snake_case_   =	torch.Generator(device=_UpperCamelCase							).manual_seed(_UpperCamelCase							)
									snake_case_   =	np.random.RandomState(_UpperCamelCase							).standard_normal((1, 8, 1_2_8, 1_6)							)
									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 				snake_case__(     self :				Any							)      ->Union[str, Any]:
									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_   =	2_5
									snake_case_   =	audioldm_pipe(**_UpperCamelCase							).audios[0]

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

									snake_case_   =	audio[7_7_2_3_0:7_7_2_4_0]
									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 				snake_case__(     self :				List[Any]							)      ->str:
									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_1_9_2_0

									snake_case_   =	audio[2_7_7_8_0:2_7_7_9_0]
									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


from __future__ import annotations


def   __SCREAMING_SNAKE_CASE					(SCREAMING_SNAKE_CASE__       ):
						return [ord(SCREAMING_SNAKE_CASE__       ) - 96 for elem in plain]


def   __SCREAMING_SNAKE_CASE					(SCREAMING_SNAKE_CASE__       ):
						return "".join(chr(elem + 96       ) for elem in encoded       )


def   __SCREAMING_SNAKE_CASE					():
						snake_case_   =	encode(input('''-> '''       ).strip().lower()       )
						print('''Encoded: '''  ,		SCREAMING_SNAKE_CASE__       )
						print('''Decoded:'''  ,		decode(SCREAMING_SNAKE_CASE__       )       )


if __name__ == "__main__":
					main()