



'''simple docstring'''





import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
  import os

  # The slow tests are often failing with OOM error on GPU
  # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
  # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
  a__    =	"""platform"""

  import jax
  import jax.numpy as jnp

  from transformers import BlenderbotTokenizer
  from transformers.models.blenderbot.modeling_flax_blenderbot import (
      FlaxBlenderbotForConditionalGeneration,
      FlaxBlenderbotModel,
      shift_tokens_right,
  )



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	str				,       SCREAMING_SNAKE_CASE__       :	Dict				,       SCREAMING_SNAKE_CASE__       :	List[Any]=None				,       SCREAMING_SNAKE_CASE__       :	List[Any]=None				,       SCREAMING_SNAKE_CASE__       :	Optional[Any]=None				,       SCREAMING_SNAKE_CASE__       :	Tuple=None				,       SCREAMING_SNAKE_CASE__       :	List[Any]=None				,       SCREAMING_SNAKE_CASE__       :	Optional[int]=None				,       )  ->     Dict:
 if attention_mask is None:
  _snake_case    :  str					=					np.where(input_ids != config.pad_token_id				,       1				,       0				)
 if decoder_attention_mask is None:
  _snake_case    :  Optional[Any]					=					np.where(decoder_input_ids != config.pad_token_id				,       1				,       0				)
 if head_mask is None:
  _snake_case    :  Any					=					np.ones((config.encoder_layers, config.encoder_attention_heads)				)
 if decoder_head_mask is None:
  _snake_case    :  Any					=					np.ones((config.decoder_layers, config.decoder_attention_heads)				)
 if cross_attn_head_mask is None:
  _snake_case    :  str					=					np.ones((config.decoder_layers, config.decoder_attention_heads)				)
 return {
     "input_ids": input_ids,
     "decoder_input_ids": decoder_input_ids,
     "attention_mask": attention_mask,
     "decoder_attention_mask": attention_mask,
 }







