




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

    '''simple docstring'''






    UpperCAmelCase__				    = AudioLDMPipeline
    UpperCAmelCase__				    = TEXT_TO_AUDIO_PARAMS
    UpperCAmelCase__				    = TEXT_TO_AUDIO_BATCH_PARAMS
    UpperCAmelCase__				    = frozenset(
        [
            '''num_inference_steps''',
            '''num_waveforms_per_prompt''',
            '''generator''',
            '''latents''',
            '''output_type''',
            '''return_dict''',
            '''callback''',
            '''callback_steps''',
        ]	)


    def 			SCREAMING_SNAKE_CASE				(     self :      int)		->List[str]:





           '''simple docstring'''

           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=UpperCAmelCase__   ,					)
           A__ 					=       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)
           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=1_000   ,					projection_dim=32   ,					)
           A__ 					=       ClapTextModelWithProjection(UpperCAmelCase__)
           A__ 					=       RobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-random-roberta'''   ,					model_max_length=77)

           A__ 					=       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=UpperCAmelCase__   ,					)

           A__ 					=       SpeechTaHifiGan(UpperCAmelCase__)

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


    def 			SCREAMING_SNAKE_CASE				(     self :      Dict   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      Any=0)		->List[str]:





           '''simple docstring'''

           if str(UpperCAmelCase__).startswith('''mps'''):
                  A__ 					=       torch.manual_seed(UpperCAmelCase__)
           else:
                  A__ 					=       torch.Generator(device=UpperCAmelCase__).manual_seed(UpperCAmelCase__)
           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 :      Tuple)		->List[Any]:





           '''simple docstring'''

           A__ 					=       '''cpu'''  # ensure determinism for the device-dependent torch.Generator

           A__ 					=       self.get_dummy_components()
           A__ 					=       AudioLDMPipeline(**UpperCAmelCase__)
           A__ 					=       audioldm_pipe.to(UpperCAmelCase__)
           audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

           A__ 					=       self.get_dummy_inputs(UpperCAmelCase__)
           A__ 					=       audioldm_pipe(**UpperCAmelCase__)
           A__ 					=       output.audios[0]

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

           A__ 					=       audio[:10]
           A__ 					=       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 			SCREAMING_SNAKE_CASE				(     self :      Union[str, Any])		->Optional[Any]:





           '''simple docstring'''

           A__ 					=       self.get_dummy_components()
           A__ 					=       AudioLDMPipeline(**UpperCAmelCase__)
           A__ 					=       audioldm_pipe.to(UpperCAmelCase__)
           A__ 					=       audioldm_pipe.to(UpperCAmelCase__)
           audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

           A__ 					=       self.get_dummy_inputs(UpperCAmelCase__)
           A__ 					=       3 * [inputs['''prompt''']]

           # forward
           A__ 					=       audioldm_pipe(**UpperCAmelCase__)
           A__ 					=       output.audios[0]

           A__ 					=       self.get_dummy_inputs(UpperCAmelCase__)
           A__ 					=       3 * [inputs.pop('''prompt''')]

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

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

           A__ 					=       prompt_embeds

           # forward
           A__ 					=       audioldm_pipe(**UpperCAmelCase__)
           A__ 					=       output.audios[0]

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


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[int])		->Any:





           '''simple docstring'''

           A__ 					=       self.get_dummy_components()
           A__ 					=       AudioLDMPipeline(**UpperCAmelCase__)
           A__ 					=       audioldm_pipe.to(UpperCAmelCase__)
           A__ 					=       audioldm_pipe.to(UpperCAmelCase__)
           audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

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

           # forward
           A__ 					=       audioldm_pipe(**UpperCAmelCase__)
           A__ 					=       output.audios[0]

           A__ 					=       self.get_dummy_inputs(UpperCAmelCase__)
           A__ 					=       3 * [inputs.pop('''prompt''')]

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

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

                  embeds.append(UpperCAmelCase__)

           A__					,				A__ 					=       embeds

           # forward
           A__ 					=       audioldm_pipe(**UpperCAmelCase__)
           A__ 					=       output.audios[0]

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


    def 			SCREAMING_SNAKE_CASE				(     self :      Union[str, Any])		->Tuple:





           '''simple docstring'''

           A__ 					=       '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           A__ 					=       self.get_dummy_components()
           A__ 					=       PNDMScheduler(skip_prk_steps=UpperCAmelCase__)
           A__ 					=       AudioLDMPipeline(**UpperCAmelCase__)
           A__ 					=       audioldm_pipe.to(UpperCAmelCase__)
           audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

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

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

           A__ 					=       audio[:10]
           A__ 					=       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 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->Tuple:





           '''simple docstring'''

           A__ 					=       '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           A__ 					=       self.get_dummy_components()
           A__ 					=       PNDMScheduler(skip_prk_steps=UpperCAmelCase__)
           A__ 					=       AudioLDMPipeline(**UpperCAmelCase__)
           A__ 					=       audioldm_pipe.to(UpperCAmelCase__)
           audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

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

           # test num_waveforms_per_prompt=1 (default)
           A__ 					=       audioldm_pipe(UpperCAmelCase__   ,					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(UpperCAmelCase__   ,					num_inference_steps=2   ,					num_waveforms_per_prompt=UpperCAmelCase__).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=UpperCAmelCase__).audios

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


    def 			SCREAMING_SNAKE_CASE				(     self :      List[Any])		->Tuple:





           '''simple docstring'''

           A__ 					=       '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           A__ 					=       self.get_dummy_components()
           A__ 					=       AudioLDMPipeline(**UpperCAmelCase__)
           A__ 					=       audioldm_pipe.to(UpperCAmelCase__)
           audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase__)
           A__ 					=       audioldm_pipe.vocoder.config.sampling_rate

           A__ 					=       self.get_dummy_inputs(UpperCAmelCase__)
           A__ 					=       audioldm_pipe(audio_length_in_s=0.016   ,					**UpperCAmelCase__)
           A__ 					=       output.audios[0]

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

           A__ 					=       audioldm_pipe(audio_length_in_s=0.032   ,					**UpperCAmelCase__)
           A__ 					=       output.audios[0]

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


    def 			SCREAMING_SNAKE_CASE				(     self :      str)		->Dict:





           '''simple docstring'''

           A__ 					=       self.get_dummy_components()
           A__ 					=       AudioLDMPipeline(**UpperCAmelCase__)
           A__ 					=       audioldm_pipe.to(UpperCAmelCase__)
           audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

           A__ 					=       ['''hey''']

           A__ 					=       audioldm_pipe(UpperCAmelCase__   ,					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(UpperCAmelCase__).to(UpperCAmelCase__)
           A__ 					=       audioldm_pipe(UpperCAmelCase__   ,					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[Any])		->int:





           '''simple docstring'''

           self._test_attention_slicing_forward_pass(test_mean_pixel_difference=UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[int])		->List[str]:





           '''simple docstring'''

           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 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Optional[int]:





           '''simple docstring'''

           self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=UpperCAmelCase__)


@slow
class 							UpperCamelCase_       (					unittest.TestCase	):

    '''simple docstring'''






    def 			SCREAMING_SNAKE_CASE				(     self :      int)		->int:





           '''simple docstring'''

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


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str]   ,					UpperCAmelCase__ :      Tuple   ,					UpperCAmelCase__ :      int="cpu"   ,					UpperCAmelCase__ :      List[Any]=torch.floataa   ,					UpperCAmelCase__ :      str=0)		->str:





           '''simple docstring'''

           A__ 					=       torch.Generator(device=UpperCAmelCase__).manual_seed(UpperCAmelCase__)
           A__ 					=       np.random.RandomState(UpperCAmelCase__).standard_normal((1, 8, 128, 16))
           A__ 					=       torch.from_numpy(UpperCAmelCase__).to(device=UpperCAmelCase__   ,					dtype=UpperCAmelCase__)
           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 :      Dict)		->Union[str, Any]:





           '''simple docstring'''

           A__ 					=       AudioLDMPipeline.from_pretrained('''cvssp/audioldm''')
           A__ 					=       audioldm_pipe.to(UpperCAmelCase__)
           audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

           A__ 					=       self.get_inputs(UpperCAmelCase__)
           A__ 					=       25
           A__ 					=       audioldm_pipe(**UpperCAmelCase__).audios[0]

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

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





    def 			SCREAMING_SNAKE_CASE				(     self :      int)		->Optional[Any]:





           '''simple docstring'''

           A__ 					=       AudioLDMPipeline.from_pretrained('''cvssp/audioldm''')
           A__ 					=       LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config)
           A__ 					=       audioldm_pipe.to(UpperCAmelCase__)
           audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase__)

           A__ 					=       self.get_inputs(UpperCAmelCase__)
           A__ 					=       audioldm_pipe(**UpperCAmelCase__).audios[0]

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

           A__ 					=       audio[27_780:27_790]
           A__ 					=       np.array([-0.2131, -0.0873, -0.0124, -0.0189, 0.0569, 0.1373, 0.1883, 0.2886, 0.3297, 0.2212])
           A__ 					=       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.
__lowerCAmelCase        =							abspath(join(dirname(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       __UpperCamelCase					(      lowercase_       :  Dict   ):
       """simple docstring"""




       from transformers.testing_utils import pytest_addoption_shared

       pytest_addoption_shared(lowercase_   )


def       __UpperCamelCase					(      lowercase_       :  List[Any]   ):
       """simple docstring"""




       from transformers.testing_utils import pytest_terminal_summary_main

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




