





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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          torch.manual_seed(0)
          __A						: Tuple      =       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=_UpperCAmelCase						,	)
          __A						: Union[str, Any]      =       DDIMScheduler(
              beta_start=0.00085						,	beta_end=0.012						,	beta_schedule='scaled_linear'						,	clip_sample=_UpperCAmelCase						,	set_alpha_to_one=_UpperCAmelCase						,	)
          torch.manual_seed(0)
          __A						: Dict      =       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						: Tuple      =       ClapTextConfig(
              bos_token_id=0						,	eos_token_id=2						,	hidden_size=32						,	intermediate_size=37						,	layer_norm_eps=1e-0_5						,	num_attention_heads=4						,	num_hidden_layers=5						,	pad_token_id=1						,	vocab_size=1000						,	projection_dim=32						,	)
          __A						: int      =       ClapTextModelWithProjection(_UpperCAmelCase)
          __A						: Tuple      =       RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta'						,	model_max_length=77)

          __A						: Tuple      =       SpeechTaHifiGanConfig(
              model_in_dim=8						,	sampling_rate=1_6000						,	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=_UpperCAmelCase						,	)

          __A						: List[str]      =       SpeechTaHifiGan(_UpperCAmelCase)

          __A						: Union[str, Any]      =       {
              'unet': unet,
              'scheduler': scheduler,
              'vae': vae,
              'text_encoder': text_encoder,
              'tokenizer': tokenizer,
              'vocoder': vocoder,
          }
          return components


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase=0):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: str      =       'cpu'  # ensure determinism for the device-dependent torch.Generator

          __A						: Dict      =       self.get_dummy_components()
          __A						: Tuple      =       AudioLDMPipeline(**_UpperCAmelCase)
          __A						: Union[str, Any]      =       audioldm_pipe.to(_UpperCAmelCase)
          audioldm_pipe.set_progress_bar_config(disable=_UpperCAmelCase)

          __A						: Optional[int]      =       self.get_dummy_inputs(_UpperCAmelCase)
          __A						: Any      =       audioldm_pipe(**_UpperCAmelCase)
          __A						: Optional[Any]      =       output.audios[0]

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

          __A						: List[str]      =       audio[:10]
          __A						: Tuple      =       np.array(
              [-0.0050, 0.0050, -0.0060, 0.0033, -0.0026, 0.0033, -0.0027, 0.0033, -0.0028, 0.0033])

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Any      =       self.get_dummy_components()
          __A						: str      =       AudioLDMPipeline(**_UpperCAmelCase)
          __A						: Dict      =       audioldm_pipe.to(_UpperCAmelCase)
          __A						: Any      =       audioldm_pipe.to(_UpperCAmelCase)
          audioldm_pipe.set_progress_bar_config(disable=_UpperCAmelCase)

          __A						: Union[str, Any]      =       self.get_dummy_inputs(_UpperCAmelCase)
          __A						: Optional[Any]      =       3 * [inputs['prompt']]

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

          __A						: int      =       self.get_dummy_inputs(_UpperCAmelCase)
          __A						: Any      =       3 * [inputs.pop('prompt')]

          __A						: List[str]      =       audioldm_pipe.tokenizer(
              _UpperCAmelCase						,	padding='max_length'						,	max_length=audioldm_pipe.tokenizer.model_max_length						,	truncation=_UpperCAmelCase						,	return_tensors='pt'						,	)
          __A						: Union[str, Any]      =       text_inputs['input_ids'].to(_UpperCAmelCase)

          __A						: Any      =       audioldm_pipe.text_encoder(
              _UpperCAmelCase						,	)
          __A						: str      =       prompt_embeds.text_embeds
          # additional L_2 normalization over each hidden-state
          __A						: List[str]      =       F.normalize(_UpperCAmelCase						,	dim=-1)

          __A						: Dict      =       prompt_embeds

          # forward
          __A						: Tuple      =       audioldm_pipe(**_UpperCAmelCase)
          __A						: Any      =       output.audios[0]

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[str]      =       self.get_dummy_components()
          __A						: Union[str, Any]      =       AudioLDMPipeline(**_UpperCAmelCase)
          __A						: Dict      =       audioldm_pipe.to(_UpperCAmelCase)
          __A						: str      =       audioldm_pipe.to(_UpperCAmelCase)
          audioldm_pipe.set_progress_bar_config(disable=_UpperCAmelCase)

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

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

          __A						: Union[str, Any]      =       self.get_dummy_inputs(_UpperCAmelCase)
          __A						: Any      =       3 * [inputs.pop('prompt')]

          __A						: List[Any]      =       []
          for p in [prompt, negative_prompt]:
               __A						: List[Any]      =       audioldm_pipe.tokenizer(
                   _UpperCAmelCase						,	padding='max_length'						,	max_length=audioldm_pipe.tokenizer.model_max_length						,	truncation=_UpperCAmelCase						,	return_tensors='pt'						,	)
               __A						: Optional[int]      =       text_inputs['input_ids'].to(_UpperCAmelCase)

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

               embeds.append(_UpperCAmelCase)

          __A ,__A						: List[str]      =       embeds

          # forward
          __A						: List[Any]      =       audioldm_pipe(**_UpperCAmelCase)
          __A						: int      =       output.audios[0]

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

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

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

          __A						: Union[str, Any]      =       audio[:10]
          __A						: Any      =       np.array(
              [-0.0051, 0.0050, -0.0060, 0.0034, -0.0026, 0.0033, -0.0027, 0.0033, -0.0028, 0.0032])

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: List[str]      =       'cpu'  # ensure determinism for the device-dependent torch.Generator
          __A						: Tuple      =       self.get_dummy_components()
          __A						: int      =       PNDMScheduler(skip_prk_steps=_UpperCAmelCase)
          __A						: Union[str, Any]      =       AudioLDMPipeline(**_UpperCAmelCase)
          __A						: Tuple      =       audioldm_pipe.to(_UpperCAmelCase)
          audioldm_pipe.set_progress_bar_config(disable=_UpperCAmelCase)

          __A						: Tuple      =       'A hammer hitting a wooden surface'

          # test num_waveforms_per_prompt=1 (default)
          __A						: Dict      =       audioldm_pipe(_UpperCAmelCase						,	num_inference_steps=2).audios

          assert audios.shape == (1, 256)

          # test num_waveforms_per_prompt=1 (default) for batch of prompts
          __A						: Optional[Any]      =       2
          __A						: List[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
          __A						: Tuple      =       2
          __A						: Any      =       audioldm_pipe(_UpperCAmelCase						,	num_inference_steps=2						,	num_waveforms_per_prompt=_UpperCAmelCase).audios

          assert audios.shape == (num_waveforms_per_prompt, 256)

          # test num_waveforms_per_prompt for batch of prompts
          __A						: Any      =       2
          __A						: List[str]      =       audioldm_pipe(
              [prompt] * batch_size						,	num_inference_steps=2						,	num_waveforms_per_prompt=_UpperCAmelCase).audios

          assert audios.shape == (batch_size * num_waveforms_per_prompt, 256)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Dict      =       'cpu'  # ensure determinism for the device-dependent torch.Generator
          __A						: int      =       self.get_dummy_components()
          __A						: int      =       AudioLDMPipeline(**_UpperCAmelCase)
          __A						: str      =       audioldm_pipe.to(_UpperCAmelCase)
          audioldm_pipe.set_progress_bar_config(disable=_UpperCAmelCase)
          __A						: int      =       audioldm_pipe.vocoder.config.sampling_rate

          __A						: Any      =       self.get_dummy_inputs(_UpperCAmelCase)
          __A						: List[Any]      =       audioldm_pipe(audio_length_in_s=0.016						,	**_UpperCAmelCase)
          __A						: int      =       output.audios[0]

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

          __A						: List[str]      =       audioldm_pipe(audio_length_in_s=0.032						,	**_UpperCAmelCase)
          __A						: Any      =       output.audios[0]

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Tuple      =       self.get_dummy_components()
          __A						: str      =       AudioLDMPipeline(**_UpperCAmelCase)
          __A						: Any      =       audioldm_pipe.to(_UpperCAmelCase)
          audioldm_pipe.set_progress_bar_config(disable=_UpperCAmelCase)

          __A						: Dict      =       ['hey']

          __A						: Dict      =       audioldm_pipe(_UpperCAmelCase						,	num_inference_steps=1)
          __A						: Tuple      =       output.audios.shape
          assert audio_shape == (1, 256)

          __A						: int      =       audioldm_pipe.vocoder.config
          config.model_in_dim *= 2
          __A						: Any      =       SpeechTaHifiGan(_UpperCAmelCase).to(_UpperCAmelCase)
          __A						: Dict      =       audioldm_pipe(_UpperCAmelCase						,	num_inference_steps=1)
          __A						: 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):

          '''simple docstring'''






          self._test_attention_slicing_forward_pass(test_mean_pixel_difference=_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          self._test_inference_batch_single_identical(test_mean_pixel_difference=_UpperCAmelCase)


     @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):

          '''simple docstring'''






          self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=_UpperCAmelCase)