class 		snake_case						:



   '''simple docstring'''






   def __init__(					self							:			Union[str, Any]				,		lowerCAmelCase							:			List[Any]				,		lowerCAmelCase							:			List[Any]=13				,		lowerCAmelCase							:			int=7				,		lowerCAmelCase							:			Union[str, Any]=True				,		lowerCAmelCase							:			List[str]=False				,		lowerCAmelCase							:			Any=99				,		lowerCAmelCase							:			List[Any]=16				,		lowerCAmelCase							:			Union[str, Any]=2				,		lowerCAmelCase							:			Optional[Any]=4				,		lowerCAmelCase							:			Optional[int]=4				,		lowerCAmelCase							:			Dict="gelu"				,		lowerCAmelCase							:			List[Any]=0.1				,		lowerCAmelCase							:			Union[str, Any]=0.1				,		lowerCAmelCase							:			Optional[Any]=32				,		lowerCAmelCase							:			Dict=2				,		lowerCAmelCase							:			str=1				,		lowerCAmelCase							:			str=0				,		lowerCAmelCase							:			List[str]=0.02				,		)				->				Optional[Any]:

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					parent
    _snake_case    :  Optional[Any]					=					batch_size
    _snake_case    :  Union[str, Any]					=					seq_length
    _snake_case    :  List[Any]					=					is_training
    _snake_case    :  Tuple					=					use_labels
    _snake_case    :  List[str]					=					vocab_size
    _snake_case    :  int					=					hidden_size
    _snake_case    :  List[str]					=					num_hidden_layers
    _snake_case    :  int					=					num_attention_heads
    _snake_case    :  int					=					intermediate_size
    _snake_case    :  Tuple					=					hidden_act
    _snake_case    :  Dict					=					hidden_dropout_prob
    _snake_case    :  int					=					attention_probs_dropout_prob
    _snake_case    :  int					=					max_position_embeddings
    _snake_case    :  str					=					eos_token_id
    _snake_case    :  int					=					pad_token_id
    _snake_case    :  List[Any]					=					bos_token_id
    _snake_case    :  int					=					initializer_range

   def      UpperCamelCase_  (					self							:			Optional[Any])				->				int:

    """simple docstring"""







    _snake_case    :  Any					=					np.clip(ids_tensor([self.batch_size, self.seq_length - 1]				,		self.vocab_size)				,		3				,		self.vocab_size)
    _snake_case    :  Optional[int]					=					np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)				,		dtype=np.intaa))				,		-1)

    _snake_case    :  List[str]					=					shift_tokens_right(lowerCAmelCase				,		1				,		2)

    _snake_case    :  Tuple					=					BlenderbotConfig(
        vocab_size=self.vocab_size				,		d_model=self.hidden_size				,		encoder_layers=self.num_hidden_layers				,		decoder_layers=self.num_hidden_layers				,		encoder_attention_heads=self.num_attention_heads				,		decoder_attention_heads=self.num_attention_heads				,		encoder_ffn_dim=self.intermediate_size				,		decoder_ffn_dim=self.intermediate_size				,		dropout=self.hidden_dropout_prob				,		attention_dropout=self.attention_probs_dropout_prob				,		max_position_embeddings=self.max_position_embeddings				,		eos_token_id=self.eos_token_id				,		bos_token_id=self.bos_token_id				,		pad_token_id=self.pad_token_id				,		initializer_range=self.initializer_range				,		use_cache=lowerCAmelCase				,		)
    _snake_case    :  str					=					prepare_blenderbot_inputs_dict(lowerCAmelCase				,		lowerCAmelCase				,		lowerCAmelCase)
    return config, inputs_dict

   def      UpperCamelCase_  (					self							:			Optional[int])				->				Optional[Any]:

    """simple docstring"""







    _snake_case    :  int					=					self.prepare_config_and_inputs()
    return config, inputs_dict

   def      UpperCamelCase_  (					self							:			Tuple				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			int				,		lowerCAmelCase							:			Optional[Any])				->				List[Any]:

    """simple docstring"""







    _snake_case    :  Optional[int]					=					20
    _snake_case    :  List[str]					=					model_class_name(lowerCAmelCase)

    _snake_case    :  Optional[Any]					=					model.encode(inputs_dict["""input_ids"""])

    _snake_case    :  List[str]					=					(
        inputs_dict["""decoder_input_ids"""],
        inputs_dict["""decoder_attention_mask"""],
    )

    _snake_case    :  Tuple					=					model.init_cache(decoder_input_ids.shape[0]				,		lowerCAmelCase				,		lowerCAmelCase)
    _snake_case    :  Optional[int]					=					jnp.ones((decoder_input_ids.shape[0], max_decoder_length)				,		dtype="""i4""")

    _snake_case    :  Dict					=					jnp.broadcast_to(
        jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :]				,		(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,		)
    _snake_case    :  str					=					model.decode(
        decoder_input_ids[:, :-1]				,		lowerCAmelCase				,		decoder_attention_mask=lowerCAmelCase				,		past_key_values=lowerCAmelCase				,		decoder_position_ids=lowerCAmelCase				,		)

    _snake_case    :  Optional[Any]					=					jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,		dtype="""i4""")
    _snake_case    :  Dict					=					model.decode(
        decoder_input_ids[:, -1:]				,		lowerCAmelCase				,		decoder_attention_mask=lowerCAmelCase				,		past_key_values=outputs_cache.past_key_values				,		decoder_position_ids=lowerCAmelCase				,		)

    _snake_case    :  Any					=					model.decode(lowerCAmelCase				,		lowerCAmelCase)

    _snake_case    :  List[str]					=					np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
    self.parent.assertTrue(diff < 1E-3				,		msg=F'''Max diff is {diff}''')


   def      UpperCamelCase_  (					self							:			Dict				,		lowerCAmelCase							:			Optional[int]				,		lowerCAmelCase							:			Dict				,		lowerCAmelCase							:			Optional[int])				->				int:

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					20
    _snake_case    :  Any					=					model_class_name(lowerCAmelCase)

    _snake_case    :  Tuple					=					model.encode(inputs_dict["""input_ids"""])

    _snake_case    :  str					=					(
        inputs_dict["""decoder_input_ids"""],
        inputs_dict["""decoder_attention_mask"""],
    )

    _snake_case    :  int					=					jnp.concatenate(
        [
            decoder_attention_mask,
            jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])),
        ]				,		axis=-1				,		)

    _snake_case    :  List[Any]					=					model.init_cache(decoder_input_ids.shape[0]				,		lowerCAmelCase				,		lowerCAmelCase)
    _snake_case    :  Optional[Any]					=					jnp.broadcast_to(
        jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :]				,		(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)				,		)

    _snake_case    :  Union[str, Any]					=					model.decode(
        decoder_input_ids[:, :-1]				,		lowerCAmelCase				,		decoder_attention_mask=lowerCAmelCase				,		past_key_values=lowerCAmelCase				,		decoder_position_ids=lowerCAmelCase				,		)
    _snake_case    :  Tuple					=					jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]				,		dtype="""i4""")
    _snake_case    :  List[Any]					=					model.decode(
        decoder_input_ids[:, -1:]				,		lowerCAmelCase				,		past_key_values=outputs_cache.past_key_values				,		decoder_attention_mask=lowerCAmelCase				,		decoder_position_ids=lowerCAmelCase				,		)

    _snake_case    :  Union[str, Any]					=					model.decode(lowerCAmelCase				,		lowerCAmelCase				,		decoder_attention_mask=lowerCAmelCase)

    _snake_case    :  str					=					np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
    self.parent.assertTrue(diff < 1E-3				,		msg=F'''Max diff is {diff}''')







