






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



	'''simple docstring'''




	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 				_UpperCamelCase     (			self   )			->       Union[str, Any]:
								torch.manual_seed(0   )
								snake_case_           =				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=a			,     )
								snake_case_           =				DDIMScheduler(
								    beta_start=0.00_085			,     beta_end=0.012			,     beta_schedule='scaled_linear'			,     clip_sample=a			,     set_alpha_to_one=a			,     )
								torch.manual_seed(0   )
								snake_case_           =				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   )
								snake_case_           =				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			,     )
								snake_case_           =				ClapTextModelWithProjection(a   )
								snake_case_           =				RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta'			,     model_max_length=77   )

								snake_case_           =				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=a			,     )

								snake_case_           =				SpeechTaHifiGan(a   )

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




	def 				_UpperCamelCase     (			self			,     a			,     a=0   )			->       Any:
								if str(a   ).startswith('mps'   ):
															snake_case_           =				torch.manual_seed(a   )
								else:
															snake_case_           =				torch.Generator(device=a   ).manual_seed(a   )
								snake_case_           =				{
								    'prompt': 'A hammer hitting a wooden surface',
								    'generator': generator,
								    'num_inference_steps': 2,
								    'guidance_scale': 6.0,
								}
								return inputs




	def 				_UpperCamelCase     (			self   )			->       Optional[int]:
								snake_case_           =				'cpu'  # ensure determinism for the device-dependent torch.Generator

								snake_case_           =				self.get_dummy_components()
								snake_case_           =				AudioLDMPipeline(**a   )
								snake_case_           =				audioldm_pipe.to(a   )
								audioldm_pipe.set_progress_bar_config(disable=a   )

								snake_case_           =				self.get_dummy_inputs(a   )
								snake_case_           =				audioldm_pipe(**a   )
								snake_case_           =				output.audios[0]

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

								snake_case_           =				audio[:10]
								snake_case_           =				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 				_UpperCamelCase     (			self   )			->       Dict:
								snake_case_           =				self.get_dummy_components()
								snake_case_           =				AudioLDMPipeline(**a   )
								snake_case_           =				audioldm_pipe.to(a   )
								snake_case_           =				audioldm_pipe.to(a   )
								audioldm_pipe.set_progress_bar_config(disable=a   )

								snake_case_           =				self.get_dummy_inputs(a   )
								snake_case_           =				3 * [inputs['prompt']]

								# forward
								snake_case_           =				audioldm_pipe(**a   )
								snake_case_           =				output.audios[0]

								snake_case_           =				self.get_dummy_inputs(a   )
								snake_case_           =				3 * [inputs.pop('prompt'   )]

								snake_case_           =				audioldm_pipe.tokenizer(
								    a			,     padding='max_length'			,     max_length=audioldm_pipe.tokenizer.model_max_length			,     truncation=a			,     return_tensors='pt'			,     )
								snake_case_           =				text_inputs['input_ids'].to(a   )

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

								snake_case_           =				prompt_embeds

								# forward
								snake_case_           =				audioldm_pipe(**a   )
								snake_case_           =				output.audios[0]

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




	def 				_UpperCamelCase     (			self   )			->       List[str]:
								snake_case_           =				self.get_dummy_components()
								snake_case_           =				AudioLDMPipeline(**a   )
								snake_case_           =				audioldm_pipe.to(a   )
								snake_case_           =				audioldm_pipe.to(a   )
								audioldm_pipe.set_progress_bar_config(disable=a   )

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

								# forward
								snake_case_           =				audioldm_pipe(**a   )
								snake_case_           =				output.audios[0]

								snake_case_           =				self.get_dummy_inputs(a   )
								snake_case_           =				3 * [inputs.pop('prompt'   )]

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

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

															embeds.append(a   )

								snake_case_     ,						snake_case_           =				embeds

								# forward
								snake_case_           =				audioldm_pipe(**a   )
								snake_case_           =				output.audios[0]

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




	def 				_UpperCamelCase     (			self   )			->       List[Any]:
								snake_case_           =				'cpu'  # ensure determinism for the device-dependent torch.Generator
								snake_case_           =				self.get_dummy_components()
								snake_case_           =				PNDMScheduler(skip_prk_steps=a   )
								snake_case_           =				AudioLDMPipeline(**a   )
								snake_case_           =				audioldm_pipe.to(a   )
								audioldm_pipe.set_progress_bar_config(disable=a   )

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

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

								snake_case_           =				audio[:10]
								snake_case_           =				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 				_UpperCamelCase     (			self   )			->       List[str]:
								snake_case_           =				'cpu'  # ensure determinism for the device-dependent torch.Generator
								snake_case_           =				self.get_dummy_components()
								snake_case_           =				PNDMScheduler(skip_prk_steps=a   )
								snake_case_           =				AudioLDMPipeline(**a   )
								snake_case_           =				audioldm_pipe.to(a   )
								audioldm_pipe.set_progress_bar_config(disable=a   )

								snake_case_           =				'A hammer hitting a wooden surface'

								# test num_waveforms_per_prompt=1 (default)
								snake_case_           =				audioldm_pipe(a			,     num_inference_steps=2   ).audios

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

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

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

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

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




	def 				_UpperCamelCase     (			self   )			->       int:
								snake_case_           =				'cpu'  # ensure determinism for the device-dependent torch.Generator
								snake_case_           =				self.get_dummy_components()
								snake_case_           =				AudioLDMPipeline(**a   )
								snake_case_           =				audioldm_pipe.to(a   )
								audioldm_pipe.set_progress_bar_config(disable=a   )
								snake_case_           =				audioldm_pipe.vocoder.config.sampling_rate

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

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

								snake_case_           =				audioldm_pipe(audio_length_in_s=0.032			,     **a   )
								snake_case_           =				output.audios[0]

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




	def 				_UpperCamelCase     (			self   )			->       List[Any]:
								snake_case_           =				self.get_dummy_components()
								snake_case_           =				AudioLDMPipeline(**a   )
								snake_case_           =				audioldm_pipe.to(a   )
								audioldm_pipe.set_progress_bar_config(disable=a   )

								snake_case_           =				['hey']

								snake_case_           =				audioldm_pipe(a			,     num_inference_steps=1   )
								snake_case_           =				output.audios.shape
								assert audio_shape == (1, 2_56)

								snake_case_           =				audioldm_pipe.vocoder.config
								config.model_in_dim *= 2
								snake_case_           =				SpeechTaHifiGan(a   ).to(a   )
								snake_case_           =				audioldm_pipe(a			,     num_inference_steps=1   )
								snake_case_           =				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   )			->       int:
								self._test_attention_slicing_forward_pass(test_mean_pixel_difference=a   )




	def 				_UpperCamelCase     (			self   )			->       Any:
								self._test_inference_batch_single_identical(test_mean_pixel_difference=a   )






	@unittest.skipIf(
	    torch_device != 'cuda' or not is_xformers_available()			,     reason='XFormers attention is only available with CUDA and `xformers` installed'			,     )
	def 				_UpperCamelCase     (			self   )			->       int:
								self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=a   )






