


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




			def 	lowerCamelCase_       (    self				:		Optional[Any]    ):







						"""simple docstring"""







						torch.manual_seed(0    )
						UpperCamelCase								=       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=__lowercase							,							)
						UpperCamelCase								=       DDIMScheduler(
						    beta_start=0.00_085							,							beta_end=0.012							,							beta_schedule="""scaled_linear"""							,							clip_sample=__lowercase							,							set_alpha_to_one=__lowercase							,							)
						torch.manual_seed(0    )
						UpperCamelCase								=       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    )
						UpperCamelCase								=       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							,							)
						UpperCamelCase								=       ClapTextModelWithProjection(__lowercase    )
						UpperCamelCase								=       RobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-random-roberta"""							,							model_max_length=7_7    )

						UpperCamelCase								=       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=__lowercase							,							)

						UpperCamelCase								=       SpeechTaHifiGan(__lowercase    )

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




			def 	lowerCamelCase_       (    self				:		Any							,							__magic_name__				:		Union[str, Any]							,							__magic_name__				:		Any=0    ):







						"""simple docstring"""







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




			def 	lowerCamelCase_       (    self				:		Optional[Any]    ):







						"""simple docstring"""







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

						UpperCamelCase								=       self.get_dummy_components()
						UpperCamelCase								=       AudioLDMPipeline(**__lowercase    )
						UpperCamelCase								=       audioldm_pipe.to(__lowercase    )
						audioldm_pipe.set_progress_bar_config(disable=__lowercase    )

						UpperCamelCase								=       self.get_dummy_inputs(__lowercase    )
						UpperCamelCase								=       audioldm_pipe(**__lowercase    )
						UpperCamelCase								=       output.audios[0]

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

						UpperCamelCase								=       audio[:1_0]
						UpperCamelCase								=       np.array(
						    [-0.0_050, 0.0_050, -0.0_060, 0.0_033, -0.0_026, 0.0_033, -0.0_027, 0.0_033, -0.0_028, 0.0_033]    )

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




			def 	lowerCamelCase_       (    self				:		Dict    ):







						"""simple docstring"""







						UpperCamelCase								=       self.get_dummy_components()
						UpperCamelCase								=       AudioLDMPipeline(**__lowercase    )
						UpperCamelCase								=       audioldm_pipe.to(__lowercase    )
						UpperCamelCase								=       audioldm_pipe.to(__lowercase    )
						audioldm_pipe.set_progress_bar_config(disable=__lowercase    )

						UpperCamelCase								=       self.get_dummy_inputs(__lowercase    )
						UpperCamelCase								=       3 * [inputs['prompt']]

						# forward
						UpperCamelCase								=       audioldm_pipe(**__lowercase    )
						UpperCamelCase								=       output.audios[0]

						UpperCamelCase								=       self.get_dummy_inputs(__lowercase    )
						UpperCamelCase								=       3 * [inputs.pop("""prompt"""    )]

						UpperCamelCase								=       audioldm_pipe.tokenizer(
						    __lowercase							,							padding="""max_length"""							,							max_length=audioldm_pipe.tokenizer.model_max_length							,							truncation=__lowercase							,							return_tensors="""pt"""							,							)
						UpperCamelCase								=       text_inputs['input_ids'].to(__lowercase    )

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

						UpperCamelCase								=       prompt_embeds

						# forward
						UpperCamelCase								=       audioldm_pipe(**__lowercase    )
						UpperCamelCase								=       output.audios[0]

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




			def 	lowerCamelCase_       (    self				:		int    ):







						"""simple docstring"""







						UpperCamelCase								=       self.get_dummy_components()
						UpperCamelCase								=       AudioLDMPipeline(**__lowercase    )
						UpperCamelCase								=       audioldm_pipe.to(__lowercase    )
						UpperCamelCase								=       audioldm_pipe.to(__lowercase    )
						audioldm_pipe.set_progress_bar_config(disable=__lowercase    )

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

						# forward
						UpperCamelCase								=       audioldm_pipe(**__lowercase    )
						UpperCamelCase								=       output.audios[0]

						UpperCamelCase								=       self.get_dummy_inputs(__lowercase    )
						UpperCamelCase								=       3 * [inputs.pop("""prompt"""    )]

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

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

									embeds.append(__lowercase    )

						UpperCamelCase								=       embeds

						# forward
						UpperCamelCase								=       audioldm_pipe(**__lowercase    )
						UpperCamelCase								=       output.audios[0]

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




			def 	lowerCamelCase_       (    self				:		Union[str, Any]    ):







						"""simple docstring"""







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

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

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

						UpperCamelCase								=       audio[:1_0]
						UpperCamelCase								=       np.array(
						    [-0.0_051, 0.0_050, -0.0_060, 0.0_034, -0.0_026, 0.0_033, -0.0_027, 0.0_033, -0.0_028, 0.0_032]    )

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




			def 	lowerCamelCase_       (    self				:		str    ):







						"""simple docstring"""







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

						UpperCamelCase								=       'A hammer hitting a wooden surface'

						# test num_waveforms_per_prompt=1 (default)
						UpperCamelCase								=       audioldm_pipe(__lowercase							,							num_inference_steps=2    ).audios

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

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

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

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

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




			def 	lowerCamelCase_       (    self				:		Any    ):







						"""simple docstring"""







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

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

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

						UpperCamelCase								=       audioldm_pipe(audio_length_in_s=0.032							,							**__lowercase    )
						UpperCamelCase								=       output.audios[0]

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




			def 	lowerCamelCase_       (    self				:		Union[str, Any]    ):







						"""simple docstring"""







						UpperCamelCase								=       self.get_dummy_components()
						UpperCamelCase								=       AudioLDMPipeline(**__lowercase    )
						UpperCamelCase								=       audioldm_pipe.to(__lowercase    )
						audioldm_pipe.set_progress_bar_config(disable=__lowercase    )

						UpperCamelCase								=       ['hey']

						UpperCamelCase								=       audioldm_pipe(__lowercase							,							num_inference_steps=1    )
						UpperCamelCase								=       output.audios.shape
						assert audio_shape == (1, 2_5_6)

						UpperCamelCase								=       audioldm_pipe.vocoder.config
						config.model_in_dim *= 2
						UpperCamelCase								=       SpeechTaHifiGan(__lowercase    ).to(__lowercase    )
						UpperCamelCase								=       audioldm_pipe(__lowercase							,							num_inference_steps=1    )
						UpperCamelCase								=       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 	lowerCamelCase_       (    self				:		Tuple    ):







						"""simple docstring"""







						self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__lowercase    )




			def 	lowerCamelCase_       (    self				:		str    ):







						"""simple docstring"""







						self._test_inference_batch_single_identical(test_mean_pixel_difference=__lowercase    )





			@unittest.skipIf(
			    torch_device != """cuda""" or not is_xformers_available()							,							reason="""XFormers attention is only available with CUDA and `xformers` installed"""							,							)
			def 	lowerCamelCase_       (    self				:		Union[str, Any]    ):







						"""simple docstring"""







						self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__lowercase    )




@slow
class  UpperCAmelCase       (							unittest.TestCase     ):




			def 	lowerCamelCase_       (    self				:		Dict    ):







						"""simple docstring"""







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




			def 	lowerCamelCase_       (    self				:		Tuple							,							__magic_name__				:		str							,							__magic_name__				:		Union[str, Any]="cpu"							,							__magic_name__				:		List[Any]=torch.floataa							,							__magic_name__				:		Tuple=0    ):







						"""simple docstring"""







						UpperCamelCase								=       torch.Generator(device=__lowercase    ).manual_seed(__lowercase    )
						UpperCamelCase								=       np.random.RandomState(__lowercase    ).standard_normal((1, 8, 1_2_8, 1_6)    )
						UpperCamelCase								=       torch.from_numpy(__lowercase    ).to(device=__lowercase							,							dtype=__lowercase    )
						UpperCamelCase								=       {
						    'prompt': 'A hammer hitting a wooden surface',
						    'latents': latents,
						    'generator': generator,
						    'num_inference_steps': 3,
						    'guidance_scale': 2.5,
						}
						return inputs




			def 	lowerCamelCase_       (    self				:		str    ):







						"""simple docstring"""







						UpperCamelCase								=       AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""    )
						UpperCamelCase								=       audioldm_pipe.to(__lowercase    )
						audioldm_pipe.set_progress_bar_config(disable=__lowercase    )

						UpperCamelCase								=       self.get_inputs(__lowercase    )
						UpperCamelCase								=       2_5
						UpperCamelCase								=       audioldm_pipe(**__lowercase    ).audios[0]

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

						UpperCamelCase								=       audio[7_7_2_3_0:7_7_2_4_0]
						UpperCamelCase								=       np.array(
						    [-0.4_884, -0.4_607, 0.0_023, 0.5_007, 0.5_896, 0.5_151, 0.3_813, -0.0_208, -0.3_687, -0.4_315]    )
						UpperCamelCase								=       np.abs(expected_slice - audio_slice    ).max()
						assert max_diff < 1e-2





			def 	lowerCamelCase_       (    self				:		Any    ):







						"""simple docstring"""







						UpperCamelCase								=       AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""    )
						UpperCamelCase								=       LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config    )
						UpperCamelCase								=       audioldm_pipe.to(__lowercase    )
						audioldm_pipe.set_progress_bar_config(disable=__lowercase    )

						UpperCamelCase								=       self.get_inputs(__lowercase    )
						UpperCamelCase								=       audioldm_pipe(**__lowercase    ).audios[0]

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

						UpperCamelCase								=       audio[2_7_7_8_0:2_7_7_9_0]
						UpperCamelCase								=       np.array([-0.2_131, -0.0_873, -0.0_124, -0.0_189, 0.0_569, 0.1_373, 0.1_883, 0.2_886, 0.3_297, 0.2_212]    )
						UpperCamelCase								=       np.abs(expected_slice - audio_slice    ).max()
						assert max_diff < 3e-2








'''simple docstring'''
import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch

@require_torch
class     _lowerCamelCase       :




 '''simple docstring'''


 @property
 def 			snake_case__ (   self    ):

  """simple docstring"""







  return self.get_dummy_input()


 @property
 def 			snake_case__ (   self    ):

  """simple docstring"""







  if self.block_type == "down":
   return (4, 32, 16, 16)
  elif self.block_type == "mid":
   return (4, 32, 32, 32)
  elif self.block_type == "up":
   return (4, 32, 64, 64)

  raise ValueError(F"""'{self.block_type}' is not a supported block_type. Set it to 'up', 'mid', or 'down'."""    )


 def 			snake_case__ (   self   ,   __lowercase=True   ,   __lowercase=False   ,   __lowercase=False   ,   __lowercase=False   ,   ):

  """simple docstring"""







  __A	:		List[str]	  =      4
  __A	:		List[Any]	  =      32
  __A	:		Optional[Any]	  =      (32, 32)

  __A	:		List[str]	  =      torch.manual_seed(0    )
  __A	:		Union[str, Any]	  =      torch.device(__lowercase    )
  __A	:		Optional[int]	  =      (batch_size, num_channels) + sizes
  __A	:		str	  =      randn_tensor(__lowercase   ,   generator=__lowercase   ,   device=__lowercase    )
  __A	:		int	  =      {'hidden_states': hidden_states}

  if include_temb:
   __A	:		Optional[Any]	  =      128
   __A	:		Any	  =      randn_tensor((batch_size, temb_channels)   ,   generator=__lowercase   ,   device=__lowercase    )

  if include_res_hidden_states_tuple:
   __A	:		Any	  =      torch.manual_seed(1    )
   __A	:		List[Any]	  =      (randn_tensor(__lowercase   ,   generator=__lowercase   ,   device=__lowercase    ),)

  if include_encoder_hidden_states:
   __A	:		str	  =      floats_tensor((batch_size, 32, 32)    ).to(__lowercase    )

  if include_skip_sample:
   __A	:		List[Any]	  =      randn_tensor(((batch_size, 3) + sizes)   ,   generator=__lowercase   ,   device=__lowercase    )

  return dummy_input


 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A	:		Tuple	  =      {
      'in_channels': 32,
      'out_channels': 32,
      'temb_channels': 128,
  }
  if self.block_type == "up":
   __A	:		int	  =      32

  if self.block_type == "mid":
   init_dict.pop('out_channels'    )

  __A	:		Optional[int]	  =      self.dummy_input
  return init_dict, inputs_dict


 def 			snake_case__ (   self   ,   __lowercase    ):

  """simple docstring"""







  __A  ,__A	:		List[Any]	  =      self.prepare_init_args_and_inputs_for_common()
  __A	:		int	  =      self.block_class(**__lowercase    )
  unet_block.to(__lowercase    )
  unet_block.eval()

  with torch.no_grad():
   __A	:		List[Any]	  =      unet_block(**__lowercase    )

  if isinstance(__lowercase   ,   __lowercase    ):
   __A	:		Tuple	  =      output[0]

  self.assertEqual(output.shape   ,   self.output_shape    )

  __A	:		Any	  =      output[0, -1, -3:, -3:]
  __A	:		Dict	  =      torch.tensor(__lowercase    ).to(__lowercase    )
  assert torch_all_close(output_slice.flatten()   ,   __lowercase   ,   atol=5E-3    )







 @unittest.skipIf(torch_device == 'mps'   ,   'Training is not supported in mps'    )
 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A  ,__A	:		str	  =      self.prepare_init_args_and_inputs_for_common()
  __A	:		Optional[Any]	  =      self.block_class(**__lowercase    )
  model.to(__lowercase    )
  model.train()
  __A	:		Any	  =      model(**__lowercase    )

  if isinstance(__lowercase   ,   __lowercase    ):
   __A	:		Any	  =      output[0]

  __A	:		int	  =      torch.device(__lowercase    )
  __A	:		Dict	  =      randn_tensor(output.shape   ,   device=__lowercase    )
  __A	:		Dict	  =      torch.nn.functional.mse_loss(__lowercase   ,   __lowercase    )
  loss.backward()






