



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 (     lowerCamelCase__       ,						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 							snake_case   (							self		):



									"""simple docstring"""





									torch.manual_seed(0		)
									_lowerCAmelCase						=    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=_snake_case						,				)
									_lowerCAmelCase						=    DDIMScheduler(
									    beta_start=0.0_0085						,				beta_end=0.012						,				beta_schedule="""scaled_linear"""						,				clip_sample=_snake_case						,				set_alpha_to_one=_snake_case						,				)
									torch.manual_seed(0		)
									_lowerCAmelCase						=    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		)
									_lowerCAmelCase						=    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						,				)
									_lowerCAmelCase						=    ClapTextModelWithProjection(_snake_case		)
									_lowerCAmelCase						=    RobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-random-roberta"""						,				model_max_length=77		)

									_lowerCAmelCase						=    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=_snake_case						,				)

									_lowerCAmelCase						=    SpeechTaHifiGan(_snake_case		)

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




				def 							snake_case   (							self						,				_snake_case						,				_snake_case=0		):



									"""simple docstring"""





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




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    """cpu"""  # ensure determinism for the device-dependent torch.Generator

									_lowerCAmelCase						=    self.get_dummy_components()
									_lowerCAmelCase						=    AudioLDMPipeline(**_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe.to(_snake_case		)
									audioldm_pipe.set_progress_bar_config(disable=_snake_case		)

									_lowerCAmelCase						=    self.get_dummy_inputs(_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe(**_snake_case		)
									_lowerCAmelCase						=    output.audios[0]

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

									_lowerCAmelCase						=    audio[:10]
									_lowerCAmelCase						=    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 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.get_dummy_components()
									_lowerCAmelCase						=    AudioLDMPipeline(**_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe.to(_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe.to(_snake_case		)
									audioldm_pipe.set_progress_bar_config(disable=_snake_case		)

									_lowerCAmelCase						=    self.get_dummy_inputs(_snake_case		)
									_lowerCAmelCase						=    3 * [inputs["""prompt"""]]

									# forward
									_lowerCAmelCase						=    audioldm_pipe(**_snake_case		)
									_lowerCAmelCase						=    output.audios[0]

									_lowerCAmelCase						=    self.get_dummy_inputs(_snake_case		)
									_lowerCAmelCase						=    3 * [inputs.pop("""prompt"""		)]

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

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

									_lowerCAmelCase						=    prompt_embeds

									# forward
									_lowerCAmelCase						=    audioldm_pipe(**_snake_case		)
									_lowerCAmelCase						=    output.audios[0]

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




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.get_dummy_components()
									_lowerCAmelCase						=    AudioLDMPipeline(**_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe.to(_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe.to(_snake_case		)
									audioldm_pipe.set_progress_bar_config(disable=_snake_case		)

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

									# forward
									_lowerCAmelCase						=    audioldm_pipe(**_snake_case		)
									_lowerCAmelCase						=    output.audios[0]

									_lowerCAmelCase						=    self.get_dummy_inputs(_snake_case		)
									_lowerCAmelCase						=    3 * [inputs.pop("""prompt"""		)]

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

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

														embeds.append(_snake_case		)

									_lowerCAmelCase	,	_lowerCAmelCase						=    embeds

									# forward
									_lowerCAmelCase						=    audioldm_pipe(**_snake_case		)
									_lowerCAmelCase						=    output.audios[0]

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




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    """cpu"""  # ensure determinism for the device-dependent torch.Generator
									_lowerCAmelCase						=    self.get_dummy_components()
									_lowerCAmelCase						=    PNDMScheduler(skip_prk_steps=_snake_case		)
									_lowerCAmelCase						=    AudioLDMPipeline(**_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe.to(_snake_case		)
									audioldm_pipe.set_progress_bar_config(disable=_snake_case		)

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

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

									_lowerCAmelCase						=    audio[:10]
									_lowerCAmelCase						=    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 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    """cpu"""  # ensure determinism for the device-dependent torch.Generator
									_lowerCAmelCase						=    self.get_dummy_components()
									_lowerCAmelCase						=    PNDMScheduler(skip_prk_steps=_snake_case		)
									_lowerCAmelCase						=    AudioLDMPipeline(**_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe.to(_snake_case		)
									audioldm_pipe.set_progress_bar_config(disable=_snake_case		)

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

									# test num_waveforms_per_prompt=1 (default)
									_lowerCAmelCase						=    audioldm_pipe(_snake_case						,				num_inference_steps=2		).audios

									assert audios.shape == (1, 256)

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

									assert audios.shape == (batch_size, 256)

									# test num_waveforms_per_prompt for single prompt
									_lowerCAmelCase						=    2
									_lowerCAmelCase						=    audioldm_pipe(_snake_case						,				num_inference_steps=2						,				num_waveforms_per_prompt=_snake_case		).audios

									assert audios.shape == (num_waveforms_per_prompt, 256)

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

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




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    """cpu"""  # ensure determinism for the device-dependent torch.Generator
									_lowerCAmelCase						=    self.get_dummy_components()
									_lowerCAmelCase						=    AudioLDMPipeline(**_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe.to(_snake_case		)
									audioldm_pipe.set_progress_bar_config(disable=_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe.vocoder.config.sampling_rate

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

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

									_lowerCAmelCase						=    audioldm_pipe(audio_length_in_s=0.032						,				**_snake_case		)
									_lowerCAmelCase						=    output.audios[0]

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




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    self.get_dummy_components()
									_lowerCAmelCase						=    AudioLDMPipeline(**_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe.to(_snake_case		)
									audioldm_pipe.set_progress_bar_config(disable=_snake_case		)

									_lowerCAmelCase						=    ["""hey"""]

									_lowerCAmelCase						=    audioldm_pipe(_snake_case						,				num_inference_steps=1		)
									_lowerCAmelCase						=    output.audios.shape
									assert audio_shape == (1, 256)

									_lowerCAmelCase						=    audioldm_pipe.vocoder.config
									config.model_in_dim *= 2
									_lowerCAmelCase						=    SpeechTaHifiGan(_snake_case		).to(_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe(_snake_case						,				num_inference_steps=1		)
									_lowerCAmelCase						=    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		):



									"""simple docstring"""





									self._test_attention_slicing_forward_pass(test_mean_pixel_difference=_snake_case		)




				def 							snake_case   (							self		):



									"""simple docstring"""





									self._test_inference_batch_single_identical(test_mean_pixel_difference=_snake_case		)




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



									"""simple docstring"""





									self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=_snake_case		)





@slow
class 					__lowerCAmelCase (     unittest.TestCase       ):




				def 							snake_case   (							self		):



									"""simple docstring"""





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




				def 							snake_case   (							self						,				_snake_case						,				_snake_case="cpu"						,				_snake_case=torch.floataa						,				_snake_case=0		):



									"""simple docstring"""





									_lowerCAmelCase						=    torch.Generator(device=_snake_case		).manual_seed(_snake_case		)
									_lowerCAmelCase						=    np.random.RandomState(_snake_case		).standard_normal((1, 8, 128, 16)		)
									_lowerCAmelCase						=    torch.from_numpy(_snake_case		).to(device=_snake_case						,				dtype=_snake_case		)
									_lowerCAmelCase						=    {
									    """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		):



									"""simple docstring"""





									_lowerCAmelCase						=    AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""		)
									_lowerCAmelCase						=    audioldm_pipe.to(_snake_case		)
									audioldm_pipe.set_progress_bar_config(disable=_snake_case		)

									_lowerCAmelCase						=    self.get_inputs(_snake_case		)
									_lowerCAmelCase						=    25
									_lowerCAmelCase						=    audioldm_pipe(**_snake_case		).audios[0]

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

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




				def 							snake_case   (							self		):



									"""simple docstring"""





									_lowerCAmelCase						=    AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""		)
									_lowerCAmelCase						=    LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config		)
									_lowerCAmelCase						=    audioldm_pipe.to(_snake_case		)
									audioldm_pipe.set_progress_bar_config(disable=_snake_case		)

									_lowerCAmelCase						=    self.get_inputs(_snake_case		)
									_lowerCAmelCase						=    audioldm_pipe(**_snake_case		).audios[0]

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

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










