


"""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    lowerCamelCase__						(		_a			,   unittest.TestCase    ):
      a :     Tuple											=		DanceDiffusionPipeline
      a :     Optional[Any]											=		UNCONDITIONAL_AUDIO_GENERATION_PARAMS
      a :     str											=		PipelineTesterMixin.required_optional_params - {
          """callback""",
          """latents""",
          """callback_steps""",
          """output_type""",
          """num_images_per_prompt""",
      }
      a :     List[Any]											=		UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
      a :     Union[str, Any]											=		False
      a :     str											=		False
      def SCREAMING_SNAKE_CASE_	(		self	:	str							):






       '''simple docstring'''


       torch.manual_seed(0							)
       __lowercase									=	UNetaDModel(
           block_out_channels=(3_2, 3_2, 6_4)      ,       extra_in_channels=1_6      ,       sample_size=5_1_2      ,       sample_rate=1_6_0_0_0      ,       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""")      ,       )
       __lowercase									=	IPNDMScheduler()

       __lowercase									=	{
           """unet""": unet,
           """scheduler""": scheduler,
       }
       return components
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]      ,       A_	:	Tuple      ,       A_	:	List[str]=0							):






       '''simple docstring'''


       if str(A_							).startswith("""mps"""							):
        __lowercase									=	torch.manual_seed(A_							)
       else:
        __lowercase									=	torch.Generator(device=A_							).manual_seed(A_							)
       __lowercase									=	{
           """batch_size""": 1,
           """generator""": generator,
           """num_inference_steps""": 4,
       }
       return inputs
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]							):






       '''simple docstring'''


       __lowercase									=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
       __lowercase									=	self.get_dummy_components()
       __lowercase									=	DanceDiffusionPipeline(**A_							)
       __lowercase									=	pipe.to(A_							)
       pipe.set_progress_bar_config(disable=A_							)

       __lowercase									=	self.get_dummy_inputs(A_							)
       __lowercase									=	pipe(**A_							)
       __lowercase									=	output.audios

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

       assert audio.shape == (1, 2, components["unet"].sample_size)
       __lowercase									=	np.array([-0.72_65, 1.00_00, -0.83_88, 0.11_75, 0.94_98, -1.00_00]							)
       assert np.abs(audio_slice.flatten() - expected_slice							).max() < 1e-2
      @skip_mps
      def SCREAMING_SNAKE_CASE_	(		self	:	Union[str, Any]							):






       '''simple docstring'''


       return super().test_save_load_local()
      @skip_mps
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


       return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3							)
      @skip_mps
      def SCREAMING_SNAKE_CASE_	(		self	:	Tuple							):






       '''simple docstring'''


       return super().test_save_load_optional_components()
      @skip_mps
      def SCREAMING_SNAKE_CASE_	(		self	:	Any							):






       '''simple docstring'''


       return super().test_attention_slicing_forward_pass()






      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


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






@slow
@require_torch_gpu
class    lowerCamelCase__						(		unittest.TestCase    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]							):






       '''simple docstring'''


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	torch_device

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

       __lowercase									=	torch.manual_seed(0							)
       __lowercase									=	pipe(generator=A_      ,       num_inference_steps=1_0_0      ,       audio_length_in_s=4.0_96							)
       __lowercase									=	output.audios

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

       assert audio.shape == (1, 2, pipe.unet.sample_size)
       __lowercase									=	np.array([-0.01_92, -0.02_31, -0.03_18, -0.00_59, 0.00_02, -0.00_20]							)

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






      def SCREAMING_SNAKE_CASE_	(		self	:	List[Any]							):






       '''simple docstring'''


       __lowercase									=	torch_device

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

       __lowercase									=	torch.manual_seed(0							)
       __lowercase									=	pipe(generator=A_      ,       num_inference_steps=1_0_0      ,       audio_length_in_s=4.0_96							)
       __lowercase									=	output.audios

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

       assert audio.shape == (1, 2, pipe.unet.sample_size)
       __lowercase									=	np.array([-0.03_67, -0.04_88, -0.07_71, -0.05_25, -0.04_44, -0.03_41]							)

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




"""simple docstring"""

from collections.abc import Callable

import numpy as np


def lowerCAmelCase_      (				UpperCamelCase__	: Callable				,   UpperCamelCase__	: float				,   UpperCamelCase__	: float				,   UpperCamelCase__	: float				,   UpperCamelCase__	: float ):






 """simple docstring"""







 __lowercase									=	int(np.ceil((x_end - xa) / step_size ) )
 __lowercase									=	np.zeros((n + 1,) )
 __lowercase									=	ya
 __lowercase									=	xa

 for k in range(UpperCamelCase__ ):
  __lowercase									=	y[k] + step_size * ode_func(UpperCamelCase__				,   y[k] )
  x += step_size

 return y


if __name__ == "__main__":
       import doctest

       doctest.testmod()