@slow
class  SCREAMING_SNAKE_CASE    (unittest.TestCase							):


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase="cpu"						,	_UpperCAmelCase=torch.floataa						,	_UpperCAmelCase=0):

          '''simple docstring'''






          __A						: int      =       torch.Generator(device=_UpperCAmelCase).manual_seed(_UpperCAmelCase)
          __A						: Dict      =       np.random.RandomState(_UpperCAmelCase).standard_normal((1, 8, 128, 16))
          __A						: Optional[int]      =       torch.from_numpy(_UpperCAmelCase).to(device=_UpperCAmelCase						,	dtype=_UpperCAmelCase)
          __A						: Dict      =       {
              '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):

          '''simple docstring'''






          __A						: int      =       AudioLDMPipeline.from_pretrained('cvssp/audioldm')
          __A						: Dict      =       audioldm_pipe.to(_UpperCAmelCase)
          audioldm_pipe.set_progress_bar_config(disable=_UpperCAmelCase)

          __A						: List[Any]      =       self.get_inputs(_UpperCAmelCase)
          __A						: Dict      =       25
          __A						: Union[str, Any]      =       audioldm_pipe(**_UpperCAmelCase).audios[0]

          assert audio.ndim == 1
          assert len(_UpperCAmelCase) == 8_1920

          __A						: Dict      =       audio[7_7230:7_7240]
          __A						: List[str]      =       np.array(
              [-0.4884, -0.4607, 0.0023, 0.5007, 0.5896, 0.5151, 0.3813, -0.0208, -0.3687, -0.4315])
          __A						: Union[str, Any]      =       np.abs(expected_slice - audio_slice).max()
          assert max_diff < 1e-2


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Tuple      =       AudioLDMPipeline.from_pretrained('cvssp/audioldm')
          __A						: Tuple      =       LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config)
          __A						: List[str]      =       audioldm_pipe.to(_UpperCAmelCase)
          audioldm_pipe.set_progress_bar_config(disable=_UpperCAmelCase)

          __A						: Optional[Any]      =       self.get_inputs(_UpperCAmelCase)
          __A						: int      =       audioldm_pipe(**_UpperCAmelCase).audios[0]

          assert audio.ndim == 1
          assert len(_UpperCAmelCase) == 8_1920

          __A						: Union[str, Any]      =       audio[2_7780:2_7790]
          __A						: Optional[Any]      =       np.array([-0.2131, -0.0873, -0.0124, -0.0189, 0.0569, 0.1373, 0.1883, 0.2886, 0.3297, 0.2212])
          __A						: List[str]      =       np.abs(expected_slice - audio_slice).max()
          assert max_diff < 3e-2






