





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 				A_				(      UpperCAmelCase					,	unittest.TestCase	):







       """simple docstring"""





       SCREAMING_SNAKE_CASE_      : Dict      =  DanceDiffusionPipeline
       SCREAMING_SNAKE_CASE_      : Union[str, Any]      =  UNCONDITIONAL_AUDIO_GENERATION_PARAMS
       SCREAMING_SNAKE_CASE_      : 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_      : List[str]      =  False
       SCREAMING_SNAKE_CASE_      : Tuple      =  False







       def    __UpperCAmelCase (    self	:	Tuple  )						->					int:
         torch.manual_seed(0  )
         _lowercase							= UNetaDModel(
             block_out_channels=(32, 32, 64)						,extra_in_channels=16						,sample_size=512						,sample_rate=1_6000						,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    __UpperCAmelCase (    self	:	Any						,__A	:	Union[str, Any]						,__A	:	Optional[Any]=0  )						->					int:
         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    __UpperCAmelCase (    self	:	int  )						->					Optional[int]:
         _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.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	:	Tuple  )						->					Any:
         return super().test_save_load_local()







       @skip_mps
       def    __UpperCAmelCase (    self	:	Optional[Any]  )						->					Dict:
         return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3  )







       @skip_mps
       def    __UpperCAmelCase (    self	:	Tuple  )						->					Tuple:
         return super().test_save_load_optional_components()







       @skip_mps
       def    __UpperCAmelCase (    self	:	Any  )						->					Optional[Any]:
         return super().test_attention_slicing_forward_pass()







       def    __UpperCAmelCase (    self	:	str  )						->					List[Any]:
         super().test_inference_batch_single_identical(expected_max_diff=3e-3  )





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







       """simple docstring"""





       def    __UpperCAmelCase (    self	:	Dict  )						->					Optional[Any]:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()







       def    __UpperCAmelCase (    self	:	Any  )						->					List[str]:
         _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=100						,audio_length_in_s=4.096  )
         _lowercase							= output.audios

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

         assert audio.shape == (1, 2, pipe.unet.sample_size)
         _lowercase							= 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	:	Tuple  )						->					List[str]:
         _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=100						,audio_length_in_s=4.096  )
         _lowercase							= output.audios

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

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






snake_case 				=       8.3_144_598

def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :float			,							snake_case__  :float     )      ->						float:
  if temperature < 0:
    raise Exception('Temperature cannot be less than 0 K'     )
  if molar_mass <= 0:
    raise Exception('Molar mass cannot be less than or equal to 0 kg/mol'     )
  else:
    return (3 * UNIVERSAL_GAS_CONSTANT * temperature / molar_mass) ** 0.5


if __name__ == "__main__":
      import doctest

      # run doctest
      doctest.testmod()

      # example
      snake_case 				=       3_0_0
      snake_case 				=       2_8
      snake_case 				=       rms_speed_of_molecule(temperature, molar_mass)
      print(F"""Vrms of Nitrogen gas at 300 K is {vrms} m/s""")