



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





					'''simple docstring'''



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

					def 			snake_case__      (    self   ):
									torch.manual_seed(0   )
									_lowerCamelCase        =				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=lowerCamelCase__				,       )
									_lowerCamelCase        =				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   )
									_lowerCamelCase        =				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   )
									_lowerCamelCase        =				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				,       )
									_lowerCamelCase        =				ClapTextModelWithProjection(lowerCamelCase__   )
									_lowerCamelCase        =				RobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-random-roberta'''				,       model_max_length=7_7   )

									_lowerCamelCase        =				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=lowerCamelCase__				,       )

									_lowerCamelCase        =				SpeechTaHifiGan(lowerCamelCase__   )

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

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__=0   ):
									if str(lowerCamelCase__   ).startswith('''mps'''   ):
													_lowerCamelCase        =				torch.manual_seed(lowerCamelCase__   )
									else:
													_lowerCamelCase        =				torch.Generator(device=lowerCamelCase__   ).manual_seed(lowerCamelCase__   )
									_lowerCamelCase        =				{
									    '''prompt''': '''A hammer hitting a wooden surface''',
									    '''generator''': generator,
									    '''num_inference_steps''': 2,
									    '''guidance_scale''': 6.0,
									}
									return inputs

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator

									_lowerCamelCase        =				self.get_dummy_components()
									_lowerCamelCase        =				AudioLDMPipeline(**lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe.to(lowerCamelCase__   )
									audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase__   )

									_lowerCamelCase        =				self.get_dummy_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe(**lowerCamelCase__   )
									_lowerCamelCase        =				output.audios[0]

									assert audio.ndim == 1
									assert len(lowerCamelCase__   ) == 2_5_6

									_lowerCamelCase        =				audio[:1_0]
									_lowerCamelCase        =				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   ):
									_lowerCamelCase        =				self.get_dummy_components()
									_lowerCamelCase        =				AudioLDMPipeline(**lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe.to(lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe.to(lowerCamelCase__   )
									audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase__   )

									_lowerCamelCase        =				self.get_dummy_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				3 * [inputs['''prompt''']]

									# forward
									_lowerCamelCase        =				audioldm_pipe(**lowerCamelCase__   )
									_lowerCamelCase        =				output.audios[0]

									_lowerCamelCase        =				self.get_dummy_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				3 * [inputs.pop('''prompt'''   )]

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

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

									_lowerCamelCase        =				prompt_embeds

									# forward
									_lowerCamelCase        =				audioldm_pipe(**lowerCamelCase__   )
									_lowerCamelCase        =				output.audios[0]

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

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.get_dummy_components()
									_lowerCamelCase        =				AudioLDMPipeline(**lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe.to(lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe.to(lowerCamelCase__   )
									audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase__   )

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

									# forward
									_lowerCamelCase        =				audioldm_pipe(**lowerCamelCase__   )
									_lowerCamelCase        =				output.audios[0]

									_lowerCamelCase        =				self.get_dummy_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				3 * [inputs.pop('''prompt'''   )]

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

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

													embeds.append(lowerCamelCase__   )

									_lowerCamelCase  ,     _lowerCamelCase        =				embeds

									# forward
									_lowerCamelCase        =				audioldm_pipe(**lowerCamelCase__   )
									_lowerCamelCase        =				output.audios[0]

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

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									_lowerCamelCase        =				self.get_dummy_components()
									_lowerCamelCase        =				PNDMScheduler(skip_prk_steps=lowerCamelCase__   )
									_lowerCamelCase        =				AudioLDMPipeline(**lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe.to(lowerCamelCase__   )
									audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase__   )

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

									assert audio.ndim == 1
									assert len(lowerCamelCase__   ) == 2_5_6

									_lowerCamelCase        =				audio[:1_0]
									_lowerCamelCase        =				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   ):
									_lowerCamelCase        =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									_lowerCamelCase        =				self.get_dummy_components()
									_lowerCamelCase        =				PNDMScheduler(skip_prk_steps=lowerCamelCase__   )
									_lowerCamelCase        =				AudioLDMPipeline(**lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe.to(lowerCamelCase__   )
									audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase__   )

									_lowerCamelCase        =				'''A hammer hitting a wooden surface'''

									# test num_waveforms_per_prompt=1 (default)
									_lowerCamelCase        =				audioldm_pipe(lowerCamelCase__				,       num_inference_steps=2   ).audios

									assert audios.shape == (1, 2_5_6)

									# test num_waveforms_per_prompt=1 (default) for batch of prompts
									_lowerCamelCase        =				2
									_lowerCamelCase        =				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
									_lowerCamelCase        =				2
									_lowerCamelCase        =				audioldm_pipe(lowerCamelCase__				,       num_inference_steps=2				,       num_waveforms_per_prompt=lowerCamelCase__   ).audios

									assert audios.shape == (num_waveforms_per_prompt, 2_5_6)

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

									assert audios.shape == (batch_size * num_waveforms_per_prompt, 2_5_6)

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									_lowerCamelCase        =				self.get_dummy_components()
									_lowerCamelCase        =				AudioLDMPipeline(**lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe.to(lowerCamelCase__   )
									audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe.vocoder.config.sampling_rate

									_lowerCamelCase        =				self.get_dummy_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe(audio_length_in_s=0.0_1_6				,       **lowerCamelCase__   )
									_lowerCamelCase        =				output.audios[0]

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

									_lowerCamelCase        =				audioldm_pipe(audio_length_in_s=0.0_3_2				,       **lowerCamelCase__   )
									_lowerCamelCase        =				output.audios[0]

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

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.get_dummy_components()
									_lowerCamelCase        =				AudioLDMPipeline(**lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe.to(lowerCamelCase__   )
									audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase__   )

									_lowerCamelCase        =				['''hey''']

									_lowerCamelCase        =				audioldm_pipe(lowerCamelCase__				,       num_inference_steps=1   )
									_lowerCamelCase        =				output.audios.shape
									assert audio_shape == (1, 2_5_6)

									_lowerCamelCase        =				audioldm_pipe.vocoder.config
									config.model_in_dim *= 2
									_lowerCamelCase        =				SpeechTaHifiGan(lowerCamelCase__   ).to(lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe(lowerCamelCase__				,       num_inference_steps=1   )
									_lowerCamelCase        =				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   ):
									self._test_attention_slicing_forward_pass(test_mean_pixel_difference=lowerCamelCase__   )

					def 			snake_case__      (    self   ):
									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   ):
									self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowerCamelCase__   )


@slow
class 							lowerCamelCase_(   unittest.TestCase    ):





					'''simple docstring'''



					def 			snake_case__      (    self   ):
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__="cpu"				,       lowerCamelCase__=torch.floataa				,       lowerCamelCase__=0   ):
									_lowerCamelCase        =				torch.Generator(device=lowerCamelCase__   ).manual_seed(lowerCamelCase__   )
									_lowerCamelCase        =				np.random.RandomState(lowerCamelCase__   ).standard_normal((1, 8, 1_2_8, 1_6)   )
									_lowerCamelCase        =				torch.from_numpy(lowerCamelCase__   ).to(device=lowerCamelCase__				,       dtype=lowerCamelCase__   )
									_lowerCamelCase        =				{
									    '''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   ):
									_lowerCamelCase        =				AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''   )
									_lowerCamelCase        =				audioldm_pipe.to(lowerCamelCase__   )
									audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase__   )

									_lowerCamelCase        =				self.get_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				2_5
									_lowerCamelCase        =				audioldm_pipe(**lowerCamelCase__   ).audios[0]

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

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





					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''   )
									_lowerCamelCase        =				LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config   )
									_lowerCamelCase        =				audioldm_pipe.to(lowerCamelCase__   )
									audioldm_pipe.set_progress_bar_config(disable=lowerCamelCase__   )

									_lowerCamelCase        =				self.get_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				audioldm_pipe(**lowerCamelCase__   ).audios[0]

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

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





"""simple docstring"""
# tests directory-specific settings - this file is run automatically
# by pytest before any tests are run

import sys
import warnings
from os.path import abspath, dirname, join


# allow having multiple repository checkouts and not needing to remember to rerun
# 'pip install -e .[dev]' when switching between checkouts and running tests.
__SCREAMING_SNAKE_CASE  :						Optional[Any]    							=      abspath(join(dirname(dirname(__file__)), '''src'''))
sys.path.insert(1, git_repo_path)

# silence FutureWarning warnings in tests since often we can't act on them until
# they become normal warnings - i.e. the tests still need to test the current functionality
warnings.simplefilter(action='''ignore''', category=FutureWarning)



def        lowerCAmelCase_(   lowercase_     :							List[Any]							)	->       Optional[Any]:
				from diffusers.utils.testing_utils import pytest_addoption_shared

				pytest_addoption_shared(lowercase_							)




def        lowerCAmelCase_(   lowercase_     :							List[str]							)	->       List[str]:
				from diffusers.utils.testing_utils import pytest_terminal_summary_main

				_lowerCamelCase        =				terminalreporter.config.getoption('''--make-reports'''							)
				if make_reports:
								pytest_terminal_summary_main(lowercase_						,   id=lowercase_							)