import argparse
import re

import torch
from CLAP import create_model

from transformers import AutoFeatureExtractor, ClapConfig, ClapModel


__UpperCAmelCase       :   Any				       =							{
    "text_branch": "text_model",
    "audio_branch": "audio_model.audio_encoder",
    "attn": "attention.self",
    "self.proj": "output.dense",
    "attention.self_mask": "attn_mask",
    "mlp.fc1": "intermediate.dense",
    "mlp.fc2": "output.dense",
    "norm1": "layernorm_before",
    "norm2": "layernorm_after",
    "bn0": "batch_norm",
}

__UpperCAmelCase       :   List[str]				       =							AutoFeatureExtractor.from_pretrained("laion/clap-htsat-unfused", truncation="rand_trunc")







def 			A__							(    SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__=False)   ->    Union[str, Any]:
  __snake_case  ,   __snake_case:						int   =						create_model(
      """HTSAT-tiny"""					,					"""roberta"""					,					SCREAMING_SNAKE_CASE__					,					precision="""fp32"""					,					device="""cuda:0""" if torch.cuda.is_available() else """cpu"""					,					enable_fusion=SCREAMING_SNAKE_CASE__					,					fusion_type="""aff_2d""" if enable_fusion else None					,					)
  return model, model_cfg







def 			A__							(    SCREAMING_SNAKE_CASE__)   ->    Any:
  __snake_case:						Optional[Any]   =						{}

  __snake_case:						int   =						r""".*sequential.(\d+).*"""
  __snake_case:						List[str]   =						r""".*_projection.(\d+).*"""

  for key, value in state_dict.items():
    # check if any key needs to be modified
    for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
      if key_to_modify in key:
        __snake_case:						Tuple   =						key.replace(SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__)

    if re.match(SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__):
      # replace sequential layers with list
      __snake_case:						Optional[int]   =						re.match(SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__).group(1)

      __snake_case:						str   =						key.replace(F'''sequential.{sequential_layer}.'''					,					F'''layers.{int(SCREAMING_SNAKE_CASE__)//3}.linear.''')
    elif re.match(SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__):
      __snake_case:						Any   =						int(re.match(SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__).group(1))

      # Because in CLAP they use `nn.Sequential`...
      __snake_case:						Dict   =						1 if projecton_layer == 0 else 2

      __snake_case:						Any   =						key.replace(F'''_projection.{projecton_layer}.'''					,					F'''_projection.linear{transformers_projection_layer}.''')

    if "audio" and "qkv" in key:
      # split qkv into query key and value
      __snake_case:						List[str]   =						value
      __snake_case:						Optional[Any]   =						mixed_qkv.size(0) // 3

      __snake_case:						Union[str, Any]   =						mixed_qkv[:qkv_dim]
      __snake_case:						Dict   =						mixed_qkv[qkv_dim : qkv_dim * 2]
      __snake_case:						int   =						mixed_qkv[qkv_dim * 2 :]

      __snake_case:						Optional[Any]   =						query_layer
      __snake_case:						str   =						key_layer
      __snake_case:						int   =						value_layer
    else:
      __snake_case:						Dict   =						value

  return model_state_dict







