



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			(							lowerCamelCase__			,      unittest.TestCase					):
							__lowerCamelCase	:	int    =AudioLDMPipeline
							__lowerCamelCase	:	List[str]    =TEXT_TO_AUDIO_PARAMS
							__lowerCamelCase	:	int    =TEXT_TO_AUDIO_BATCH_PARAMS
							__lowerCamelCase	:	List[Any]    =frozenset(
							    [
							        'num_inference_steps',
							        'num_waveforms_per_prompt',
							        'generator',
							        'latents',
							        'output_type',
							        'return_dict',
							        'callback',
							        'callback_steps',
							    ]					)





							def 						UpperCamelCase_					(   self		:		List[str]	):

											'''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=__lowercase      ,				)
											__a       =		DDIMScheduler(
											    beta_start=0.00085      ,				beta_end=0.012      ,				beta_schedule="""scaled_linear"""      ,				clip_sample=__lowercase      ,				set_alpha_to_one=__lowercase      ,				)
											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=1000      ,				projection_dim=32      ,				)
											__a       =		ClapTextModelWithProjection(__lowercase	)
											__a       =		RobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-random-roberta"""      ,				model_max_length=77	)

											__a       =		SpeechTaHifiGanConfig(
											    model_in_dim=8      ,				sampling_rate=16000      ,				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      ,				)

											__a       =		SpeechTaHifiGan(__lowercase	)

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





							def 						UpperCamelCase_					(   self		:		int      ,				__lowercase		:		Optional[Any]      ,				__lowercase		:		Union[str, Any]=0	):

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





							def 						UpperCamelCase_					(   self		:		str	):

											'''simple docstring'''
											__a       =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator

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

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

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

											__a       =		audio[:10]
											__a       =		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 						UpperCamelCase_					(   self		:		Dict	):

											'''simple docstring'''
											__a       =		self.get_dummy_components()
											__a       =		AudioLDMPipeline(**__lowercase	)
											__a       =		audioldm_pipe.to(__lowercase	)
											__a       =		audioldm_pipe.to(__lowercase	)
											audioldm_pipe.set_progress_bar_config(disable=__lowercase	)

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

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

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

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

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

											__a       =		prompt_embeds

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

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





							def 						UpperCamelCase_					(   self		:		List[str]	):

											'''simple docstring'''
											__a       =		self.get_dummy_components()
											__a       =		AudioLDMPipeline(**__lowercase	)
											__a       =		audioldm_pipe.to(__lowercase	)
											__a       =		audioldm_pipe.to(__lowercase	)
											audioldm_pipe.set_progress_bar_config(disable=__lowercase	)

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

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

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

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

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

															embeds.append(__lowercase	)

											__a     ,					__a       =		embeds

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

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





							def 						UpperCamelCase_					(   self		:		int	):

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

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

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

											__a       =		audio[:10]
											__a       =		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 						UpperCamelCase_					(   self		:		List[Any]	):

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

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

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

											assert audios.shape == (1, 256)

											# 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, 256)

											# test num_waveforms_per_prompt for single prompt
											__a       =		2
											__a       =		audioldm_pipe(__lowercase      ,				num_inference_steps=2      ,				num_waveforms_per_prompt=__lowercase	).audios

											assert audios.shape == (num_waveforms_per_prompt, 256)

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

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





							def 						UpperCamelCase_					(   self		:		Dict	):

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

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

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

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

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





							def 						UpperCamelCase_					(   self		:		str	):

											'''simple docstring'''
											__a       =		self.get_dummy_components()
											__a       =		AudioLDMPipeline(**__lowercase	)
											__a       =		audioldm_pipe.to(__lowercase	)
											audioldm_pipe.set_progress_bar_config(disable=__lowercase	)

											__a       =		["""hey"""]

											__a       =		audioldm_pipe(__lowercase      ,				num_inference_steps=1	)
											__a       =		output.audios.shape
											assert audio_shape == (1, 256)

											__a       =		audioldm_pipe.vocoder.config
											config.model_in_dim *= 2
											__a       =		SpeechTaHifiGan(__lowercase	).to(__lowercase	)
											__a       =		audioldm_pipe(__lowercase      ,				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, 256)





							def 						UpperCamelCase_					(   self		:		str	):

											'''simple docstring'''
											self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__lowercase	)





							def 						UpperCamelCase_					(   self		:		Optional[Any]	):

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

											'''simple docstring'''
											self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__lowercase	)






@slow
class    SCREAMING_SNAKE_CASE			(							unittest.TestCase					):





							def 						UpperCamelCase_					(   self		:		Optional[Any]	):

											'''simple docstring'''
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()





							def 						UpperCamelCase_					(   self		:		List[Any]      ,				__lowercase		:		Tuple      ,				__lowercase		:		List[Any]="cpu"      ,				__lowercase		:		Dict=torch.floataa      ,				__lowercase		:		int=0	):

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





							def 						UpperCamelCase_					(   self		:		List[Any]	):

											'''simple docstring'''
											__a       =		AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""	)
											__a       =		audioldm_pipe.to(__lowercase	)
											audioldm_pipe.set_progress_bar_config(disable=__lowercase	)

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

											assert audio.ndim == 1
											assert len(__lowercase	) == 81920

											__a       =		audio[77230:77240]
											__a       =		np.array(
											    [-0.4884, -0.4607, 0.0023, 0.5007, 0.5896, 0.5151, 0.3813, -0.0208, -0.3687, -0.4315]	)
											__a       =		np.abs(expected_slice - audio_slice	).max()
											assert max_diff < 1E-2







							def 						UpperCamelCase_					(   self		:		Any	):

											'''simple docstring'''
											__a       =		AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""	)
											__a       =		LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config	)
											__a       =		audioldm_pipe.to(__lowercase	)
											audioldm_pipe.set_progress_bar_config(disable=__lowercase	)

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

											assert audio.ndim == 1
											assert len(__lowercase	) == 81920

											__a       =		audio[27780:27790]
											__a       =		np.array([-0.2131, -0.0873, -0.0124, -0.0189, 0.0569, 0.1373, 0.1883, 0.2886, 0.3297, 0.2212]	)
											__a       =		np.abs(expected_slice - audio_slice	).max()
											assert max_diff < 3E-2





import unittest
from dataclasses import dataclass

import pytest

from accelerate.commands.config.config_args import SageMakerConfig
from accelerate.utils import ComputeEnvironment
from accelerate.utils.launch import _convert_nargs_to_dict






@dataclass
class    SCREAMING_SNAKE_CASE			(							lowerCamelCase__					):
							__lowerCamelCase	:	List[Any]    =ComputeEnvironment.AMAZON_SAGEMAKER
							__lowerCamelCase	:	Optional[int]    =True
							__lowerCamelCase	:	List[Any]    ='ml.p3.2xlarge'
							__lowerCamelCase	:	Optional[int]    ='accelerate_sagemaker_execution_role'
							__lowerCamelCase	:	Optional[Any]    ='hf-sm'
							__lowerCamelCase	:	Any    ='us-east-1'
							__lowerCamelCase	:	Any    =1
							__lowerCamelCase	:	List[str]    ='accelerate-sagemaker-1'
							__lowerCamelCase	:	Union[str, Any]    ='1.6'
							__lowerCamelCase	:	str    ='4.4'
							__lowerCamelCase	:	List[str]    ='train.py'
							__lowerCamelCase	:	List[Any]    =[
							    '--model_name_or_path',
							    'bert',
							    '--do_train',
							    'False',
							    '--epochs',
							    '3',
							    '--learning_rate',
							    '5e-5',
							    '--max_steps',
							    '50.5',
							]
							__lowerCamelCase	:	Tuple    =[
							    '--model_name_or_path',
							    'bert',
							    '--do_train',
							    '--do_test',
							    'False',
							    '--do_predict',
							    '--epochs',
							    '3',
							    '--learning_rate',
							    '5e-5',
							    '--max_steps',
							    '50.5',
							]






class    SCREAMING_SNAKE_CASE			(							unittest.TestCase					):





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

											'''simple docstring'''
											# If no defaults are changed, `to_kwargs` returns an empty dict.
											__a       =		_convert_nargs_to_dict(MockLaunchConfig.success_training_script_args	)
											assert isinstance(converted_args["""model_name_or_path"""]      ,				__lowercase	)
											assert isinstance(converted_args["""do_train"""]      ,				__lowercase	)
											assert isinstance(converted_args["""epochs"""]      ,				__lowercase	)
											assert isinstance(converted_args["""learning_rate"""]      ,				__lowercase	)
											assert isinstance(converted_args["""max_steps"""]      ,				__lowercase	)

											with pytest.raises(__lowercase	):
															_convert_nargs_to_dict(MockLaunchConfig.fail_training_script_args	)
