

import argparse
import json
import os
from collections import OrderedDict

import torch

from transformers import LukeConfig, LukeForMaskedLM, MLukeTokenizer, XLMRobertaTokenizer
from transformers.tokenization_utils_base import AddedToken


@torch.no_grad()
def        __lowerCAmelCase    (       _UpperCamelCase		:  Tuple		,							_UpperCamelCase		:  List[Any]		,							_UpperCamelCase		:  Union[str, Any]		,							_UpperCamelCase		:  Optional[Any]		,							_UpperCamelCase		:  List[str]	)		->							int:
 '''simple docstring'''


 with open(_UpperCamelCase	) as metadata_file:
  SCREAMING_SNAKE_CASE								=			json.load(_UpperCamelCase	)
 SCREAMING_SNAKE_CASE								=			LukeConfig(use_entity_aware_attention=_UpperCamelCase		,							**metadata['model_config']	)

 # Load in the weights from the checkpoint_path
 SCREAMING_SNAKE_CASE								=			torch.load(_UpperCamelCase		,							map_location='cpu'	)['module']

 # Load the entity vocab file
 SCREAMING_SNAKE_CASE								=			load_original_entity_vocab(_UpperCamelCase	)
 # add an entry for [MASK2]
 SCREAMING_SNAKE_CASE								=			max(entity_vocab.values()	) + 1
 config.entity_vocab_size += 1

 SCREAMING_SNAKE_CASE								=			XLMRobertaTokenizer.from_pretrained(metadata['model_config']['bert_model_name']	)

 # Add special tokens to the token vocabulary for downstream tasks
 SCREAMING_SNAKE_CASE								=			AddedToken('<ent>'		,							lstrip=_UpperCamelCase		,							rstrip=_UpperCamelCase	)
 SCREAMING_SNAKE_CASE								=			AddedToken('<ent2>'		,							lstrip=_UpperCamelCase		,							rstrip=_UpperCamelCase	)
 tokenizer.add_special_tokens({'additional_special_tokens': [entity_token_a, entity_token_a]}	)
 config.vocab_size += 2

 print(f"""Saving tokenizer to {pytorch_dump_folder_path}"""	)
 tokenizer.save_pretrained(_UpperCamelCase	)
 with open(os.path.join(_UpperCamelCase		,							'tokenizer_config.json'	)		,							'r'	) as f:
  SCREAMING_SNAKE_CASE								=			json.load(_UpperCamelCase	)
 SCREAMING_SNAKE_CASE								=			'MLukeTokenizer'
 with open(os.path.join(_UpperCamelCase		,							'tokenizer_config.json'	)		,							'w'	) as f:
  json.dump(_UpperCamelCase		,							_UpperCamelCase	)

 with open(os.path.join(_UpperCamelCase		,							MLukeTokenizer.vocab_files_names['entity_vocab_file']	)		,							'w'	) as f:
  json.dump(_UpperCamelCase		,							_UpperCamelCase	)

 SCREAMING_SNAKE_CASE								=			MLukeTokenizer.from_pretrained(_UpperCamelCase	)

 # Initialize the embeddings of the special tokens
 SCREAMING_SNAKE_CASE								=			tokenizer.convert_tokens_to_ids(['@']	)[0]
 SCREAMING_SNAKE_CASE								=			tokenizer.convert_tokens_to_ids(['#']	)[0]

 SCREAMING_SNAKE_CASE								=			state_dict['embeddings.word_embeddings.weight']
 SCREAMING_SNAKE_CASE								=			word_emb[ent_init_index].unsqueeze(0	)
 SCREAMING_SNAKE_CASE								=			word_emb[enta_init_index].unsqueeze(0	)
 SCREAMING_SNAKE_CASE								=			torch.cat([word_emb, ent_emb, enta_emb]	)
 # add special tokens for 'entity_predictions.bias'
 for bias_name in ["lm_head.decoder.bias", "lm_head.bias"]:
  SCREAMING_SNAKE_CASE								=			state_dict[bias_name]
  SCREAMING_SNAKE_CASE								=			decoder_bias[ent_init_index].unsqueeze(0	)
  SCREAMING_SNAKE_CASE								=			decoder_bias[enta_init_index].unsqueeze(0	)
  SCREAMING_SNAKE_CASE								=			torch.cat([decoder_bias, ent_decoder_bias, enta_decoder_bias]	)

 # Initialize the query layers of the entity-aware self-attention mechanism
 for layer_index in range(config.num_hidden_layers	):
  for matrix_name in ["query.weight", "query.bias"]:
   SCREAMING_SNAKE_CASE								=			f"""encoder.layer.{layer_index}.attention.self."""
   SCREAMING_SNAKE_CASE								=			state_dict[prefix + matrix_name]
   SCREAMING_SNAKE_CASE								=			state_dict[prefix + matrix_name]
   SCREAMING_SNAKE_CASE								=			state_dict[prefix + matrix_name]

    # Initialize the embedding of the [MASK2] entity using that of the [MASK] entity for downstream tasks
 SCREAMING_SNAKE_CASE								=			state_dict['entity_embeddings.entity_embeddings.weight']
 SCREAMING_SNAKE_CASE								=			entity_emb[entity_vocab['[MASK]']].unsqueeze(0	)
 SCREAMING_SNAKE_CASE								=			torch.cat([entity_emb, entity_mask_emb]	)
 # add [MASK2] for 'entity_predictions.bias'
 SCREAMING_SNAKE_CASE								=			state_dict['entity_predictions.bias']
 SCREAMING_SNAKE_CASE								=			entity_prediction_bias[entity_vocab['[MASK]']].unsqueeze(0	)
 SCREAMING_SNAKE_CASE								=			torch.cat([entity_prediction_bias, entity_mask_bias]	)

 SCREAMING_SNAKE_CASE								=			LukeForMaskedLM(config=_UpperCamelCase	).eval()

 state_dict.pop('entity_predictions.decoder.weight'	)
 state_dict.pop('lm_head.decoder.weight'	)
 state_dict.pop('lm_head.decoder.bias'	)
 SCREAMING_SNAKE_CASE								=			OrderedDict()
 for key, value in state_dict.items():
  if not (key.startswith('lm_head'	) or key.startswith('entity_predictions'	)):
   SCREAMING_SNAKE_CASE								=			state_dict[key]
  else:
   SCREAMING_SNAKE_CASE								=			state_dict[key]

 SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			model.load_state_dict(_UpperCamelCase		,							strict=_UpperCamelCase	)

 if set(_UpperCamelCase	) != {"luke.embeddings.position_ids"}:
  raise ValueError(f"""Unexpected unexpected_keys: {unexpected_keys}"""	)
 if set(_UpperCamelCase	) != {
     "lm_head.decoder.weight",
     "lm_head.decoder.bias",
     "entity_predictions.decoder.weight",
 }:
  raise ValueError(f"""Unexpected missing_keys: {missing_keys}"""	)

 model.tie_weights()
 assert (model.luke.embeddings.word_embeddings.weight == model.lm_head.decoder.weight).all()
 assert (model.luke.entity_embeddings.entity_embeddings.weight == model.entity_predictions.decoder.weight).all()

 # Check outputs
 SCREAMING_SNAKE_CASE								=			MLukeTokenizer.from_pretrained(_UpperCamelCase		,							task='entity_classification'	)

 SCREAMING_SNAKE_CASE								=			'ISO 639-3 uses the code fas for the dialects spoken across Iran and アフガニスタン (Afghanistan).'
 SCREAMING_SNAKE_CASE								=			(0, 9)
 SCREAMING_SNAKE_CASE								=			tokenizer(_UpperCamelCase		,							entity_spans=[span]		,							return_tensors='pt'	)

 SCREAMING_SNAKE_CASE								=			model(**_UpperCamelCase	)

 # Verify word hidden states
 if model_size == "large":
  raise NotImplementedError
 else:  # base
  SCREAMING_SNAKE_CASE								=			torch.Size((1, 33, 7_68)	)
  SCREAMING_SNAKE_CASE								=			torch.tensor([[0.08_92, 0.05_96, -0.28_19], [0.01_34, 0.11_99, 0.05_73], [-0.01_69, 0.09_27, 0.06_44]]	)

 if not (outputs.last_hidden_state.shape == expected_shape):
  raise ValueError(
      f"""Outputs.last_hidden_state.shape is {outputs.last_hidden_state.shape}, Expected shape is {expected_shape}"""	)
 if not torch.allclose(outputs.last_hidden_state[0, :3, :3]		,							_UpperCamelCase		,							atol=1e-4	):
  raise ValueError

 # Verify entity hidden states
 if model_size == "large":
  raise NotImplementedError
 else:  # base
  SCREAMING_SNAKE_CASE								=			torch.Size((1, 1, 7_68)	)
  SCREAMING_SNAKE_CASE								=			torch.tensor([[-0.14_82, 0.06_09, 0.03_22]]	)

 if not (outputs.entity_last_hidden_state.shape == expected_shape):
  raise ValueError(
      f"""Outputs.entity_last_hidden_state.shape is {outputs.entity_last_hidden_state.shape}, Expected shape is"""
      f""" {expected_shape}"""	)
 if not torch.allclose(outputs.entity_last_hidden_state[0, :3, :3]		,							_UpperCamelCase		,							atol=1e-4	):
  raise ValueError

 # Verify masked word/entity prediction
 SCREAMING_SNAKE_CASE								=			MLukeTokenizer.from_pretrained(_UpperCamelCase	)
 SCREAMING_SNAKE_CASE								=			'Tokyo is the capital of <mask>.'
 SCREAMING_SNAKE_CASE								=			(24, 30)
 SCREAMING_SNAKE_CASE								=			tokenizer(_UpperCamelCase		,							entity_spans=[span]		,							return_tensors='pt'	)

 SCREAMING_SNAKE_CASE								=			model(**_UpperCamelCase	)

 SCREAMING_SNAKE_CASE								=			encoding['input_ids'][0].tolist()
 SCREAMING_SNAKE_CASE								=			input_ids.index(tokenizer.convert_tokens_to_ids('<mask>'	)	)
 SCREAMING_SNAKE_CASE								=			outputs.logits[0][mask_position_id].argmax(dim=-1	)
 assert "Japan" == tokenizer.decode(_UpperCamelCase	)

 SCREAMING_SNAKE_CASE								=			outputs.entity_logits[0][0].argmax().item()
 SCREAMING_SNAKE_CASE								=			[
     entity for entity, entity_id in tokenizer.entity_vocab.items() if entity_id == predicted_entity_id
 ]
 assert [e for e in multilingual_predicted_entities if e.startswith('en:'	)][0] == "en:Japan"

 # Finally, save our PyTorch model and tokenizer
 print('Saving PyTorch model to {}'.format(_UpperCamelCase	)	)
 model.save_pretrained(_UpperCamelCase	)


