
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       _lowerCAmelCase  (		__magic_name__			,	unittest.TestCase							):

 """simple docstring"""
 SCREAMING_SNAKE_CASE_				:					Optional[int]       						=DanceDiffusionPipeline
 SCREAMING_SNAKE_CASE_				:					List[str]       						=UNCONDITIONAL_AUDIO_GENERATION_PARAMS
 SCREAMING_SNAKE_CASE_				:					List[Any]       						=PipelineTesterMixin.required_optional_params - {
     "callback",
     "latents",
     "callback_steps",
     "output_type",
     "num_images_per_prompt",
 }
 SCREAMING_SNAKE_CASE_				:					List[Any]       						=UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
 SCREAMING_SNAKE_CASE_				:					Dict       						=False
 SCREAMING_SNAKE_CASE_				:					List[str]       						=False

 def      __lowerCAmelCase  (						self			:			Optional[int]   ):






      """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=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')					, )
      UpperCamelCase						    =					IPNDMScheduler()

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

 def      __lowerCAmelCase  (						self			:			Tuple					, SCREAMING_SNAKE_CASE__			:			int					, SCREAMING_SNAKE_CASE__			:			Optional[int]=0   ):






      """simple docstring"""



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

 def      __lowerCAmelCase  (						self			:			Dict   ):






      """simple docstring"""



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

      UpperCamelCase						    =					self.get_dummy_inputs(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					pipe(**SCREAMING_SNAKE_CASE__   )
      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      __lowerCAmelCase  (						self			:			List[str]   ):






      """simple docstring"""



      return super().test_save_load_local()

 @skip_mps
 def      __lowerCAmelCase  (						self			:			Optional[int]   ):






      """simple docstring"""



      return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3   )

 @skip_mps
 def      __lowerCAmelCase  (						self			:			str   ):






      """simple docstring"""



      return super().test_save_load_optional_components()

 @skip_mps
 def      __lowerCAmelCase  (						self			:			Dict   ):






      """simple docstring"""



      return super().test_attention_slicing_forward_pass()






 def      __lowerCAmelCase  (						self			:			Tuple   ):






      """simple docstring"""



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




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

 """simple docstring"""

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






      """simple docstring"""



      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()

 def      __lowerCAmelCase  (						self			:			str   ):






      """simple docstring"""



      UpperCamelCase						    =					torch_device

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

      UpperCamelCase						    =					torch.manual_seed(0   )
      UpperCamelCase						    =					pipe(generator=SCREAMING_SNAKE_CASE__					, 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      __lowerCAmelCase  (						self			:			Optional[Any]   ):






      """simple docstring"""



      UpperCamelCase						    =					torch_device

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

      UpperCamelCase						    =					torch.manual_seed(0   )
      UpperCamelCase						    =					pipe(generator=SCREAMING_SNAKE_CASE__					, 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




import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest
class       _lowerCAmelCase  (		__magic_name__							):

 """simple docstring"""
 SCREAMING_SNAKE_CASE_				:					Tuple       						=(DEISMultistepScheduler,)
 SCREAMING_SNAKE_CASE_				:					List[str]       						=(("num_inference_steps", 25),)

 def      __lowerCAmelCase  (						self			:			Tuple					, **SCREAMING_SNAKE_CASE__			:			List[str]   ):






      """simple docstring"""



      UpperCamelCase						    =					{
          'num_train_timesteps': 10_00,
          'beta_start': 0.0001,
          'beta_end': 0.02,
          'beta_schedule': 'linear',
          'solver_order': 2,
      }

      config.update(**SCREAMING_SNAKE_CASE__   )
      return config

 def      __lowerCAmelCase  (						self			:			Any					, SCREAMING_SNAKE_CASE__			:			Dict=0					, **SCREAMING_SNAKE_CASE__			:			List[str]   ):






      """simple docstring"""



      UpperCamelCase						    =					dict(self.forward_default_kwargs   )
      UpperCamelCase						    =					kwargs.pop('num_inference_steps'					, SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					self.dummy_sample
      UpperCamelCase						    =					0.1 * sample
      UpperCamelCase						    =					[residual + 0.2, residual + 0.15, residual + 0.10]

      for scheduler_class in self.scheduler_classes:
           UpperCamelCase						    =					self.get_scheduler_config(**SCREAMING_SNAKE_CASE__   )
           UpperCamelCase						    =					scheduler_class(**SCREAMING_SNAKE_CASE__   )
           scheduler.set_timesteps(SCREAMING_SNAKE_CASE__   )
           # copy over dummy past residuals
           UpperCamelCase						    =					dummy_past_residuals[: scheduler.config.solver_order]

           with tempfile.TemporaryDirectory() as tmpdirname:
                scheduler.save_config(SCREAMING_SNAKE_CASE__   )
                UpperCamelCase						    =					scheduler_class.from_pretrained(SCREAMING_SNAKE_CASE__   )
                new_scheduler.set_timesteps(SCREAMING_SNAKE_CASE__   )
                # copy over dummy past residuals
                UpperCamelCase						    =					dummy_past_residuals[: new_scheduler.config.solver_order]

           UpperCamelCase    ,		UpperCamelCase						    =					sample, sample
           for t in range(SCREAMING_SNAKE_CASE__					, time_step + scheduler.config.solver_order + 1   ):
                UpperCamelCase						    =					scheduler.step(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, **SCREAMING_SNAKE_CASE__   ).prev_sample
                UpperCamelCase						    =					new_scheduler.step(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, **SCREAMING_SNAKE_CASE__   ).prev_sample

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

 def      __lowerCAmelCase  (						self			:			Optional[int]   ):






      """simple docstring"""



      pass

 def      __lowerCAmelCase  (						self			:			Optional[int]					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=0					, **SCREAMING_SNAKE_CASE__			:			int   ):






      """simple docstring"""



      UpperCamelCase						    =					dict(self.forward_default_kwargs   )
      UpperCamelCase						    =					kwargs.pop('num_inference_steps'					, SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					self.dummy_sample
      UpperCamelCase						    =					0.1 * sample
      UpperCamelCase						    =					[residual + 0.2, residual + 0.15, residual + 0.10]

      for scheduler_class in self.scheduler_classes:
           UpperCamelCase						    =					self.get_scheduler_config()
           UpperCamelCase						    =					scheduler_class(**SCREAMING_SNAKE_CASE__   )
           scheduler.set_timesteps(SCREAMING_SNAKE_CASE__   )

           # copy over dummy past residuals (must be after setting timesteps)
           UpperCamelCase						    =					dummy_past_residuals[: scheduler.config.solver_order]

           with tempfile.TemporaryDirectory() as tmpdirname:
                scheduler.save_config(SCREAMING_SNAKE_CASE__   )
                UpperCamelCase						    =					scheduler_class.from_pretrained(SCREAMING_SNAKE_CASE__   )
                # copy over dummy past residuals
                new_scheduler.set_timesteps(SCREAMING_SNAKE_CASE__   )

                # copy over dummy past residual (must be after setting timesteps)
                UpperCamelCase						    =					dummy_past_residuals[: new_scheduler.config.solver_order]

           UpperCamelCase						    =					scheduler.step(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, **SCREAMING_SNAKE_CASE__   ).prev_sample
           UpperCamelCase						    =					new_scheduler.step(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, **SCREAMING_SNAKE_CASE__   ).prev_sample

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

 def      __lowerCAmelCase  (						self			:			Dict					, SCREAMING_SNAKE_CASE__			:			List[Any]=None					, **SCREAMING_SNAKE_CASE__			:			Union[str, Any]   ):






      """simple docstring"""



      if scheduler is None:
           UpperCamelCase						    =					self.scheduler_classes[0]
           UpperCamelCase						    =					self.get_scheduler_config(**SCREAMING_SNAKE_CASE__   )
           UpperCamelCase						    =					scheduler_class(**SCREAMING_SNAKE_CASE__   )

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

      UpperCamelCase						    =					10
      UpperCamelCase						    =					self.dummy_model()
      UpperCamelCase						    =					self.dummy_sample_deter
      scheduler.set_timesteps(SCREAMING_SNAKE_CASE__   )

      for i, t in enumerate(scheduler.timesteps   ):
           UpperCamelCase						    =					model(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )
           UpperCamelCase						    =					scheduler.step(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   ).prev_sample

      return sample

 def      __lowerCAmelCase  (						self			:			List[str]   ):






      """simple docstring"""



      UpperCamelCase						    =					dict(self.forward_default_kwargs   )

      UpperCamelCase						    =					kwargs.pop('num_inference_steps'					, SCREAMING_SNAKE_CASE__   )

      for scheduler_class in self.scheduler_classes:
           UpperCamelCase						    =					self.get_scheduler_config()
           UpperCamelCase						    =					scheduler_class(**SCREAMING_SNAKE_CASE__   )

           UpperCamelCase						    =					self.dummy_sample
           UpperCamelCase						    =					0.1 * sample

           if num_inference_steps is not None and hasattr(SCREAMING_SNAKE_CASE__					, 'set_timesteps'   ):
                scheduler.set_timesteps(SCREAMING_SNAKE_CASE__   )
           elif num_inference_steps is not None and not hasattr(SCREAMING_SNAKE_CASE__					, 'set_timesteps'   ):
                UpperCamelCase						    =					num_inference_steps

           # copy over dummy past residuals (must be done after set_timesteps)
           UpperCamelCase						    =					[residual + 0.2, residual + 0.15, residual + 0.10]
           UpperCamelCase						    =					dummy_past_residuals[: scheduler.config.solver_order]

           UpperCamelCase						    =					scheduler.timesteps[5]
           UpperCamelCase						    =					scheduler.timesteps[6]

           UpperCamelCase						    =					scheduler.step(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, **SCREAMING_SNAKE_CASE__   ).prev_sample
           UpperCamelCase						    =					scheduler.step(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, **SCREAMING_SNAKE_CASE__   ).prev_sample

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

 def      __lowerCAmelCase  (						self			:			List[Any]   ):






      """simple docstring"""



      UpperCamelCase						    =					DEISMultistepScheduler(**self.get_scheduler_config()   )
      UpperCamelCase						    =					self.full_loop(scheduler=SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					torch.mean(torch.abs(SCREAMING_SNAKE_CASE__   )   )

      assert abs(result_mean.item() - 0.23916   ) < 1e-3

      UpperCamelCase						    =					DPMSolverSinglestepScheduler.from_config(scheduler.config   )
      UpperCamelCase						    =					DPMSolverMultistepScheduler.from_config(scheduler.config   )
      UpperCamelCase						    =					UniPCMultistepScheduler.from_config(scheduler.config   )
      UpperCamelCase						    =					DEISMultistepScheduler.from_config(scheduler.config   )

      UpperCamelCase						    =					self.full_loop(scheduler=SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					torch.mean(torch.abs(SCREAMING_SNAKE_CASE__   )   )

      assert abs(result_mean.item() - 0.23916   ) < 1e-3

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






      """simple docstring"""



      for timesteps in [25, 50, 1_00, 9_99, 10_00]:
           self.check_over_configs(num_train_timesteps=SCREAMING_SNAKE_CASE__   )

 def      __lowerCAmelCase  (						self			:			str   ):






      """simple docstring"""



      self.check_over_configs(thresholding=SCREAMING_SNAKE_CASE__   )
      for order in [1, 2, 3]:
           for solver_type in ["logrho"]:
                for threshold in [0.5, 1.0, 2.0]:
                     for prediction_type in ["epsilon", "sample"]:
                          self.check_over_configs(
                              thresholding=SCREAMING_SNAKE_CASE__					, prediction_type=SCREAMING_SNAKE_CASE__					, sample_max_value=SCREAMING_SNAKE_CASE__					, algorithm_type='deis'					, solver_order=SCREAMING_SNAKE_CASE__					, solver_type=SCREAMING_SNAKE_CASE__					, )

 def      __lowerCAmelCase  (						self			:			List[Any]   ):






      """simple docstring"""



      for prediction_type in ["epsilon", "v_prediction"]:
           self.check_over_configs(prediction_type=SCREAMING_SNAKE_CASE__   )

 def      __lowerCAmelCase  (						self			:			Optional[int]   ):






      """simple docstring"""



      for algorithm_type in ["deis"]:
           for solver_type in ["logrho"]:
                for order in [1, 2, 3]:
                     for prediction_type in ["epsilon", "sample"]:
                          self.check_over_configs(
                              solver_order=SCREAMING_SNAKE_CASE__					, solver_type=SCREAMING_SNAKE_CASE__					, prediction_type=SCREAMING_SNAKE_CASE__					, algorithm_type=SCREAMING_SNAKE_CASE__					, )
                          UpperCamelCase						    =					self.full_loop(
                              solver_order=SCREAMING_SNAKE_CASE__					, solver_type=SCREAMING_SNAKE_CASE__					, prediction_type=SCREAMING_SNAKE_CASE__					, algorithm_type=SCREAMING_SNAKE_CASE__					, )
                          assert not torch.isnan(SCREAMING_SNAKE_CASE__   ).any(), "Samples have nan numbers"

 def      __lowerCAmelCase  (						self			:			Optional[Any]   ):






      """simple docstring"""



      self.check_over_configs(lower_order_final=SCREAMING_SNAKE_CASE__   )
      self.check_over_configs(lower_order_final=SCREAMING_SNAKE_CASE__   )

 def      __lowerCAmelCase  (						self			:			Any   ):






      """simple docstring"""



      for num_inference_steps in [1, 2, 3, 5, 10, 50, 1_00, 9_99, 10_00]:
           self.check_over_forward(num_inference_steps=SCREAMING_SNAKE_CASE__					, time_step=0   )

 def      __lowerCAmelCase  (						self			:			Optional[Any]   ):






      """simple docstring"""



      UpperCamelCase						    =					self.full_loop()
      UpperCamelCase						    =					torch.mean(torch.abs(SCREAMING_SNAKE_CASE__   )   )

      assert abs(result_mean.item() - 0.23916   ) < 1e-3

 def      __lowerCAmelCase  (						self			:			Dict   ):






      """simple docstring"""



      UpperCamelCase						    =					self.full_loop(prediction_type='v_prediction'   )
      UpperCamelCase						    =					torch.mean(torch.abs(SCREAMING_SNAKE_CASE__   )   )

      assert abs(result_mean.item() - 0.091   ) < 1e-3






 def      __lowerCAmelCase  (						self			:			Tuple   ):






      """simple docstring"""



      UpperCamelCase						    =					self.scheduler_classes[0]
      UpperCamelCase						    =					self.get_scheduler_config(thresholding=SCREAMING_SNAKE_CASE__					, dynamic_thresholding_ratio=0   )
      UpperCamelCase						    =					scheduler_class(**SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					10
      UpperCamelCase						    =					self.dummy_model()
      UpperCamelCase						    =					self.dummy_sample_deter.half()
      scheduler.set_timesteps(SCREAMING_SNAKE_CASE__   )

      for i, t in enumerate(scheduler.timesteps   ):
           UpperCamelCase						    =					model(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )
           UpperCamelCase						    =					scheduler.step(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   ).prev_sample

      assert sample.dtype == torch.floataa


