



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 	lowerCamelCase_    ( _A						,unittest.TestCase						):






						'''simple docstring'''
						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 			SCREAMING_SNAKE_CASE__		(	self      :			List[Any] )      ->				Optional[Any]:
									torch.manual_seed(0 )
									A							:					List[str]              =						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							:					List[str]              =						DDIMScheduler(
									    beta_start=0.00085	,	beta_end=0.012	,	beta_schedule="scaled_linear"	,	clip_sample=__lowerCamelCase	,	set_alpha_to_one=__lowerCamelCase	,	)
									torch.manual_seed(0 )
									A							:					Optional[int]              =						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							:					List[str]              =						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							:					Any              =						ClapTextModelWithProjection(__lowerCamelCase )
									A							:					List[str]              =						RobertaTokenizer.from_pretrained("hf-internal-testing/tiny-random-roberta"	,	model_max_length=77 )

									A							:					List[str]              =						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							:					Optional[int]              =						SpeechTaHifiGan(__lowerCamelCase )

									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      :			Optional[Any]	,	__lowerCamelCase      :			Optional[Any]	,	__lowerCamelCase      :			Optional[Any]=0 )      ->				int:
									if str(__lowerCamelCase ).startswith("mps" ):
												A							:					Any              =						torch.manual_seed(__lowerCamelCase )
									else:
												A							:					Union[str, Any]              =						torch.Generator(device=__lowerCamelCase ).manual_seed(__lowerCamelCase )
									A							:					List[Any]              =						{
									    "prompt": "A hammer hitting a wooden surface",
									    "generator": generator,
									    "num_inference_steps": 2,
									    "guidance_scale": 6.0,
									}
									return inputs







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Dict )      ->				Union[str, Any]:
									A							:					int              =						"cpu"  # ensure determinism for the device-dependent torch.Generator

									A							:					Dict              =						self.get_dummy_components()
									A							:					List[str]              =						AudioLDMPipeline(**__lowerCamelCase )
									A							:					Tuple              =						audioldm_pipe.to(__lowerCamelCase )
									audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase )

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

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

									A							:					int              =						audio[:10]
									A							:					Dict              =						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      :			List[Any] )      ->				List[str]:
									A							:					List[str]              =						self.get_dummy_components()
									A							:					List[Any]              =						AudioLDMPipeline(**__lowerCamelCase )
									A							:					Tuple              =						audioldm_pipe.to(__lowerCamelCase )
									A							:					Optional[Any]              =						audioldm_pipe.to(__lowerCamelCase )
									audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase )

									A							:					str              =						self.get_dummy_inputs(__lowerCamelCase )
									A							:					Optional[Any]              =						3 * [inputs["prompt"]]

									# forward
									A							:					List[str]              =						audioldm_pipe(**__lowerCamelCase )
									A							:					Dict              =						output.audios[0]

									A							:					Tuple              =						self.get_dummy_inputs(__lowerCamelCase )
									A							:					List[Any]              =						3 * [inputs.pop("prompt" )]

									A							:					Optional[int]              =						audioldm_pipe.tokenizer(
									    __lowerCamelCase	,	padding="max_length"	,	max_length=audioldm_pipe.tokenizer.model_max_length	,	truncation=__lowerCamelCase	,	return_tensors="pt"	,	)
									A							:					str              =						text_inputs["input_ids"].to(__lowerCamelCase )

									A							:					str              =						audioldm_pipe.text_encoder(
									    __lowerCamelCase	,	)
									A							:					Tuple              =						prompt_embeds.text_embeds
									# additional L_2 normalization over each hidden-state
									A							:					Optional[int]              =						F.normalize(__lowerCamelCase	,	dim=-1 )

									A							:					Union[str, Any]              =						prompt_embeds

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

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







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Optional[int] )      ->				List[str]:
									A							:					Optional[Any]              =						self.get_dummy_components()
									A							:					List[str]              =						AudioLDMPipeline(**__lowerCamelCase )
									A							:					List[Any]              =						audioldm_pipe.to(__lowerCamelCase )
									A							:					Optional[int]              =						audioldm_pipe.to(__lowerCamelCase )
									audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase )

									A							:					Optional[Any]              =						self.get_dummy_inputs(__lowerCamelCase )
									A							:					int              =						3 * ["this is a negative prompt"]
									A							:					Dict              =						negative_prompt
									A							:					str              =						3 * [inputs["prompt"]]

									# forward
									A							:					Tuple              =						audioldm_pipe(**__lowerCamelCase )
									A							:					str              =						output.audios[0]

									A							:					List[str]              =						self.get_dummy_inputs(__lowerCamelCase )
									A							:					Dict              =						3 * [inputs.pop("prompt" )]

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

												A							:					List[Any]              =						audioldm_pipe.text_encoder(
												    __lowerCamelCase	,	)
												A							:					List[Any]              =						text_embeds.text_embeds
												# additional L_2 normalization over each hidden-state
												A							:					Dict              =						F.normalize(__lowerCamelCase	,	dim=-1 )

												embeds.append(__lowerCamelCase )

									A       , A							:					List[Any]              =						embeds

									# forward
									A							:					Union[str, Any]              =						audioldm_pipe(**__lowerCamelCase )
									A							:					Optional[Any]              =						output.audios[0]

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







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Optional[Any] )      ->				str:
									A							:					Dict              =						"cpu"  # ensure determinism for the device-dependent torch.Generator
									A							:					List[Any]              =						self.get_dummy_components()
									A							:					Optional[int]              =						PNDMScheduler(skip_prk_steps=__lowerCamelCase )
									A							:					Any              =						AudioLDMPipeline(**__lowerCamelCase )
									A							:					List[str]              =						audioldm_pipe.to(__lowerCamelCase )
									audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase )

									A							:					int              =						self.get_dummy_inputs(__lowerCamelCase )
									A							:					Any              =						"egg cracking"
									A							:					List[Any]              =						audioldm_pipe(**__lowerCamelCase	,	negative_prompt=__lowerCamelCase )
									A							:					Any              =						output.audios[0]

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

									A							:					Any              =						audio[:10]
									A							:					Tuple              =						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      :			int )      ->				Dict:
									A							:					Tuple              =						"cpu"  # ensure determinism for the device-dependent torch.Generator
									A							:					Optional[Any]              =						self.get_dummy_components()
									A							:					List[Any]              =						PNDMScheduler(skip_prk_steps=__lowerCamelCase )
									A							:					Tuple              =						AudioLDMPipeline(**__lowerCamelCase )
									A							:					Optional[Any]              =						audioldm_pipe.to(__lowerCamelCase )
									audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase )

									A							:					Union[str, Any]              =						"A hammer hitting a wooden surface"

									# test num_waveforms_per_prompt=1 (default)
									A							:					Optional[int]              =						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							:					str              =						2
									A							:					List[str]              =						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							:					Union[str, Any]              =						2
									A							:					Optional[Any]              =						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							:					Optional[int]              =						2
									A							:					Tuple              =						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 			SCREAMING_SNAKE_CASE__		(	self      :			str )      ->				Optional[int]:
									A							:					Union[str, Any]              =						"cpu"  # ensure determinism for the device-dependent torch.Generator
									A							:					int              =						self.get_dummy_components()
									A							:					List[Any]              =						AudioLDMPipeline(**__lowerCamelCase )
									A							:					Optional[Any]              =						audioldm_pipe.to(__lowerCamelCase )
									audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase )
									A							:					Any              =						audioldm_pipe.vocoder.config.sampling_rate

									A							:					Tuple              =						self.get_dummy_inputs(__lowerCamelCase )
									A							:					Optional[Any]              =						audioldm_pipe(audio_length_in_s=0.016	,	**__lowerCamelCase )
									A							:					str              =						output.audios[0]

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

									A							:					Dict              =						audioldm_pipe(audio_length_in_s=0.032	,	**__lowerCamelCase )
									A							:					Dict              =						output.audios[0]

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







						def 			SCREAMING_SNAKE_CASE__		(	self      :			Any )      ->				Union[str, Any]:
									A							:					Any              =						self.get_dummy_components()
									A							:					List[str]              =						AudioLDMPipeline(**__lowerCamelCase )
									A							:					Any              =						audioldm_pipe.to(__lowerCamelCase )
									audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase )

									A							:					Tuple              =						["hey"]

									A							:					Tuple              =						audioldm_pipe(__lowerCamelCase	,	num_inference_steps=1 )
									A							:					Optional[int]              =						output.audios.shape
									assert audio_shape == (1, 2_56)

									A							:					Tuple              =						audioldm_pipe.vocoder.config
									config.model_in_dim *= 2
									A							:					Tuple              =						SpeechTaHifiGan(__lowerCamelCase ).to(__lowerCamelCase )
									A							:					Any              =						audioldm_pipe(__lowerCamelCase	,	num_inference_steps=1 )
									A							:					Any              =						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 			SCREAMING_SNAKE_CASE__		(	self      :			List[str] )      ->				int:
									self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__lowerCamelCase )







						def 			SCREAMING_SNAKE_CASE__		(	self      :			List[Any] )      ->				Optional[Any]:
									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 			SCREAMING_SNAKE_CASE__		(	self      :			Dict )      ->				Optional[Any]:
									self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__lowerCamelCase )







