



'''simple docstring'''
import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest






class 							lowercase__     (							lowercase       ):
	lowercase__      =      (PNDMScheduler,)
	lowercase__      =      (("""num_inference_steps""", 50),)





	def      UpperCamelCase_  (			self :    List[Any]			,**lowerCamelCase__ :    Optional[int]			):







							'''simple docstring'''
							_UpperCamelCase :		Union[str, Any]			    =       {
							    'num_train_timesteps': 1000,
							    'beta_start': 0.0_0_0_1,
							    'beta_end': 0.0_2,
							    'beta_schedule': 'linear',
							}

							config.update(**lowerCamelCase__			)
							return config





	def      UpperCamelCase_  (			self :    Tuple			,lowerCamelCase__ :    Any=0			,**lowerCamelCase__ :    str			):







							'''simple docstring'''
							_UpperCamelCase :		int			    =       dict(self.forward_default_kwargs			)
							_UpperCamelCase :		Dict			    =       kwargs.pop('num_inference_steps'			,lowerCamelCase__			)
							_UpperCamelCase :		Any			    =       self.dummy_sample
							_UpperCamelCase :		int			    =       0.1 * sample
							_UpperCamelCase :		Optional[int]			    =       [residual + 0.2, residual + 0.1_5, residual + 0.1, residual + 0.0_5]

							for scheduler_class in self.scheduler_classes:
													_UpperCamelCase :		Dict			    =       self.get_scheduler_config(**lowerCamelCase__			)
													_UpperCamelCase :		Dict			    =       scheduler_class(**lowerCamelCase__			)
													scheduler.set_timesteps(lowerCamelCase__			)
													# copy over dummy past residuals
													_UpperCamelCase :		Optional[Any]			    =       dummy_past_residuals[:]

													with tempfile.TemporaryDirectory() as tmpdirname:
																			scheduler.save_config(lowerCamelCase__			)
																			_UpperCamelCase :		Dict			    =       scheduler_class.from_pretrained(lowerCamelCase__			)
																			new_scheduler.set_timesteps(lowerCamelCase__			)
																			# copy over dummy past residuals
																			_UpperCamelCase :		Optional[int]			    =       dummy_past_residuals[:]

													_UpperCamelCase :		int			    =       scheduler.step_prk(lowerCamelCase__			,lowerCamelCase__			,lowerCamelCase__			,**lowerCamelCase__			).prev_sample
													_UpperCamelCase :		Any			    =       new_scheduler.step_prk(lowerCamelCase__			,lowerCamelCase__			,lowerCamelCase__			,**lowerCamelCase__			).prev_sample

													assert torch.sum(torch.abs(output - new_output			)			) < 1E-5, "Scheduler outputs are not identical"

													_UpperCamelCase :		Optional[int]			    =       scheduler.step_plms(lowerCamelCase__			,lowerCamelCase__			,lowerCamelCase__			,**lowerCamelCase__			).prev_sample
													_UpperCamelCase :		Tuple			    =       new_scheduler.step_plms(lowerCamelCase__			,lowerCamelCase__			,lowerCamelCase__			,**lowerCamelCase__			).prev_sample

													assert torch.sum(torch.abs(output - new_output			)			) < 1E-5, "Scheduler outputs are not identical"





	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							pass





	def      UpperCamelCase_  (			self :    List[str]			,lowerCamelCase__ :    str=0			,**lowerCamelCase__ :    Tuple			):







							'''simple docstring'''
							_UpperCamelCase :		Optional[Any]			    =       dict(self.forward_default_kwargs			)
							_UpperCamelCase :		Optional[int]			    =       kwargs.pop('num_inference_steps'			,lowerCamelCase__			)
							_UpperCamelCase :		List[Any]			    =       self.dummy_sample
							_UpperCamelCase :		List[str]			    =       0.1 * sample
							_UpperCamelCase :		int			    =       [residual + 0.2, residual + 0.1_5, residual + 0.1, residual + 0.0_5]

							for scheduler_class in self.scheduler_classes:
													_UpperCamelCase :		Union[str, Any]			    =       self.get_scheduler_config()
													_UpperCamelCase :		List[str]			    =       scheduler_class(**lowerCamelCase__			)
													scheduler.set_timesteps(lowerCamelCase__			)

													# copy over dummy past residuals (must be after setting timesteps)
													_UpperCamelCase :		int			    =       dummy_past_residuals[:]

													with tempfile.TemporaryDirectory() as tmpdirname:
																			scheduler.save_config(lowerCamelCase__			)
																			_UpperCamelCase :		Optional[Any]			    =       scheduler_class.from_pretrained(lowerCamelCase__			)
																			# copy over dummy past residuals
																			new_scheduler.set_timesteps(lowerCamelCase__			)

																			# copy over dummy past residual (must be after setting timesteps)
																			_UpperCamelCase :		Tuple			    =       dummy_past_residuals[:]

													_UpperCamelCase :		List[str]			    =       scheduler.step_prk(lowerCamelCase__			,lowerCamelCase__			,lowerCamelCase__			,**lowerCamelCase__			).prev_sample
													_UpperCamelCase :		int			    =       new_scheduler.step_prk(lowerCamelCase__			,lowerCamelCase__			,lowerCamelCase__			,**lowerCamelCase__			).prev_sample

													assert torch.sum(torch.abs(output - new_output			)			) < 1E-5, "Scheduler outputs are not identical"

													_UpperCamelCase :		Tuple			    =       scheduler.step_plms(lowerCamelCase__			,lowerCamelCase__			,lowerCamelCase__			,**lowerCamelCase__			).prev_sample
													_UpperCamelCase :		List[Any]			    =       new_scheduler.step_plms(lowerCamelCase__			,lowerCamelCase__			,lowerCamelCase__			,**lowerCamelCase__			).prev_sample

													assert torch.sum(torch.abs(output - new_output			)			) < 1E-5, "Scheduler outputs are not identical"





	def      UpperCamelCase_  (			self :    Optional[int]			,**lowerCamelCase__ :    Any			):







							'''simple docstring'''
							_UpperCamelCase :		Dict			    =       self.scheduler_classes[0]
							_UpperCamelCase :		int			    =       self.get_scheduler_config(**lowerCamelCase__			)
							_UpperCamelCase :		Any			    =       scheduler_class(**lowerCamelCase__			)

							_UpperCamelCase :		Optional[Any]			    =       10
							_UpperCamelCase :		Any			    =       self.dummy_model()
							_UpperCamelCase :		List[str]			    =       self.dummy_sample_deter
							scheduler.set_timesteps(lowerCamelCase__			)

							for i, t in enumerate(scheduler.prk_timesteps			):
													_UpperCamelCase :		List[Any]			    =       model(lowerCamelCase__			,lowerCamelCase__			)
													_UpperCamelCase :		List[str]			    =       scheduler.step_prk(lowerCamelCase__			,lowerCamelCase__			,lowerCamelCase__			).prev_sample

							for i, t in enumerate(scheduler.plms_timesteps			):
													_UpperCamelCase :		Union[str, Any]			    =       model(lowerCamelCase__			,lowerCamelCase__			)
													_UpperCamelCase :		List[Any]			    =       scheduler.step_plms(lowerCamelCase__			,lowerCamelCase__			,lowerCamelCase__			).prev_sample

							return sample





	def      UpperCamelCase_  (			self :    Union[str, Any]			):







							'''simple docstring'''
							_UpperCamelCase :		int			    =       dict(self.forward_default_kwargs			)

							_UpperCamelCase :		Dict			    =       kwargs.pop('num_inference_steps'			,lowerCamelCase__			)

							for scheduler_class in self.scheduler_classes:
													_UpperCamelCase :		Optional[int]			    =       self.get_scheduler_config()
													_UpperCamelCase :		int			    =       scheduler_class(**lowerCamelCase__			)

													_UpperCamelCase :		Union[str, Any]			    =       self.dummy_sample
													_UpperCamelCase :		List[str]			    =       0.1 * sample

													if num_inference_steps is not None and hasattr(lowerCamelCase__			,'set_timesteps'			):
																			scheduler.set_timesteps(lowerCamelCase__			)
													elif num_inference_steps is not None and not hasattr(lowerCamelCase__			,'set_timesteps'			):
																			_UpperCamelCase :		Optional[int]			    =       num_inference_steps

													# copy over dummy past residuals (must be done after set_timesteps)
													_UpperCamelCase :		Optional[int]			    =       [residual + 0.2, residual + 0.1_5, residual + 0.1, residual + 0.0_5]
													_UpperCamelCase :		int			    =       dummy_past_residuals[:]

													_UpperCamelCase :		Union[str, Any]			    =       scheduler.step_prk(lowerCamelCase__			,0			,lowerCamelCase__			,**lowerCamelCase__			).prev_sample
													_UpperCamelCase :		Optional[int]			    =       scheduler.step_prk(lowerCamelCase__			,1			,lowerCamelCase__			,**lowerCamelCase__			).prev_sample

													self.assertEqual(output_a.shape			,sample.shape			)
													self.assertEqual(output_a.shape			,output_a.shape			)

													_UpperCamelCase :		Optional[Any]			    =       scheduler.step_plms(lowerCamelCase__			,0			,lowerCamelCase__			,**lowerCamelCase__			).prev_sample
													_UpperCamelCase :		Optional[Any]			    =       scheduler.step_plms(lowerCamelCase__			,1			,lowerCamelCase__			,**lowerCamelCase__			).prev_sample

													self.assertEqual(output_a.shape			,sample.shape			)
													self.assertEqual(output_a.shape			,output_a.shape			)





	def      UpperCamelCase_  (			self :    List[str]			):







							'''simple docstring'''
							for timesteps in [100, 1000]:
													self.check_over_configs(num_train_timesteps=lowerCamelCase__			)





	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							for steps_offset in [0, 1]:
													self.check_over_configs(steps_offset=lowerCamelCase__			)

							_UpperCamelCase :		Union[str, Any]			    =       self.scheduler_classes[0]
							_UpperCamelCase :		Tuple			    =       self.get_scheduler_config(steps_offset=1			)
							_UpperCamelCase :		Tuple			    =       scheduler_class(**lowerCamelCase__			)
							scheduler.set_timesteps(10			)
							assert torch.equal(
							    scheduler.timesteps			,torch.LongTensor(
							        [901, 851, 851, 801, 801, 751, 751, 701, 701, 651, 651, 601, 601, 501, 401, 301, 201, 101, 1]			)			,)





	def      UpperCamelCase_  (			self :    Tuple			):







							'''simple docstring'''
							for beta_start, beta_end in zip([0.0_0_0_1, 0.0_0_1]			,[0.0_0_2, 0.0_2]			):
													self.check_over_configs(beta_start=lowerCamelCase__			,beta_end=lowerCamelCase__			)





	def      UpperCamelCase_  (			self :    Dict			):







							'''simple docstring'''
							for schedule in ["linear", "squaredcos_cap_v2"]:
													self.check_over_configs(beta_schedule=lowerCamelCase__			)





	def      UpperCamelCase_  (			self :    Dict			):







							'''simple docstring'''
							for prediction_type in ["epsilon", "v_prediction"]:
													self.check_over_configs(prediction_type=lowerCamelCase__			)





	def      UpperCamelCase_  (			self :    Union[str, Any]			):







							'''simple docstring'''
							for t in [1, 5, 10]:
													self.check_over_forward(time_step=lowerCamelCase__			)





	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							for t, num_inference_steps in zip([1, 5, 10]			,[10, 50, 100]			):
													self.check_over_forward(num_inference_steps=lowerCamelCase__			)





	def      UpperCamelCase_  (			self :    List[str]			):







							'''simple docstring'''
							# earlier version of set_timesteps() caused an error indexing alpha's with inference steps as power of 3
							_UpperCamelCase :		Tuple			    =       27

							for scheduler_class in self.scheduler_classes:
													_UpperCamelCase :		List[Any]			    =       self.dummy_sample
													_UpperCamelCase :		Any			    =       0.1 * sample

													_UpperCamelCase :		Union[str, Any]			    =       self.get_scheduler_config()
													_UpperCamelCase :		Union[str, Any]			    =       scheduler_class(**lowerCamelCase__			)

													scheduler.set_timesteps(lowerCamelCase__			)

													# before power of 3 fix, would error on first step, so we only need to do two
													for i, t in enumerate(scheduler.prk_timesteps[:2]			):
																			_UpperCamelCase :		Dict			    =       scheduler.step_prk(lowerCamelCase__			,lowerCamelCase__			,lowerCamelCase__			).prev_sample





	def      UpperCamelCase_  (			self :    Union[str, Any]			):







							'''simple docstring'''
							with self.assertRaises(lowerCamelCase__			):
													_UpperCamelCase :		Any			    =       self.scheduler_classes[0]
													_UpperCamelCase :		str			    =       self.get_scheduler_config()
													_UpperCamelCase :		List[Any]			    =       scheduler_class(**lowerCamelCase__			)

													scheduler.step_plms(self.dummy_sample			,1			,self.dummy_sample			).prev_sample





	def      UpperCamelCase_  (			self :    Tuple			):







							'''simple docstring'''
							_UpperCamelCase :		str			    =       self.full_loop()
							_UpperCamelCase :		Dict			    =       torch.sum(torch.abs(lowerCamelCase__			)			)
							_UpperCamelCase :		int			    =       torch.mean(torch.abs(lowerCamelCase__			)			)

							assert abs(result_sum.item() - 1_9_8.1_3_1_8			) < 1E-2
							assert abs(result_mean.item() - 0.2_5_8_0			) < 1E-3





	def      UpperCamelCase_  (			self :    Union[str, Any]			):







							'''simple docstring'''
							_UpperCamelCase :		List[Any]			    =       self.full_loop(prediction_type='v_prediction'			)
							_UpperCamelCase :		Tuple			    =       torch.sum(torch.abs(lowerCamelCase__			)			)
							_UpperCamelCase :		Dict			    =       torch.mean(torch.abs(lowerCamelCase__			)			)

							assert abs(result_sum.item() - 6_7.3_9_8_6			) < 1E-2
							assert abs(result_mean.item() - 0.0_8_7_8			) < 1E-3





	def      UpperCamelCase_  (			self :    Optional[Any]			):







							'''simple docstring'''
							# We specify different beta, so that the first alpha is 0.99
							_UpperCamelCase :		str			    =       self.full_loop(set_alpha_to_one=lowerCamelCase__			,beta_start=0.0_1			)
							_UpperCamelCase :		Optional[int]			    =       torch.sum(torch.abs(lowerCamelCase__			)			)
							_UpperCamelCase :		int			    =       torch.mean(torch.abs(lowerCamelCase__			)			)

							assert abs(result_sum.item() - 2_3_0.0_3_9_9			) < 1E-2
							assert abs(result_mean.item() - 0.2_9_9_5			) < 1E-3





	def      UpperCamelCase_  (			self :    List[Any]			):







							'''simple docstring'''
							# We specify different beta, so that the first alpha is 0.99
							_UpperCamelCase :		Dict			    =       self.full_loop(set_alpha_to_one=lowerCamelCase__			,beta_start=0.0_1			)
							_UpperCamelCase :		Tuple			    =       torch.sum(torch.abs(lowerCamelCase__			)			)
							_UpperCamelCase :		List[Any]			    =       torch.mean(torch.abs(lowerCamelCase__			)			)

							assert abs(result_sum.item() - 1_8_6.9_4_8_2			) < 1E-2
							assert abs(result_mean.item() - 0.2_4_3_4			) < 1E-3







