






'''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 					_snake_case						(			__lowerCamelCase       ,    unittest.TestCase ):
	_A          =   AudioLDMPipeline
	_A          =   TEXT_TO_AUDIO_PARAMS
	_A          =   TEXT_TO_AUDIO_BATCH_PARAMS
	_A          =   frozenset(
	    [
	        'num_inference_steps',
	        'num_waveforms_per_prompt',
	        'generator',
	        'latents',
	        'output_type',
	        'return_dict',
	        'callback',
	        'callback_steps',
	    ] )






	def    lowerCAmelCase_							(		self       )						->    Dict:
					torch.manual_seed(0       )
					snake_case__      :str         =   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_					,)
					snake_case__      :Optional[Any]         =   DDIMScheduler(
					    beta_start=0.00085					,beta_end=0.012					,beta_schedule="scaled_linear"					,clip_sample=a_					,set_alpha_to_one=a_					,)
					torch.manual_seed(0       )
					snake_case__      :str         =   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       )
					snake_case__      :Dict         =   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=1_000					,projection_dim=32					,)
					snake_case__      :str         =   ClapTextModelWithProjection(a_       )
					snake_case__      :Any         =   RobertaTokenizer.from_pretrained("hf-internal-testing/tiny-random-roberta"					,model_max_length=77       )

					snake_case__      :Optional[Any]         =   SpeechTaHifiGanConfig(
					    model_in_dim=8					,sampling_rate=16_000					,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_					,)

					snake_case__      :str         =   SpeechTaHifiGan(a_       )

					snake_case__      :Tuple         =   {
					    "unet": unet,
					    "scheduler": scheduler,
					    "vae": vae,
					    "text_encoder": text_encoder,
					    "tokenizer": tokenizer,
					    "vocoder": vocoder,
					}
					return components






	def    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase=0       )						->    int:
					if str(a_       ).startswith("mps"       ):
									snake_case__      :str         =   torch.manual_seed(a_       )
					else:
									snake_case__      :List[Any]         =   torch.Generator(device=a_       ).manual_seed(a_       )
					snake_case__      :int         =   {
					    "prompt": "A hammer hitting a wooden surface",
					    "generator": generator,
					    "num_inference_steps": 2,
					    "guidance_scale": 6.0,
					}
					return inputs






	def    lowerCAmelCase_							(		self       )						->    Optional[Any]:
					snake_case__      :Optional[int]         =   "cpu"  # ensure determinism for the device-dependent torch.Generator

					snake_case__      :List[Any]         =   self.get_dummy_components()
					snake_case__      :Union[str, Any]         =   AudioLDMPipeline(**a_       )
					snake_case__      :Any         =   audioldm_pipe.to(a_       )
					audioldm_pipe.set_progress_bar_config(disable=a_       )

					snake_case__      :Optional[int]         =   self.get_dummy_inputs(a_       )
					snake_case__      :Optional[int]         =   audioldm_pipe(**a_       )
					snake_case__      :List[str]         =   output.audios[0]

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

					snake_case__      :Union[str, Any]         =   audio[:10]
					snake_case__      :Union[str, Any]         =   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    lowerCAmelCase_							(		self       )						->    int:
					snake_case__      :List[str]         =   self.get_dummy_components()
					snake_case__      :Optional[Any]         =   AudioLDMPipeline(**a_       )
					snake_case__      :List[Any]         =   audioldm_pipe.to(a_       )
					snake_case__      :Dict         =   audioldm_pipe.to(a_       )
					audioldm_pipe.set_progress_bar_config(disable=a_       )

					snake_case__      :List[Any]         =   self.get_dummy_inputs(a_       )
					snake_case__      :Tuple         =   3 * [inputs["prompt"]]

					# forward
					snake_case__      :Tuple         =   audioldm_pipe(**a_       )
					snake_case__      :Optional[int]         =   output.audios[0]

					snake_case__      :List[Any]         =   self.get_dummy_inputs(a_       )
					snake_case__      :Optional[Any]         =   3 * [inputs.pop("prompt"       )]

					snake_case__      :Optional[Any]         =   audioldm_pipe.tokenizer(
					    a_					,padding="max_length"					,max_length=audioldm_pipe.tokenizer.model_max_length					,truncation=a_					,return_tensors="pt"					,)
					snake_case__      :Optional[Any]         =   text_inputs["input_ids"].to(a_       )

					snake_case__      :Union[str, Any]         =   audioldm_pipe.text_encoder(
					    a_					,)
					snake_case__      :str         =   prompt_embeds.text_embeds
					# additional L_2 normalization over each hidden-state
					snake_case__      :Optional[Any]         =   F.normalize(a_					,dim=-1       )

					snake_case__      :Optional[Any]         =   prompt_embeds

					# forward
					snake_case__      :List[Any]         =   audioldm_pipe(**a_       )
					snake_case__      :Tuple         =   output.audios[0]

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






	def    lowerCAmelCase_							(		self       )						->    Tuple:
					snake_case__      :List[Any]         =   self.get_dummy_components()
					snake_case__      :Tuple         =   AudioLDMPipeline(**a_       )
					snake_case__      :Dict         =   audioldm_pipe.to(a_       )
					snake_case__      :Any         =   audioldm_pipe.to(a_       )
					audioldm_pipe.set_progress_bar_config(disable=a_       )

					snake_case__      :Union[str, Any]         =   self.get_dummy_inputs(a_       )
					snake_case__      :Dict         =   3 * ["this is a negative prompt"]
					snake_case__      :Optional[int]         =   negative_prompt
					snake_case__      :List[str]         =   3 * [inputs["prompt"]]

					# forward
					snake_case__      :Optional[Any]         =   audioldm_pipe(**a_       )
					snake_case__      :Dict         =   output.audios[0]

					snake_case__      :Dict         =   self.get_dummy_inputs(a_       )
					snake_case__      :int         =   3 * [inputs.pop("prompt"       )]

					snake_case__      :Dict         =   []
					for p in [prompt, negative_prompt]:
									snake_case__      :Dict         =   audioldm_pipe.tokenizer(
									    a_					,padding="max_length"					,max_length=audioldm_pipe.tokenizer.model_max_length					,truncation=a_					,return_tensors="pt"					,)
									snake_case__      :str         =   text_inputs["input_ids"].to(a_       )

									snake_case__      :Optional[Any]         =   audioldm_pipe.text_encoder(
									    a_					,)
									snake_case__      :int         =   text_embeds.text_embeds
									# additional L_2 normalization over each hidden-state
									snake_case__      :List[str]         =   F.normalize(a_					,dim=-1       )

									embeds.append(a_       )

					snake_case__      :List[str]         =   embeds

					# forward
					snake_case__      :Dict         =   audioldm_pipe(**a_       )
					snake_case__      :Optional[Any]         =   output.audios[0]

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






	def    lowerCAmelCase_							(		self       )						->    Optional[int]:
					snake_case__      :int         =   "cpu"  # ensure determinism for the device-dependent torch.Generator
					snake_case__      :int         =   self.get_dummy_components()
					snake_case__      :Any         =   PNDMScheduler(skip_prk_steps=a_       )
					snake_case__      :List[Any]         =   AudioLDMPipeline(**a_       )
					snake_case__      :Union[str, Any]         =   audioldm_pipe.to(a_       )
					audioldm_pipe.set_progress_bar_config(disable=a_       )

					snake_case__      :Union[str, Any]         =   self.get_dummy_inputs(a_       )
					snake_case__      :List[str]         =   "egg cracking"
					snake_case__      :Tuple         =   audioldm_pipe(**a_					,negative_prompt=a_       )
					snake_case__      :Tuple         =   output.audios[0]

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

					snake_case__      :Optional[Any]         =   audio[:10]
					snake_case__      :List[str]         =   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    lowerCAmelCase_							(		self       )						->    str:
					snake_case__      :Tuple         =   "cpu"  # ensure determinism for the device-dependent torch.Generator
					snake_case__      :Optional[Any]         =   self.get_dummy_components()
					snake_case__      :int         =   PNDMScheduler(skip_prk_steps=a_       )
					snake_case__      :int         =   AudioLDMPipeline(**a_       )
					snake_case__      :Dict         =   audioldm_pipe.to(a_       )
					audioldm_pipe.set_progress_bar_config(disable=a_       )

					snake_case__      :List[str]         =   "A hammer hitting a wooden surface"

					# test num_waveforms_per_prompt=1 (default)
					snake_case__      :Any         =   audioldm_pipe(a_					,num_inference_steps=2       ).audios

					assert audios.shape == (1, 256)

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

					assert audios.shape == (batch_size, 256)

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

					snake_case__      :Union[str, Any]         =   self.get_dummy_inputs(a_       )
					snake_case__      :Union[str, Any]         =   audioldm_pipe(audio_length_in_s=0.016					,**a_       )
					snake_case__      :Union[str, Any]         =   output.audios[0]

					assert audio.ndim == 1
					assert len(a_       ) / vocoder_sampling_rate == 0.016

					snake_case__      :Dict         =   audioldm_pipe(audio_length_in_s=0.032					,**a_       )
					snake_case__      :List[Any]         =   output.audios[0]

					assert audio.ndim == 1
					assert len(a_       ) / vocoder_sampling_rate == 0.032






	def    lowerCAmelCase_							(		self       )						->    Optional[Any]:
					snake_case__      :List[str]         =   self.get_dummy_components()
					snake_case__      :Tuple         =   AudioLDMPipeline(**a_       )
					snake_case__      :Tuple         =   audioldm_pipe.to(a_       )
					audioldm_pipe.set_progress_bar_config(disable=a_       )

					snake_case__      :Dict         =   ["hey"]

					snake_case__      :int         =   audioldm_pipe(a_					,num_inference_steps=1       )
					snake_case__      :str         =   output.audios.shape
					assert audio_shape == (1, 256)

					snake_case__      :Optional[int]         =   audioldm_pipe.vocoder.config
					config.model_in_dim *= 2
					snake_case__      :Tuple         =   SpeechTaHifiGan(a_       ).to(a_       )
					snake_case__      :Optional[int]         =   audioldm_pipe(a_					,num_inference_steps=1       )
					snake_case__      :Any         =   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       )						->    str:
					self._test_attention_slicing_forward_pass(test_mean_pixel_difference=a_       )






	def    lowerCAmelCase_							(		self       )						->    int:
					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       )						->    Dict:
					self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=a_       )




@slow
class 					_snake_case						(			unittest.TestCase ):






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






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






	def    lowerCAmelCase_							(		self       )						->    str:
					snake_case__      :str         =   AudioLDMPipeline.from_pretrained("cvssp/audioldm"       )
					snake_case__      :Any         =   audioldm_pipe.to(a_       )
					audioldm_pipe.set_progress_bar_config(disable=a_       )

					snake_case__      :Tuple         =   self.get_inputs(a_       )
					snake_case__      :Dict         =   25
					snake_case__      :Optional[Any]         =   audioldm_pipe(**a_       ).audios[0]

					assert audio.ndim == 1
					assert len(a_       ) == 81_920

					snake_case__      :Tuple         =   audio[77_230:77_240]
					snake_case__      :str         =   np.array(
					    [-0.4884, -0.4607, 0.0023, 0.5007, 0.5896, 0.5151, 0.3813, -0.0208, -0.3687, -0.4315]       )
					snake_case__      :int         =   np.abs(expected_slice - audio_slice       ).max()
					assert max_diff < 1E-2






	def    lowerCAmelCase_							(		self       )						->    Dict:
					snake_case__      :int         =   AudioLDMPipeline.from_pretrained("cvssp/audioldm"       )
					snake_case__      :Tuple         =   LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config       )
					snake_case__      :List[str]         =   audioldm_pipe.to(a_       )
					audioldm_pipe.set_progress_bar_config(disable=a_       )

					snake_case__      :int         =   self.get_inputs(a_       )
					snake_case__      :str         =   audioldm_pipe(**a_       ).audios[0]

					assert audio.ndim == 1
					assert len(a_       ) == 81_920

					snake_case__      :Optional[int]         =   audio[27_780:27_790]
					snake_case__      :Union[str, Any]         =   np.array([-0.2131, -0.0873, -0.0124, -0.0189, 0.0569, 0.1373, 0.1883, 0.2886, 0.3297, 0.2212]       )
					snake_case__      :str         =   np.abs(expected_slice - audio_slice       ).max()
					assert max_diff < 3E-2







from typing import Any





def 					lowercase_ (    __snake_case     :							list	,	__snake_case     :							list	,	__snake_case     :							dict	,	__snake_case     :							dict	,	__snake_case     :							dict	,	)     -> list:





				'''simple docstring'''


				_validation(
				    __snake_case	,	__snake_case	,	__snake_case	,	__snake_case	,	__snake_case	,	)
				# Creates data structures and fill initial step
				snake_case__      :dict         =   {}
				snake_case__      :dict         =   {}
				for state in states_space:
								snake_case__      :List[Any]         =   observations_space[0]
								snake_case__      :str         =   (
								    initial_probabilities[state] * emission_probabilities[state][observation]
								)
								snake_case__      :str         =   None

				# Fills the data structure with the probabilities of
				# different transitions and pointers to previous states
				for o in range(1	,	len(__snake_case      )      ):
								snake_case__      :Any         =   observations_space[o]
								snake_case__      :Tuple         =   observations_space[o - 1]
								for state in states_space:
												# Calculates the argmax for probability function
												snake_case__      :Tuple         =   ""
												snake_case__      :Union[str, Any]         =   -1
												for k_state in states_space:
																snake_case__      :int         =   (
																    probabilities[(k_state, prior_observation)]
																    * transition_probabilities[k_state][state]
																    * emission_probabilities[state][observation]
																)
																if probability > max_probability:
																				snake_case__      :str         =   probability
																				snake_case__      :Tuple         =   k_state

            # Update probabilities and pointers dicts
												snake_case__      :List[str]         =   (
												    probabilities[(arg_max, prior_observation)]
												    * transition_probabilities[arg_max][state]
												    * emission_probabilities[state][observation]
												)

												snake_case__      :List[str]         =   arg_max

    # The final observation
				snake_case__      :str         =   observations_space[len(__snake_case      ) - 1]

				# argmax for given final observation
				snake_case__      :Optional[int]         =   ""
				snake_case__      :List[str]         =   -1
				for k_state in states_space:
								snake_case__      :List[str]         =   probabilities[(k_state, final_observation)]
								if probability > max_probability:
												snake_case__      :List[str]         =   probability
												snake_case__      :int         =   k_state
				snake_case__      :Any         =   arg_max

				# Process pointers backwards
				snake_case__      :int         =   last_state
				snake_case__      :List[str]         =   []
				for o in range(len(__snake_case      ) - 1	,	-1	,	-1      ):
								result.append(__snake_case      )
								snake_case__      :List[str]         =   pointers[previous, observations_space[o]]
				result.reverse()

				return result





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	)     -> None:





				'''simple docstring'''


				_validate_not_empty(
				    __snake_case	,	__snake_case	,	__snake_case	,	__snake_case	,	__snake_case	,	)
				_validate_lists(__snake_case	,	__snake_case      )
				_validate_dicts(
				    __snake_case	,	__snake_case	,	__snake_case      )





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	)     -> None:





				'''simple docstring'''


				if not all(
				    [
				        observations_space,
				        states_space,
				        initial_probabilities,
				        transition_probabilities,
				        emission_probabilities,
				    ]      ):
								raise ValueError("There's an empty parameter"      )





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							Any      )     -> None:





				'''simple docstring'''


				_validate_list(__snake_case	,	"observations_space"      )
				_validate_list(__snake_case	,	"states_space"      )





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							str      )     -> None:





				'''simple docstring'''


				if not isinstance(_object	,	__snake_case      ):
								snake_case__      :Optional[int]         =   F'{var_name} must be a list'
								raise ValueError(__snake_case      )
				else:
								for x in _object:
												if not isinstance(__snake_case	,	__snake_case      ):
																snake_case__      :Any         =   F'{var_name} must be a list of strings'
																raise ValueError(__snake_case      )





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							Any	,	)     -> None:





				'''simple docstring'''


				_validate_dict(__snake_case	,	"initial_probabilities"	,	__snake_case      )
				_validate_nested_dict(__snake_case	,	"transition_probabilities"      )
				_validate_nested_dict(__snake_case	,	"emission_probabilities"      )





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							str      )     -> None:





				'''simple docstring'''


				_validate_dict(_object	,	__snake_case	,	__snake_case      )
				for x in _object.values():
								_validate_dict(__snake_case	,	__snake_case	,	__snake_case	,	__snake_case      )





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							str	,	__snake_case     :							type	,	__snake_case     :							bool = False      )     -> None:





				'''simple docstring'''


				if not isinstance(_object	,	__snake_case      ):
								snake_case__      :str         =   F'{var_name} must be a dict'
								raise ValueError(__snake_case      )
				if not all(isinstance(__snake_case	,	__snake_case      ) for x in _object      ):
								snake_case__      :List[Any]         =   F'{var_name} all keys must be strings'
								raise ValueError(__snake_case      )
				if not all(isinstance(__snake_case	,	__snake_case      ) for x in _object.values()      ):
								snake_case__      :Optional[int]         =   "nested dictionary " if nested else ""
								snake_case__      :int         =   F'{var_name} {nested_text}all values must be {value_type.__name__}'
								raise ValueError(__snake_case      )


if __name__ == "__main__":
	from doctest import testmod

	testmod()