'''simple docstring'''




import json
from typing import Iterator, List, Union

from tokenizers import AddedToken, Regex, Tokenizer, decoders, normalizers, pre_tokenizers, trainers
from tokenizers.implementations.base_tokenizer import BaseTokenizer
from tokenizers.models import Unigram
from tokenizers.processors import TemplateProcessing


class  SCREAMING_SNAKE_CASE    (a__							):


     def __init__(						self						,	_UpperCAmelCase = "▁"						,	_UpperCAmelCase = True						,	_UpperCAmelCase = "<unk>"						,	_UpperCAmelCase = "</s>"						,	_UpperCAmelCase = "<pad>"						,	):

          '''simple docstring'''






          __A						: Dict      =       {
              'pad': {'id': 0, 'token': pad_token},
              'eos': {'id': 1, 'token': eos_token},
              'unk': {'id': 2, 'token': unk_token},
          }

          __A						: List[Any]      =       [None] * len(self.special_tokens)
          for token_dict in self.special_tokens.values():
               __A						: List[str]      =       token_dict['token']

          __A						: str      =       Tokenizer(Unigram())

          __A						: Dict      =       normalizers.Sequence(
              [
                  normalizers.Nmt(),
                  normalizers.NFKC(),
                  normalizers.Replace(Regex(' {2,}')						,	' '),
                  normalizers.Lowercase(),
              ])
          __A						: Union[str, Any]      =       pre_tokenizers.Sequence(
              [
                  pre_tokenizers.Metaspace(replacement=_UpperCAmelCase						,	add_prefix_space=_UpperCAmelCase),
                  pre_tokenizers.Digits(individual_digits=_UpperCAmelCase),
                  pre_tokenizers.Punctuation(),
              ])
          __A						: Any      =       decoders.Metaspace(replacement=_UpperCAmelCase						,	add_prefix_space=_UpperCAmelCase)

          __A						: Dict      =       TemplateProcessing(
              single=F'$A {self.special_tokens["eos"]["token"]}'						,	special_tokens=[(self.special_tokens['eos']['token'], self.special_tokens['eos']['id'])]						,	)

          __A						: Any      =       {
              'model': 'SentencePieceUnigram',
              'replacement': replacement,
              'add_prefix_space': add_prefix_space,
          }

          super().__init__(_UpperCAmelCase						,	_UpperCAmelCase)


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase = 8000						,	_UpperCAmelCase = True						,	):

          '''simple docstring'''






          __A						: str      =       trainers.UnigramTrainer(
              vocab_size=_UpperCAmelCase						,	special_tokens=self.special_tokens_list						,	show_progress=_UpperCAmelCase						,	)

          if isinstance(_UpperCAmelCase						,	_UpperCAmelCase):
               __A						: Union[str, Any]      =       [files]
          self._tokenizer.train(_UpperCAmelCase						,	trainer=_UpperCAmelCase)

          self.add_unk_id()


     def 							SCREAMING_SNAKE_CASE  (						self						,	_UpperCAmelCase						,	_UpperCAmelCase = 8000						,	_UpperCAmelCase = True						,	):

          '''simple docstring'''






          __A						: Dict      =       trainers.UnigramTrainer(
              vocab_size=_UpperCAmelCase						,	special_tokens=self.special_tokens_list						,	show_progress=_UpperCAmelCase						,	)

          self._tokenizer.train_from_iterator(_UpperCAmelCase						,	trainer=_UpperCAmelCase)

          self.add_unk_id()


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						: Any      =       json.loads(self._tokenizer.to_str())

          __A						: Union[str, Any]      =       self.special_tokens['unk']['id']

          __A						: str      =       Tokenizer.from_str(json.dumps(_UpperCAmelCase))