





import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotSmallConfig, is_flax_available
from transformers.testing_utils import 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
     _UpperCamelCase							=			'''platform'''

     import jax
     import jax.numpy as jnp

     from transformers.models.blenderbot_small.modeling_flax_blenderbot_small import (
         FlaxBlenderbotSmallForConditionalGeneration,
         FlaxBlenderbotSmallModel,
         shift_tokens_right,
     )







def    lowerCAmelCase__(						lowercase							:     Tuple ,  lowercase							:     Any ,  lowercase							:     Optional[int]=None ,  lowercase							:     Tuple=None ,  lowercase							:     List[Any]=None ,  lowercase							:     List[str]=None ,  lowercase							:     Any=None ,  lowercase							:     List[Any]=None ,  )			->     Optional[Any]:
  if attention_mask is None:
    __snake_case :					List[Any]          =							np.where(input_ids != config.pad_token_id ,  1 ,  0					)
  if decoder_attention_mask is None:
    __snake_case :					Any          =							np.where(decoder_input_ids != config.pad_token_id ,  1 ,  0					)
  if head_mask is None:
    __snake_case :					Union[str, Any]          =							np.ones((config.encoder_layers, config.encoder_attention_heads)					)
  if decoder_head_mask is None:
    __snake_case :					List[str]          =							np.ones((config.decoder_layers, config.decoder_attention_heads)					)
  if cross_attn_head_mask is None:
    __snake_case :					Optional[Any]          =							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 					_lowerCamelCase			:



      """simple docstring"""






      def __init__(			self					,    UpperCAmelCase					,    UpperCAmelCase=13					,    UpperCAmelCase=7					,    UpperCAmelCase=True					,    UpperCAmelCase=False					,    UpperCAmelCase=99					,    UpperCAmelCase=16					,    UpperCAmelCase=2					,    UpperCAmelCase=4					,    UpperCAmelCase=4					,    UpperCAmelCase="gelu"					,    UpperCAmelCase=0.1					,    UpperCAmelCase=0.1					,    UpperCAmelCase=32					,    UpperCAmelCase=2					,    UpperCAmelCase=1					,    UpperCAmelCase=0					,    UpperCAmelCase=0.02					,    )							-> int:






        '''simple docstring'''






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







      def     UpperCAmelCase		(			self )							-> List[str]:






        '''simple docstring'''






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

        __snake_case :					List[Any]          =							shift_tokens_right(__SCREAMING_SNAKE_CASE					,    1					,    2 )

        __snake_case :					str          =							BlenderbotSmallConfig(
            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=__SCREAMING_SNAKE_CASE					,    )
        __snake_case :					Optional[Any]          =							prepare_blenderbot_inputs_dict(__SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE )
        return config, inputs_dict







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






        '''simple docstring'''






        __snake_case :					List[Any]          =							self.prepare_config_and_inputs()
        return config, inputs_dict







      def     UpperCAmelCase		(			self					,    UpperCAmelCase					,    UpperCAmelCase					,    UpperCAmelCase )							-> str:






        '''simple docstring'''






        __snake_case :					Optional[Any]          =							20
        __snake_case :					int          =							model_class_name(__SCREAMING_SNAKE_CASE )

        __snake_case :					List[str]          =							model.encode(inputs_dict["input_ids"] )

        __snake_case :					Optional[Any]          =							(
            inputs_dict['''decoder_input_ids'''],
            inputs_dict['''decoder_attention_mask'''],
        )

        __snake_case :					Any          =							model.init_cache(decoder_input_ids.shape[0]					,    __SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE )
        __snake_case :					Union[str, Any]          =							jnp.ones((decoder_input_ids.shape[0], max_decoder_length)					,    dtype="i4" )

        __snake_case :					Tuple          =							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 :					List[Any]          =							model.decode(
            decoder_input_ids[:, :-1]					,    __SCREAMING_SNAKE_CASE					,    decoder_attention_mask=__SCREAMING_SNAKE_CASE					,    past_key_values=__SCREAMING_SNAKE_CASE					,    decoder_position_ids=__SCREAMING_SNAKE_CASE					,    )

        __snake_case :					Union[str, 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:]					,    __SCREAMING_SNAKE_CASE					,    decoder_attention_mask=__SCREAMING_SNAKE_CASE					,    past_key_values=outputs_cache.past_key_values					,    decoder_position_ids=__SCREAMING_SNAKE_CASE					,    )

        __snake_case :					int          =							model.decode(__SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE )

        __snake_case :					Dict          =							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					,    UpperCAmelCase					,    UpperCAmelCase					,    UpperCAmelCase )							-> List[str]:






        '''simple docstring'''






        __snake_case :					List[Any]          =							20
        __snake_case :					str          =							model_class_name(__SCREAMING_SNAKE_CASE )

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

        __snake_case :					Union[str, Any]          =							(
            inputs_dict['''decoder_input_ids'''],
            inputs_dict['''decoder_attention_mask'''],
        )

        __snake_case :					Union[str, Any]          =							jnp.concatenate(
            [
                decoder_attention_mask,
                jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1]) ),
            ]					,    axis=-1					,    )

        __snake_case :					int          =							model.init_cache(decoder_input_ids.shape[0]					,    __SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE )
        __snake_case :					str          =							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 :					Tuple          =							model.decode(
            decoder_input_ids[:, :-1]					,    __SCREAMING_SNAKE_CASE					,    decoder_attention_mask=__SCREAMING_SNAKE_CASE					,    past_key_values=__SCREAMING_SNAKE_CASE					,    decoder_position_ids=__SCREAMING_SNAKE_CASE					,    )
        __snake_case :					List[str]          =							jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]					,    dtype="i4" )
        __snake_case :					Any          =							model.decode(
            decoder_input_ids[:, -1:]					,    __SCREAMING_SNAKE_CASE					,    past_key_values=outputs_cache.past_key_values					,    decoder_attention_mask=__SCREAMING_SNAKE_CASE					,    decoder_position_ids=__SCREAMING_SNAKE_CASE					,    )

        __snake_case :					int          =							model.decode(__SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE					,    decoder_attention_mask=__SCREAMING_SNAKE_CASE )

        __snake_case :					Union[str, Any]          =							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 					_lowerCamelCase			(     unittest.TestCase     ):



      """simple docstring"""






      UpperCAmelCase_    :   Union[str, Any]      =99







      def     UpperCAmelCase		(			self )							-> List[str]:






        '''simple docstring'''






        __snake_case :					int          =							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 :					List[str]          =							input_ids.shape[0]
        __snake_case :					Any          =							BlenderbotSmallConfig(
            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 )							-> Dict:






        '''simple docstring'''






        __snake_case :					Optional[int]          =							self._get_config_and_data()
        __snake_case :					Dict          =							FlaxBlenderbotSmallForConditionalGeneration(__SCREAMING_SNAKE_CASE )
        __snake_case :					Union[str, Any]          =							lm_model(input_ids=__SCREAMING_SNAKE_CASE )
        __snake_case :					Optional[Any]          =							(batch_size, input_ids.shape[1], config.vocab_size)
        self.assertEqual(outputs["logits"].shape					,    __SCREAMING_SNAKE_CASE )







      def     UpperCAmelCase		(			self )							-> List[str]:






        '''simple docstring'''






        __snake_case :					int          =							BlenderbotSmallConfig(
            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 :					int          =							FlaxBlenderbotSmallForConditionalGeneration(__SCREAMING_SNAKE_CASE )
        __snake_case :					Any          =							np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]					,    dtype=np.intaa )
        __snake_case :					Tuple          =							np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]					,    dtype=np.intaa )
        __snake_case :					Tuple          =							lm_model(input_ids=__SCREAMING_SNAKE_CASE					,    decoder_input_ids=__SCREAMING_SNAKE_CASE )
        __snake_case :					List[str]          =							(*summary.shape, config.vocab_size)
        self.assertEqual(outputs["logits"].shape					,    __SCREAMING_SNAKE_CASE )







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






        '''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 :					str          =							shift_tokens_right(__SCREAMING_SNAKE_CASE					,    1					,    2 )
        __snake_case :					List[str]          =							np.equal(__SCREAMING_SNAKE_CASE					,    1 ).astype(np.floataa ).sum()
        __snake_case :					Union[str, Any]          =							np.equal(__SCREAMING_SNAKE_CASE					,    1 ).astype(np.floataa ).sum()
        self.assertEqual(shifted.shape					,    input_ids.shape )
        self.assertEqual(__SCREAMING_SNAKE_CASE					,    n_pad_before - 1 )
        self.assertTrue(np.equal(shifted[:, 0]					,    2 ).all() )




@require_flax
class 					_lowerCamelCase			(     UpperCamelCase_			, unittest.TestCase			, UpperCamelCase_     ):



      """simple docstring"""






      UpperCAmelCase_    :   Optional[int]      =True
      UpperCAmelCase_    :   Dict      =(
          (
              FlaxBlenderbotSmallModel,
              FlaxBlenderbotSmallForConditionalGeneration,
          )
          if is_flax_available()
          else ()
      )
      UpperCAmelCase_    :   List[str]      =(FlaxBlenderbotSmallForConditionalGeneration,) if is_flax_available() else ()







      def     UpperCAmelCase		(			self )							-> int:






        '''simple docstring'''






        __snake_case :					Any          =							FlaxBlenderbotSmallModelTester(self )







      def     UpperCAmelCase		(			self )							-> str:






        '''simple docstring'''






        __snake_case :					Optional[Any]          =							self.model_tester.prepare_config_and_inputs()
        for model_class in self.all_model_classes:
          self.model_tester.check_use_cache_forward(__SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE )







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






        '''simple docstring'''






        __snake_case :					Dict          =							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(__SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE )







      def     UpperCAmelCase		(			self )							-> List[str]:






        '''simple docstring'''






        __snake_case :					Union[str, Any]          =							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 :					List[Any]          =							self._prepare_for_class(__SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE )
            __snake_case :					str          =							model_class(__SCREAMING_SNAKE_CASE )

            @jax.jit
            def encode_jitted(UpperCAmelCase					,    UpperCAmelCase=None					,    **UpperCAmelCase ):
              return model.encode(input_ids=__SCREAMING_SNAKE_CASE					,    attention_mask=__SCREAMING_SNAKE_CASE )

            with self.subTest("JIT Enabled" ):
              __snake_case :					Optional[Any]          =							encode_jitted(**__SCREAMING_SNAKE_CASE ).to_tuple()

            with self.subTest("JIT Disabled" ):
              with jax.disable_jit():
                __snake_case :					Any          =							encode_jitted(**__SCREAMING_SNAKE_CASE ).to_tuple()

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







      def     UpperCAmelCase		(			self )							-> Optional[Any]:






        '''simple docstring'''






        __snake_case :					Any          =							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 :					List[str]          =							model_class(__SCREAMING_SNAKE_CASE )
            __snake_case :					List[str]          =							model.encode(inputs_dict["input_ids"]					,    inputs_dict["attention_mask"] )

            __snake_case :					Any          =							{
                '''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(UpperCAmelCase					,    UpperCAmelCase					,    UpperCAmelCase ):
              return model.decode(
                  decoder_input_ids=__SCREAMING_SNAKE_CASE					,    decoder_attention_mask=__SCREAMING_SNAKE_CASE					,    encoder_outputs=__SCREAMING_SNAKE_CASE					,    )

            with self.subTest("JIT Enabled" ):
              __snake_case :					List[Any]          =							decode_jitted(**__SCREAMING_SNAKE_CASE ).to_tuple()

            with self.subTest("JIT Disabled" ):
              with jax.disable_jit():
                __snake_case :					List[Any]          =							decode_jitted(**__SCREAMING_SNAKE_CASE ).to_tuple()

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







      @slow
      def     UpperCAmelCase		(			self )							-> List[str]:






        '''simple docstring'''






        for model_class_name in self.all_model_classes:
          __snake_case :					Optional[int]          =							model_class_name.from_pretrained("facebook/blenderbot_small-90M" )
          # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
          __snake_case :					List[str]          =							np.ones((1, 1) ) * model.config.eos_token_id
          __snake_case :					Optional[Any]          =							model(__SCREAMING_SNAKE_CASE )
          self.assertIsNotNone(__SCREAMING_SNAKE_CASE )