@slow
class 	lowerCamelCase_    ( unittest.TestCase						):






						'''simple docstring'''







						def 			SCREAMING_SNAKE_CASE__		(	self      :			int )      ->				int:
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()







						def 			SCREAMING_SNAKE_CASE__		(	self      :			List[str]	,	__lowerCamelCase      :			str	,	__lowerCamelCase      :			str="cpu"	,	__lowerCamelCase      :			int=torch.floataa	,	__lowerCamelCase      :			Optional[Any]=0 )      ->				Any:
									A							:					List[str]              =						torch.Generator(device=__lowerCamelCase ).manual_seed(__lowerCamelCase )
									A							:					Optional[Any]              =						np.random.RandomState(__lowerCamelCase ).standard_normal((1, 8, 1_28, 16) )
									A							:					Optional[Any]              =						torch.from_numpy(__lowerCamelCase ).to(device=__lowerCamelCase	,	dtype=__lowerCamelCase )
									A							:					int              =						{
									    "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      :			Optional[Any] )      ->				List[str]:
									A							:					int              =						AudioLDMPipeline.from_pretrained("cvssp/audioldm" )
									A							:					List[Any]              =						audioldm_pipe.to(__lowerCamelCase )
									audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase )

									A							:					Optional[Any]              =						self.get_inputs(__lowerCamelCase )
									A							:					int              =						25
									A							:					Dict              =						audioldm_pipe(**__lowerCamelCase ).audios[0]

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

									A							:					str              =						audio[7_72_30:7_72_40]
									A							:					Tuple              =						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      :			str )      ->				Tuple:
									A							:					List[str]              =						AudioLDMPipeline.from_pretrained("cvssp/audioldm" )
									A							:					str              =						LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config )
									A							:					str              =						audioldm_pipe.to(__lowerCamelCase )
									audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase )

									A							:					List[str]              =						self.get_inputs(__lowerCamelCase )
									A							:					str              =						audioldm_pipe(**__lowerCamelCase ).audios[0]

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

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




