
"""simple docstring"""


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 					lowerCAmelCase_					(   _lowercase		,       unittest.TestCase     ):



						'''simple docstring'''


						_lowerCamelCase:     str   =     AudioLDMPipeline
						_lowerCamelCase:     Optional[int]   =     TEXT_TO_AUDIO_PARAMS
						_lowerCamelCase:     Optional[int]   =     TEXT_TO_AUDIO_BATCH_PARAMS
						_lowerCamelCase:     Optional[int]   =     frozenset(
						    [
						        '''num_inference_steps''',
						        '''num_waveforms_per_prompt''',
						        '''generator''',
						        '''latents''',
						        '''output_type''',
						        '''return_dict''',
						        '''callback''',
						        '''callback_steps''',
						    ]     )


						def     _SCREAMING_SNAKE_CASE   (			self  :   Union[str, Any]   )  -> Any:
							torch.manual_seed(0   )
							A     = 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_	,)
							A     = 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   )
							A     = 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   )
							A     = 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	,)
							A     = ClapTextModelWithProjection(A_   )
							A     = RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta'	,model_max_length=77   )

							A     = 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_	,)

							A     = SpeechTaHifiGan(A_   )

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


						def     _SCREAMING_SNAKE_CASE   (			self  :   Union[str, Any]	,A_  :   Any	,A_  :   Dict=0   )  -> str:
							if str(A_   ).startswith('mps'   ):
								A     = torch.manual_seed(A_   )
							else:
								A     = torch.Generator(device=A_   ).manual_seed(A_   )
							A     = {
							    'prompt': 'A hammer hitting a wooden surface',
							    'generator': generator,
							    'num_inference_steps': 2,
							    'guidance_scale': 6.0,
							}
							return inputs


						def     _SCREAMING_SNAKE_CASE   (			self  :   Any   )  -> Dict:
							A     = 'cpu'  # ensure determinism for the device-dependent torch.Generator

							A     = self.get_dummy_components()
							A     = AudioLDMPipeline(**A_   )
							A     = audioldm_pipe.to(A_   )
							audioldm_pipe.set_progress_bar_config(disable=A_   )

							A     = self.get_dummy_inputs(A_   )
							A     = audioldm_pipe(**A_   )
							A     = output.audios[0]

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

							A     = audio[:10]
							A     = 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     _SCREAMING_SNAKE_CASE   (			self  :   Union[str, Any]   )  -> Tuple:
							A     = self.get_dummy_components()
							A     = AudioLDMPipeline(**A_   )
							A     = audioldm_pipe.to(A_   )
							A     = audioldm_pipe.to(A_   )
							audioldm_pipe.set_progress_bar_config(disable=A_   )

							A     = self.get_dummy_inputs(A_   )
							A     = 3 * [inputs['prompt']]

							# forward
							A     = audioldm_pipe(**A_   )
							A     = output.audios[0]

							A     = self.get_dummy_inputs(A_   )
							A     = 3 * [inputs.pop('prompt'   )]

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

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

							A     = prompt_embeds

							# forward
							A     = audioldm_pipe(**A_   )
							A     = output.audios[0]

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


						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[int]   )  -> Any:
							A     = self.get_dummy_components()
							A     = AudioLDMPipeline(**A_   )
							A     = audioldm_pipe.to(A_   )
							A     = audioldm_pipe.to(A_   )
							audioldm_pipe.set_progress_bar_config(disable=A_   )

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

							# forward
							A     = audioldm_pipe(**A_   )
							A     = output.audios[0]

							A     = self.get_dummy_inputs(A_   )
							A     = 3 * [inputs.pop('prompt'   )]

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

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

								embeds.append(A_   )

							A ,   A     = embeds

							# forward
							A     = audioldm_pipe(**A_   )
							A     = output.audios[0]

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


						def     _SCREAMING_SNAKE_CASE   (			self  :   str   )  -> int:
							A     = 'cpu'  # ensure determinism for the device-dependent torch.Generator
							A     = self.get_dummy_components()
							A     = PNDMScheduler(skip_prk_steps=A_   )
							A     = AudioLDMPipeline(**A_   )
							A     = audioldm_pipe.to(A_   )
							audioldm_pipe.set_progress_bar_config(disable=A_   )

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

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

							A     = audio[:10]
							A     = 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     _SCREAMING_SNAKE_CASE   (			self  :   Any   )  -> List[str]:
							A     = 'cpu'  # ensure determinism for the device-dependent torch.Generator
							A     = self.get_dummy_components()
							A     = PNDMScheduler(skip_prk_steps=A_   )
							A     = AudioLDMPipeline(**A_   )
							A     = audioldm_pipe.to(A_   )
							audioldm_pipe.set_progress_bar_config(disable=A_   )

							A     = 'A hammer hitting a wooden surface'

							# test num_waveforms_per_prompt=1 (default)
							A     = audioldm_pipe(A_	,num_inference_steps=2   ).audios

							assert audios.shape == (1, 256)

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

							assert audios.shape == (batch_size, 256)

							# test num_waveforms_per_prompt for single prompt
							A     = 2
							A     = 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
							A     = 2
							A     = 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     _SCREAMING_SNAKE_CASE   (			self  :   List[Any]   )  -> Optional[Any]:
							A     = 'cpu'  # ensure determinism for the device-dependent torch.Generator
							A     = self.get_dummy_components()
							A     = AudioLDMPipeline(**A_   )
							A     = audioldm_pipe.to(A_   )
							audioldm_pipe.set_progress_bar_config(disable=A_   )
							A     = audioldm_pipe.vocoder.config.sampling_rate

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

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

							A     = audioldm_pipe(audio_length_in_s=0.0_32	,**A_   )
							A     = output.audios[0]

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


						def     _SCREAMING_SNAKE_CASE   (			self  :   Dict   )  -> List[Any]:
							A     = self.get_dummy_components()
							A     = AudioLDMPipeline(**A_   )
							A     = audioldm_pipe.to(A_   )
							audioldm_pipe.set_progress_bar_config(disable=A_   )

							A     = ['hey']

							A     = audioldm_pipe(A_	,num_inference_steps=1   )
							A     = output.audios.shape
							assert audio_shape == (1, 256)

							A     = audioldm_pipe.vocoder.config
							config.model_in_dim *= 2
							A     = SpeechTaHifiGan(A_   ).to(A_   )
							A     = audioldm_pipe(A_	,num_inference_steps=1   )
							A     = 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     _SCREAMING_SNAKE_CASE   (			self  :   List[str]   )  -> int:
							self._test_attention_slicing_forward_pass(test_mean_pixel_difference=A_   )


						def     _SCREAMING_SNAKE_CASE   (			self  :   Any   )  -> List[str]:
							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     _SCREAMING_SNAKE_CASE   (			self  :   int   )  -> List[str]:
							self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=A_   )