@require_flax
class 		snake_case						(	unittest.TestCase  ):



   '''simple docstring'''






   snake_case_					: str          =		99

   def      UpperCamelCase_  (					self							:			List[Any])				->				Any:

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					np.array(
        [
            [71, 82, 18, 33, 46, 91, 2],
            [68, 34, 26, 58, 30, 82, 2],
            [5, 97, 17, 39, 94, 40, 2],
            [76, 83, 94, 25, 70, 78, 2],
            [87, 59, 41, 35, 48, 66, 2],
            [55, 13, 16, 58, 5, 2, 1],  # note padding
            [64, 27, 31, 51, 12, 75, 2],
            [52, 64, 86, 17, 83, 39, 2],
            [48, 61, 9, 24, 71, 82, 2],
            [26, 1, 60, 48, 22, 13, 2],
            [21, 5, 62, 28, 14, 76, 2],
            [45, 98, 37, 86, 59, 48, 2],
            [70, 70, 50, 9, 28, 0, 2],
        ]				,		dtype=np.intaa				,		)

    _snake_case    :  Dict					=					input_ids.shape[0]
    _snake_case    :  Dict					=					BlenderbotConfig(
        vocab_size=self.vocab_size				,		d_model=24				,		encoder_layers=2				,		decoder_layers=2				,		encoder_attention_heads=2				,		decoder_attention_heads=2				,		encoder_ffn_dim=32				,		decoder_ffn_dim=32				,		max_position_embeddings=48				,		eos_token_id=2				,		pad_token_id=1				,		bos_token_id=0				,		)
    return config, input_ids, batch_size

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

    """simple docstring"""







    _snake_case    :  str					=					self._get_config_and_data()
    _snake_case    :  Tuple					=					FlaxBlenderbotForConditionalGeneration(lowerCAmelCase)
    _snake_case    :  Dict					=					lm_model(input_ids=lowerCAmelCase)
    _snake_case    :  List[str]					=					(batch_size, input_ids.shape[1], config.vocab_size)
    self.assertEqual(outputs["""logits"""].shape				,		lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			List[str])				->				Dict:

    """simple docstring"""







    _snake_case    :  Tuple					=					BlenderbotConfig(
        vocab_size=self.vocab_size				,		d_model=14				,		encoder_layers=2				,		decoder_layers=2				,		encoder_attention_heads=2				,		decoder_attention_heads=2				,		encoder_ffn_dim=8				,		decoder_ffn_dim=8				,		max_position_embeddings=48				,		)
    _snake_case    :  str					=					FlaxBlenderbotForConditionalGeneration(lowerCAmelCase)
    _snake_case    :  int					=					np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]				,		dtype=np.intaa)
    _snake_case    :  Union[str, Any]					=					np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]				,		dtype=np.intaa)
    _snake_case    :  Dict					=					lm_model(input_ids=lowerCAmelCase				,		decoder_input_ids=lowerCAmelCase)
    _snake_case    :  Union[str, Any]					=					(*summary.shape, config.vocab_size)
    self.assertEqual(outputs["""logits"""].shape				,		lowerCAmelCase)


   def      UpperCamelCase_  (					self							:			List[str])				->				Tuple:

    """simple docstring"""







    _snake_case    :  Any					=					np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]				,		dtype=np.intaa)
    _snake_case    :  List[str]					=					shift_tokens_right(lowerCAmelCase				,		1				,		2)
    _snake_case    :  Optional[Any]					=					np.equal(lowerCAmelCase				,		1).astype(np.floataa).sum()
    _snake_case    :  int					=					np.equal(lowerCAmelCase				,		1).astype(np.floataa).sum()
    self.assertEqual(shifted.shape				,		input_ids.shape)
    self.assertEqual(lowerCAmelCase				,		n_pad_before - 1)
    self.assertTrue(np.equal(shifted[:, 0]				,		2).all())







