


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






		'''simple docstring'''

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			torch.manual_seed(0					)
			__snake_case	:						List[Any]						=				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				,       )
			__snake_case	:						Tuple						=				DDIMScheduler(
			    beta_start=0.0_0085				,       beta_end=0.012				,       beta_schedule='''scaled_linear'''				,       clip_sample=_UpperCamelCase				,       set_alpha_to_one=_UpperCamelCase				,       )
			torch.manual_seed(0					)
			__snake_case	:						Optional[int]						=				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=10_00				,       projection_dim=32				,       )
			__snake_case	:						str						=				ClapTextModelWithProjection(_UpperCamelCase					)
			__snake_case	:						str						=				RobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-random-roberta'''				,       model_max_length=77					)

			__snake_case	:						List[Any]						=				SpeechTaHifiGanConfig(
			    model_in_dim=8				,       sampling_rate=1_60_00				,       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				,       )

			__snake_case	:						Tuple						=				SpeechTaHifiGan(_UpperCamelCase					)

			__snake_case	:						List[str]						=				{
			    '''unet''': unet,
			    '''scheduler''': scheduler,
			    '''vae''': vae,
			    '''text_encoder''': text_encoder,
			    '''tokenizer''': tokenizer,
			    '''vocoder''': vocoder,
			}
			return components



		def 				SCREAMING_SNAKE_CASE  (self				,       a_				,       a_=0					):





			'''simple docstring'''

			if str(_UpperCamelCase					).startswith('''mps'''					):
				__snake_case	:						int						=				torch.manual_seed(_UpperCamelCase					)
			else:
				__snake_case	:						Tuple						=				torch.Generator(device=_UpperCamelCase					).manual_seed(_UpperCamelCase					)
			__snake_case	:						Dict						=				{
			    '''prompt''': '''A hammer hitting a wooden surface''',
			    '''generator''': generator,
			    '''num_inference_steps''': 2,
			    '''guidance_scale''': 6.0,
			}
			return inputs



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						str						=				'''cpu'''  # ensure determinism for the device-dependent torch.Generator

			__snake_case	:						List[Any]						=				self.get_dummy_components()
			__snake_case	:						int						=				AudioLDMPipeline(**_UpperCamelCase					)
			__snake_case	:						Optional[Any]						=				audioldm_pipe.to(_UpperCamelCase					)
			audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase					)

			__snake_case	:						Optional[int]						=				self.get_dummy_inputs(_UpperCamelCase					)
			__snake_case	:						int						=				audioldm_pipe(**_UpperCamelCase					)
			__snake_case	:						Union[str, Any]						=				output.audios[0]

			assert audio.ndim == 1
			assert len(_UpperCamelCase					) == 2_56

			__snake_case	:						str						=				audio[:10]
			__snake_case	:						int						=				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					):





			'''simple docstring'''

			__snake_case	:						Optional[Any]						=				self.get_dummy_components()
			__snake_case	:						Dict						=				AudioLDMPipeline(**_UpperCamelCase					)
			__snake_case	:						List[Any]						=				audioldm_pipe.to(_UpperCamelCase					)
			__snake_case	:						Tuple						=				audioldm_pipe.to(_UpperCamelCase					)
			audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase					)

			__snake_case	:						List[Any]						=				self.get_dummy_inputs(_UpperCamelCase					)
			__snake_case	:						Tuple						=				3 * [inputs['''prompt''']]

			# forward
			__snake_case	:						str						=				audioldm_pipe(**_UpperCamelCase					)
			__snake_case	:						Optional[Any]						=				output.audios[0]

			__snake_case	:						List[str]						=				self.get_dummy_inputs(_UpperCamelCase					)
			__snake_case	:						List[str]						=				3 * [inputs.pop('''prompt'''					)]

			__snake_case	:						List[Any]						=				audioldm_pipe.tokenizer(
			    _UpperCamelCase				,       padding='''max_length'''				,       max_length=audioldm_pipe.tokenizer.model_max_length				,       truncation=_UpperCamelCase				,       return_tensors='''pt'''				,       )
			__snake_case	:						int						=				text_inputs['''input_ids'''].to(_UpperCamelCase					)

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

			__snake_case	:						Tuple						=				prompt_embeds

			# forward
			__snake_case	:						int						=				audioldm_pipe(**_UpperCamelCase					)
			__snake_case	:						int						=				output.audios[0]

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

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

			__snake_case	:						str						=				self.get_dummy_inputs(_UpperCamelCase					)
			__snake_case	:						Union[str, Any]						=				3 * ['''this is a negative prompt''']
			__snake_case	:						str						=				negative_prompt
			__snake_case	:						Optional[int]						=				3 * [inputs['''prompt''']]

			# forward
			__snake_case	:						Optional[int]						=				audioldm_pipe(**_UpperCamelCase					)
			__snake_case	:						int						=				output.audios[0]

			__snake_case	:						List[str]						=				self.get_dummy_inputs(_UpperCamelCase					)
			__snake_case	:						Tuple						=				3 * [inputs.pop('''prompt'''					)]

			__snake_case	:						List[str]						=				[]
			for p in [prompt, negative_prompt]:
				__snake_case	:						Optional[int]						=				audioldm_pipe.tokenizer(
				    _UpperCamelCase				,       padding='''max_length'''				,       max_length=audioldm_pipe.tokenizer.model_max_length				,       truncation=_UpperCamelCase				,       return_tensors='''pt'''				,       )
				__snake_case	:						List[str]						=				text_inputs['''input_ids'''].to(_UpperCamelCase					)

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

				embeds.append(_UpperCamelCase					)

			__snake_case	,							__snake_case	:						List[Any]						=				embeds

			# forward
			__snake_case	:						Dict						=				audioldm_pipe(**_UpperCamelCase					)
			__snake_case	:						Dict						=				output.audios[0]

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

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

			__snake_case	:						Optional[int]						=				self.get_dummy_inputs(_UpperCamelCase					)
			__snake_case	:						List[Any]						=				'''egg cracking'''
			__snake_case	:						List[Any]						=				audioldm_pipe(**_UpperCamelCase				,       negative_prompt=_UpperCamelCase					)
			__snake_case	:						Any						=				output.audios[0]

			assert audio.ndim == 1
			assert len(_UpperCamelCase					) == 2_56

			__snake_case	:						List[str]						=				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 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

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

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

			# test num_waveforms_per_prompt=1 (default)
			__snake_case	:						List[Any]						=				audioldm_pipe(_UpperCamelCase				,       num_inference_steps=2					).audios

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

			# test num_waveforms_per_prompt=1 (default) for batch of prompts
			__snake_case	:						Optional[int]						=				2
			__snake_case	:						Union[str, Any]						=				audioldm_pipe([prompt] * batch_size				,       num_inference_steps=2					).audios

			assert audios.shape == (batch_size, 2_56)

			# test num_waveforms_per_prompt for single prompt
			__snake_case	:						Tuple						=				2
			__snake_case	:						Optional[Any]						=				audioldm_pipe(_UpperCamelCase				,       num_inference_steps=2				,       num_waveforms_per_prompt=_UpperCamelCase					).audios

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

			# test num_waveforms_per_prompt for batch of prompts
			__snake_case	:						Any						=				2
			__snake_case	:						Optional[int]						=				audioldm_pipe(
			    [prompt] * batch_size				,       num_inference_steps=2				,       num_waveforms_per_prompt=_UpperCamelCase					).audios

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

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

			__snake_case	:						Optional[Any]						=				self.get_dummy_inputs(_UpperCamelCase					)
			__snake_case	:						str						=				audioldm_pipe(audio_length_in_s=0.016				,       **_UpperCamelCase					)
			__snake_case	:						Optional[Any]						=				output.audios[0]

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

			__snake_case	:						Optional[Any]						=				audioldm_pipe(audio_length_in_s=0.032				,       **_UpperCamelCase					)
			__snake_case	:						Any						=				output.audios[0]

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



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						List[str]						=				self.get_dummy_components()
			__snake_case	:						str						=				AudioLDMPipeline(**_UpperCamelCase					)
			__snake_case	:						Any						=				audioldm_pipe.to(_UpperCamelCase					)
			audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase					)

			__snake_case	:						Optional[int]						=				['''hey''']

			__snake_case	:						Optional[Any]						=				audioldm_pipe(_UpperCamelCase				,       num_inference_steps=1					)
			__snake_case	:						Optional[Any]						=				output.audios.shape
			assert audio_shape == (1, 2_56)

			__snake_case	:						Optional[int]						=				audioldm_pipe.vocoder.config
			config.model_in_dim *= 2
			__snake_case	:						Optional[Any]						=				SpeechTaHifiGan(_UpperCamelCase					).to(_UpperCamelCase					)
			__snake_case	:						str						=				audioldm_pipe(_UpperCamelCase				,       num_inference_steps=1					)
			__snake_case	:						List[str]						=				output.audios.shape
			# waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
			assert audio_shape == (1, 2_56)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			self._test_attention_slicing_forward_pass(test_mean_pixel_difference=_UpperCamelCase					)



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''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					):





			'''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					):





			'''simple docstring'''

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



		def 				SCREAMING_SNAKE_CASE  (self				,       a_				,       a_="cpu"				,       a_=torch.floataa				,       a_=0					):





			'''simple docstring'''

			__snake_case	:						List[str]						=				torch.Generator(device=_UpperCamelCase					).manual_seed(_UpperCamelCase					)
			__snake_case	:						Optional[int]						=				np.random.RandomState(_UpperCamelCase					).standard_normal((1, 8, 1_28, 16)					)
			__snake_case	:						int						=				torch.from_numpy(_UpperCamelCase					).to(device=_UpperCamelCase				,       dtype=_UpperCamelCase					)
			__snake_case	:						Tuple						=				{
			    '''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					):





			'''simple docstring'''

			__snake_case	:						int						=				AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''					)
			__snake_case	:						Optional[Any]						=				audioldm_pipe.to(_UpperCamelCase					)
			audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase					)

			__snake_case	:						Dict						=				self.get_inputs(_UpperCamelCase					)
			__snake_case	:						Dict						=				25
			__snake_case	:						str						=				audioldm_pipe(**_UpperCamelCase					).audios[0]

			assert audio.ndim == 1
			assert len(_UpperCamelCase					) == 8_19_20

			__snake_case	:						List[str]						=				audio[7_72_30:7_72_40]
			__snake_case	:						Any						=				np.array(
			    [-0.4884, -0.4607, 0.0023, 0.5007, 0.5896, 0.5151, 0.3813, -0.0208, -0.3687, -0.4315]					)
			__snake_case	:						Optional[int]						=				np.abs(expected_slice - audio_slice					).max()
			assert max_diff < 1E-2



		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

			__snake_case	:						Optional[Any]						=				AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''					)
			__snake_case	:						List[str]						=				LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config					)
			__snake_case	:						Dict						=				audioldm_pipe.to(_UpperCamelCase					)
			audioldm_pipe.set_progress_bar_config(disable=_UpperCamelCase					)

			__snake_case	:						int						=				self.get_inputs(_UpperCamelCase					)
			__snake_case	:						Tuple						=				audioldm_pipe(**_UpperCamelCase					).audios[0]

			assert audio.ndim == 1
			assert len(_UpperCamelCase					) == 8_19_20

			__snake_case	:						Tuple						=				audio[2_77_80:2_77_90]
			__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	:						List[str]						=				np.abs(expected_slice - audio_slice					).max()
			assert max_diff < 3E-2








from typing import Optional

import pyspark

from .. import Features, NamedSplit
from ..download import DownloadMode
from ..packaged_modules.spark.spark import Spark
from .abc import AbstractDatasetReader





class 	snake_case_   (		__A					):







			'''simple docstring'''

			def __init__(     self :				int							,   _UpperCamelCase :				pyspark.sql.DataFrame							,   _UpperCamelCase :				Optional[NamedSplit] = None							,   _UpperCamelCase :				Optional[Features] = None							,   _UpperCamelCase :				bool = True							,   _UpperCamelCase :				str = None							,   _UpperCamelCase :				bool = False							,   _UpperCamelCase :				str = None							,   _UpperCamelCase :				bool = True							,   _UpperCamelCase :				str = "arrow"							,   **_UpperCamelCase :				Tuple							,   )      ->str:
									super().__init__(
									    split=_UpperCamelCase							,   features=_UpperCamelCase							,   cache_dir=_UpperCamelCase							,   keep_in_memory=_UpperCamelCase							,   streaming=_UpperCamelCase							,   **_UpperCamelCase							,   )
									snake_case_   =	load_from_cache_file
									snake_case_   =	file_format
									snake_case_   =	Spark(
									    df=_UpperCamelCase							,   features=_UpperCamelCase							,   cache_dir=_UpperCamelCase							,   working_dir=_UpperCamelCase							,   **_UpperCamelCase							,   )







			def 				snake_case__(     self :				int							)      ->Tuple:
									if self.streaming:
															return self.builder.as_streaming_dataset(split=self.split							)
									snake_case_   =	None if self._load_from_cache_file else DownloadMode.FORCE_REDOWNLOAD
									self.builder.download_and_prepare(
									    download_mode=_UpperCamelCase							,   file_format=self._file_format							,   )
									return self.builder.as_dataset(split=self.split							)