@slow
class 					lowerCAmelCase_					(   unittest.TestCase     ):



						'''simple docstring'''


						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[Any]   )  -> List[str]:
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[str]	,A_  :   List[Any]	,A_  :   str="cpu"	,A_  :   List[str]=torch.floataa	,A_  :   str=0   )  -> List[Any]:
							A     = torch.Generator(device=A_   ).manual_seed(A_   )
							A     = np.random.RandomState(A_   ).standard_normal((1, 8, 128, 16)   )
							A     = torch.from_numpy(A_   ).to(device=A_	,dtype=A_   )
							A     = {
							    'prompt': 'A hammer hitting a wooden surface',
							    'latents': latents,
							    'generator': generator,
							    'num_inference_steps': 3,
							    'guidance_scale': 2.5,
							}
							return inputs


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[str]   )  -> int:
							A     = AudioLDMPipeline.from_pretrained('cvssp/audioldm'   )
							A     = audioldm_pipe.to(A_   )
							audioldm_pipe.set_progress_bar_config(disable=A_   )

							A     = self.get_inputs(A_   )
							A     = 25
							A     = audioldm_pipe(**A_   ).audios[0]

							assert audio.ndim == 1
							assert len(A_   ) == 8_1920

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




						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[int]   )  -> List[Any]:
							A     = AudioLDMPipeline.from_pretrained('cvssp/audioldm'   )
							A     = LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config   )
							A     = audioldm_pipe.to(A_   )
							audioldm_pipe.set_progress_bar_config(disable=A_   )

							A     = self.get_inputs(A_   )
							A     = audioldm_pipe(**A_   ).audios[0]

							assert audio.ndim == 1
							assert len(A_   ) == 8_1920

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

