'''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     __magic_name__						(    lowerCAmelCase_,     unittest.TestCase):
			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				:  str    ):
										torch.manual_seed(0    )
										lowercase_					:			Dict										=       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				,			)
										lowercase_					:			Dict										=       DDIMScheduler(
										    beta_start=0.0_00_85				,			beta_end=0.0_12				,			beta_schedule="""scaled_linear"""				,			clip_sample=__a				,			set_alpha_to_one=__a				,			)
										torch.manual_seed(0    )
										lowercase_					:			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    )
										lowercase_					:			Any										=       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				,			)
										lowercase_					:			List[str]										=       ClapTextModelWithProjection(__a    )
										lowercase_					:			Any										=       RobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-random-roberta"""				,			model_max_length=77    )

										lowercase_					:			str										=       SpeechTaHifiGanConfig(
										    model_in_dim=8				,			sampling_rate=16000				,			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				,			)

										lowercase_					:			Optional[int]										=       SpeechTaHifiGan(__a    )

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







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]				,			lowercase_				:  str				,			lowercase_				:  Any=0    ):





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







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[str]    ):
										lowercase_					:			List[str]										=       'cpu'  # ensure determinism for the device-dependent torch.Generator

										lowercase_					:			Tuple										=       self.get_dummy_components()
										lowercase_					:			List[str]										=       AudioLDMPipeline(**__a    )
										lowercase_					:			int										=       audioldm_pipe.to(__a    )
										audioldm_pipe.set_progress_bar_config(disable=__a    )

										lowercase_					:			List[Any]										=       self.get_dummy_inputs(__a    )
										lowercase_					:			Any										=       audioldm_pipe(**__a    )
										lowercase_					:			Optional[Any]										=       output.audios[0]

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

										lowercase_					:			Dict										=       audio[:10]
										lowercase_					:			Union[str, Any]										=       np.array(
										    [-0.00_50, 0.00_50, -0.00_60, 0.00_33, -0.00_26, 0.00_33, -0.00_27, 0.00_33, -0.00_28, 0.00_33]    )

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







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										lowercase_					:			str										=       self.get_dummy_components()
										lowercase_					:			Optional[Any]										=       AudioLDMPipeline(**__a    )
										lowercase_					:			Any										=       audioldm_pipe.to(__a    )
										lowercase_					:			List[Any]										=       audioldm_pipe.to(__a    )
										audioldm_pipe.set_progress_bar_config(disable=__a    )

										lowercase_					:			Any										=       self.get_dummy_inputs(__a    )
										lowercase_					:			Optional[int]										=       3 * [inputs['prompt']]

										# forward
										lowercase_					:			List[Any]										=       audioldm_pipe(**__a    )
										lowercase_					:			str										=       output.audios[0]

										lowercase_					:			Any										=       self.get_dummy_inputs(__a    )
										lowercase_					:			Optional[Any]										=       3 * [inputs.pop("""prompt"""    )]

										lowercase_					:			int										=       audioldm_pipe.tokenizer(
										    __a				,			padding="""max_length"""				,			max_length=audioldm_pipe.tokenizer.model_max_length				,			truncation=__a				,			return_tensors="""pt"""				,			)
										lowercase_					:			str										=       text_inputs['input_ids'].to(__a    )

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

										lowercase_					:			Optional[int]										=       prompt_embeds

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

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







			def 	SCREAMING_SNAKE_CASE_     (	self				:  int    ):
										lowercase_					:			Any										=       self.get_dummy_components()
										lowercase_					:			Optional[int]										=       AudioLDMPipeline(**__a    )
										lowercase_					:			Union[str, Any]										=       audioldm_pipe.to(__a    )
										lowercase_					:			Union[str, Any]										=       audioldm_pipe.to(__a    )
										audioldm_pipe.set_progress_bar_config(disable=__a    )

										lowercase_					:			str										=       self.get_dummy_inputs(__a    )
										lowercase_					:			List[str]										=       3 * ['this is a negative prompt']
										lowercase_					:			int										=       negative_prompt
										lowercase_					:			List[Any]										=       3 * [inputs['prompt']]

										# forward
										lowercase_					:			Dict										=       audioldm_pipe(**__a    )
										lowercase_					:			str										=       output.audios[0]

										lowercase_					:			List[Any]										=       self.get_dummy_inputs(__a    )
										lowercase_					:			List[str]										=       3 * [inputs.pop("""prompt"""    )]

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

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

																	embeds.append(__a    )

										lowercase_					:			Union[str, Any]										=       embeds

										# forward
										lowercase_					:			List[str]										=       audioldm_pipe(**__a    )
										lowercase_					:			Optional[Any]										=       output.audios[0]

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







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

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

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

										lowercase_					:			Tuple										=       audio[:10]
										lowercase_					:			List[Any]										=       np.array(
										    [-0.00_51, 0.00_50, -0.00_60, 0.00_34, -0.00_26, 0.00_33, -0.00_27, 0.00_33, -0.00_28, 0.00_32]    )

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







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Any    ):
										lowercase_					:			int										=       'cpu'  # ensure determinism for the device-dependent torch.Generator
										lowercase_					:			int										=       self.get_dummy_components()
										lowercase_					:			Dict										=       PNDMScheduler(skip_prk_steps=__a    )
										lowercase_					:			Any										=       AudioLDMPipeline(**__a    )
										lowercase_					:			Dict										=       audioldm_pipe.to(__a    )
										audioldm_pipe.set_progress_bar_config(disable=__a    )

										lowercase_					:			List[str]										=       'A hammer hitting a wooden surface'

										# test num_waveforms_per_prompt=1 (default)
										lowercase_					:			str										=       audioldm_pipe(__a				,			num_inference_steps=2    ).audios

										assert audios.shape == (1, 256)

										# test num_waveforms_per_prompt=1 (default) for batch of prompts
										lowercase_					:			Tuple										=       2
										lowercase_					:			Optional[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
										lowercase_					:			Tuple										=       2
										lowercase_					:			Dict										=       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
										lowercase_					:			Optional[Any]										=       2
										lowercase_					:			List[str]										=       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 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]    ):
										lowercase_					:			List[Any]										=       'cpu'  # ensure determinism for the device-dependent torch.Generator
										lowercase_					:			int										=       self.get_dummy_components()
										lowercase_					:			Optional[Any]										=       AudioLDMPipeline(**__a    )
										lowercase_					:			Any										=       audioldm_pipe.to(__a    )
										audioldm_pipe.set_progress_bar_config(disable=__a    )
										lowercase_					:			Dict										=       audioldm_pipe.vocoder.config.sampling_rate

										lowercase_					:			Union[str, Any]										=       self.get_dummy_inputs(__a    )
										lowercase_					:			str										=       audioldm_pipe(audio_length_in_s=0.0_16				,			**__a    )
										lowercase_					:			Optional[Any]										=       output.audios[0]

										assert audio.ndim == 1
										assert len(__a    ) / vocoder_sampling_rate == 0.0_16

										lowercase_					:			Tuple										=       audioldm_pipe(audio_length_in_s=0.0_32				,			**__a    )
										lowercase_					:			Optional[Any]										=       output.audios[0]

										assert audio.ndim == 1
										assert len(__a    ) / vocoder_sampling_rate == 0.0_32







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple    ):
										lowercase_					:			Optional[Any]										=       self.get_dummy_components()
										lowercase_					:			int										=       AudioLDMPipeline(**__a    )
										lowercase_					:			Optional[Any]										=       audioldm_pipe.to(__a    )
										audioldm_pipe.set_progress_bar_config(disable=__a    )

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

										lowercase_					:			Optional[Any]										=       audioldm_pipe(__a				,			num_inference_steps=1    )
										lowercase_					:			int										=       output.audios.shape
										assert audio_shape == (1, 256)

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







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Any    ):
										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 	SCREAMING_SNAKE_CASE_     (	self				:  Any    ):
										self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__a    )





@slow
class     __magic_name__						(    unittest.TestCase):







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[str]    ):
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]				,			lowercase_				:  Union[str, Any]				,			lowercase_				:  Any="cpu"				,			lowercase_				:  Union[str, Any]=torch.floataa				,			lowercase_				:  int=0    ):
										lowercase_					:			Tuple										=       torch.Generator(device=__a    ).manual_seed(__a    )
										lowercase_					:			List[str]										=       np.random.RandomState(__a    ).standard_normal((1, 8, 128, 16)    )
										lowercase_					:			Tuple										=       torch.from_numpy(__a    ).to(device=__a				,			dtype=__a    )
										lowercase_					:			Optional[Any]										=       {
										    '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				:  List[Any]    ):
										lowercase_					:			List[str]										=       AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""    )
										lowercase_					:			int										=       audioldm_pipe.to(__a    )
										audioldm_pipe.set_progress_bar_config(disable=__a    )

										lowercase_					:			Optional[Any]										=       self.get_inputs(__a    )
										lowercase_					:			Optional[Any]										=       25
										lowercase_					:			List[Any]										=       audioldm_pipe(**__a    ).audios[0]

										assert audio.ndim == 1
										assert len(__a    ) == 81920

										lowercase_					:			List[str]										=       audio[77230:77240]
										lowercase_					:			str										=       np.array(
										    [-0.48_84, -0.46_07, 0.00_23, 0.50_07, 0.58_96, 0.51_51, 0.38_13, -0.02_08, -0.36_87, -0.43_15]    )
										lowercase_					:			Dict										=       np.abs(expected_slice - audio_slice    ).max()
										assert max_diff < 1E-2







			def 	SCREAMING_SNAKE_CASE_     (	self				:  int    ):
										lowercase_					:			Any										=       AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""    )
										lowercase_					:			Optional[int]										=       LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config    )
										lowercase_					:			Optional[Any]										=       audioldm_pipe.to(__a    )
										audioldm_pipe.set_progress_bar_config(disable=__a    )

										lowercase_					:			Dict										=       self.get_inputs(__a    )
										lowercase_					:			Tuple										=       audioldm_pipe(**__a    ).audios[0]

										assert audio.ndim == 1
										assert len(__a    ) == 81920

										lowercase_					:			int										=       audio[27780:27790]
										lowercase_					:			str										=       np.array([-0.21_31, -0.08_73, -0.01_24, -0.01_89, 0.05_69, 0.13_73, 0.18_83, 0.28_86, 0.32_97, 0.22_12]    )
										lowercase_					:			Optional[Any]										=       np.abs(expected_slice - audio_slice    ).max()
										assert max_diff < 3E-2









'''simple docstring'''



from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


__lowercase	:						Dict       =	logging.get_logger(__name__)

__lowercase	:						Optional[Any]       =	{
    'google/umt5-small': 'https://huggingface.co/google/umt5-small/resolve/main/config.json',
    # See all umt5 models at https://huggingface.co/models?filter=umt5
}





class 							__UpperCamelCase				(		lowerCAmelCase_					):
	A_				=    "umt5"
	A_				=    ["past_key_values"]




	def __init__(   self						,				__a=25_0112						,				__a=512						,				__a=64						,				__a=1024						,				__a=8						,				__a=None						,				__a=6						,				__a=32						,				__a=128						,				__a=0.1						,				__a=1E-6						,				__a=1.0						,				__a="gated-gelu"						,				__a=True						,				__a=True						,				__a="T5Tokenizer"						,				__a=True						,				__a=0						,				__a=1						,				__a=0						,				**__a						,				):

						'''simple docstring'''





						super().__init__(
						    is_encoder_decoder=__a						,				tokenizer_class=__a						,				tie_word_embeddings=__a						,				pad_token_id=__a						,				eos_token_id=__a						,				decoder_start_token_id=__a						,				**__a						,				)
						__a			:       Any					    =  vocab_size
						__a			:       Any					    =  d_model
						__a			:       str					    =  d_kv
						__a			:       Dict					    =  d_ff
						__a			:       Union[str, Any]					    =  num_layers
						__a			:       int					    =  (
						    num_decoder_layers if num_decoder_layers is not None else self.num_layers
						)  # default = symmetry
						__a			:       Optional[int]					    =  num_heads
						__a			:       Tuple					    =  relative_attention_num_buckets
						__a			:       Optional[Any]					    =  relative_attention_max_distance
						__a			:       Optional[int]					    =  dropout_rate
						__a			:       List[Any]					    =  layer_norm_epsilon
						__a			:       int					    =  initializer_factor
						__a			:       Union[str, Any]					    =  feed_forward_proj
						__a			:       Any					    =  use_cache

						__a			:       List[Any]					    =  self.feed_forward_proj.split('-'							)
						__a			:       Dict					    =  act_info[-1]
						__a			:       Dict					    =  act_info[0] == 'gated'

						if len(__a							) > 1 and act_info[0] != "gated" or len(__a							) > 2:
											raise ValueError(
											    f"""`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."""
											    'Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. '
											    '\'gated-gelu\' or \'relu\''							)

						if feed_forward_proj == "gated-gelu":
											__a			:       Optional[int]					    =  'gelu_new'




	@property
	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





						return self.d_model




	@property
	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





						return self.num_heads





	@property
	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





						return self.num_layers





class 							__UpperCamelCase				(		lowerCAmelCase_					):




	@property
	# Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.inputs
	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





						__a			:       str					    =  {
						    'input_ids': {0: 'batch', 1: 'encoder_sequence'},
						    'attention_mask': {0: 'batch', 1: 'encoder_sequence'},
						}
						if self.use_past:
											__a			:       Dict					    =  'past_encoder_sequence + sequence'
											__a			:       Tuple					    =  {0: 'batch'}
											__a			:       Tuple					    =  {0: 'batch', 1: 'past_decoder_sequence + sequence'}
						else:
											__a			:       List[Any]					    =  {0: 'batch', 1: 'decoder_sequence'}
											__a			:       int					    =  {0: 'batch', 1: 'decoder_sequence'}

						if self.use_past:
											self.fill_with_past_key_values_(__a						,				direction='inputs'							)

						return common_inputs




	@property
	# Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.default_onnx_opset
	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





						return 13





	@property
	def       __UpperCAmelCase							(   self							):

						'''simple docstring'''





						return 5E-4



