






"""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_       (UpperCamelCase__				,						unittest.TestCase     ):
	__magic_name__	      =				AudioLDMPipeline
	__magic_name__	      =				TEXT_TO_AUDIO_PARAMS
	__magic_name__	      =				TEXT_TO_AUDIO_BATCH_PARAMS
	__magic_name__	      =				frozenset(
	    [
	        '''num_inference_steps''',
	        '''num_waveforms_per_prompt''',
	        '''generator''',
	        '''latents''',
	        '''output_type''',
	        '''return_dict''',
	        '''callback''',
	        '''callback_steps''',
	    ]     )





	def      _SCREAMING_SNAKE_CASE      (      self					: List[Any]						)						->			Any:
				torch.manual_seed(0						)
				UpperCAmelCase_     : 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=lowerCAmelCase__	,	)
				UpperCAmelCase_     : Optional[int]			     =					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						)
				UpperCAmelCase_     : 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						)
				UpperCAmelCase_     : int			     =					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	,	)
				UpperCAmelCase_     : str			     =					ClapTextModelWithProjection(lowerCAmelCase__						)
				UpperCAmelCase_     : int			     =					RobertaTokenizer.from_pretrained("hf-internal-testing/tiny-random-roberta"	,	model_max_length=77						)

				UpperCAmelCase_     : Any			     =					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=lowerCAmelCase__	,	)

				UpperCAmelCase_     : Dict			     =					SpeechTaHifiGan(lowerCAmelCase__						)

				UpperCAmelCase_     : Dict			     =					{
				    "unet": unet,
				    "scheduler": scheduler,
				    "vae": vae,
				    "text_encoder": text_encoder,
				    "tokenizer": tokenizer,
				    "vocoder": vocoder,
				}
				return components





	def      _SCREAMING_SNAKE_CASE      (      self					: List[Any]	,	lowerCAmelCase_					: Optional[int]	,	lowerCAmelCase_					: str=0						)						->			Optional[int]:




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





	def      _SCREAMING_SNAKE_CASE      (      self					: Any						)						->			List[Any]:
				UpperCAmelCase_     : Tuple			     =					"cpu"  # ensure determinism for the device-dependent torch.Generator

				UpperCAmelCase_     : List[Any]			     =					self.get_dummy_components()
				UpperCAmelCase_     : List[Any]			     =					AudioLDMPipeline(**lowerCAmelCase__						)
				UpperCAmelCase_     : Tuple			     =					audioldm_pipe.to(lowerCAmelCase__						)
				audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase__						)

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

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

				UpperCAmelCase_     : Optional[Any]			     =					audio[:10]
				UpperCAmelCase_     : Optional[int]			     =					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      _SCREAMING_SNAKE_CASE      (      self					: Optional[int]						)						->			int:
				UpperCAmelCase_     : List[str]			     =					self.get_dummy_components()
				UpperCAmelCase_     : Union[str, Any]			     =					AudioLDMPipeline(**lowerCAmelCase__						)
				UpperCAmelCase_     : Tuple			     =					audioldm_pipe.to(lowerCAmelCase__						)
				UpperCAmelCase_     : List[str]			     =					audioldm_pipe.to(lowerCAmelCase__						)
				audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase__						)

				UpperCAmelCase_     : Optional[int]			     =					self.get_dummy_inputs(lowerCAmelCase__						)
				UpperCAmelCase_     : Tuple			     =					3 * [inputs["prompt"]]

				# forward
				UpperCAmelCase_     : List[Any]			     =					audioldm_pipe(**lowerCAmelCase__						)
				UpperCAmelCase_     : Optional[int]			     =					output.audios[0]

				UpperCAmelCase_     : int			     =					self.get_dummy_inputs(lowerCAmelCase__						)
				UpperCAmelCase_     : List[Any]			     =					3 * [inputs.pop("prompt"						)]

				UpperCAmelCase_     : str			     =					audioldm_pipe.tokenizer(
				    lowerCAmelCase__	,	padding="max_length"	,	max_length=audioldm_pipe.tokenizer.model_max_length	,	truncation=lowerCAmelCase__	,	return_tensors="pt"	,	)
				UpperCAmelCase_     : Any			     =					text_inputs["input_ids"].to(lowerCAmelCase__						)

				UpperCAmelCase_     : int			     =					audioldm_pipe.text_encoder(
				    lowerCAmelCase__	,	)
				UpperCAmelCase_     : Any			     =					prompt_embeds.text_embeds
				# additional L_2 normalization over each hidden-state
				UpperCAmelCase_     : Any			     =					F.normalize(lowerCAmelCase__	,	dim=-1						)

				UpperCAmelCase_     : List[Any]			     =					prompt_embeds

				# forward
				UpperCAmelCase_     : int			     =					audioldm_pipe(**lowerCAmelCase__						)
				UpperCAmelCase_     : Optional[Any]			     =					output.audios[0]

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





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[Any]						)						->			Tuple:
				UpperCAmelCase_     : Any			     =					self.get_dummy_components()
				UpperCAmelCase_     : Union[str, Any]			     =					AudioLDMPipeline(**lowerCAmelCase__						)
				UpperCAmelCase_     : Optional[Any]			     =					audioldm_pipe.to(lowerCAmelCase__						)
				UpperCAmelCase_     : str			     =					audioldm_pipe.to(lowerCAmelCase__						)
				audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase__						)

				UpperCAmelCase_     : Any			     =					self.get_dummy_inputs(lowerCAmelCase__						)
				UpperCAmelCase_     : Optional[int]			     =					3 * ["this is a negative prompt"]
				UpperCAmelCase_     : int			     =					negative_prompt
				UpperCAmelCase_     : List[str]			     =					3 * [inputs["prompt"]]

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

				UpperCAmelCase_     : int			     =					self.get_dummy_inputs(lowerCAmelCase__						)
				UpperCAmelCase_     : Dict			     =					3 * [inputs.pop("prompt"						)]

				UpperCAmelCase_     : Tuple			     =					[]
				for p in [prompt, negative_prompt]:
							UpperCAmelCase_     : Any			     =					audioldm_pipe.tokenizer(
							    lowerCAmelCase__	,	padding="max_length"	,	max_length=audioldm_pipe.tokenizer.model_max_length	,	truncation=lowerCAmelCase__	,	return_tensors="pt"	,	)
							UpperCAmelCase_     : Dict			     =					text_inputs["input_ids"].to(lowerCAmelCase__						)

							UpperCAmelCase_     : int			     =					audioldm_pipe.text_encoder(
							    lowerCAmelCase__	,	)
							UpperCAmelCase_     : str			     =					text_embeds.text_embeds
							# additional L_2 normalization over each hidden-state
							UpperCAmelCase_     : int			     =					F.normalize(lowerCAmelCase__	,	dim=-1						)

							embeds.append(lowerCAmelCase__						)

				UpperCAmelCase_     : str			     =					embeds

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

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





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[int]						)						->			Any:
				UpperCAmelCase_     : Tuple			     =					"cpu"  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase_     : Tuple			     =					self.get_dummy_components()
				UpperCAmelCase_     : str			     =					PNDMScheduler(skip_prk_steps=lowerCAmelCase__						)
				UpperCAmelCase_     : List[Any]			     =					AudioLDMPipeline(**lowerCAmelCase__						)
				UpperCAmelCase_     : Tuple			     =					audioldm_pipe.to(lowerCAmelCase__						)
				audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase__						)

				UpperCAmelCase_     : List[Any]			     =					self.get_dummy_inputs(lowerCAmelCase__						)
				UpperCAmelCase_     : int			     =					"egg cracking"
				UpperCAmelCase_     : str			     =					audioldm_pipe(**lowerCAmelCase__	,	negative_prompt=lowerCAmelCase__						)
				UpperCAmelCase_     : List[str]			     =					output.audios[0]

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

				UpperCAmelCase_     : Optional[Any]			     =					audio[:10]
				UpperCAmelCase_     : List[str]			     =					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      _SCREAMING_SNAKE_CASE      (      self					: Optional[Any]						)						->			str:
				UpperCAmelCase_     : Union[str, Any]			     =					"cpu"  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase_     : List[Any]			     =					self.get_dummy_components()
				UpperCAmelCase_     : str			     =					PNDMScheduler(skip_prk_steps=lowerCAmelCase__						)
				UpperCAmelCase_     : str			     =					AudioLDMPipeline(**lowerCAmelCase__						)
				UpperCAmelCase_     : Any			     =					audioldm_pipe.to(lowerCAmelCase__						)
				audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase__						)

				UpperCAmelCase_     : int			     =					"A hammer hitting a wooden surface"

				# test num_waveforms_per_prompt=1 (default)
				UpperCAmelCase_     : int			     =					audioldm_pipe(lowerCAmelCase__	,	num_inference_steps=2						).audios

				assert audios.shape == (1, 256)

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

				UpperCAmelCase_     : str			     =					self.get_dummy_inputs(lowerCAmelCase__						)
				UpperCAmelCase_     : Union[str, Any]			     =					audioldm_pipe(audio_length_in_s=0.0_1_6	,	**lowerCAmelCase__						)
				UpperCAmelCase_     : Optional[int]			     =					output.audios[0]

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

				UpperCAmelCase_     : Dict			     =					audioldm_pipe(audio_length_in_s=0.0_3_2	,	**lowerCAmelCase__						)
				UpperCAmelCase_     : List[str]			     =					output.audios[0]

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





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[int]						)						->			Union[str, Any]:
				UpperCAmelCase_     : List[str]			     =					self.get_dummy_components()
				UpperCAmelCase_     : Tuple			     =					AudioLDMPipeline(**lowerCAmelCase__						)
				UpperCAmelCase_     : List[str]			     =					audioldm_pipe.to(lowerCAmelCase__						)
				audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase__						)

				UpperCAmelCase_     : str			     =					["hey"]

				UpperCAmelCase_     : int			     =					audioldm_pipe(lowerCAmelCase__	,	num_inference_steps=1						)
				UpperCAmelCase_     : str			     =					output.audios.shape
				assert audio_shape == (1, 256)

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





	def      _SCREAMING_SNAKE_CASE      (      self					: int						)						->			Union[str, Any]:
				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      _SCREAMING_SNAKE_CASE      (      self					: Union[str, Any]						)						->			str:
				self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowerCAmelCase__						)







@slow
class 			UpperCamelCase_       (unittest.TestCase     ):





	def      _SCREAMING_SNAKE_CASE      (      self					: Optional[Any]						)						->			Optional[int]:
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()





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

				UpperCAmelCase_     : Any			     =					self.get_inputs(lowerCAmelCase__						)
				UpperCAmelCase_     : Optional[Any]			     =					25
				UpperCAmelCase_     : Union[str, Any]			     =					audioldm_pipe(**lowerCAmelCase__						).audios[0]

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

				UpperCAmelCase_     : Union[str, Any]			     =					audio[77_230:77_240]
				UpperCAmelCase_     : Union[str, 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]						)
				UpperCAmelCase_     : Optional[Any]			     =					np.abs(expected_slice - audio_slice						).max()
				assert max_diff < 1e-2





	def      _SCREAMING_SNAKE_CASE      (      self					: List[Any]						)						->			Dict:
				UpperCAmelCase_     : int			     =					AudioLDMPipeline.from_pretrained("cvssp/audioldm"						)
				UpperCAmelCase_     : Optional[int]			     =					LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config						)
				UpperCAmelCase_     : str			     =					audioldm_pipe.to(lowerCAmelCase__						)
				audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase__						)

				UpperCAmelCase_     : Any			     =					self.get_inputs(lowerCAmelCase__						)
				UpperCAmelCase_     : List[Any]			     =					audioldm_pipe(**lowerCAmelCase__						).audios[0]

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

				UpperCAmelCase_     : int			     =					audio[27_780:27_790]
				UpperCAmelCase_     : Optional[int]			     =					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]						)
				UpperCAmelCase_     : Any			     =					np.abs(expected_slice - audio_slice						).max()
				assert max_diff < 3e-2










import torch
from transformers import PreTrainedModel, XLMRobertaConfig, XLMRobertaModel







class 				__lowerCAmelCase       (    UpperCamelCase__):
      _lowercase				:					Tuple						=			"""M-CLIP"""







      def __init__(    self		,      lowerCAmelCase__=1_0_2_4		,      lowerCAmelCase__=7_6_8		,      **lowerCAmelCase__							)   ->						Any:



        '''simple docstring'''




        a__	:					int										=transformerDimSize
        a__	:					Dict										=imageDimSize
        super().__init__(**lowerCAmelCase__							)







class 				__lowerCAmelCase       (    UpperCamelCase__):
      _lowercase				:					Optional[Any]						=			MCLIPConfig







      def __init__(    self		,      lowerCAmelCase__		,      *lowerCAmelCase__		,      **lowerCAmelCase__							)   ->						List[str]:



        '''simple docstring'''




        super().__init__(lowerCAmelCase__		,      *lowerCAmelCase__		,      **lowerCAmelCase__							)
        a__	:					Tuple										=XLMRobertaModel(lowerCAmelCase__							)
        a__	:					List[str]										=torch.nn.Linear(
            in_features=config.transformerDimensions		,      out_features=config.numDims							)







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					Optional[Any]										=self.transformer(input_ids=lowerCAmelCase__		,      attention_mask=lowerCAmelCase__							)[0]
        a__	:					int										=(embs * attention_mask.unsqueeze(2							)).sum(dim=1							) / attention_mask.sum(dim=1							)[:, None]
        return self.LinearTransformation(lowerCAmelCase__							), embs