'''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 							lowercase__     (							lowercase						, 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 :		str			    =       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=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 :		int			    =       IPNDMScheduler()

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





	def      UpperCamelCase_  (			self :    Tuple			,lowerCamelCase__ :    Tuple			,lowerCamelCase__ :    int=0			):







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





	def      UpperCamelCase_  (			self :    int			):







							'''simple docstring'''
							_UpperCamelCase :		Union[str, Any]			    =       'cpu'  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase :		List[str]			    =       self.get_dummy_components()
							_UpperCamelCase :		int			    =       DanceDiffusionPipeline(**lowerCamelCase__			)
							_UpperCamelCase :		List[str]			    =       pipe.to(lowerCamelCase__			)
							pipe.set_progress_bar_config(disable=lowerCamelCase__			)

							_UpperCamelCase :		Union[str, Any]			    =       self.get_dummy_inputs(lowerCamelCase__			)
							_UpperCamelCase :		List[Any]			    =       pipe(**lowerCamelCase__			)
							_UpperCamelCase :		Union[str, Any]			    =       output.audios

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

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







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





	@skip_mps
	def      UpperCamelCase_  (			self :    Dict			):







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





	@skip_mps
	def      UpperCamelCase_  (			self :    str			):







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





	@skip_mps
	def      UpperCamelCase_  (			self :    List[Any]			):







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





	def      UpperCamelCase_  (			self :    List[str]			):







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






@slow
@require_torch_gpu
class 							lowercase__     (							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 :    int			):







							'''simple docstring'''
							_UpperCamelCase :		Union[str, Any]			    =       torch_device

							_UpperCamelCase :		Dict			    =       DanceDiffusionPipeline.from_pretrained('harmonai/maestro-150k'			)
							_UpperCamelCase :		str			    =       pipe.to(lowerCamelCase__			)
							pipe.set_progress_bar_config(disable=lowerCamelCase__			)

							_UpperCamelCase :		int			    =       torch.manual_seed(0			)
							_UpperCamelCase :		Optional[Any]			    =       pipe(generator=lowerCamelCase__			,num_inference_steps=100			,audio_length_in_s=4.0_9_6			)
							_UpperCamelCase :		Optional[int]			    =       output.audios

							_UpperCamelCase :		int			    =       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      UpperCamelCase_  (			self :    Union[str, Any]			):







							'''simple docstring'''
							_UpperCamelCase :		List[Any]			    =       torch_device

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

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

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

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


