





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     A     (SCREAMING_SNAKE_CASE  ,	unittest.TestCase       ):
			'''simple docstring'''




			__lowerCamelCase      :			Union[str, Any]							     =       AudioLDMPipeline
			__lowerCamelCase      :			Optional[int]							     =       TEXT_TO_AUDIO_PARAMS
			__lowerCamelCase      :			Dict							     =       TEXT_TO_AUDIO_BATCH_PARAMS
			__lowerCamelCase      :			Union[str, Any]							     =       frozenset(
			    [
			        '''num_inference_steps''',
			        '''num_waveforms_per_prompt''',
			        '''generator''',
			        '''latents''',
			        '''output_type''',
			        '''return_dict''',
			        '''callback''',
			        '''callback_steps''',
			    ]       )






			def   a_			(     self  :       List[Any]  )      ->      Dict:







					"""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=__lowerCAmelCase				,      )
					A__     =  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  )
					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=10_00				,      projection_dim=32				,      )
					A__     =  ClapTextModelWithProjection(__lowerCAmelCase  )
					A__     =  RobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-random-roberta"""				,      model_max_length=77  )

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

					A__     =  SpeechTaHifiGan(__lowerCAmelCase  )

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






			def   a_			(     self  :       Optional[int]				,      __lowerCAmelCase  :       Tuple				,      __lowerCAmelCase  :       str=0  )      ->      Union[str, Any]:







					"""simple docstring"""






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






			def   a_			(     self  :       Union[str, Any]  )      ->      Dict:







					"""simple docstring"""






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

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

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

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

					A__     =  audio[:10]
					A__     =  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   a_			(     self  :       Optional[int]  )      ->      Any:







					"""simple docstring"""






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

					A__     =  self.get_dummy_inputs(__lowerCAmelCase  )
					A__     =  3 * [inputs["""prompt"""]]

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

					A__     =  self.get_dummy_inputs(__lowerCAmelCase  )
					A__     =  3 * [inputs.pop("""prompt"""  )]

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

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

					A__     =  prompt_embeds

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

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






			def   a_			(     self  :       Any  )      ->      Dict:







					"""simple docstring"""






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

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

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

					A__     =  self.get_dummy_inputs(__lowerCAmelCase  )
					A__     =  3 * [inputs.pop("""prompt"""  )]

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

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

							embeds.append(__lowerCAmelCase  )

					A__ ,     A__     =  embeds

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

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






			def   a_			(     self  :       int  )      ->      List[Any]:







					"""simple docstring"""






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

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

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

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







					"""simple docstring"""






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

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

					# test num_waveforms_per_prompt=1 (default)
					A__     =  audioldm_pipe(__lowerCAmelCase				,      num_inference_steps=2  ).audios

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

					# 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, 2_56)

					# test num_waveforms_per_prompt for single prompt
					A__     =  2
					A__     =  audioldm_pipe(__lowerCAmelCase				,      num_inference_steps=2				,      num_waveforms_per_prompt=__lowerCAmelCase  ).audios

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

					# 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=__lowerCAmelCase  ).audios

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






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







					"""simple docstring"""






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

					A__     =  self.get_dummy_inputs(__lowerCAmelCase  )
					A__     =  audioldm_pipe(audio_length_in_s=0.0_1_6				,      **__lowerCAmelCase  )
					A__     =  output.audios[0]

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

					A__     =  audioldm_pipe(audio_length_in_s=0.0_3_2				,      **__lowerCAmelCase  )
					A__     =  output.audios[0]

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






			def   a_			(     self  :       List[str]  )      ->      Dict:







					"""simple docstring"""






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

					A__     =  ["""hey"""]

					A__     =  audioldm_pipe(__lowerCAmelCase				,      num_inference_steps=1  )
					A__     =  output.audios.shape
					assert audio_shape == (1, 2_56)

					A__     =  audioldm_pipe.vocoder.config
					config.model_in_dim *= 2
					A__     =  SpeechTaHifiGan(__lowerCAmelCase  ).to(__lowerCAmelCase  )
					A__     =  audioldm_pipe(__lowerCAmelCase				,      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, 2_56)






			def   a_			(     self  :       List[Any]  )      ->      int:







					"""simple docstring"""






					self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__lowerCAmelCase  )






			def   a_			(     self  :       Tuple  )      ->      Dict:







					"""simple docstring"""






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







					"""simple docstring"""






					self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__lowerCAmelCase  )