@slow
class 					UpperCamelCase_ (     unittest.TestCase							):



	'''simple docstring'''




	def 				_UpperCamelCase     (			self   )			->       str:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()




	def 				_UpperCamelCase     (			self			,     a			,     a="cpu"			,     a=torch.floataa			,     a=0   )			->       int:
								snake_case_           =				torch.Generator(device=a   ).manual_seed(a   )
								snake_case_           =				np.random.RandomState(a   ).standard_normal((1, 8, 1_28, 16)   )
								snake_case_           =				torch.from_numpy(a   ).to(device=a			,     dtype=a   )
								snake_case_           =				{
								    'prompt': 'A hammer hitting a wooden surface',
								    'latents': latents,
								    'generator': generator,
								    'num_inference_steps': 3,
								    'guidance_scale': 2.5,
								}
								return inputs




	def 				_UpperCamelCase     (			self   )			->       str:
								snake_case_           =				AudioLDMPipeline.from_pretrained('cvssp/audioldm'   )
								snake_case_           =				audioldm_pipe.to(a   )
								audioldm_pipe.set_progress_bar_config(disable=a   )

								snake_case_           =				self.get_inputs(a   )
								snake_case_           =				25
								snake_case_           =				audioldm_pipe(**a   ).audios[0]

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

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






	def 				_UpperCamelCase     (			self   )			->       List[str]:
								snake_case_           =				AudioLDMPipeline.from_pretrained('cvssp/audioldm'   )
								snake_case_           =				LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config   )
								snake_case_           =				audioldm_pipe.to(a   )
								audioldm_pipe.set_progress_bar_config(disable=a   )

								snake_case_           =				self.get_inputs(a   )
								snake_case_           =				audioldm_pipe(**a   ).audios[0]

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

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