import argparse

import torch
from huggingface_hub import hf_hub_download

from transformers import AutoTokenizer, RobertaPreLayerNormConfig, RobertaPreLayerNormForMaskedLM
from transformers.utils import logging


logging.set_verbosity_info()
__SCREAMING_SNAKE_CASE           =  logging.get_logger(__name__)




def 	UpperCAmelCase   (						_lowerCamelCase ,    _lowerCamelCase    ):
			A							:					Optional[Any]              =						RobertaPreLayerNormConfig.from_pretrained(
			    _lowerCamelCase ,    architectures=["RobertaPreLayerNormForMaskedLM"]    )

			# convert state_dict
			A							:					List[Any]              =						torch.load(hf_hub_download(repo_id=_lowerCamelCase ,    filename="pytorch_model.bin"    )    )
			A							:					Union[str, Any]              =						{}
			for tensor_key, tensor_value in original_state_dict.items():
						# The transformer implementation gives the model a unique name, rather than overwiriting 'roberta'
						if tensor_key.startswith("roberta."    ):
									A							:					int              =						"roberta_prelayernorm." + tensor_key[len("roberta."    ) :]

						# The original implementation contains weights which are not used, remove them from the state_dict
						if tensor_key.endswith(".self.LayerNorm.weight"    ) or tensor_key.endswith(".self.LayerNorm.bias"    ):
									continue

						A							:					Any              =						tensor_value

			A							:					Optional[int]              =						RobertaPreLayerNormForMaskedLM.from_pretrained(
			    pretrained_model_name_or_path=_lowerCamelCase ,    config=_lowerCamelCase ,    state_dict=_lowerCamelCase    )
			model.save_pretrained(_lowerCamelCase    )

			# convert tokenizer
			A							:					Optional[Any]              =						AutoTokenizer.from_pretrained(_lowerCamelCase    )
			tokenizer.save_pretrained(_lowerCamelCase    )


if __name__ == "__main__":
							__SCREAMING_SNAKE_CASE           =  argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    """--checkpoint-repo""",
							    default=None,
							    type=str,
							    required=True,
							    help="""Path the official PyTorch dump, e.g. 'andreasmadsen/efficient_mlm_m0.40'.""",
							)
							parser.add_argument(
							    """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
							)
							__SCREAMING_SNAKE_CASE           =  parser.parse_args()
							convert_roberta_prelayernorm_checkpoint_to_pytorch(args.checkpoint_repo, args.pytorch_dump_folder_path)