

"""simple docstring"""






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      _UpperCAmelCase      (		__a     ,  unittest.TestCase):
			__a			:							Dict					       =				DanceDiffusionPipeline
			__a			:							Any					       =				UNCONDITIONAL_AUDIO_GENERATION_PARAMS
			__a			:							List[Any]					       =				PipelineTesterMixin.required_optional_params - {
			    '''callback''',
			    '''latents''',
			    '''callback_steps''',
			    '''output_type''',
			    '''num_images_per_prompt''',
			}
			__a			:							List[Any]					       =				UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
			__a			:							List[str]					       =				False
			__a			:							Union[str, Any]					       =				False
			def 				__snake_case			(		self  )     ->     List[Any]:






				'''simple docstring'''







				torch.manual_seed(0  )
				_UpperCAmelCase     :				Tuple          =       UNetaDModel(
				    block_out_channels=(32, 32, 64)   ,		extra_in_channels=16   ,		sample_size=5_12   ,		sample_rate=1_60_00   ,		in_channels=2   ,		out_channels=2   ,		flip_sin_to_cos=__lowerCAmelCase   ,		use_timestep_embedding=__lowerCAmelCase   ,		time_embedding_type="""fourier"""   ,		mid_block_type="""UNetMidBlock1D"""   ,		down_block_types=("""DownBlock1DNoSkip""", """DownBlock1D""", """AttnDownBlock1D""")   ,		up_block_types=("""AttnUpBlock1D""", """UpBlock1D""", """UpBlock1DNoSkip""")   ,		)
				_UpperCAmelCase     :				List[Any]          =       IPNDMScheduler()

				_UpperCAmelCase     :				Union[str, Any]          =       {
				    """unet""": unet,
				    """scheduler""": scheduler,
				}
				return components
			def 				__snake_case			(		self   ,		_A   ,		_A=0  )     ->     List[Any]:






				'''simple docstring'''







				if str(__lowerCAmelCase  ).startswith("""mps"""  ):
					_UpperCAmelCase     :				List[str]          =       torch.manual_seed(__lowerCAmelCase  )
				else:
					_UpperCAmelCase     :				Optional[Any]          =       torch.Generator(device=__lowerCAmelCase  ).manual_seed(__lowerCAmelCase  )
				_UpperCAmelCase     :				Optional[Any]          =       {
				    """batch_size""": 1,
				    """generator""": generator,
				    """num_inference_steps""": 4,
				}
				return inputs
			def 				__snake_case			(		self  )     ->     str:






				'''simple docstring'''







				_UpperCAmelCase     :				Optional[int]          =       """cpu"""  # ensure determinism for the device-dependent torch.Generator
				_UpperCAmelCase     :				int          =       self.get_dummy_components()
				_UpperCAmelCase     :				List[str]          =       DanceDiffusionPipeline(**__lowerCAmelCase  )
				_UpperCAmelCase     :				str          =       pipe.to(__lowerCAmelCase  )
				pipe.set_progress_bar_config(disable=__lowerCAmelCase  )

				_UpperCAmelCase     :				Dict          =       self.get_dummy_inputs(__lowerCAmelCase  )
				_UpperCAmelCase     :				List[Any]          =       pipe(**__lowerCAmelCase  )
				_UpperCAmelCase     :				Union[str, Any]          =       output.audios

				_UpperCAmelCase     :				int          =       audio[0, -3:, -3:]

				assert audio.shape == (1, 2, components["unet"].sample_size)
				_UpperCAmelCase     :				Optional[int]          =       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 				__snake_case			(		self  )     ->     Any:






				'''simple docstring'''







				return super().test_save_load_local()
			@skip_mps
			def 				__snake_case			(		self  )     ->     Optional[Any]:






				'''simple docstring'''







				return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3  )
			@skip_mps
			def 				__snake_case			(		self  )     ->     Any:






				'''simple docstring'''







				return super().test_save_load_optional_components()
			@skip_mps
			def 				__snake_case			(		self  )     ->     Union[str, Any]:






				'''simple docstring'''







				return super().test_attention_slicing_forward_pass()
			def 				__snake_case			(		self  )     ->     Optional[int]:






				'''simple docstring'''







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






@slow
@require_torch_gpu
class      _UpperCAmelCase      (		unittest.TestCase):
			def 				__snake_case			(		self  )     ->     List[Any]:






				'''simple docstring'''







				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()
			def 				__snake_case			(		self  )     ->     List[str]:






				'''simple docstring'''







				_UpperCAmelCase     :				str          =       torch_device

				_UpperCAmelCase     :				List[Any]          =       DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k"""  )
				_UpperCAmelCase     :				str          =       pipe.to(__lowerCAmelCase  )
				pipe.set_progress_bar_config(disable=__lowerCAmelCase  )

				_UpperCAmelCase     :				str          =       torch.manual_seed(0  )
				_UpperCAmelCase     :				Union[str, Any]          =       pipe(generator=__lowerCAmelCase   ,		num_inference_steps=1_00   ,		audio_length_in_s=4.096  )
				_UpperCAmelCase     :				Optional[Any]          =       output.audios

				_UpperCAmelCase     :				List[str]          =       audio[0, -3:, -3:]

				assert audio.shape == (1, 2, pipe.unet.sample_size)
				_UpperCAmelCase     :				Optional[Any]          =       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 				__snake_case			(		self  )     ->     Optional[int]:






				'''simple docstring'''







				_UpperCAmelCase     :				Tuple          =       torch_device

				_UpperCAmelCase     :				Any          =       DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k"""   ,		torch_dtype=torch.floataa  )
				_UpperCAmelCase     :				Tuple          =       pipe.to(__lowerCAmelCase  )
				pipe.set_progress_bar_config(disable=__lowerCAmelCase  )

				_UpperCAmelCase     :				Tuple          =       torch.manual_seed(0  )
				_UpperCAmelCase     :				List[str]          =       pipe(generator=__lowerCAmelCase   ,		num_inference_steps=1_00   ,		audio_length_in_s=4.096  )
				_UpperCAmelCase     :				str          =       output.audios

				_UpperCAmelCase     :				List[str]          =       audio[0, -3:, -3:]

				assert audio.shape == (1, 2, pipe.unet.sample_size)
				_UpperCAmelCase     :				List[str]          =       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












import numpy as np







def   __lowerCamelCase   (				__a  :np.array       )						->							np.array:






		"""simple docstring"""

		return 1 / (1 + np.exp(-vector       ))







def   __lowerCamelCase   (				__a  :np.array       )						->							np.array:






		"""simple docstring"""

		return vector * sigmoid(1.702 * vector       )


if __name__ == "__main__":
					import doctest

					doctest.testmod()