import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
	from .tokenization_rembert import RemBertTokenizer
else:
	lowercase    				=   None

lowercase    				=   logging.get_logger(__name__)
lowercase    				=   {"vocab_file": "sentencepiece.model", "tokenizer_file": "tokenizer.json"}

lowercase    				=   {
    "vocab_file": {
        "google/rembert": "https://huggingface.co/google/rembert/resolve/main/sentencepiece.model",
    },
    "tokenizer_file": {
        "google/rembert": "https://huggingface.co/google/rembert/resolve/main/tokenizer.json",
    },
}

lowercase    				=   {
    "google/rembert": 256,
}

lowercase    				=   "▁"




class 					UpperCamelCase_ (     snake_case_							):



	'''simple docstring'''




	lowerCAmelCase						=					VOCAB_FILES_NAMES
	lowerCAmelCase						=					PRETRAINED_VOCAB_FILES_MAP
	lowerCAmelCase						=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	lowerCAmelCase						=					RemBertTokenizer




	def __init__(			self			,     a=None			,     a=None			,     a=True			,     a=True			,     a=False			,     a="[CLS]"			,     a="[SEP]"			,     a="<unk>"			,     a="[SEP]"			,     a="<pad>"			,     a="[CLS]"			,     a="[MASK]"			,     **a			,     )			->       Union[str, Any]:
								# Mask token behave like a normal word, i.e. include the space before it
								snake_case_           =				AddedToken(a			,     lstrip=a			,     rstrip=a   ) if isinstance(a			,     a   ) else mask_token

								super().__init__(
								    a			,     tokenizer_file=a			,     do_lower_case=a			,     remove_space=a			,     keep_accents=a			,     bos_token=a			,     eos_token=a			,     unk_token=a			,     sep_token=a			,     pad_token=a			,     cls_token=a			,     mask_token=a			,     **a			,     )

								snake_case_           =				do_lower_case
								snake_case_           =				remove_space
								snake_case_           =				keep_accents
								snake_case_           =				vocab_file
								snake_case_           =				False if not self.vocab_file else True




	def 				_UpperCamelCase     (			self			,     a			,     a = None   )			->       List[int]:
								snake_case_           =				[self.sep_token_id]
								snake_case_           =				[self.cls_token_id]
								if token_ids_a is None:
															return cls + token_ids_a + sep
								return cls + token_ids_a + sep + token_ids_a + sep




	def 				_UpperCamelCase     (			self			,     a			,     a = None			,     a = False   )			->       List[int]:

								if already_has_special_tokens:
															if token_ids_a is not None:
																						raise ValueError(
																						    'You should not supply a second sequence if the provided sequence of '
																						    'ids is already formatted with special tokens for the model.'   )
															return [1 if x in [self.sep_token_id, self.cls_token_id] else 0 for x in token_ids_a]

								if token_ids_a is not None:
															return [1] + ([0] * len(a   )) + [1] + ([0] * len(a   )) + [1]
								return [1] + ([0] * len(a   )) + [1]




	def 				_UpperCamelCase     (			self			,     a			,     a = None   )			->       List[int]:
								snake_case_           =				[self.sep_token_id]
								snake_case_           =				[self.cls_token_id]

								if token_ids_a is None:
															return len(cls + token_ids_a + sep   ) * [0]
								return len(cls + token_ids_a + sep   ) * [0] + len(token_ids_a + sep   ) * [1]






	def 				_UpperCamelCase     (			self			,     a			,     a = None   )			->       Tuple[str]:
								if not os.path.isdir(a   ):
															logger.error('Vocabulary path ({}) should be a directory'.format(a   )   )
															return
								snake_case_           =				os.path.join(
								    a			,     (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']   )

								if os.path.abspath(self.vocab_file   ) != os.path.abspath(a   ):
															copyfile(self.vocab_file			,     a   )

								return (out_vocab_file,)





