


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




       """simple docstring"""
       snake_case    =       AudioLDMPipeline
       snake_case    =       TEXT_TO_AUDIO_PARAMS
       snake_case    =       TEXT_TO_AUDIO_BATCH_PARAMS
       snake_case    =       frozenset(
           [
               "num_inference_steps",
               "num_waveforms_per_prompt",
               "generator",
               "latents",
               "output_type",
               "return_dict",
               "callback",
               "callback_steps",
           ]     )



       def 			_snake_case			( self						)->str:




         '''simple docstring'''





         torch.manual_seed(0						)
         A_ :							Optional[int]							=   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=lowerCamelCase_							,			)
         A_ :							Dict							=   DDIMScheduler(
             beta_start=0.0_0_0_8_5							,			beta_end=0.0_1_2							,			beta_schedule='''scaled_linear'''							,			clip_sample=lowerCamelCase_							,			set_alpha_to_one=lowerCamelCase_							,			)
         torch.manual_seed(0						)
         A_ :							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						)
         A_ :							str							=   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_ :							Optional[Any]							=   ClapTextModelWithProjection(lowerCamelCase_						)
         A_ :							Dict							=   RobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-random-roberta'''							,			model_max_length=77						)

         A_ :							List[str]							=   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=lowerCamelCase_							,			)

         A_ :							Tuple							=   SpeechTaHifiGan(lowerCamelCase_						)

         A_ :							Optional[int]							=   {
             'unet': unet,
             'scheduler': scheduler,
             'vae': vae,
             'text_encoder': text_encoder,
             'tokenizer': tokenizer,
             'vocoder': vocoder,
         }
         return components



       def 			_snake_case			( self							,			_SCREAMING_SNAKE_CASE							,			_SCREAMING_SNAKE_CASE=0						)->str:




         '''simple docstring'''





         if str(lowerCamelCase_						).startswith('''mps'''						):
           A_ :							int							=   torch.manual_seed(lowerCamelCase_						)
         else:
           A_ :							List[Any]							=   torch.Generator(device=lowerCamelCase_						).manual_seed(lowerCamelCase_						)
         A_ :							Optional[int]							=   {
             'prompt': 'A hammer hitting a wooden surface',
             'generator': generator,
             'num_inference_steps': 2,
             'guidance_scale': 6.0,
         }
         return inputs



       def 			_snake_case			( self						)->Optional[Any]:




         '''simple docstring'''





         A_ :							int							=   'cpu'  # ensure determinism for the device-dependent torch.Generator

         A_ :							Optional[int]							=   self.get_dummy_components()
         A_ :							str							=   AudioLDMPipeline(**lowerCamelCase_						)
         A_ :							Optional[int]							=   audioldm_pipe.to(lowerCamelCase_						)
         audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_						)

         A_ :							int							=   self.get_dummy_inputs(lowerCamelCase_						)
         A_ :							List[Any]							=   audioldm_pipe(**lowerCamelCase_						)
         A_ :							Tuple							=   output.audios[0]

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

         A_ :							Optional[Any]							=   audio[:10]
         A_ :							List[str]							=   np.array(
             [-0.0_0_5_0, 0.0_0_5_0, -0.0_0_6_0, 0.0_0_3_3, -0.0_0_2_6, 0.0_0_3_3, -0.0_0_2_7, 0.0_0_3_3, -0.0_0_2_8, 0.0_0_3_3]						)

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



       def 			_snake_case			( self						)->Dict:




         '''simple docstring'''





         A_ :							Dict							=   self.get_dummy_components()
         A_ :							Tuple							=   AudioLDMPipeline(**lowerCamelCase_						)
         A_ :							Union[str, Any]							=   audioldm_pipe.to(lowerCamelCase_						)
         A_ :							List[Any]							=   audioldm_pipe.to(lowerCamelCase_						)
         audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_						)

         A_ :							Optional[Any]							=   self.get_dummy_inputs(lowerCamelCase_						)
         A_ :							Dict							=   3 * [inputs['prompt']]

         # forward
         A_ :							List[Any]							=   audioldm_pipe(**lowerCamelCase_						)
         A_ :							Optional[Any]							=   output.audios[0]

         A_ :							Tuple							=   self.get_dummy_inputs(lowerCamelCase_						)
         A_ :							Optional[Any]							=   3 * [inputs.pop('''prompt'''						)]

         A_ :							Any							=   audioldm_pipe.tokenizer(
             lowerCamelCase_							,			padding='''max_length'''							,			max_length=audioldm_pipe.tokenizer.model_max_length							,			truncation=lowerCamelCase_							,			return_tensors='''pt'''							,			)
         A_ :							Union[str, Any]							=   text_inputs['input_ids'].to(lowerCamelCase_						)

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

         A_ :							Tuple							=   prompt_embeds

         # forward
         A_ :							str							=   audioldm_pipe(**lowerCamelCase_						)
         A_ :							Any							=   output.audios[0]

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



       def 			_snake_case			( self						)->Tuple:




         '''simple docstring'''





         A_ :							Any							=   self.get_dummy_components()
         A_ :							Optional[int]							=   AudioLDMPipeline(**lowerCamelCase_						)
         A_ :							Any							=   audioldm_pipe.to(lowerCamelCase_						)
         A_ :							List[str]							=   audioldm_pipe.to(lowerCamelCase_						)
         audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_						)

         A_ :							Any							=   self.get_dummy_inputs(lowerCamelCase_						)
         A_ :							Tuple							=   3 * ['this is a negative prompt']
         A_ :							Optional[int]							=   negative_prompt
         A_ :							Optional[int]							=   3 * [inputs['prompt']]

         # forward
         A_ :							Optional[int]							=   audioldm_pipe(**lowerCamelCase_						)
         A_ :							Union[str, Any]							=   output.audios[0]

         A_ :							str							=   self.get_dummy_inputs(lowerCamelCase_						)
         A_ :							Tuple							=   3 * [inputs.pop('''prompt'''						)]

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

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

           embeds.append(lowerCamelCase_						)

         A_ :							Optional[int]							=   embeds

         # forward
         A_ :							int							=   audioldm_pipe(**lowerCamelCase_						)
         A_ :							str							=   output.audios[0]

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



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




         '''simple docstring'''





         A_ :							Dict							=   'cpu'  # ensure determinism for the device-dependent torch.Generator
         A_ :							Optional[int]							=   self.get_dummy_components()
         A_ :							List[Any]							=   PNDMScheduler(skip_prk_steps=lowerCamelCase_						)
         A_ :							int							=   AudioLDMPipeline(**lowerCamelCase_						)
         A_ :							Union[str, Any]							=   audioldm_pipe.to(lowerCamelCase_						)
         audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_						)

         A_ :							Any							=   self.get_dummy_inputs(lowerCamelCase_						)
         A_ :							List[str]							=   'egg cracking'
         A_ :							str							=   audioldm_pipe(**lowerCamelCase_							,			negative_prompt=lowerCamelCase_						)
         A_ :							str							=   output.audios[0]

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

         A_ :							Tuple							=   audio[:10]
         A_ :							Any							=   np.array(
             [-0.0_0_5_1, 0.0_0_5_0, -0.0_0_6_0, 0.0_0_3_4, -0.0_0_2_6, 0.0_0_3_3, -0.0_0_2_7, 0.0_0_3_3, -0.0_0_2_8, 0.0_0_3_2]						)

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



       def 			_snake_case			( self						)->int:




         '''simple docstring'''





         A_ :							List[Any]							=   'cpu'  # ensure determinism for the device-dependent torch.Generator
         A_ :							List[Any]							=   self.get_dummy_components()
         A_ :							Optional[Any]							=   PNDMScheduler(skip_prk_steps=lowerCamelCase_						)
         A_ :							Dict							=   AudioLDMPipeline(**lowerCamelCase_						)
         A_ :							Dict							=   audioldm_pipe.to(lowerCamelCase_						)
         audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_						)

         A_ :							Dict							=   'A hammer hitting a wooden surface'

         # test num_waveforms_per_prompt=1 (default)
         A_ :							Tuple							=   audioldm_pipe(lowerCamelCase_							,			num_inference_steps=2						).audios

         assert audios.shape == (1, 256)

         # test num_waveforms_per_prompt=1 (default) for batch of prompts
         A_ :							List[str]							=   2
         A_ :							Optional[int]							=   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_ :							Union[str, Any]							=   2
         A_ :							List[Any]							=   audioldm_pipe(lowerCamelCase_							,			num_inference_steps=2							,			num_waveforms_per_prompt=lowerCamelCase_						).audios

         assert audios.shape == (num_waveforms_per_prompt, 256)

         # test num_waveforms_per_prompt for batch of prompts
         A_ :							Tuple							=   2
         A_ :							Any							=   audioldm_pipe(
             [prompt] * batch_size							,			num_inference_steps=2							,			num_waveforms_per_prompt=lowerCamelCase_						).audios

         assert audios.shape == (batch_size * num_waveforms_per_prompt, 256)



       def 			_snake_case			( self						)->Any:




         '''simple docstring'''





         A_ :							Tuple							=   'cpu'  # ensure determinism for the device-dependent torch.Generator
         A_ :							Optional[Any]							=   self.get_dummy_components()
         A_ :							int							=   AudioLDMPipeline(**lowerCamelCase_						)
         A_ :							Optional[int]							=   audioldm_pipe.to(lowerCamelCase_						)
         audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_						)
         A_ :							Optional[Any]							=   audioldm_pipe.vocoder.config.sampling_rate

         A_ :							List[str]							=   self.get_dummy_inputs(lowerCamelCase_						)
         A_ :							str							=   audioldm_pipe(audio_length_in_s=0.0_1_6							,			**lowerCamelCase_						)
         A_ :							List[Any]							=   output.audios[0]

         assert audio.ndim == 1
         assert len(lowerCamelCase_						) / vocoder_sampling_rate == 0.0_1_6

         A_ :							int							=   audioldm_pipe(audio_length_in_s=0.0_3_2							,			**lowerCamelCase_						)
         A_ :							List[str]							=   output.audios[0]

         assert audio.ndim == 1
         assert len(lowerCamelCase_						) / vocoder_sampling_rate == 0.0_3_2



       def 			_snake_case			( self						)->Any:




         '''simple docstring'''





         A_ :							Any							=   self.get_dummy_components()
         A_ :							List[Any]							=   AudioLDMPipeline(**lowerCamelCase_						)
         A_ :							Any							=   audioldm_pipe.to(lowerCamelCase_						)
         audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_						)

         A_ :							Optional[Any]							=   ['hey']

         A_ :							List[str]							=   audioldm_pipe(lowerCamelCase_							,			num_inference_steps=1						)
         A_ :							Tuple							=   output.audios.shape
         assert audio_shape == (1, 256)

         A_ :							Optional[Any]							=   audioldm_pipe.vocoder.config
         config.model_in_dim *= 2
         A_ :							Tuple							=   SpeechTaHifiGan(lowerCamelCase_						).to(lowerCamelCase_						)
         A_ :							Optional[int]							=   audioldm_pipe(lowerCamelCase_							,			num_inference_steps=1						)
         A_ :							Dict							=   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 			_snake_case			( self						)->Tuple:




         '''simple docstring'''





         self._test_attention_slicing_forward_pass(test_mean_pixel_difference=lowerCamelCase_						)



       def 			_snake_case			( self						)->Any:




         '''simple docstring'''





         self._test_inference_batch_single_identical(test_mean_pixel_difference=lowerCamelCase_						)







       @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						)->Any:




         '''simple docstring'''





         self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowerCamelCase_						)






@slow
class     _lowerCamelCase					(   unittest.TestCase     ):




       """simple docstring"""



       def 			_snake_case			( self						)->Dict:




         '''simple docstring'''





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



       def 			_snake_case			( self							,			_SCREAMING_SNAKE_CASE							,			_SCREAMING_SNAKE_CASE="cpu"							,			_SCREAMING_SNAKE_CASE=torch.floataa							,			_SCREAMING_SNAKE_CASE=0						)->Any:




         '''simple docstring'''





         A_ :							Optional[int]							=   torch.Generator(device=lowerCamelCase_						).manual_seed(lowerCamelCase_						)
         A_ :							List[Any]							=   np.random.RandomState(lowerCamelCase_						).standard_normal((1, 8, 128, 16)						)
         A_ :							Optional[int]							=   torch.from_numpy(lowerCamelCase_						).to(device=lowerCamelCase_							,			dtype=lowerCamelCase_						)
         A_ :							int							=   {
             '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						)->List[Any]:




         '''simple docstring'''





         A_ :							str							=   AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''						)
         A_ :							Tuple							=   audioldm_pipe.to(lowerCamelCase_						)
         audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_						)

         A_ :							List[str]							=   self.get_inputs(lowerCamelCase_						)
         A_ :							Union[str, Any]							=   25
         A_ :							str							=   audioldm_pipe(**lowerCamelCase_						).audios[0]

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

         A_ :							Optional[int]							=   audio[7_7230:7_7240]
         A_ :							List[Any]							=   np.array(
             [-0.4_8_8_4, -0.4_6_0_7, 0.0_0_2_3, 0.5_0_0_7, 0.5_8_9_6, 0.5_1_5_1, 0.3_8_1_3, -0.0_2_0_8, -0.3_6_8_7, -0.4_3_1_5]						)
         A_ :							Optional[Any]							=   np.abs(expected_slice - audio_slice						).max()
         assert max_diff < 1e-2







       def 			_snake_case			( self						)->Dict:




         '''simple docstring'''





         A_ :							str							=   AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''						)
         A_ :							Dict							=   LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config						)
         A_ :							str							=   audioldm_pipe.to(lowerCamelCase_						)
         audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase_						)

         A_ :							int							=   self.get_inputs(lowerCamelCase_						)
         A_ :							Optional[int]							=   audioldm_pipe(**lowerCamelCase_						).audios[0]

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

         A_ :							Any							=   audio[2_7780:2_7790]
         A_ :							Dict							=   np.array([-0.2_1_3_1, -0.0_8_7_3, -0.0_1_2_4, -0.0_1_8_9, 0.0_5_6_9, 0.1_3_7_3, 0.1_8_8_3, 0.2_8_8_6, 0.3_2_9_7, 0.2_2_1_2]						)
         A_ :							Tuple							=   np.abs(expected_slice - audio_slice						).max()
         assert max_diff < 3e-2











__UpperCamelCase:	str = {0: [2, 3], 1: [0], 2: [1], 3: [4], 4: []}

__UpperCamelCase:	Tuple = {0: [1, 2, 3], 1: [2], 2: [0], 3: [4], 4: [5], 5: [3]}


def       SCREAMING_SNAKE_CASE__  (       _lowercase     : dict[int, list[int]]   ,    _lowercase     : int   ,    _lowercase     : list[bool]			)   ->							list[int]:




		'''simple docstring'''
		lowercase__					:      str      =		True
		lowercase__					:      int      =		[]

		for neighbour in graph[vert]:
				if not visited[neighbour]:
						order += topology_sort(_lowercase   ,    _lowercase   ,    _lowercase			)

		order.append(_lowercase			)

		return order


def       SCREAMING_SNAKE_CASE__  (       _lowercase     : dict[int, list[int]]   ,    _lowercase     : int   ,    _lowercase     : list[bool]			)   ->							list[int]:




		'''simple docstring'''
		lowercase__					:      Union[str, Any]      =		True
		lowercase__					:      Optional[Any]      =		[vert]

		for neighbour in reversed_graph[vert]:
				if not visited[neighbour]:
						component += find_components(_lowercase   ,    _lowercase   ,    _lowercase			)

		return component





def       SCREAMING_SNAKE_CASE__  (       _lowercase     : dict[int, list[int]]			)   ->							list[list[int]]:




		'''simple docstring'''
		lowercase__					:      Optional[int]      =		len(_lowercase			) * [False]
		lowercase__					:      dict[int, list[int]]      =		{vert: [] for vert in range(len(_lowercase			)			)}

		for vert, neighbours in graph.items():
				for neighbour in neighbours:
						reversed_graph[neighbour].append(_lowercase			)

		lowercase__					:      Dict      =		[]
		for i, was_visited in enumerate(_lowercase			):
				if not was_visited:
						order += topology_sort(_lowercase   ,    _lowercase   ,    _lowercase			)

		lowercase__					:      Dict      =		[]
		lowercase__					:      Optional[Any]      =		len(_lowercase			) * [False]

		for i in range(len(_lowercase			)			):
				lowercase__					:      List[Any]      =		order[len(_lowercase			) - i - 1]
				if not visited[vert]:
						lowercase__					:      Tuple      =		find_components(_lowercase   ,    _lowercase   ,    _lowercase			)
						components_list.append(_lowercase			)

		return components_list


