






"""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	):
	__magic_name__       =							DanceDiffusionPipeline
	__magic_name__       =							UNCONDITIONAL_AUDIO_GENERATION_PARAMS
	__magic_name__       =							PipelineTesterMixin.required_optional_params - {
	    '''callback''',
	    '''latents''',
	    '''callback_steps''',
	    '''output_type''',
	    '''num_images_per_prompt''',
	}
	__magic_name__       =							UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
	__magic_name__       =							False
	__magic_name__       =							False





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[Any]		)		->     Optional[int]:
				torch.manual_seed(0		)
				UpperCAmelCase_	:					Dict     	=   UNetaDModel(
				    block_out_channels=(32, 32, 64)      ,      extra_in_channels=16      ,      sample_size=512      ,      sample_rate=16_000      ,      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_	:					str     	=   IPNDMScheduler()

				UpperCAmelCase_	:					Any     	=   {
				    "unet": unet,
				    "scheduler": scheduler,
				}
				return components





	def 	_SCREAMING_SNAKE_CASE ( self   :					Dict      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					Optional[int]=0		)		->     str:
				if str(lowerCAmelCase_		).startswith("mps"		):
							UpperCAmelCase_	:					Optional[int]     	=   torch.manual_seed(lowerCAmelCase_		)
				else:
							UpperCAmelCase_	:					Optional[int]     	=   torch.Generator(device=lowerCAmelCase_		).manual_seed(lowerCAmelCase_		)
				UpperCAmelCase_	:					str     	=   {
				    "batch_size": 1,
				    "generator": generator,
				    "num_inference_steps": 4,
				}
				return inputs





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]		)		->     Union[str, Any]:
				UpperCAmelCase_	:					Any     	=   "cpu"  # ensure determinism for the device-dependent torch.Generator
				UpperCAmelCase_	:					Optional[int]     	=   self.get_dummy_components()
				UpperCAmelCase_	:					List[str]     	=   DanceDiffusionPipeline(**lowerCAmelCase_		)
				UpperCAmelCase_	:					Union[str, Any]     	=   pipe.to(lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_	:					str     	=   self.get_dummy_inputs(lowerCAmelCase_		)
				UpperCAmelCase_	:					Optional[Any]     	=   pipe(**lowerCAmelCase_		)
				UpperCAmelCase_	:					Dict     	=   output.audios

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

				assert audio.shape == (1, 2, components["unet"].sample_size)
				UpperCAmelCase_	:					Any     	=   np.array([-0.7_2_6_5, 1.0_0_0_0, -0.8_3_8_8, 0.1_1_7_5, 0.9_4_9_8, -1.0_0_0_0]		)
				assert np.abs(audio_slice.flatten() - expected_slice		).max() < 1e-2





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





	@skip_mps
	def 	_SCREAMING_SNAKE_CASE ( self   :					str		)		->     Dict:
				return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3		)





	@skip_mps
	def 	_SCREAMING_SNAKE_CASE ( self   :					Dict		)		->     str:
				return super().test_save_load_optional_components()





	@skip_mps
	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]		)		->     Optional[int]:
				return super().test_attention_slicing_forward_pass()





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]		)		->     Union[str, Any]:
				super().test_inference_batch_single_identical(expected_max_diff=3e-3		)







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





	def 	_SCREAMING_SNAKE_CASE ( self   :					int		)		->     List[Any]:
				# clean up the VRAM after each test
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]		)		->     Any:
				UpperCAmelCase_	:					Dict     	=   torch_device

				UpperCAmelCase_	:					Any     	=   DanceDiffusionPipeline.from_pretrained("harmonai/maestro-150k"		)
				UpperCAmelCase_	:					Optional[Any]     	=   pipe.to(lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_	:					Union[str, Any]     	=   torch.manual_seed(0		)
				UpperCAmelCase_	:					Dict     	=   pipe(generator=lowerCAmelCase_      ,      num_inference_steps=100      ,      audio_length_in_s=4.0_9_6		)
				UpperCAmelCase_	:					Any     	=   output.audios

				UpperCAmelCase_	:					Union[str, Any]     	=   audio[0, -3:, -3:]

				assert audio.shape == (1, 2, pipe.unet.sample_size)
				UpperCAmelCase_	:					Tuple     	=   np.array([-0.0_1_9_2, -0.0_2_3_1, -0.0_3_1_8, -0.0_0_5_9, 0.0_0_0_2, -0.0_0_2_0]		)

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





	def 	_SCREAMING_SNAKE_CASE ( self   :					str		)		->     Optional[Any]:
				UpperCAmelCase_	:					Any     	=   torch_device

				UpperCAmelCase_	:					str     	=   DanceDiffusionPipeline.from_pretrained("harmonai/maestro-150k"      ,      torch_dtype=torch.floataa		)
				UpperCAmelCase_	:					List[Any]     	=   pipe.to(lowerCAmelCase_		)
				pipe.set_progress_bar_config(disable=lowerCAmelCase_		)

				UpperCAmelCase_	:					str     	=   torch.manual_seed(0		)
				UpperCAmelCase_	:					Dict     	=   pipe(generator=lowerCAmelCase_      ,      num_inference_steps=100      ,      audio_length_in_s=4.0_9_6		)
				UpperCAmelCase_	:					int     	=   output.audios

				UpperCAmelCase_	:					Union[str, Any]     	=   audio[0, -3:, -3:]

				assert audio.shape == (1, 2, pipe.unet.sample_size)
				UpperCAmelCase_	:					int     	=   np.array([-0.0_3_6_7, -0.0_4_8_8, -0.0_7_7_1, -0.0_5_2_5, -0.0_4_4_4, -0.0_3_4_1]		)

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











"""simple docstring"""
import os
from typing import Dict, List, Tuple, TypeVar, Union


lowerCamelCase_							      = TypeVar('''T''')

lowerCamelCase_							      = Union[List[T], Tuple[T, ...]]
lowerCamelCase_							      = Union[T, List[T], Dict[str, T]]
lowerCamelCase_							      = Union[str, bytes, os.PathLike]



