



"""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        _UpperCAmelCase       (			lowerCAmelCase__    ,     unittest.TestCase):
 _lowerCAmelCase   :			Optional[int]							=		AudioLDMPipeline
 _lowerCAmelCase   :			List[Any]							=		TEXT_TO_AUDIO_PARAMS
 _lowerCAmelCase   :			int							=		TEXT_TO_AUDIO_BATCH_PARAMS
 _lowerCAmelCase   :			Optional[Any]							=		frozenset(
     [
         """num_inference_steps""",
         """num_waveforms_per_prompt""",
         """generator""",
         """latents""",
         """output_type""",
         """return_dict""",
         """callback""",
         """callback_steps""",
     ])

 def 						_snake_case					(     self		:       List[Any]				):
     torch.manual_seed(0				)
     snake_case_ :     Dict				=	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_  ,     )
     snake_case_ :     List[Any]				=	DDIMScheduler(
         beta_start=0.0_00_85  ,     beta_end=0.0_12  ,     beta_schedule='''scaled_linear'''  ,     clip_sample=lowercase_  ,     set_alpha_to_one=lowercase_  ,     )
     torch.manual_seed(0				)
     snake_case_ :     List[Any]				=	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_ :     Dict				=	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  ,     )
     snake_case_ :     Union[str, Any]				=	ClapTextModelWithProjection(lowercase_				)
     snake_case_ :     Dict				=	RobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-random-roberta'''  ,     model_max_length=77				)

     snake_case_ :     int				=	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_  ,     )

     snake_case_ :     List[str]				=	SpeechTaHifiGan(lowercase_				)

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

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

 def 						_snake_case					(     self		:       Tuple				):
     snake_case_ :     int				=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator

     snake_case_ :     Tuple				=	self.get_dummy_components()
     snake_case_ :     Tuple				=	AudioLDMPipeline(**lowercase_				)
     snake_case_ :     Any				=	audioldm_pipe.to(lowercase_				)
     audioldm_pipe.set_progress_bar_config(disable=lowercase_				)

     snake_case_ :     List[Any]				=	self.get_dummy_inputs(lowercase_				)
     snake_case_ :     List[str]				=	audioldm_pipe(**lowercase_				)
     snake_case_ :     Any				=	output.audios[0]

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

     snake_case_ :     List[Any]				=	audio[:10]
     snake_case_ :     Tuple				=	np.array(
         [-0.00_50, 0.00_50, -0.00_60, 0.00_33, -0.00_26, 0.00_33, -0.00_27, 0.00_33, -0.00_28, 0.00_33]				)

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

 def 						_snake_case					(     self		:       Optional[Any]				):
     snake_case_ :     Tuple				=	self.get_dummy_components()
     snake_case_ :     str				=	AudioLDMPipeline(**lowercase_				)
     snake_case_ :     Optional[int]				=	audioldm_pipe.to(lowercase_				)
     snake_case_ :     Optional[int]				=	audioldm_pipe.to(lowercase_				)
     audioldm_pipe.set_progress_bar_config(disable=lowercase_				)

     snake_case_ :     str				=	self.get_dummy_inputs(lowercase_				)
     snake_case_ :     Optional[int]				=	3 * [inputs['''prompt''']]

     # forward
     snake_case_ :     List[Any]				=	audioldm_pipe(**lowercase_				)
     snake_case_ :     Optional[int]				=	output.audios[0]

     snake_case_ :     Tuple				=	self.get_dummy_inputs(lowercase_				)
     snake_case_ :     str				=	3 * [inputs.pop('''prompt'''				)]

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

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

     snake_case_ :     Tuple				=	prompt_embeds

     # forward
     snake_case_ :     Dict				=	audioldm_pipe(**lowercase_				)
     snake_case_ :     List[Any]				=	output.audios[0]

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

 def 						_snake_case					(     self		:       List[Any]				):
     snake_case_ :     Dict				=	self.get_dummy_components()
     snake_case_ :     int				=	AudioLDMPipeline(**lowercase_				)
     snake_case_ :     Tuple				=	audioldm_pipe.to(lowercase_				)
     snake_case_ :     Union[str, Any]				=	audioldm_pipe.to(lowercase_				)
     audioldm_pipe.set_progress_bar_config(disable=lowercase_				)

     snake_case_ :     Tuple				=	self.get_dummy_inputs(lowercase_				)
     snake_case_ :     int				=	3 * ['''this is a negative prompt''']
     snake_case_ :     int				=	negative_prompt
     snake_case_ :     Union[str, Any]				=	3 * [inputs['''prompt''']]

     # forward
     snake_case_ :     Union[str, Any]				=	audioldm_pipe(**lowercase_				)
     snake_case_ :     Any				=	output.audios[0]

     snake_case_ :     Optional[int]				=	self.get_dummy_inputs(lowercase_				)
     snake_case_ :     Any				=	3 * [inputs.pop('''prompt'''				)]

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

         snake_case_ :     Optional[Any]				=	audioldm_pipe.text_encoder(
             lowercase_  ,     )
         snake_case_ :     Tuple				=	text_embeds.text_embeds
         # additional L_2 normalization over each hidden-state
         snake_case_ :     List[str]				=	F.normalize(lowercase_  ,     dim=-1				)

         embeds.append(lowercase_				)

     snake_case_,	snake_case_ :     str				=	embeds

     # forward
     snake_case_ :     Optional[Any]				=	audioldm_pipe(**lowercase_				)
     snake_case_ :     int				=	output.audios[0]

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

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

     snake_case_ :     Union[str, Any]				=	self.get_dummy_inputs(lowercase_				)
     snake_case_ :     Tuple				=	'''egg cracking'''
     snake_case_ :     Optional[int]				=	audioldm_pipe(**lowercase_  ,     negative_prompt=lowercase_				)
     snake_case_ :     Any				=	output.audios[0]

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

     snake_case_ :     int				=	audio[:10]
     snake_case_ :     Union[str, Any]				=	np.array(
         [-0.00_51, 0.00_50, -0.00_60, 0.00_34, -0.00_26, 0.00_33, -0.00_27, 0.00_33, -0.00_28, 0.00_32]				)

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

 def 						_snake_case					(     self		:       List[str]				):
     snake_case_ :     List[Any]				=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
     snake_case_ :     Union[str, Any]				=	self.get_dummy_components()
     snake_case_ :     Dict				=	PNDMScheduler(skip_prk_steps=lowercase_				)
     snake_case_ :     Any				=	AudioLDMPipeline(**lowercase_				)
     snake_case_ :     Any				=	audioldm_pipe.to(lowercase_				)
     audioldm_pipe.set_progress_bar_config(disable=lowercase_				)

     snake_case_ :     int				=	'''A hammer hitting a wooden surface'''

     # test num_waveforms_per_prompt=1 (default)
     snake_case_ :     Optional[Any]				=	audioldm_pipe(lowercase_  ,     num_inference_steps=2				).audios

     assert audios.shape == (1, 256)

     # test num_waveforms_per_prompt=1 (default) for batch of prompts
     snake_case_ :     int				=	2
     snake_case_ :     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
     snake_case_ :     List[str]				=	2
     snake_case_ :     Union[str, Any]				=	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
     snake_case_ :     List[str]				=	2
     snake_case_ :     Dict				=	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 						_snake_case					(     self		:       int				):
     snake_case_ :     int				=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
     snake_case_ :     Dict				=	self.get_dummy_components()
     snake_case_ :     Optional[Any]				=	AudioLDMPipeline(**lowercase_				)
     snake_case_ :     str				=	audioldm_pipe.to(lowercase_				)
     audioldm_pipe.set_progress_bar_config(disable=lowercase_				)
     snake_case_ :     Optional[Any]				=	audioldm_pipe.vocoder.config.sampling_rate

     snake_case_ :     Any				=	self.get_dummy_inputs(lowercase_				)
     snake_case_ :     Tuple				=	audioldm_pipe(audio_length_in_s=0.0_16  ,     **lowercase_				)
     snake_case_ :     Optional[Any]				=	output.audios[0]

     assert audio.ndim == 1
     assert len(lowercase_				) / vocoder_sampling_rate == 0.0_16

     snake_case_ :     List[str]				=	audioldm_pipe(audio_length_in_s=0.0_32  ,     **lowercase_				)
     snake_case_ :     Tuple				=	output.audios[0]

     assert audio.ndim == 1
     assert len(lowercase_				) / vocoder_sampling_rate == 0.0_32

 def 						_snake_case					(     self		:       Any				):
     snake_case_ :     List[str]				=	self.get_dummy_components()
     snake_case_ :     str				=	AudioLDMPipeline(**lowercase_				)
     snake_case_ :     Optional[int]				=	audioldm_pipe.to(lowercase_				)
     audioldm_pipe.set_progress_bar_config(disable=lowercase_				)

     snake_case_ :     int				=	['''hey''']

     snake_case_ :     int				=	audioldm_pipe(lowercase_  ,     num_inference_steps=1				)
     snake_case_ :     Optional[Any]				=	output.audios.shape
     assert audio_shape == (1, 256)

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

 def 						_snake_case					(     self		:       Optional[int]				):
     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 						_snake_case					(     self		:       Optional[Any]				):
     self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowercase_				)





@slow
class        _UpperCAmelCase       (			unittest.TestCase):

 def 						_snake_case					(     self		:       Any				):
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()

 def 						_snake_case					(     self		:       List[str]  ,     lowercase_		:       List[str]  ,     lowercase_		:       Optional[Any]="cpu"  ,     lowercase_		:       Optional[Any]=torch.floataa  ,     lowercase_		:       Any=0				):
     snake_case_ :     str				=	torch.Generator(device=lowercase_				).manual_seed(lowercase_				)
     snake_case_ :     List[Any]				=	np.random.RandomState(lowercase_				).standard_normal((1, 8, 128, 16)				)
     snake_case_ :     Dict				=	torch.from_numpy(lowercase_				).to(device=lowercase_  ,     dtype=lowercase_				)
     snake_case_ :     Union[str, Any]				=	{
         '''prompt''': '''A hammer hitting a wooden surface''',
         '''latents''': latents,
         '''generator''': generator,
         '''num_inference_steps''': 3,
         '''guidance_scale''': 2.5,
     }
     return inputs

 def 						_snake_case					(     self		:       List[Any]				):
     snake_case_ :     Optional[int]				=	AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''				)
     snake_case_ :     List[str]				=	audioldm_pipe.to(lowercase_				)
     audioldm_pipe.set_progress_bar_config(disable=lowercase_				)

     snake_case_ :     Optional[Any]				=	self.get_inputs(lowercase_				)
     snake_case_ :     Tuple				=	25
     snake_case_ :     int				=	audioldm_pipe(**lowercase_				).audios[0]

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

     snake_case_ :     Tuple				=	audio[77230:77240]
     snake_case_ :     Any				=	np.array(
         [-0.48_84, -0.46_07, 0.00_23, 0.50_07, 0.58_96, 0.51_51, 0.38_13, -0.02_08, -0.36_87, -0.43_15]				)
     snake_case_ :     Union[str, Any]				=	np.abs(expected_slice - audio_slice				).max()
     assert max_diff < 1E-2






 def 						_snake_case					(     self		:       Union[str, Any]				):
     snake_case_ :     Optional[int]				=	AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''				)
     snake_case_ :     Dict				=	LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config				)
     snake_case_ :     List[Any]				=	audioldm_pipe.to(lowercase_				)
     audioldm_pipe.set_progress_bar_config(disable=lowercase_				)

     snake_case_ :     str				=	self.get_inputs(lowercase_				)
     snake_case_ :     Tuple				=	audioldm_pipe(**lowercase_				).audios[0]

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

     snake_case_ :     Dict				=	audio[27780:27790]
     snake_case_ :     Optional[Any]				=	np.array([-0.21_31, -0.08_73, -0.01_24, -0.01_89, 0.05_69, 0.13_73, 0.18_83, 0.28_86, 0.32_97, 0.22_12]				)
     snake_case_ :     int				=	np.abs(expected_slice - audio_slice				).max()
     assert max_diff < 3E-2










"""simple docstring"""

import argparse
import json

import numpy
import torch

from transformers.models.xlm.tokenization_xlm import VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()




def __lowercase			(							_a		,			_a    ):
    # Load checkpoint
    snake_case_ :     Optional[Any]				=	torch.load(_a		,			map_location='''cpu'''    )

    snake_case_ :     Union[str, Any]				=	chkpt['''model''']

    # We have the base model one level deeper than the original XLM repository
    snake_case_ :     Dict				=	{}
    for k, v in state_dict.items():
        if "pred_layer" in k:
            snake_case_ :     Union[str, Any]				=	v
        else:
            snake_case_ :     Dict				=	v

    snake_case_ :     Union[str, Any]				=	chkpt['''params''']
    snake_case_ :     int				=	{n: v for n, v in config.items() if not isinstance(_a		,			(torch.FloatTensor, numpy.ndarray)    )}

    snake_case_ :     int				=	chkpt['''dico_word2id''']
    snake_case_ :     str				=	{s + '''</w>''' if s.find('''@@'''    ) == -1 and i > 13 else s.replace('''@@'''		,			''''''    ): i for s, i in vocab.items()}

    # Save pytorch-model
    snake_case_ :     Union[str, Any]				=	pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
    snake_case_ :     Union[str, Any]				=	pytorch_dump_folder_path + '''/''' + CONFIG_NAME
    snake_case_ :     Any				=	pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''vocab_file''']

    print(f"Save PyTorch model to {pytorch_weights_dump_path}"    )
    torch.save(_a		,			_a    )

    print(f"Save configuration file to {pytorch_config_dump_path}"    )
    with open(_a		,			'''w'''		,			encoding='''utf-8'''    ) as f:
        f.write(json.dumps(_a		,			indent=2    ) + '''\n'''    )

    print(f"Save vocab file to {pytorch_config_dump_path}"    )
    with open(_a		,			'''w'''		,			encoding='''utf-8'''    ) as f:
        f.write(json.dumps(_a		,			indent=2    ) + '''\n'''    )


if __name__ == "__main__":
     lowercase__    :  Optional[Any]							    =	argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '''--xlm_checkpoint_path''', default=None, type=str, required=True, help='''Path the official PyTorch dump.'''
     )
     parser.add_argument(
         '''--pytorch_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
     )
     lowercase__    :  List[str]							    =	parser.parse_args()
     convert_xlm_checkpoint_to_pytorch(args.xlm_checkpoint_path, args.pytorch_dump_folder_path)