@require_flax
class 		snake_case						(	SCREAMING_SNAKE_CASE_					,unittest.TestCase					,SCREAMING_SNAKE_CASE_  ):



   '''simple docstring'''






   snake_case_					: Optional[int]          =		True
   snake_case_					: Tuple          =		(
       (
           FlaxBlenderbotModel,
           FlaxBlenderbotForConditionalGeneration,
       )
       if is_flax_available()
       else ()
   )
   snake_case_					: int          =		(FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()

   def      UpperCamelCase_  (					self							:			List[Any])				->				List[Any]:

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					FlaxBlenderbotModelTester(self)

   def      UpperCamelCase_  (					self							:			Any)				->				Any:

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					self.model_tester.prepare_config_and_inputs()
    for model_class in self.all_model_classes:
     self.model_tester.check_use_cache_forward(lowerCAmelCase				,		lowerCAmelCase				,		lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			Dict)				->				List[Any]:

    """simple docstring"""







    _snake_case    :  Optional[int]					=					self.model_tester.prepare_config_and_inputs()
    for model_class in self.all_model_classes:
     self.model_tester.check_use_cache_forward_with_attn_mask(lowerCAmelCase				,		lowerCAmelCase				,		lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			Any)				->				Tuple:

    """simple docstring"""







    _snake_case    :  Tuple					=					self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
     with self.subTest(model_class.__name__):
      _snake_case    :  str					=					self._prepare_for_class(lowerCAmelCase				,		lowerCAmelCase)
      _snake_case    :  int					=					model_class(lowerCAmelCase)

      @jax.jit
      def encode_jitted(lowerCAmelCase							:			List[Any]				,		lowerCAmelCase							:			Dict=None				,		**lowerCAmelCase							:			List[Any]):
       return model.encode(input_ids=lowerCAmelCase				,		attention_mask=lowerCAmelCase)

      with self.subTest("""JIT Enabled"""):
       _snake_case    :  Any					=					encode_jitted(**lowerCAmelCase).to_tuple()

      with self.subTest("""JIT Disabled"""):
       with jax.disable_jit():
        _snake_case    :  str					=					encode_jitted(**lowerCAmelCase).to_tuple()

      self.assertEqual(len(lowerCAmelCase)				,		len(lowerCAmelCase))
      for jitted_output, output in zip(lowerCAmelCase				,		lowerCAmelCase):
       self.assertEqual(jitted_output.shape				,		output.shape)

   def      UpperCamelCase_  (					self							:			List[Any])				->				Dict:

    """simple docstring"""







    _snake_case    :  str					=					self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
     with self.subTest(model_class.__name__):
      _snake_case    :  int					=					model_class(lowerCAmelCase)
      _snake_case    :  List[Any]					=					model.encode(inputs_dict["""input_ids"""]				,		inputs_dict["""attention_mask"""])

      _snake_case    :  List[str]					=					{
          """decoder_input_ids""": inputs_dict["""decoder_input_ids"""],
          """decoder_attention_mask""": inputs_dict["""decoder_attention_mask"""],
          """encoder_outputs""": encoder_outputs,
      }

      @jax.jit
      def decode_jitted(lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			List[Any]):
       return model.decode(
           decoder_input_ids=lowerCAmelCase				,		decoder_attention_mask=lowerCAmelCase				,		encoder_outputs=lowerCAmelCase				,		)

      with self.subTest("""JIT Enabled"""):
       _snake_case    :  Optional[Any]					=					decode_jitted(**lowerCAmelCase).to_tuple()

      with self.subTest("""JIT Disabled"""):
       with jax.disable_jit():
        _snake_case    :  int					=					decode_jitted(**lowerCAmelCase).to_tuple()

      self.assertEqual(len(lowerCAmelCase)				,		len(lowerCAmelCase))
      for jitted_output, output in zip(lowerCAmelCase				,		lowerCAmelCase):
       self.assertEqual(jitted_output.shape				,		output.shape)

   @slow
   def      UpperCamelCase_  (					self							:			Tuple)				->				Dict:

    """simple docstring"""







    for model_class_name in self.all_model_classes:
     _snake_case    :  Union[str, Any]					=					model_class_name.from_pretrained("""facebook/blenderbot-400M-distill""")
     # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
     _snake_case    :  Dict					=					np.ones((1, 1)) * model.config.eos_token_id
     _snake_case    :  Union[str, Any]					=					model(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)


   @unittest.skipUnless(jax_device != """cpu"""				,		"""3B test too slow on CPU.""")
   @slow
   def      UpperCamelCase_  (					self							:			Tuple)				->				Dict:

    """simple docstring"""







    _snake_case    :  Dict					=					{"""num_beams""": 1, """early_stopping""": True, """min_length""": 15, """max_length""": 25}
    _snake_case    :  Tuple					=					{"""skip_special_tokens""": True, """clean_up_tokenization_spaces""": True}

    _snake_case    :  Tuple					=					FlaxBlenderbotForConditionalGeneration.from_pretrained("""facebook/blenderbot-3B"""				,		from_pt=lowerCAmelCase)
    _snake_case    :  Any					=					BlenderbotTokenizer.from_pretrained("""facebook/blenderbot-3B""")

    _snake_case    :  Tuple					=					["""Sam"""]
    _snake_case    :  Any					=					tokenizer(lowerCAmelCase				,		return_tensors="""jax""")

    _snake_case    :  Tuple					=					model.generate(**lowerCAmelCase				,		**lowerCAmelCase)
    _snake_case    :  Union[str, Any]					=					"""Sam is a great name. It means \"sun\" in Gaelic."""

    _snake_case    :  Optional[int]					=					tokenizer.batch_decode(lowerCAmelCase				,		**lowerCAmelCase)
    assert generated_txt[0].strip() == tgt_text





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 		snake_case						(	SCREAMING_SNAKE_CASE_					,unittest.TestCase  ):



   '''simple docstring'''






   snake_case_					: str          =		AudioLDMPipeline
   snake_case_					: Any          =		TEXT_TO_AUDIO_PARAMS
   snake_case_					: Tuple          =		TEXT_TO_AUDIO_BATCH_PARAMS
   snake_case_					: Dict          =		frozenset(
       [
           """num_inference_steps""",
           """num_waveforms_per_prompt""",
           """generator""",
           """latents""",
           """output_type""",
           """return_dict""",
           """callback""",
           """callback_steps""",
       ]  )

   def      UpperCamelCase_  (					self							:			Any)				->				Any:

    """simple docstring"""







    torch.manual_seed(0)
    _snake_case    :  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				,		)
    _snake_case    :  List[Any]					=					DDIMScheduler(
        beta_start=0.00_085				,		beta_end=0.012				,		beta_schedule="""scaled_linear"""				,		clip_sample=lowerCAmelCase				,		set_alpha_to_one=lowerCAmelCase				,		)
    torch.manual_seed(0)
    _snake_case    :  str					=					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    :  Any					=					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    :  Any					=					ClapTextModelWithProjection(lowerCAmelCase)
    _snake_case    :  List[str]					=					RobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-random-roberta"""				,		model_max_length=77)

    _snake_case    :  List[Any]					=					SpeechTaHifiGanConfig(
        model_in_dim=8				,		sampling_rate=1_6000				,		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				,		)

    _snake_case    :  Tuple					=					SpeechTaHifiGan(lowerCAmelCase)

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

   def      UpperCamelCase_  (					self							:			Optional[Any]				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			Dict=0)				->				int:

    """simple docstring"""







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

   def      UpperCamelCase_  (					self							:			List[Any])				->				str:

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator

    _snake_case    :  Dict					=					self.get_dummy_components()
    _snake_case    :  Optional[Any]					=					AudioLDMPipeline(**lowerCAmelCase)
    _snake_case    :  Optional[int]					=					audioldm_pipe.to(lowerCAmelCase)
    audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase)

    _snake_case    :  str					=					self.get_dummy_inputs(lowerCAmelCase)
    _snake_case    :  Tuple					=					audioldm_pipe(**lowerCAmelCase)
    _snake_case    :  int					=					output.audios[0]

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

    _snake_case    :  Union[str, Any]					=					audio[:10]
    _snake_case    :  Union[str, Any]					=					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							:			List[str])				->				Optional[int]:

    """simple docstring"""







    _snake_case    :  Dict					=					self.get_dummy_components()
    _snake_case    :  Any					=					AudioLDMPipeline(**lowerCAmelCase)
    _snake_case    :  int					=					audioldm_pipe.to(lowerCAmelCase)
    _snake_case    :  str					=					audioldm_pipe.to(lowerCAmelCase)
    audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase)

    _snake_case    :  str					=					self.get_dummy_inputs(lowerCAmelCase)
    _snake_case    :  List[Any]					=					3 * [inputs["""prompt"""]]

    # forward
    _snake_case    :  str					=					audioldm_pipe(**lowerCAmelCase)
    _snake_case    :  Dict					=					output.audios[0]

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

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

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

    _snake_case    :  List[Any]					=					prompt_embeds

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

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

   def      UpperCamelCase_  (					self							:			Optional[Any])				->				Dict:

    """simple docstring"""







    _snake_case    :  List[Any]					=					self.get_dummy_components()
    _snake_case    :  Optional[Any]					=					AudioLDMPipeline(**lowerCAmelCase)
    _snake_case    :  Tuple					=					audioldm_pipe.to(lowerCAmelCase)
    _snake_case    :  List[Any]					=					audioldm_pipe.to(lowerCAmelCase)
    audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase)

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

    # forward
    _snake_case    :  Dict					=					audioldm_pipe(**lowerCAmelCase)
    _snake_case    :  Union[str, Any]					=					output.audios[0]

    _snake_case    :  str					=					self.get_dummy_inputs(lowerCAmelCase)
    _snake_case    :  Union[str, Any]					=					3 * [inputs.pop("""prompt""")]

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

     _snake_case    :  int					=					audioldm_pipe.text_encoder(
         lowerCAmelCase				,		)
     _snake_case    :  Optional[Any]					=					text_embeds.text_embeds
     # additional L_2 normalization over each hidden-state
     _snake_case    :  str					=					F.normalize(lowerCAmelCase				,		dim=-1)

     embeds.append(lowerCAmelCase)

    _snake_case	,    _snake_case    :  int					=					embeds

    # forward
    _snake_case    :  List[str]					=					audioldm_pipe(**lowerCAmelCase)
    _snake_case    :  Optional[Any]					=					output.audios[0]

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

   def      UpperCamelCase_  (					self							:			Optional[Any])				->				List[Any]:

    """simple docstring"""







    _snake_case    :  Optional[int]					=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
    _snake_case    :  List[Any]					=					self.get_dummy_components()
    _snake_case    :  int					=					PNDMScheduler(skip_prk_steps=lowerCAmelCase)
    _snake_case    :  Tuple					=					AudioLDMPipeline(**lowerCAmelCase)
    _snake_case    :  str					=					audioldm_pipe.to(lowerCAmelCase)
    audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase)

    _snake_case    :  List[str]					=					self.get_dummy_inputs(lowerCAmelCase)
    _snake_case    :  List[Any]					=					"""egg cracking"""
    _snake_case    :  Optional[Any]					=					audioldm_pipe(**lowerCAmelCase				,		negative_prompt=lowerCAmelCase)
    _snake_case    :  List[str]					=					output.audios[0]

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

    _snake_case    :  int					=					audio[:10]
    _snake_case    :  Any					=					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							:			Optional[int])				->				Optional[int]:

    """simple docstring"""







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

    _snake_case    :  List[str]					=					"""A hammer hitting a wooden surface"""

    # test num_waveforms_per_prompt=1 (default)
    _snake_case    :  List[Any]					=					audioldm_pipe(lowerCAmelCase				,		num_inference_steps=2).audios

    assert audios.shape == (1, 256)

    # test num_waveforms_per_prompt=1 (default) for batch of prompts
    _snake_case    :  Optional[Any]					=					2
    _snake_case    :  Optional[int]					=					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[Any]					=					2
    _snake_case    :  Tuple					=					audioldm_pipe(lowerCAmelCase				,		num_inference_steps=2				,		num_waveforms_per_prompt=lowerCAmelCase).audios

    assert audios.shape == (num_waveforms_per_prompt, 256)

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

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

   def      UpperCamelCase_  (					self							:			Optional[int])				->				int:

    """simple docstring"""







    _snake_case    :  Dict					=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
    _snake_case    :  Optional[int]					=					self.get_dummy_components()
    _snake_case    :  Any					=					AudioLDMPipeline(**lowerCAmelCase)
    _snake_case    :  Any					=					audioldm_pipe.to(lowerCAmelCase)
    audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase)
    _snake_case    :  str					=					audioldm_pipe.vocoder.config.sampling_rate

    _snake_case    :  List[str]					=					self.get_dummy_inputs(lowerCAmelCase)
    _snake_case    :  Optional[int]					=					audioldm_pipe(audio_length_in_s=0.016				,		**lowerCAmelCase)
    _snake_case    :  str					=					output.audios[0]

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

    _snake_case    :  int					=					audioldm_pipe(audio_length_in_s=0.032				,		**lowerCAmelCase)
    _snake_case    :  List[Any]					=					output.audios[0]

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

   def      UpperCamelCase_  (					self							:			int)				->				List[Any]:

    """simple docstring"""







    _snake_case    :  str					=					self.get_dummy_components()
    _snake_case    :  List[str]					=					AudioLDMPipeline(**lowerCAmelCase)
    _snake_case    :  str					=					audioldm_pipe.to(lowerCAmelCase)
    audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase)

    _snake_case    :  List[str]					=					["""hey"""]

    _snake_case    :  int					=					audioldm_pipe(lowerCAmelCase				,		num_inference_steps=1)
    _snake_case    :  Optional[int]					=					output.audios.shape
    assert audio_shape == (1, 256)

    _snake_case    :  Dict					=					audioldm_pipe.vocoder.config
    config.model_in_dim *= 2
    _snake_case    :  str					=					SpeechTaHifiGan(lowerCAmelCase).to(lowerCAmelCase)
    _snake_case    :  List[str]					=					audioldm_pipe(lowerCAmelCase				,		num_inference_steps=1)
    _snake_case    :  Dict					=					output.audios.shape
    # waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
    assert audio_shape == (1, 256)

   def      UpperCamelCase_  (					self							:			Dict)				->				Dict:

    """simple docstring"""







    self._test_attention_slicing_forward_pass(test_mean_pixel_difference=lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			str)				->				Optional[int]:

    """simple docstring"""







    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      UpperCamelCase_  (					self							:			Dict)				->				Union[str, Any]:

    """simple docstring"""







    self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowerCAmelCase)







@slow
class 		snake_case						(	unittest.TestCase  ):



   '''simple docstring'''






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

    """simple docstring"""







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

   def      UpperCamelCase_  (					self							:			Any				,		lowerCAmelCase							:			Dict				,		lowerCAmelCase							:			Optional[Any]="cpu"				,		lowerCAmelCase							:			Any=torch.floataa				,		lowerCAmelCase							:			Optional[Any]=0)				->				Tuple:

    """simple docstring"""







    _snake_case    :  Any					=					torch.Generator(device=lowerCAmelCase).manual_seed(lowerCAmelCase)
    _snake_case    :  Optional[Any]					=					np.random.RandomState(lowerCAmelCase).standard_normal((1, 8, 128, 16))
    _snake_case    :  List[str]					=					torch.from_numpy(lowerCAmelCase).to(device=lowerCAmelCase				,		dtype=lowerCAmelCase)
    _snake_case    :  Optional[int]					=					{
        """prompt""": """A hammer hitting a wooden surface""",
        """latents""": latents,
        """generator""": generator,
        """num_inference_steps""": 3,
        """guidance_scale""": 2.5,
    }
    return inputs

   def      UpperCamelCase_  (					self							:			Any)				->				List[str]:

    """simple docstring"""







    _snake_case    :  Dict					=					AudioLDMPipeline.from_pretrained("""cvssp/audioldm""")
    _snake_case    :  Optional[int]					=					audioldm_pipe.to(lowerCAmelCase)
    audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase)

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

    assert audio.ndim == 1
    assert len(lowerCAmelCase) == 8_1920

    _snake_case    :  Any					=					audio[7_7230:7_7240]
    _snake_case    :  str					=					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])
    _snake_case    :  List[str]					=					np.abs(expected_slice - audio_slice).max()
    assert max_diff < 1E-2


   def      UpperCamelCase_  (					self							:			Dict)				->				List[Any]:

    """simple docstring"""







    _snake_case    :  Dict					=					AudioLDMPipeline.from_pretrained("""cvssp/audioldm""")
    _snake_case    :  Any					=					LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config)
    _snake_case    :  Optional[int]					=					audioldm_pipe.to(lowerCAmelCase)
    audioldm_pipe.set_progress_bar_config(disable=lowerCAmelCase)

    _snake_case    :  Union[str, Any]					=					self.get_inputs(lowerCAmelCase)
    _snake_case    :  List[str]					=					audioldm_pipe(**lowerCAmelCase).audios[0]

    assert audio.ndim == 1
    assert len(lowerCAmelCase) == 8_1920

    _snake_case    :  Optional[Any]					=					audio[2_7780:2_7790]
    _snake_case    :  int					=					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])
    _snake_case    :  List[Any]					=					np.abs(expected_slice - audio_slice).max()
    assert max_diff < 3E-2