@slow
class     A     (unittest.TestCase       ):
			'''simple docstring'''




			def   a_			(     self  :       Optional[Any]  )      ->      List[Any]:







					"""simple docstring"""






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






			def   a_			(     self  :       List[str]				,      __lowerCAmelCase  :       Optional[int]				,      __lowerCAmelCase  :       List[str]="cpu"				,      __lowerCAmelCase  :       Optional[Any]=torch.floataa				,      __lowerCAmelCase  :       Optional[int]=0  )      ->      Union[str, Any]:







					"""simple docstring"""






					A__     =  torch.Generator(device=__lowerCAmelCase  ).manual_seed(__lowerCAmelCase  )
					A__     =  np.random.RandomState(__lowerCAmelCase  ).standard_normal((1, 8, 1_28, 16)  )
					A__     =  torch.from_numpy(__lowerCAmelCase  ).to(device=__lowerCAmelCase				,      dtype=__lowerCAmelCase  )
					A__     =  {
					    """prompt""": """A hammer hitting a wooden surface""",
					    """latents""": latents,
					    """generator""": generator,
					    """num_inference_steps""": 3,
					    """guidance_scale""": 2.5,
					}
					return inputs






			def   a_			(     self  :       Any  )      ->      str:







					"""simple docstring"""






					A__     =  AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""  )
					A__     =  audioldm_pipe.to(__lowerCAmelCase  )
					audioldm_pipe.set_progress_bar_config(disable=__lowerCAmelCase  )

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

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

					A__     =  audio[7_72_30:7_72_40]
					A__     =  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]  )
					A__     =  np.abs(expected_slice - audio_slice  ).max()
					assert max_diff < 1e-2







			def   a_			(     self  :       Optional[int]  )      ->      int:







					"""simple docstring"""






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

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

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

					A__     =  audio[2_77_80:2_77_90]
					A__     =  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]  )
					A__     =  np.abs(expected_slice - audio_slice  ).max()
					assert max_diff < 3e-2










import multiprocessing
from typing import TYPE_CHECKING, Optional, Union

from .. import Dataset, Features, config
from ..formatting import query_table
from ..packaged_modules.sql.sql import Sql
from ..utils import logging
from .abc import AbstractDatasetInputStream


if TYPE_CHECKING:
					import sqlitea

					import sqlalchemy


class     A     (SCREAMING_SNAKE_CASE       ):
			'''simple docstring'''




			def __init__(     self  :       Any				,      __lowerCAmelCase  :       Union[str, "sqlalchemy.sql.Selectable"]				,      __lowerCAmelCase  :       Union[str, "sqlalchemy.engine.Connection", "sqlalchemy.engine.Engine", "sqlite3.Connection"]				,      __lowerCAmelCase  :       Optional[Features] = None				,      __lowerCAmelCase  :       str = None				,      __lowerCAmelCase  :       bool = False				,      **__lowerCAmelCase  :       Dict				,      )      ->      Any:







					"""simple docstring"""






					super().__init__(features=__lowerCAmelCase				,      cache_dir=__lowerCAmelCase				,      keep_in_memory=__lowerCAmelCase				,      **__lowerCAmelCase  )
					A__     =  Sql(
					    cache_dir=__lowerCAmelCase				,      features=__lowerCAmelCase				,      sql=__lowerCAmelCase				,      con=__lowerCAmelCase				,      **__lowerCAmelCase				,      )







			def   a_			(     self  :       Optional[int]  )      ->      Optional[int]:







					"""simple docstring"""






					A__     =  None
					A__     =  None
					A__     =  None
					A__     =  None

					self.builder.download_and_prepare(
					    download_config=__lowerCAmelCase				,      download_mode=__lowerCAmelCase				,      verification_mode=__lowerCAmelCase				,      base_path=__lowerCAmelCase				,      )

					# Build dataset for splits
					A__     =  self.builder.as_dataset(
					    split="""train"""				,      verification_mode=__lowerCAmelCase				,      in_memory=self.keep_in_memory  )
					return dataset


