


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






  def       _UpperCAmelCase	(       self   :       str			):




         """simple docstring"""







         torch.manual_seed(0			)
         UpperCAmelCase__					    =  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=_lowercase		,    )
         UpperCAmelCase__					    =  DDIMScheduler(
             beta_start=0.0_0_0_8_5		,    beta_end=0.0_1_2		,    beta_schedule="scaled_linear"		,    clip_sample=_lowercase		,    set_alpha_to_one=_lowercase		,    )
         torch.manual_seed(0			)
         UpperCAmelCase__					    =  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__					    =  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		,    )
         UpperCAmelCase__					    =  ClapTextModelWithProjection(_lowercase			)
         UpperCAmelCase__					    =  RobertaTokenizer.from_pretrained("hf-internal-testing/tiny-random-roberta"		,    model_max_length=77			)

         UpperCAmelCase__					    =  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=_lowercase		,    )

         UpperCAmelCase__					    =  SpeechTaHifiGan(_lowercase			)

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






  def       _UpperCAmelCase	(       self   :       Any		,    _lowercase   :       List[str]		,    _lowercase   :       int=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       _UpperCAmelCase	(       self   :       Optional[int]			):




         """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_56

         UpperCAmelCase__					    =  audio[:10]
         UpperCAmelCase__					    =  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       _UpperCAmelCase	(       self   :       List[str]			):




         """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       _UpperCAmelCase	(       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__    ,      UpperCAmelCase__					    =  embeds

         # forward
         UpperCAmelCase__					    =  audioldm_pipe(**_lowercase			)
         UpperCAmelCase__					    =  output.audios[0]

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






  def       _UpperCAmelCase	(       self   :       List[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_56

         UpperCAmelCase__					    =  audio[:10]
         UpperCAmelCase__					    =  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       _UpperCAmelCase	(       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_56)

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

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

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






  def       _UpperCAmelCase	(       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.0_1_6		,    **_lowercase			)
         UpperCAmelCase__					    =  output.audios[0]

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

         UpperCAmelCase__					    =  audioldm_pipe(audio_length_in_s=0.0_3_2		,    **_lowercase			)
         UpperCAmelCase__					    =  output.audios[0]

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






  def       _UpperCAmelCase	(       self   :       List[str]			):




         """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_56)

         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_56)






  def       _UpperCAmelCase	(       self   :       str			):




         """simple docstring"""







         self._test_attention_slicing_forward_pass(test_mean_pixel_difference=_lowercase			)






  def       _UpperCAmelCase	(       self   :       Tuple			):




         """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       _UpperCAmelCase	(       self   :       List[str]			):




         """simple docstring"""







         self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=_lowercase			)




@slow
class      lowercase__	(       unittest.TestCase		):






  def       _UpperCAmelCase	(       self   :       Any			):




         """simple docstring"""







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






  def       _UpperCAmelCase	(       self   :       Optional[int]		,    _lowercase   :       List[Any]		,    _lowercase   :       Optional[int]="cpu"		,    _lowercase   :       Any=torch.floataa		,    _lowercase   :       Optional[Any]=0			):




         """simple docstring"""







         UpperCAmelCase__					    =  torch.Generator(device=_lowercase			).manual_seed(_lowercase			)
         UpperCAmelCase__					    =  np.random.RandomState(_lowercase			).standard_normal((1, 8, 1_28, 16)			)
         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       _UpperCAmelCase	(       self   :       Optional[Any]			):




         """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__					    =  25
         UpperCAmelCase__					    =  audioldm_pipe(**_lowercase			).audios[0]

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

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






  def       _UpperCAmelCase	(       self   :       Union[str, 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_19_20

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













'''simple docstring'''



__UpperCamelCase   :							List[Any]         =   """ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"""



def  __UpperCAmelCase						(					SCREAMING_SNAKE_CASE__:						bytes     )	->      bytes:

   """simple docstring"""
   # Make sure the supplied data is a bytes-like object
   if not isinstance(SCREAMING_SNAKE_CASE__,       SCREAMING_SNAKE_CASE__     ):
      __a					  =     f"""a bytes-like object is required, not '{data.__class__.__name__}'"""
      raise TypeError(SCREAMING_SNAKE_CASE__     )

   __a					  =     ''.join(bin(SCREAMING_SNAKE_CASE__     )[2:].zfill(8     ) for byte in data     )

   __a					  =     len(SCREAMING_SNAKE_CASE__     ) % 6 != 0

   if padding_needed:
      # The padding that will be added later
      __a					  =     b'=' * ((6 - len(SCREAMING_SNAKE_CASE__     ) % 6) // 2)

      # Append binary_stream with arbitrary binary digits (0's by default) to make its
      # length a multiple of 6.
      binary_stream += "0" * (6 - len(SCREAMING_SNAKE_CASE__     ) % 6)
   else:
      __a					  =     b''

   # Encode every 6 binary digits to their corresponding Base64 character
   return (
       "".join(
           B64_CHARSET[int(binary_stream[index : index + 6],       2     )]
           for index in range(0,       len(SCREAMING_SNAKE_CASE__     ),       6     )     ).encode()
       + padding
   )



def  __UpperCAmelCase						(					SCREAMING_SNAKE_CASE__:						str     )	->      bytes:

   """simple docstring"""
   # Make sure encoded_data is either a string or a bytes-like object
   if not isinstance(SCREAMING_SNAKE_CASE__,       SCREAMING_SNAKE_CASE__     ) and not isinstance(SCREAMING_SNAKE_CASE__,       SCREAMING_SNAKE_CASE__     ):
      __a					  =     (
          'argument should be a bytes-like object or ASCII string, '
          f"""not '{encoded_data.__class__.__name__}'"""
      )
      raise TypeError(SCREAMING_SNAKE_CASE__     )

   # In case encoded_data is a bytes-like object, make sure it contains only
   # ASCII characters so we convert it to a string object
   if isinstance(SCREAMING_SNAKE_CASE__,       SCREAMING_SNAKE_CASE__     ):
      try:
         __a					  =     encoded_data.decode('utf-8'     )
      except UnicodeDecodeError:
         raise ValueError('base64 encoded data should only contain ASCII characters'     )

   __a					  =     encoded_data.count('='     )

   # Check if the encoded string contains non base64 characters
   if padding:
      assert all(
          char in B64_CHARSET for char in encoded_data[:-padding]     ), "Invalid base64 character(s) found."
   else:
      assert all(
          char in B64_CHARSET for char in encoded_data     ), "Invalid base64 character(s) found."

   # Check the padding
   assert len(SCREAMING_SNAKE_CASE__     ) % 4 == 0 and padding < 3, "Incorrect padding"

   if padding:
      # Remove padding if there is one
      __a					  =     encoded_data[:-padding]

      __a					  =     ''.join(
          bin(B64_CHARSET.index(SCREAMING_SNAKE_CASE__     )     )[2:].zfill(6     ) for char in encoded_data     )[: -padding * 2]
   else:
      __a					  =     ''.join(
          bin(B64_CHARSET.index(SCREAMING_SNAKE_CASE__     )     )[2:].zfill(6     ) for char in encoded_data     )

   __a					  =     [
       int(binary_stream[index : index + 8],       2     )
       for index in range(0,       len(SCREAMING_SNAKE_CASE__     ),       8     )
   ]

   return bytes(SCREAMING_SNAKE_CASE__     )


if __name__ == "__main__":
 import doctest

 doctest.testmod()