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					(    _A       ,     unittest.TestCase    ):
						lowercase					 =     DanceDiffusionPipeline
						lowercase					 =     UNCONDITIONAL_AUDIO_GENERATION_PARAMS
						lowercase					 =     PipelineTesterMixin.required_optional_params - {
						    'callback',
						    'latents',
						    'callback_steps',
						    'output_type',
						    'num_images_per_prompt',
						}
						lowercase					 =     UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
						lowercase					 =     False
						lowercase					 =     False

						def        __a		(	self				:      str   )							->		Dict:


									'''simple docstring'''

									torch.manual_seed(0   )
									lowercase							=		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=__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''')							,							)
									lowercase							=		IPNDMScheduler()

									lowercase							=		{
									    '''unet''': unet,
									    '''scheduler''': scheduler,
									}
									return components

						def        __a		(	self				:      Optional[int]							,							__lowerCamelCase				:      Any							,							__lowerCamelCase				:      Dict=0   )							->		int:


									'''simple docstring'''

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

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


									'''simple docstring'''

									lowercase							=		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									lowercase							=		self.get_dummy_components()
									lowercase							=		DanceDiffusionPipeline(**__lowerCamelCase   )
									lowercase							=		pipe.to(__lowerCamelCase   )
									pipe.set_progress_bar_config(disable=__lowerCamelCase   )

									lowercase							=		self.get_dummy_inputs(__lowerCamelCase   )
									lowercase							=		pipe(**__lowerCamelCase   )
									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        __a		(	self				:      str   )							->		List[str]:


									'''simple docstring'''

									return super().test_save_load_local()

						@skip_mps
						def        __a		(	self				:      int   )							->		int:


									'''simple docstring'''

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

						@skip_mps
						def        __a		(	self				:      str   )							->		Optional[int]:


									'''simple docstring'''

									return super().test_save_load_optional_components()

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


									'''simple docstring'''

									return super().test_attention_slicing_forward_pass()



						def        __a		(	self				:      Optional[Any]   )							->		List[Any]:


									'''simple docstring'''

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






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

						def        __a		(	self				:      Optional[Any]   )							->		Optional[int]:


									'''simple docstring'''

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

						def        __a		(	self				:      Tuple   )							->		Dict:


									'''simple docstring'''

									lowercase							=		torch_device

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

									lowercase							=		torch.manual_seed(0   )
									lowercase							=		pipe(generator=__lowerCamelCase							,							num_inference_steps=1_00							,							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        __a		(	self				:      List[Any]   )							->		Union[str, Any]:


									'''simple docstring'''

									lowercase							=		torch_device

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

									lowercase							=		torch.manual_seed(0   )
									lowercase							=		pipe(generator=__lowerCamelCase							,							num_inference_steps=1_00							,							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






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


A_      =      {
    "configuration_gpt_bigcode": ["GPT_BIGCODE_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTBigCodeConfig"],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						A_      =      [
						    "GPT_BIGCODE_PRETRAINED_MODEL_ARCHIVE_LIST",
						    "GPTBigCodeForSequenceClassification",
						    "GPTBigCodeForTokenClassification",
						    "GPTBigCodeForCausalLM",
						    "GPTBigCodeModel",
						    "GPTBigCodePreTrainedModel",
						]

if TYPE_CHECKING:
						from .configuration_gpt_bigcode import GPT_BIGCODE_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTBigCodeConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_gpt_bigcode import (
												    GPT_BIGCODE_PRETRAINED_MODEL_ARCHIVE_LIST,
												    GPTBigCodeForCausalLM,
												    GPTBigCodeForSequenceClassification,
												    GPTBigCodeForTokenClassification,
												    GPTBigCodeModel,
												    GPTBigCodePreTrainedModel,
												)


else:
						import sys

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