import gc
import unittest

import numpy as np
import torch

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS, UNCONDITIONAL_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class 					_lowerCamelCase			(   UpperCamelCase_						,				unittest.TestCase			):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		 =	DanceDiffusionPipeline
     SCREAMING_SNAKE_CASE_		 =	UNCONDITIONAL_AUDIO_GENERATION_PARAMS
     SCREAMING_SNAKE_CASE_		 =	PipelineTesterMixin.required_optional_params - {
         '''callback''',
         '''latents''',
         '''callback_steps''',
         '''output_type''',
         '''num_images_per_prompt''',
     }
     SCREAMING_SNAKE_CASE_		 =	UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
     SCREAMING_SNAKE_CASE_		 =	False
     SCREAMING_SNAKE_CASE_		 =	False






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Any:




           """simple docstring"""







           torch.manual_seed(0	)
           UpperCamelCase__     : List[str]											=					UNetaDModel(
               block_out_channels=(3_2, 3_2, 6_4)    ,    extra_in_channels=1_6    ,    sample_size=5_1_2    ,    sample_rate=1_6_0_0_0    ,    in_channels=2    ,    out_channels=2    ,    flip_sin_to_cos=__SCREAMING_SNAKE_CASE    ,    use_timestep_embedding=__SCREAMING_SNAKE_CASE    ,    time_embedding_type='''fourier'''    ,    mid_block_type='''UNetMidBlock1D'''    ,    down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''')    ,    up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''')    ,    )
           UpperCamelCase__     : Union[str, Any]											=					IPNDMScheduler()

           UpperCamelCase__     : List[str]											=					{
               '''unet''': unet,
               '''scheduler''': scheduler,
           }
           return components






     def  __SCREAMING_SNAKE_CASE      (				self    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE=0	)	-> Any:




           """simple docstring"""







           if str(__SCREAMING_SNAKE_CASE	).startswith('''mps'''	):
                 UpperCamelCase__     : Optional[Any]											=					torch.manual_seed(__SCREAMING_SNAKE_CASE	)
           else:
                 UpperCamelCase__     : Optional[Any]											=					torch.Generator(device=__SCREAMING_SNAKE_CASE	).manual_seed(__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : int											=					{
               '''batch_size''': 1,
               '''generator''': generator,
               '''num_inference_steps''': 4,
           }
           return inputs






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




           """simple docstring"""







           UpperCamelCase__     : str											=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
           UpperCamelCase__     : Dict											=					self.get_dummy_components()
           UpperCamelCase__     : str											=					DanceDiffusionPipeline(**__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : List[Any]											=					pipe.to(__SCREAMING_SNAKE_CASE	)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : int											=					self.get_dummy_inputs(__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : Dict											=					pipe(**__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : List[str]											=					output.audios

           UpperCamelCase__     : Dict											=					audio[0, -3:, -3:]

           assert audio.shape == (1, 2, components["unet"].sample_size)
           UpperCamelCase__     : List[Any]											=					np.array([-0.7265, 1.0000, -0.8388, 0.1175, 0.9498, -1.0000]	)
           assert np.abs(audio_slice.flatten() - expected_slice	).max() < 1e-2






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> List[Any]:




           """simple docstring"""







           return super().test_save_load_local()






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Tuple:




           """simple docstring"""







           return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3	)






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> str:




           """simple docstring"""







           return super().test_save_load_optional_components()






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Optional[int]:




           """simple docstring"""







           return super().test_attention_slicing_forward_pass()






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Tuple:




           """simple docstring"""







           super().test_inference_batch_single_identical(expected_max_diff=3e-3	)




@slow
@require_torch_gpu
class 					_lowerCamelCase			(   unittest.TestCase			):





     """simple docstring"""







     def  __SCREAMING_SNAKE_CASE      (				self	)	-> int:




           """simple docstring"""







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






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




           """simple docstring"""







           UpperCamelCase__     : Any											=					torch_device

           UpperCamelCase__     : Any											=					DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''	)
           UpperCamelCase__     : int											=					pipe.to(__SCREAMING_SNAKE_CASE	)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : int											=					torch.manual_seed(0	)
           UpperCamelCase__     : Optional[Any]											=					pipe(generator=__SCREAMING_SNAKE_CASE    ,    num_inference_steps=1_0_0    ,    audio_length_in_s=4.096	)
           UpperCamelCase__     : str											=					output.audios

           UpperCamelCase__     : List[Any]											=					audio[0, -3:, -3:]

           assert audio.shape == (1, 2, pipe.unet.sample_size)
           UpperCamelCase__     : Tuple											=					np.array([-0.0192, -0.0231, -0.0318, -0.0059, 0.0002, -0.0020]	)

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






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




           """simple docstring"""







           UpperCamelCase__     : Any											=					torch_device

           UpperCamelCase__     : Union[str, Any]											=					DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''    ,    torch_dtype=torch.floataa	)
           UpperCamelCase__     : Tuple											=					pipe.to(__SCREAMING_SNAKE_CASE	)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : int											=					torch.manual_seed(0	)
           UpperCamelCase__     : Union[str, Any]											=					pipe(generator=__SCREAMING_SNAKE_CASE    ,    num_inference_steps=1_0_0    ,    audio_length_in_s=4.096	)
           UpperCamelCase__     : List[Any]											=					output.audios

           UpperCamelCase__     : List[Any]											=					audio[0, -3:, -3:]

           assert audio.shape == (1, 2, pipe.unet.sample_size)
           UpperCamelCase__     : Optional[Any]											=					np.array([-0.0367, -0.0488, -0.0771, -0.0525, -0.0444, -0.0341]	)

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