class     A     :
			'''simple docstring'''




			def __init__(     self  :       Optional[int]				,      __lowerCAmelCase  :       Dataset				,      __lowerCAmelCase  :       str				,      __lowerCAmelCase  :       Union[str, "sqlalchemy.engine.Connection", "sqlalchemy.engine.Engine", "sqlite3.Connection"]				,      __lowerCAmelCase  :       Optional[int] = None				,      __lowerCAmelCase  :       Optional[int] = None				,      **__lowerCAmelCase  :       List[Any]				,      )      ->      Any:







					"""simple docstring"""






					if num_proc is not None and num_proc <= 0:
							raise ValueError(f'num_proc {num_proc} must be an integer > 0.'  )

					A__     =  dataset
					A__     =  name
					A__     =  con
					A__     =  batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
					A__     =  num_proc
					A__     =  to_sql_kwargs






			def   a_			(     self  :       Any  )      ->      int:







					"""simple docstring"""






					A__     =  self.to_sql_kwargs.pop("""sql"""				,      __lowerCAmelCase  )
					A__     =  self.to_sql_kwargs.pop("""con"""				,      __lowerCAmelCase  )
					A__     =  self.to_sql_kwargs.pop("""index"""				,      __lowerCAmelCase  )

					A__     =  self._write(index=__lowerCAmelCase				,      **self.to_sql_kwargs  )
					return written






			def   a_			(     self  :       Union[str, Any]				,      __lowerCAmelCase  :       Optional[Any]  )      ->      Optional[Any]:







					"""simple docstring"""






					A__ ,     A__ ,     A__     =  args
					A__     =  {**to_sql_kwargs, """if_exists""": """append"""} if offset > 0 else to_sql_kwargs
					A__     =  query_table(
					    table=self.dataset.data				,      key=slice(__lowerCAmelCase				,      offset + self.batch_size  )				,      indices=self.dataset._indices				,      )
					A__     =  batch.to_pandas()
					A__     =  df.to_sql(self.name				,      self.con				,      index=__lowerCAmelCase				,      **__lowerCAmelCase  )
					return num_rows or len(__lowerCAmelCase  )







			def   a_			(     self  :       Optional[int]				,      __lowerCAmelCase  :       Tuple				,      **__lowerCAmelCase  :       Dict  )      ->      int:







					"""simple docstring"""






					A__     =  0

					if self.num_proc is None or self.num_proc == 1:
							for offset in logging.tqdm(
							    range(0				,      len(self.dataset  )				,      self.batch_size  )				,      unit="""ba"""				,      disable=not logging.is_progress_bar_enabled()				,      desc="""Creating SQL from Arrow format"""				,      ):
									written += self._batch_sql((offset, index, to_sql_kwargs)  )
					else:
							A__ ,     A__     =  len(self.dataset  ), self.batch_size
							with multiprocessing.Pool(self.num_proc  ) as pool:
									for num_rows in logging.tqdm(
									    pool.imap(
									        self._batch_sql				,      [(offset, index, to_sql_kwargs) for offset in range(0				,      __lowerCAmelCase				,      __lowerCAmelCase  )]				,      )				,      total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size				,      unit="""ba"""				,      disable=not logging.is_progress_bar_enabled()				,      desc="""Creating SQL from Arrow format"""				,      ):
											written += num_rows

					return written



