





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 	UpperCamelCase__				(   lowerCAmelCase_							,   unittest.TestCase     ):



     '''simple docstring'''

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






     def        SCREAMING_SNAKE_CASE__						(							self		:					str )  ->	Union[str, Any]:




           '''simple docstring'''


           torch.manual_seed(0 )
           SCREAMING_SNAKE_CASE	=		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""")		,)
           SCREAMING_SNAKE_CASE	=		IPNDMScheduler()

           SCREAMING_SNAKE_CASE	=		{
               """unet""": unet,
               """scheduler""": scheduler,
           }
           return components






     def        SCREAMING_SNAKE_CASE__						(							self		:					Union[str, Any]		,lowerCamelCase__		:					str		,lowerCamelCase__		:					Union[str, Any]=0 )  ->	List[str]:




           '''simple docstring'''


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






     def        SCREAMING_SNAKE_CASE__						(							self		:					List[Any] )  ->	int:




           '''simple docstring'''


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

           SCREAMING_SNAKE_CASE	=		self.get_dummy_inputs(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		pipe(**lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		output.audios

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

           assert audio.shape == (1, 2, components["unet"].sample_size)
           SCREAMING_SNAKE_CASE	=		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        SCREAMING_SNAKE_CASE__						(							self		:					Optional[Any] )  ->	int:




           '''simple docstring'''


           return super().test_save_load_local()






     @skip_mps
     def        SCREAMING_SNAKE_CASE__						(							self		:					Any )  ->	Tuple:




           '''simple docstring'''


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






     @skip_mps
     def        SCREAMING_SNAKE_CASE__						(							self		:					Any )  ->	Dict:




           '''simple docstring'''


           return super().test_save_load_optional_components()






     @skip_mps
     def        SCREAMING_SNAKE_CASE__						(							self		:					str )  ->	List[str]:




           '''simple docstring'''


           return super().test_attention_slicing_forward_pass()






     def        SCREAMING_SNAKE_CASE__						(							self		:					Optional[Any] )  ->	Dict:




           '''simple docstring'''


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

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



     '''simple docstring'''

     def        SCREAMING_SNAKE_CASE__						(							self		:					Optional[int] )  ->	str:




           '''simple docstring'''


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






     def        SCREAMING_SNAKE_CASE__						(							self		:					int )  ->	List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		torch_device

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

           SCREAMING_SNAKE_CASE	=		torch.manual_seed(0 )
           SCREAMING_SNAKE_CASE	=		pipe(generator=lowerCamelCase__		,num_inference_steps=100		,audio_length_in_s=4.096 )
           SCREAMING_SNAKE_CASE	=		output.audios

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

           assert audio.shape == (1, 2, pipe.unet.sample_size)
           SCREAMING_SNAKE_CASE	=		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        SCREAMING_SNAKE_CASE__						(							self		:					Dict )  ->	List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		torch_device

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

           SCREAMING_SNAKE_CASE	=		torch.manual_seed(0 )
           SCREAMING_SNAKE_CASE	=		pipe(generator=lowerCamelCase__		,num_inference_steps=100		,audio_length_in_s=4.096 )
           SCREAMING_SNAKE_CASE	=		output.audios

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

           assert audio.shape == (1, 2, pipe.unet.sample_size)
           SCREAMING_SNAKE_CASE	=		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 inspect
import re
from hashlib import shaaaa
from typing import Dict, List

from .arrow import arrow
from .audiofolder import audiofolder
from .csv import csv
from .imagefolder import imagefolder
from .json import json
from .pandas import pandas
from .parquet import parquet
from .sql import sql  # noqa F401
from .text import text



def      __lowercase     (   _SCREAMING_SNAKE_CASE			)				->     str:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE	=		[]
      for line in lines:
            SCREAMING_SNAKE_CASE	=		re.sub(r"""#.*"""				,		""""""				,		_SCREAMING_SNAKE_CASE			)  # remove comments
            if line:
                  filtered_lines.append(_SCREAMING_SNAKE_CASE			)
      SCREAMING_SNAKE_CASE	=		"""\n""".join(_SCREAMING_SNAKE_CASE			)

      # Make a hash from all this code
      SCREAMING_SNAKE_CASE	=		full_str.encode("""utf-8"""			)
      return shaaaa(_SCREAMING_SNAKE_CASE			).hexdigest()


# get importable module names and hash for caching
SCREAMING_SNAKE_CASE_		     = {
    """csv""": (csv.__name__, _hash_python_lines(inspect.getsource(csv).splitlines())),
    """json""": (json.__name__, _hash_python_lines(inspect.getsource(json).splitlines())),
    """pandas""": (pandas.__name__, _hash_python_lines(inspect.getsource(pandas).splitlines())),
    """parquet""": (parquet.__name__, _hash_python_lines(inspect.getsource(parquet).splitlines())),
    """arrow""": (arrow.__name__, _hash_python_lines(inspect.getsource(arrow).splitlines())),
    """text""": (text.__name__, _hash_python_lines(inspect.getsource(text).splitlines())),
    """imagefolder""": (imagefolder.__name__, _hash_python_lines(inspect.getsource(imagefolder).splitlines())),
    """audiofolder""": (audiofolder.__name__, _hash_python_lines(inspect.getsource(audiofolder).splitlines())),
}

# Used to infer the module to use based on the data files extensions
SCREAMING_SNAKE_CASE_		     = {
    """.csv""": ("""csv""", {}),
    """.tsv""": ("""csv""", {"""sep""": """\t"""}),
    """.json""": ("""json""", {}),
    """.jsonl""": ("""json""", {}),
    """.parquet""": ("""parquet""", {}),
    """.arrow""": ("""arrow""", {}),
    """.txt""": ("""text""", {}),
}
_EXTENSION_TO_MODULE.update({ext: ("""imagefolder""", {}) for ext in imagefolder.ImageFolder.EXTENSIONS})
_EXTENSION_TO_MODULE.update({ext.upper(): ("""imagefolder""", {}) for ext in imagefolder.ImageFolder.EXTENSIONS})
_EXTENSION_TO_MODULE.update({ext: ("""audiofolder""", {}) for ext in audiofolder.AudioFolder.EXTENSIONS})
_EXTENSION_TO_MODULE.update({ext.upper(): ("""audiofolder""", {}) for ext in audiofolder.AudioFolder.EXTENSIONS})
SCREAMING_SNAKE_CASE_		     = {"""imagefolder""", """audiofolder"""}

# Used to filter data files based on extensions given a module name
SCREAMING_SNAKE_CASE_		     = {}
for _ext, (_module, _) in _EXTENSION_TO_MODULE.items():
      _MODULE_TO_EXTENSIONS.setdefault(_module, []).append(_ext)

_MODULE_TO_EXTENSIONS["imagefolder"].append(""".zip""")



_MODULE_TO_EXTENSIONS["audiofolder"].append(""".zip""")


