




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






				'''simple docstring'''
				__UpperCamelCase    						=		DanceDiffusionPipeline
				__UpperCamelCase    						=		UNCONDITIONAL_AUDIO_GENERATION_PARAMS
				__UpperCamelCase    						=		PipelineTesterMixin.required_optional_params - {
				    '''callback''',
				    '''latents''',
				    '''callback_steps''',
				    '''output_type''',
				    '''num_images_per_prompt''',
				}
				__UpperCamelCase    						=		UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
				__UpperCamelCase    						=		False
				__UpperCamelCase    						=		False





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							torch.manual_seed(0				)
							lowerCamelCase       =      UNetaDModel(
							    block_out_channels=(32, 32, 64)	,				extra_in_channels=16	,				sample_size=512	,				sample_rate=16_000	,				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""")	,				)
							lowerCamelCase       =      IPNDMScheduler()

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





				def      _lowerCAmelCase       ( self	,				_a	,				_a=0				):





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





				def      _lowerCAmelCase       ( self				):





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

							lowerCamelCase       =      self.get_dummy_inputs(_a				)
							lowerCamelCase       =      pipe(**_a				)
							lowerCamelCase       =      output.audios

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

							assert audio.shape == (1, 2, components["unet"].sample_size)
							lowerCamelCase       =      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      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							return super().test_save_load_local()





				@skip_mps
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3				)





				@skip_mps
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							return super().test_save_load_optional_components()





				@skip_mps
				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							return super().test_attention_slicing_forward_pass()





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							super().test_inference_batch_single_identical(expected_max_diff=3e-3				)






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






				'''simple docstring'''





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      torch_device

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

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

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

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





							"""simple docstring"""
							lowerCamelCase       =      torch_device

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

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

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

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









"""simple docstring"""




from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCAmelCase							:    List[str]  							=							{
    """configuration_autoformer""": [
        """AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """AutoformerConfig""",
    ],
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCAmelCase							:    Tuple  							=							[
					    """AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """AutoformerForPrediction""",
					    """AutoformerModel""",
					    """AutoformerPreTrainedModel""",
					]


if TYPE_CHECKING:
					from .configuration_autoformer import (
					    AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
					    AutoformerConfig,
					)

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_autoformer import (
										    AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
										    AutoformerForPrediction,
										    AutoformerModel,
										    AutoformerPreTrainedModel,
										)

else:
					import sys

					lowerCAmelCase							:    str  							=							_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



