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



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






							'''simple docstring'''
							torch.manual_seed(0)
							_UpperCamelCase							 =  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=__a				,		use_timestep_embedding=__a				,		time_embedding_type='''fourier'''				,		mid_block_type='''UNetMidBlock1D'''				,		down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''')				,		up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''')				,		)
							_UpperCamelCase							 =  IPNDMScheduler()

							_UpperCamelCase							 =  {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							}
							return components



		def 						UpperCAmelCase							(					self				,		__a				,		__a=0)       ->       int:






							'''simple docstring'''
							if str(__a).startswith('''mps'''):
												_UpperCamelCase							 =  torch.manual_seed(__a)
							else:
												_UpperCamelCase							 =  torch.Generator(device=__a).manual_seed(__a)
							_UpperCamelCase							 =  {
							    '''batch_size''': 1,
							    '''generator''': generator,
							    '''num_inference_steps''': 4,
							}
							return inputs



		def 						UpperCAmelCase							(					self)       ->       str:






							'''simple docstring'''
							_UpperCamelCase							 =  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase							 =  self.get_dummy_components()
							_UpperCamelCase							 =  DanceDiffusionPipeline(**__a)
							_UpperCamelCase							 =  pipe.to(__a)
							pipe.set_progress_bar_config(disable=__a)

							_UpperCamelCase							 =  self.get_dummy_inputs(__a)
							_UpperCamelCase							 =  pipe(**__a)
							_UpperCamelCase							 =  output.audios

							_UpperCamelCase							 =  audio[0, -3:, -3:]

							assert audio.shape == (1, 2, components["unet"].sample_size)
							_UpperCamelCase							 =  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 						UpperCAmelCase							(					self)       ->       str:






							'''simple docstring'''
							return super().test_save_load_local()



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






							'''simple docstring'''
							return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3)



		@skip_mps
		def 						UpperCAmelCase							(					self)       ->       int:






							'''simple docstring'''
							return super().test_save_load_optional_components()



		@skip_mps
		def 						UpperCAmelCase							(					self)       ->       Tuple:






							'''simple docstring'''
							return super().test_attention_slicing_forward_pass()




		def 						UpperCAmelCase							(					self)       ->       int:






							'''simple docstring'''
							super().test_inference_batch_single_identical(expected_max_diff=3e-3)






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



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






							'''simple docstring'''
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()



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






							'''simple docstring'''
							_UpperCamelCase							 =  torch_device

							_UpperCamelCase							 =  DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k''')
							_UpperCamelCase							 =  pipe.to(__a)
							pipe.set_progress_bar_config(disable=__a)

							_UpperCamelCase							 =  torch.manual_seed(0)
							_UpperCamelCase							 =  pipe(generator=__a				,		num_inference_steps=1_00				,		audio_length_in_s=4.096)
							_UpperCamelCase							 =  output.audios

							_UpperCamelCase							 =  audio[0, -3:, -3:]

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






							'''simple docstring'''
							_UpperCamelCase							 =  torch_device

							_UpperCamelCase							 =  DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''				,		torch_dtype=torch.floataa)
							_UpperCamelCase							 =  pipe.to(__a)
							pipe.set_progress_bar_config(disable=__a)

							_UpperCamelCase							 =  torch.manual_seed(0)
							_UpperCamelCase							 =  pipe(generator=__a				,		num_inference_steps=1_00				,		audio_length_in_s=4.096)
							_UpperCamelCase							 =  output.audios

							_UpperCamelCase							 =  audio[0, -3:, -3:]

							assert audio.shape == (1, 2, pipe.unet.sample_size)
							_UpperCamelCase							 =  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








"""simple docstring"""





from __future__ import annotations



def       lowerCamelCase__		(			__snake_case,     __snake_case,     __snake_case	)							->		int | float:


					"""simple docstring"""





					if len(__snake_case	) == 0:
										raise ValueError('''find_max() arg is an empty sequence'''	)
					if (
					    left >= len(__snake_case	)
					    or left < -len(__snake_case	)
					    or right >= len(__snake_case	)
					    or right < -len(__snake_case	)
					):
										raise IndexError('''list index out of range'''	)
					if left == right:
										return nums[left]
					_UpperCamelCase							 =  (left + right) >> 1  # the middle
					_UpperCamelCase							 =  find_max(__snake_case,     __snake_case,     __snake_case	)  # find max in range[left, mid]
					_UpperCamelCase							 =  find_max(__snake_case,     mid + 1,     __snake_case	)  # find max in range[mid + 1, right]

					return left_max if left_max >= right_max else right_max


if __name__ == "__main__":
				import doctest

				doctest.testmod(verbose=True)






