"""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		(      __lowercase , unittest.TestCase      ):
				A__ :  Tuple    				=			DanceDiffusionPipeline
				A__ :  Optional[int]    				=			UNCONDITIONAL_AUDIO_GENERATION_PARAMS
				A__ :  int    				=			PipelineTesterMixin.required_optional_params - {
				    '''callback''',
				    '''latents''',
				    '''callback_steps''',
				    '''output_type''',
				    '''num_images_per_prompt''',
				}
				A__ :  int    				=			UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
				A__ :  Dict    				=			False
				A__ :  Optional[Any]    				=			False

				def 						_a      ( self ):
										torch.manual_seed(0 )
										lowerCamelCase__   =UNetaDModel(
										    block_out_channels=(32, 32, 64)      ,					extra_in_channels=16      ,					sample_size=512      ,					sample_rate=16000      ,					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")      ,					)
										lowerCamelCase__   =IPNDMScheduler()

										lowerCamelCase__   ={
										    "unet": unet,
										    "scheduler": scheduler,
										}
										return components

				def 						_a      ( self      ,					_lowerCamelCase      ,					_lowerCamelCase=0 ):







										if str(SCREAMING_SNAKE_CASE_ ).startswith("mps" ):
																lowerCamelCase__   =torch.manual_seed(SCREAMING_SNAKE_CASE_ )
										else:
																lowerCamelCase__   =torch.Generator(device=SCREAMING_SNAKE_CASE_ ).manual_seed(SCREAMING_SNAKE_CASE_ )
										lowerCamelCase__   ={
										    "batch_size": 1,
										    "generator": generator,
										    "num_inference_steps": 4,
										}
										return inputs

				def 						_a      ( self ):
										lowerCamelCase__   ="cpu"  # ensure determinism for the device-dependent torch.Generator
										lowerCamelCase__   =self.get_dummy_components()
										lowerCamelCase__   =DanceDiffusionPipeline(**SCREAMING_SNAKE_CASE_ )
										lowerCamelCase__   =pipe.to(SCREAMING_SNAKE_CASE_ )
										pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )

										lowerCamelCase__   =self.get_dummy_inputs(SCREAMING_SNAKE_CASE_ )
										lowerCamelCase__   =pipe(**SCREAMING_SNAKE_CASE_ )
										lowerCamelCase__   =output.audios

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

										assert audio.shape == (1, 2, components["unet"].sample_size)
										lowerCamelCase__   =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 						_a      ( self ):
										return super().test_save_load_local()

				@skip_mps
				def 						_a      ( self ):
										return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3 )

				@skip_mps
				def 						_a      ( self ):
										return super().test_save_load_optional_components()

				@skip_mps
				def 						_a      ( self ):
										return super().test_attention_slicing_forward_pass()






				def 						_a      ( self ):
										super().test_inference_batch_single_identical(expected_max_diff=3E-3 )






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

				def 						_a      ( self ):
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()

				def 						_a      ( self ):
										lowerCamelCase__   =torch_device

										lowerCamelCase__   =DanceDiffusionPipeline.from_pretrained("harmonai/maestro-150k" )
										lowerCamelCase__   =pipe.to(SCREAMING_SNAKE_CASE_ )
										pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )

										lowerCamelCase__   =torch.manual_seed(0 )
										lowerCamelCase__   =pipe(generator=SCREAMING_SNAKE_CASE_      ,					num_inference_steps=100      ,					audio_length_in_s=4.0_9_6 )
										lowerCamelCase__   =output.audios

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

										assert audio.shape == (1, 2, pipe.unet.sample_size)
										lowerCamelCase__   =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 						_a      ( self ):
										lowerCamelCase__   =torch_device

										lowerCamelCase__   =DanceDiffusionPipeline.from_pretrained("harmonai/maestro-150k"      ,					torch_dtype=torch.floataa )
										lowerCamelCase__   =pipe.to(SCREAMING_SNAKE_CASE_ )
										pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )

										lowerCamelCase__   =torch.manual_seed(0 )
										lowerCamelCase__   =pipe(generator=SCREAMING_SNAKE_CASE_      ,					num_inference_steps=100      ,					audio_length_in_s=4.0_9_6 )
										lowerCamelCase__   =output.audios

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

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

from diffusers import UnCLIPScheduler

from .test_schedulers import SchedulerCommonTest





class 			snake_case  (				__lowercase ):
							UpperCAmelCase__      =   (UnCLIPScheduler,)

							def _lowercase       (self							,							**SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			{
									    '''num_train_timesteps''': 10_00,
									    '''variance_type''': '''fixed_small_log''',
									    '''clip_sample''': True,
									    '''clip_sample_range''': 1.0,
									    '''prediction_type''': '''epsilon''',
									}

									config.update(**SCREAMING_SNAKE_CASE_     )
									return config

							def _lowercase       (self     ):

									"""simple docstring"""







									for timesteps in [1, 5, 1_00, 10_00]:
											self.check_over_configs(num_train_timesteps=SCREAMING_SNAKE_CASE_     )

							def _lowercase       (self     ):

									"""simple docstring"""







									for variance in ["fixed_small_log", "learned_range"]:
											self.check_over_configs(variance_type=SCREAMING_SNAKE_CASE_     )

							def _lowercase       (self     ):

									"""simple docstring"""







									for clip_sample in [True, False]:
											self.check_over_configs(clip_sample=SCREAMING_SNAKE_CASE_     )

							def _lowercase       (self     ):

									"""simple docstring"""







									for clip_sample_range in [1, 5, 10, 20]:
											self.check_over_configs(clip_sample_range=SCREAMING_SNAKE_CASE_     )

							def _lowercase       (self     ):

									"""simple docstring"""







									for prediction_type in ["epsilon", "sample"]:
											self.check_over_configs(prediction_type=SCREAMING_SNAKE_CASE_     )

							def _lowercase       (self     ):

									"""simple docstring"""







									for time_step in [0, 5_00, 9_99]:
											for prev_timestep in [None, 5, 1_00, 2_50, 5_00, 7_50]:
													if prev_timestep is not None and prev_timestep >= time_step:
															continue

													self.check_over_forward(time_step=SCREAMING_SNAKE_CASE_							,							prev_timestep=SCREAMING_SNAKE_CASE_     )

							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.scheduler_classes[0]
									SCREAMING_SNAKE_CASE_    			=			self.get_scheduler_config(variance_type='''fixed_small_log'''     )
									SCREAMING_SNAKE_CASE_    			=			scheduler_class(**SCREAMING_SNAKE_CASE_     )

									assert torch.sum(torch.abs(scheduler._get_variance(0     ) - 1.0_0_0_0e-1_0     )     ) < 1e-5
									assert torch.sum(torch.abs(scheduler._get_variance(4_87     ) - 0.0_54_96_25     )     ) < 1e-5
									assert torch.sum(torch.abs(scheduler._get_variance(9_99     ) - 0.9_99_49_87     )     ) < 1e-5

							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.scheduler_classes[0]
									SCREAMING_SNAKE_CASE_    			=			self.get_scheduler_config(variance_type='''learned_range'''     )
									SCREAMING_SNAKE_CASE_    			=			scheduler_class(**SCREAMING_SNAKE_CASE_     )

									SCREAMING_SNAKE_CASE_    			=			0.5

									assert scheduler._get_variance(1							,							predicted_variance=SCREAMING_SNAKE_CASE_     ) - -10.1_71_27_90 < 1e-5
									assert scheduler._get_variance(4_87							,							predicted_variance=SCREAMING_SNAKE_CASE_     ) - -5.7_99_80_52 < 1e-5
									assert scheduler._get_variance(9_99							,							predicted_variance=SCREAMING_SNAKE_CASE_     ) - -0.0_01_00_11 < 1e-5

							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.scheduler_classes[0]
									SCREAMING_SNAKE_CASE_    			=			self.get_scheduler_config()
									SCREAMING_SNAKE_CASE_    			=			scheduler_class(**SCREAMING_SNAKE_CASE_     )

									SCREAMING_SNAKE_CASE_    			=			scheduler.timesteps

									SCREAMING_SNAKE_CASE_    			=			self.dummy_model()
									SCREAMING_SNAKE_CASE_    			=			self.dummy_sample_deter
									SCREAMING_SNAKE_CASE_    			=			torch.manual_seed(0     )

									for i, t in enumerate(SCREAMING_SNAKE_CASE_     ):
											# 1. predict noise residual
											SCREAMING_SNAKE_CASE_    			=			model(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     )

											# 2. predict previous mean of sample x_t-1
											SCREAMING_SNAKE_CASE_    			=			scheduler.step(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							generator=SCREAMING_SNAKE_CASE_     ).prev_sample

											SCREAMING_SNAKE_CASE_    			=			pred_prev_sample

									SCREAMING_SNAKE_CASE_    			=			torch.sum(torch.abs(SCREAMING_SNAKE_CASE_     )     )
									SCREAMING_SNAKE_CASE_    			=			torch.mean(torch.abs(SCREAMING_SNAKE_CASE_     )     )

									assert abs(result_sum.item() - 2_52.2_68_24_95     ) < 1e-2
									assert abs(result_mean.item() - 0.3_28_47_43     ) < 1e-3

							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.scheduler_classes[0]
									SCREAMING_SNAKE_CASE_    			=			self.get_scheduler_config()
									SCREAMING_SNAKE_CASE_    			=			scheduler_class(**SCREAMING_SNAKE_CASE_     )

									scheduler.set_timesteps(25     )

									SCREAMING_SNAKE_CASE_    			=			scheduler.timesteps

									SCREAMING_SNAKE_CASE_    			=			self.dummy_model()
									SCREAMING_SNAKE_CASE_    			=			self.dummy_sample_deter
									SCREAMING_SNAKE_CASE_    			=			torch.manual_seed(0     )

									for i, t in enumerate(SCREAMING_SNAKE_CASE_     ):
											# 1. predict noise residual
											SCREAMING_SNAKE_CASE_    			=			model(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     )

											if i + 1 == timesteps.shape[0]:
													SCREAMING_SNAKE_CASE_    			=			None
											else:
													SCREAMING_SNAKE_CASE_    			=			timesteps[i + 1]

											# 2. predict previous mean of sample x_t-1
											SCREAMING_SNAKE_CASE_    			=			scheduler.step(
											    SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							prev_timestep=SCREAMING_SNAKE_CASE_							,							generator=SCREAMING_SNAKE_CASE_     ).prev_sample

											SCREAMING_SNAKE_CASE_    			=			pred_prev_sample

									SCREAMING_SNAKE_CASE_    			=			torch.sum(torch.abs(SCREAMING_SNAKE_CASE_     )     )
									SCREAMING_SNAKE_CASE_    			=			torch.mean(torch.abs(SCREAMING_SNAKE_CASE_     )     )

									assert abs(result_sum.item() - 2_58.2_04_49_83     ) < 1e-2
									assert abs(result_mean.item() - 0.3_36_20_38     ) < 1e-3

							def _lowercase       (self     ):

									"""simple docstring"""







									pass







							def _lowercase       (self     ):

									"""simple docstring"""







									pass