'''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_       (		_UpperCamelCase    ,    unittest.TestCase			):







   """simple docstring"""

   __lowerCAmelCase         =   DanceDiffusionPipeline
   __lowerCAmelCase         =   UNCONDITIONAL_AUDIO_GENERATION_PARAMS
   __lowerCAmelCase         =   PipelineTesterMixin.required_optional_params - {
       "callback",
       "latents",
       "callback_steps",
       "output_type",
       "num_images_per_prompt",
   }
   __lowerCAmelCase         =   UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
   __lowerCAmelCase         =   False
   __lowerCAmelCase         =   False
   def 	__UpperCAmelCase   (      self      :			int     )  ->			Optional[Any]:
        torch.manual_seed(0     )
        _A    		=  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=UpperCamelCase__, use_timestep_embedding=UpperCamelCase__, time_embedding_type='fourier', mid_block_type='UNetMidBlock1D', down_block_types=('DownBlock1DNoSkip', 'DownBlock1D', 'AttnDownBlock1D'), up_block_types=('AttnUpBlock1D', 'UpBlock1D', 'UpBlock1DNoSkip'), )
        _A    		=  IPNDMScheduler()

        _A    		=  {
            'unet': unet,
            'scheduler': scheduler,
        }
        return components
   def 	__UpperCAmelCase   (      self      :			Dict, UpperCamelCase__      :			Optional[Any], UpperCamelCase__      :			str=0     )  ->			List[str]:
        if str(UpperCamelCase__     ).startswith('mps'     ):
             _A    		=  torch.manual_seed(UpperCamelCase__     )
        else:
             _A    		=  torch.Generator(device=UpperCamelCase__     ).manual_seed(UpperCamelCase__     )
        _A    		=  {
            'batch_size': 1,
            'generator': generator,
            'num_inference_steps': 4,
        }
        return inputs
   def 	__UpperCAmelCase   (      self      :			Optional[int]     )  ->			Union[str, Any]:
        _A    		=  'cpu'  # ensure determinism for the device-dependent torch.Generator
        _A    		=  self.get_dummy_components()
        _A    		=  DanceDiffusionPipeline(**UpperCamelCase__     )
        _A    		=  pipe.to(UpperCamelCase__     )
        pipe.set_progress_bar_config(disable=UpperCamelCase__     )

        _A    		=  self.get_dummy_inputs(UpperCamelCase__     )
        _A    		=  pipe(**UpperCamelCase__     )
        _A    		=  output.audios

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

        assert audio.shape == (1, 2, components["unet"].sample_size)
        _A    		=  np.array([-0.7_265, 1.0_000, -0.8_388, 0.1_175, 0.9_498, -1.0_000]     )
        assert np.abs(audio_slice.flatten() - expected_slice     ).max() < 1e-2
   @skip_mps
   def 	__UpperCAmelCase   (      self      :			Tuple     )  ->			Tuple:
        return super().test_save_load_local()
   @skip_mps
   def 	__UpperCAmelCase   (      self      :			Dict     )  ->			Union[str, Any]:
        return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3     )
   @skip_mps
   def 	__UpperCAmelCase   (      self      :			Union[str, Any]     )  ->			List[Any]:
        return super().test_save_load_optional_components()
   @skip_mps
   def 	__UpperCAmelCase   (      self      :			int     )  ->			Optional[int]:
        return super().test_attention_slicing_forward_pass()

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







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







   """simple docstring"""

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

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

        _A    		=  torch.manual_seed(0     )
        _A    		=  pipe(generator=UpperCamelCase__, num_inference_steps=1_00, audio_length_in_s=4.096     )
        _A    		=  output.audios

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

        assert audio.shape == (1, 2, pipe.unet.sample_size)
        _A    		=  np.array([-0.0_192, -0.0_231, -0.0_318, -0.0_059, 0.0_002, -0.0_020]     )

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

   def 	__UpperCAmelCase   (      self      :			str     )  ->			Dict:
        _A    		=  torch_device

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

        _A    		=  torch.manual_seed(0     )
        _A    		=  pipe(generator=UpperCamelCase__, num_inference_steps=1_00, audio_length_in_s=4.096     )
        _A    		=  output.audios

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

        assert audio.shape == (1, 2, pipe.unet.sample_size)
        _A    		=  np.array([-0.0_367, -0.0_488, -0.0_771, -0.0_525, -0.0_444, -0.0_341]     )

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





import importlib.metadata
import operator
import re
import sys
from typing import Optional

from packaging import version


lowercase   =	{
    '''<''': operator.lt,
    '''<=''': operator.le,
    '''==''': operator.eq,
    '''!=''': operator.ne,
    '''>=''': operator.ge,
    '''>''': operator.gt,
}



def   __lowerCAmelCase    (				UpperCAmelCase__					:   Tuple      ,					UpperCAmelCase__					:   int      ,					UpperCAmelCase__					:   Optional[int]      ,					UpperCAmelCase__					:   Optional[int]      ,					UpperCAmelCase__					:   Optional[int]      ,					UpperCAmelCase__					:   int							)	->     Union[str, Any]:
							if got_ver is None or want_ver is None:
														raise ValueError(
														    F'''Unable to compare versions for {requirement}: need={want_ver} found={got_ver}. This is unusual. Consider'''
														    F''' reinstalling {pkg}.'''							)
							if not ops[op](version.parse(UpperCAmelCase__							)      ,					version.parse(UpperCAmelCase__							)							):
														raise ImportError(
														    F'''{requirement} is required for a normal functioning of this module, but found {pkg}=={got_ver}.{hint}'''							)



def   __lowerCAmelCase    (				UpperCAmelCase__					:   str      ,					UpperCAmelCase__					:   Optional[str] = None							)	->     None:
							lowerCamelCase_					=       F'''\n{hint}''' if hint is not None else """"""

							# non-versioned check
							if re.match(R"""^[\w_\-\d]+$"""      ,					UpperCAmelCase__							):
														lowerCamelCase_	, lowerCamelCase_	, lowerCamelCase_					=       requirement, None, None
							else:
														lowerCamelCase_					=       re.findall(R"""^([^!=<>\s]+)([\s!=<>]{1,2}.+)"""      ,					UpperCAmelCase__							)
														if not match:
																					raise ValueError(
																					    """requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23, but"""
																					    F''' got {requirement}'''							)
														lowerCamelCase_	, lowerCamelCase_					=       match[0]
														lowerCamelCase_					=       want_full.split(""","""							)  # there could be multiple requirements
														lowerCamelCase_					=       {}
														for w in want_range:
																					lowerCamelCase_					=       re.findall(R"""^([\s!=<>]{1,2})(.+)"""      ,					UpperCAmelCase__							)
																					if not match:
																												raise ValueError(
																												    """requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23,"""
																												    F''' but got {requirement}'''							)
																					lowerCamelCase_	, lowerCamelCase_					=       match[0]
																					lowerCamelCase_					=       want_ver
																					if op not in ops:
																												raise ValueError(F'''{requirement}: need one of {list(ops.keys()							)}, but got {op}'''							)

    # special case
							if pkg == "python":
														lowerCamelCase_					=       """.""".join([str(UpperCAmelCase__							) for x in sys.version_info[:3]]							)
														for op, want_ver in wanted.items():
																					_compare_versions(UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__							)
														return

							# check if any version is installed
							try:
														lowerCamelCase_					=       importlib.metadata.version(UpperCAmelCase__							)
							except importlib.metadata.PackageNotFoundError:
														raise importlib.metadata.PackageNotFoundError(
														    F'''The \'{requirement}\' distribution was not found and is required by this application. {hint}'''							)

							# check that the right version is installed if version number or a range was provided
							if want_ver is not None:
														for op, want_ver in wanted.items():
																					_compare_versions(UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__							)



def   __lowerCAmelCase    (				UpperCAmelCase__					:   List[str]							)	->     Dict:
							lowerCamelCase_					=       """Try: pip install transformers -U or pip install -e '.[dev]' if you're working with git main"""
							return require_version(UpperCAmelCase__      ,					UpperCAmelCase__							)