"""simple docstring"""






from ..utils import DummyObject, requires_backends



class 			lowercase					( metaclass=lowercase__		):
			lowercase        =     ['''flax''', '''transformers''']






			def __init__(self  :	List[Any]      ,*SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,**SCREAMING_SNAKE_CASE_  :	Union[str, Any]			)       ->	str:





				"""simple docstring"""


				requires_backends(self      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	List[Any]      ,*SCREAMING_SNAKE_CASE_  :	List[Any]      ,**SCREAMING_SNAKE_CASE_  :	int			)       ->	Optional[int]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	List[Any]      ,*SCREAMING_SNAKE_CASE_  :	str      ,**SCREAMING_SNAKE_CASE_  :	Union[str, Any]			)       ->	Union[str, Any]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)





class 			lowercase					( metaclass=lowercase__		):
			lowercase        =     ['''flax''', '''transformers''']






			def __init__(self  :	int      ,*SCREAMING_SNAKE_CASE_  :	List[Any]      ,**SCREAMING_SNAKE_CASE_  :	Any			)       ->	Union[str, Any]:





				"""simple docstring"""


				requires_backends(self      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	int      ,*SCREAMING_SNAKE_CASE_  :	str      ,**SCREAMING_SNAKE_CASE_  :	int			)       ->	Optional[Any]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	List[str]      ,*SCREAMING_SNAKE_CASE_  :	List[str]      ,**SCREAMING_SNAKE_CASE_  :	Optional[Any]			)       ->	Dict:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)





class 			lowercase					( metaclass=lowercase__		):
			lowercase        =     ['''flax''', '''transformers''']






			def __init__(self  :	Tuple      ,*SCREAMING_SNAKE_CASE_  :	Optional[int]      ,**SCREAMING_SNAKE_CASE_  :	Dict			)       ->	List[str]:





				"""simple docstring"""


				requires_backends(self      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	Tuple      ,*SCREAMING_SNAKE_CASE_  :	Dict      ,**SCREAMING_SNAKE_CASE_  :	Optional[Any]			)       ->	List[Any]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	Tuple      ,*SCREAMING_SNAKE_CASE_  :	Tuple      ,**SCREAMING_SNAKE_CASE_  :	Tuple			)       ->	Union[str, Any]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)





class 			lowercase					( metaclass=lowercase__		):
			lowercase        =     ['''flax''', '''transformers''']






			def __init__(self  :	Optional[Any]      ,*SCREAMING_SNAKE_CASE_  :	Any      ,**SCREAMING_SNAKE_CASE_  :	str			)       ->	Union[str, Any]:





				"""simple docstring"""


				requires_backends(self      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	List[Any]      ,*SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,**SCREAMING_SNAKE_CASE_  :	Optional[int]			)       ->	Union[str, Any]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)






			@classmethod
			def 					UpperCAmelCase							(cls  :	List[Any]      ,*SCREAMING_SNAKE_CASE_  :	List[Any]      ,**SCREAMING_SNAKE_CASE_  :	List[Any]			)       ->	Optional[Any]:





				"""simple docstring"""


				requires_backends(cls      ,['''flax''', '''transformers''']			)


