


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    (__SCREAMING_SNAKE_CASE	,       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 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					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=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_		:     str					=					IPNDMScheduler()

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					{
										    'unet': unet,
										    'scheduler': scheduler,
										}
										return components






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__=0       ):





										"""simple docstring"""




										if str(lowerCAmelCase__       ).startswith('mps'       ):
													SCREAMING_SNAKE_CASE_		:     Dict					=					torch.manual_seed(lowerCAmelCase__       )
										else:
													SCREAMING_SNAKE_CASE_		:     Any					=					torch.Generator(device=lowerCAmelCase__       ).manual_seed(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Dict					=					{
										    'batch_size': 1,
										    'generator': generator,
										    'num_inference_steps': 4,
										}
										return inputs






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     str					=					'cpu'  # ensure determinism for the device-dependent torch.Generator
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					self.get_dummy_components()
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					DanceDiffusionPipeline(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Dict					=					pipe.to(lowerCAmelCase__       )
										pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     int					=					self.get_dummy_inputs(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					pipe(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     str					=					output.audios

										SCREAMING_SNAKE_CASE_		:     int					=					audio[0, -3:, -3:]

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





										"""simple docstring"""




										return super().test_save_load_local()






							@skip_mps
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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






							@skip_mps
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										return super().test_save_load_optional_components()






							@skip_mps
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										return super().test_attention_slicing_forward_pass()






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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







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





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					torch_device

										SCREAMING_SNAKE_CASE_		:     Any					=					DanceDiffusionPipeline.from_pretrained('harmonai/maestro-150k'       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					pipe.to(lowerCAmelCase__       )
										pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					pipe(generator=lowerCAmelCase__			,						num_inference_steps=1_0_0			,						audio_length_in_s=4.096       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					output.audios

										SCREAMING_SNAKE_CASE_		:     int					=					audio[0, -3:, -3:]

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





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					torch_device

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					DanceDiffusionPipeline.from_pretrained('harmonai/maestro-150k'			,						torch_dtype=torch.floataa       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					pipe.to(lowerCAmelCase__       )
										pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     int					=					pipe(generator=lowerCAmelCase__			,						num_inference_steps=1_0_0			,						audio_length_in_s=4.096       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					output.audios

										SCREAMING_SNAKE_CASE_		:     Tuple					=					audio[0, -3:, -3:]

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

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_lowerCAmelCase       =				{
    "configuration_timesformer": ["TIMESFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP", "TimesformerConfig"],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							_lowerCAmelCase       =				[
							    "TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
							    "TimesformerModel",
							    "TimesformerForVideoClassification",
							    "TimesformerPreTrainedModel",
							]

if TYPE_CHECKING:
							from .configuration_timesformer import TIMESFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, TimesformerConfig

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_timesformer import (
														    TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
														    TimesformerForVideoClassification,
														    TimesformerModel,
														    TimesformerPreTrainedModel,
														)

else:
							import sys

							_lowerCAmelCase       =				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