def 			A__							(    SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__=False)   ->    Optional[Any]:
  __snake_case  ,   __snake_case:						List[str]   =						init_clap(SCREAMING_SNAKE_CASE__					,					enable_fusion=SCREAMING_SNAKE_CASE__)

  clap_model.eval()
  __snake_case:						List[str]   =						clap_model.state_dict()
  __snake_case:						Optional[int]   =						rename_state_dict(SCREAMING_SNAKE_CASE__)

  __snake_case:						Any   =						ClapConfig()
  __snake_case:						Dict   =						enable_fusion
  __snake_case:						List[str]   =						ClapModel(SCREAMING_SNAKE_CASE__)

  # ignore the spectrogram embedding layer
  model.load_state_dict(SCREAMING_SNAKE_CASE__					,					strict=SCREAMING_SNAKE_CASE__)

  model.save_pretrained(SCREAMING_SNAKE_CASE__)
  transformers_config.save_pretrained(SCREAMING_SNAKE_CASE__)


if __name__ == "__main__":
      __UpperCAmelCase       :   Union[str, Any]				       =							argparse.ArgumentParser()
      parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
      parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to fairseq checkpoint")
      parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
      parser.add_argument("--enable_fusion", action="store_true", help="Whether to enable fusion or not")
      __UpperCAmelCase       :   Tuple				       =							parser.parse_args()

      convert_clap_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.enable_fusion)