def        __lowerCAmelCase    (       _UpperCamelCase		:  Any	)		->							str:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE								=			['[MASK]', '[PAD]', '[UNK]']

 SCREAMING_SNAKE_CASE								=			[json.loads(_UpperCamelCase	) for line in open(_UpperCamelCase	)]

 SCREAMING_SNAKE_CASE								=			{}
 for entry in data:
  SCREAMING_SNAKE_CASE								=			entry['id']
  for entity_name, language in entry["entities"]:
   if entity_name in SPECIAL_TOKENS:
    SCREAMING_SNAKE_CASE								=			entity_id
    break
   SCREAMING_SNAKE_CASE								=			f"""{language}:{entity_name}"""
   SCREAMING_SNAKE_CASE								=			entity_id
 return new_mapping


if __name__ == "__main__":
      a_    :						Union[str, Any]    		=      argparse.ArgumentParser()
      # Required parameters
      parser.add_argument("--checkpoint_path", type=str, help="Path to a pytorch_model.bin file.")
      parser.add_argument(
          "--metadata_path", default=None, type=str, help="Path to a metadata.json file, defining the configuration."
      )
      parser.add_argument(
          "--entity_vocab_path",
          default=None,
          type=str,
          help="Path to an entity_vocab.tsv file, containing the entity vocabulary.",
      )
      parser.add_argument(
          "--pytorch_dump_folder_path", default=None, type=str, help="Path to where to dump the output PyTorch model."
      )
      parser.add_argument(
          "--model_size", default="base", type=str, choices=["base", "large"], help="Size of the model to be converted."
      )
      a_    :						Union[str, Any]    		=      parser.parse_args()
      convert_luke_checkpoint(
          args.checkpoint_path,
          args.metadata_path,
          args.entity_vocab_path,
          args.pytorch_dump_folder_path,
          args.model_size,
      )









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 (     SCREAMING_SNAKE_CASE	,			unittest.TestCase   ):
    __UpperCamelCase         =AudioLDMPipeline
    __UpperCamelCase         =TEXT_TO_AUDIO_PARAMS
    __UpperCamelCase         =TEXT_TO_AUDIO_BATCH_PARAMS
    __UpperCamelCase         =frozenset(
        [
            "num_inference_steps",
            "num_waveforms_per_prompt",
            "generator",
            "latents",
            "output_type",
            "return_dict",
            "callback",
            "callback_steps",
        ]   )

    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






     torch.manual_seed(0    )
     SCREAMING_SNAKE_CASE								=			UNetaDConditionModel(
         block_out_channels=(3_2, 6_4)							,    layers_per_block=2							,    sample_size=3_2							,    in_channels=4							,    out_channels=4							,    down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')							,    up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')							,    cross_attention_dim=(3_2, 6_4)							,    class_embed_type='simple_projection'							,    projection_class_embeddings_input_dim=3_2							,    class_embeddings_concat=snake_case__							,    )
     SCREAMING_SNAKE_CASE								=			DDIMScheduler(
         beta_start=0.00_085							,    beta_end=0.012							,    beta_schedule='scaled_linear'							,    clip_sample=snake_case__							,    set_alpha_to_one=snake_case__							,    )
     torch.manual_seed(0    )
     SCREAMING_SNAKE_CASE								=			AutoencoderKL(
         block_out_channels=[3_2, 6_4]							,    in_channels=1							,    out_channels=1							,    down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']							,    up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']							,    latent_channels=4							,    )
     torch.manual_seed(0    )
     SCREAMING_SNAKE_CASE								=			ClapTextConfig(
         bos_token_id=0							,    eos_token_id=2							,    hidden_size=3_2							,    intermediate_size=3_7							,    layer_norm_eps=1E-05							,    num_attention_heads=4							,    num_hidden_layers=5							,    pad_token_id=1							,    vocab_size=1_0_0_0							,    projection_dim=3_2							,    )
     SCREAMING_SNAKE_CASE								=			ClapTextModelWithProjection(snake_case__    )
     SCREAMING_SNAKE_CASE								=			RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta'							,    model_max_length=7_7    )

     SCREAMING_SNAKE_CASE								=			SpeechTaHifiGanConfig(
         model_in_dim=8							,    sampling_rate=1_6_0_0_0							,    upsample_initial_channel=1_6							,    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__							,    )

     SCREAMING_SNAKE_CASE								=			SpeechTaHifiGan(snake_case__    )

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

    def UpperCamelCase			(			self :      Optional[int]							,    snake_case__ :      int							,    snake_case__ :      int=0    ):

     """simple docstring"""






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

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

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			'cpu'  # ensure determinism for the device-dependent torch.Generator

     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

     assert audio.ndim == 1
     assert len(snake_case__    ) == 2_5_6

     SCREAMING_SNAKE_CASE								=			audio[:1_0]
     SCREAMING_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    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			3 * [inputs['prompt']]

     # forward
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			3 * [inputs.pop('prompt'    )]

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

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

     SCREAMING_SNAKE_CASE								=			prompt_embeds

     # forward
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

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

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

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

     # forward
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			3 * [inputs.pop('prompt'    )]

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

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

      embeds.append(snake_case__    )

     SCREAMING_SNAKE_CASE						,	SCREAMING_SNAKE_CASE								=			embeds

     # forward
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

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

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






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

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

     assert audio.ndim == 1
     assert len(snake_case__    ) == 2_5_6

     SCREAMING_SNAKE_CASE								=			audio[:1_0]
     SCREAMING_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 :      str    ):

     """simple docstring"""






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

     SCREAMING_SNAKE_CASE								=			'A hammer hitting a wooden surface'

     # test num_waveforms_per_prompt=1 (default)
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(snake_case__							,    num_inference_steps=2    ).audios

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

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

     assert audios.shape == (batch_size, 2_5_6)

     # test num_waveforms_per_prompt for single prompt
     SCREAMING_SNAKE_CASE								=			2
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(snake_case__							,    num_inference_steps=2							,    num_waveforms_per_prompt=snake_case__    ).audios

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

     # test num_waveforms_per_prompt for batch of prompts
     SCREAMING_SNAKE_CASE								=			2
     SCREAMING_SNAKE_CASE								=			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, 2_5_6)

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

     """simple docstring"""






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

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

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

     SCREAMING_SNAKE_CASE								=			audioldm_pipe(audio_length_in_s=0.032							,    **snake_case__    )
     SCREAMING_SNAKE_CASE								=			output.audios[0]

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

    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			['hey']

     SCREAMING_SNAKE_CASE								=			audioldm_pipe(snake_case__							,    num_inference_steps=1    )
     SCREAMING_SNAKE_CASE								=			output.audios.shape
     assert audio_shape == (1, 2_5_6)

     SCREAMING_SNAKE_CASE								=			audioldm_pipe.vocoder.config
     config.model_in_dim *= 2
     SCREAMING_SNAKE_CASE								=			SpeechTaHifiGan(snake_case__    ).to(snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(snake_case__							,    num_inference_steps=1    )
     SCREAMING_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_5_6)

    def UpperCamelCase			(			self :      Tuple    ):

     """simple docstring"""






     self._test_attention_slicing_forward_pass(test_mean_pixel_difference=snake_case__    )

    def UpperCamelCase			(			self :      int    ):

     """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 UpperCamelCase			(			self :      Dict    ):

     """simple docstring"""






     self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=snake_case__    )




@slow
class      UpperCamelCase (     unittest.TestCase   ):

    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






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

    def UpperCamelCase			(			self :      int							,    snake_case__ :      int							,    snake_case__ :      Tuple="cpu"							,    snake_case__ :      List[str]=torch.floataa							,    snake_case__ :      Optional[Any]=0    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			torch.Generator(device=snake_case__    ).manual_seed(snake_case__    )
     SCREAMING_SNAKE_CASE								=			np.random.RandomState(snake_case__    ).standard_normal((1, 8, 1_2_8, 1_6)    )
     SCREAMING_SNAKE_CASE								=			torch.from_numpy(snake_case__    ).to(device=snake_case__							,    dtype=snake_case__    )
     SCREAMING_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 :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline.from_pretrained('cvssp/audioldm'    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			2_5
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    ).audios[0]

     assert audio.ndim == 1
     assert len(snake_case__    ) == 8_1_9_2_0

     SCREAMING_SNAKE_CASE								=			audio[7_7_2_3_0:7_7_2_4_0]
     SCREAMING_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]    )
     SCREAMING_SNAKE_CASE								=			np.abs(expected_slice - audio_slice    ).max()
     assert max_diff < 1E-2




    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			AudioLDMPipeline.from_pretrained('cvssp/audioldm'    )
     SCREAMING_SNAKE_CASE								=			LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe.to(snake_case__    )
     audioldm_pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			audioldm_pipe(**snake_case__    ).audios[0]

     assert audio.ndim == 1
     assert len(snake_case__    ) == 8_1_9_2_0

     SCREAMING_SNAKE_CASE								=			audio[2_7_7_8_0:2_7_7_9_0]
     SCREAMING_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]    )
     SCREAMING_SNAKE_CASE								=			np.abs(expected_slice - audio_slice    ).max()
     assert max_diff < 3E-2

