




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



	"""simple docstring"""


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






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







						"""simple docstring"""







						torch.manual_seed(0	)
						__UpperCamelCase					:	int             =      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""")					,     )
						__UpperCamelCase					:	List[Any]             =      IPNDMScheduler()

						__UpperCamelCase					:	List[str]             =      {
						    """unet""": unet,
						    """scheduler""": scheduler,
						}
						return components






	def 					lowerCamelCase__     (					self		:						List[Any]					,     lowerCAmelCase		:						int					,     lowerCAmelCase		:						int=0	)	->       Optional[Any]:







						"""simple docstring"""







						if str(lowerCAmelCase	).startswith("""mps"""	):
											__UpperCamelCase					:	Optional[int]             =      torch.manual_seed(lowerCAmelCase	)
						else:
											__UpperCamelCase					:	Any             =      torch.Generator(device=lowerCAmelCase	).manual_seed(lowerCAmelCase	)
						__UpperCamelCase					:	Optional[Any]             =      {
						    """batch_size""": 1,
						    """generator""": generator,
						    """num_inference_steps""": 4,
						}
						return inputs






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







						"""simple docstring"""







						__UpperCamelCase					:	Tuple             =      """cpu"""  # ensure determinism for the device-dependent torch.Generator
						__UpperCamelCase					:	int             =      self.get_dummy_components()
						__UpperCamelCase					:	str             =      DanceDiffusionPipeline(**lowerCAmelCase	)
						__UpperCamelCase					:	Optional[Any]             =      pipe.to(lowerCAmelCase	)
						pipe.set_progress_bar_config(disable=lowerCAmelCase	)

						__UpperCamelCase					:	Dict             =      self.get_dummy_inputs(lowerCAmelCase	)
						__UpperCamelCase					:	Any             =      pipe(**lowerCAmelCase	)
						__UpperCamelCase					:	Tuple             =      output.audios

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

						assert audio.shape == (1, 2, components["unet"].sample_size)
						__UpperCamelCase					:	int             =      np.array([-0.72_65, 1.00_00, -0.83_88, 0.11_75, 0.94_98, -1.00_00]	)
						assert np.abs(audio_slice.flatten() - expected_slice	).max() < 1E-2






	@skip_mps
	def 					lowerCamelCase__     (					self		:						Optional[int]	)	->       Dict:







						"""simple docstring"""







						return super().test_save_load_local()






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







						"""simple docstring"""







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






	@skip_mps
	def 					lowerCamelCase__     (					self		:						List[str]	)	->       Union[str, Any]:







						"""simple docstring"""







						return super().test_save_load_optional_components()






	@skip_mps
	def 					lowerCamelCase__     (					self		:						Tuple	)	->       Tuple:







						"""simple docstring"""







						return super().test_attention_slicing_forward_pass()






	def 					lowerCamelCase__     (					self		:						Any	)	->       List[Any]:







						"""simple docstring"""







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






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



	"""simple docstring"""


	def 					lowerCamelCase__     (					self		:						Dict	)	->       str:







						"""simple docstring"""







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






	def 					lowerCamelCase__     (					self		:						Tuple	)	->       List[str]:







						"""simple docstring"""







						__UpperCamelCase					:	Optional[Any]             =      torch_device

						__UpperCamelCase					:	Union[str, Any]             =      DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k"""	)
						__UpperCamelCase					:	Any             =      pipe.to(lowerCAmelCase	)
						pipe.set_progress_bar_config(disable=lowerCAmelCase	)

						__UpperCamelCase					:	Optional[int]             =      torch.manual_seed(0	)
						__UpperCamelCase					:	str             =      pipe(generator=lowerCAmelCase					,     num_inference_steps=100					,     audio_length_in_s=4.0_96	)
						__UpperCamelCase					:	Optional[int]             =      output.audios

						__UpperCamelCase					:	Dict             =      audio[0, -3:, -3:]

						assert audio.shape == (1, 2, pipe.unet.sample_size)
						__UpperCamelCase					:	int             =      np.array([-0.01_92, -0.02_31, -0.03_18, -0.00_59, 0.00_02, -0.00_20]	)

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






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







						"""simple docstring"""







						__UpperCamelCase					:	Optional[Any]             =      torch_device

						__UpperCamelCase					:	Any             =      DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k"""					,     torch_dtype=torch.floataa	)
						__UpperCamelCase					:	List[str]             =      pipe.to(lowerCAmelCase	)
						pipe.set_progress_bar_config(disable=lowerCAmelCase	)

						__UpperCamelCase					:	Any             =      torch.manual_seed(0	)
						__UpperCamelCase					:	Any             =      pipe(generator=lowerCAmelCase					,     num_inference_steps=100					,     audio_length_in_s=4.0_96	)
						__UpperCamelCase					:	List[Any]             =      output.audios

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

						assert audio.shape == (1, 2, pipe.unet.sample_size)
						__UpperCamelCase					:	Dict             =      np.array([-0.03_67, -0.04_88, -0.07_71, -0.05_25, -0.04_44, -0.03_41]	)

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












import importlib
import json
import os
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
import transformers.models.auto
from transformers.models.auto.configuration_auto import CONFIG_MAPPING, AutoConfig
from transformers.models.bert.configuration_bert import BertConfig
from transformers.models.roberta.configuration_roberta import RobertaConfig
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / '''utils'''))

from test_module.custom_configuration import CustomConfig  # noqa E402


a__		     =     get_tests_dir('''fixtures/dummy-config.json''')


class    SCREAMING_SNAKE_CASE_				(       unittest.TestCase   ):



	"""simple docstring"""


	def 					lowerCamelCase__     (					self		:						str	)	->       Dict:







						"""simple docstring"""







						__UpperCamelCase					:	Optional[Any]             =      0






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







						"""simple docstring"""







						self.assertIsNotNone(transformers.models.auto.__spec__	)
						self.assertIsNotNone(importlib.util.find_spec("""transformers.models.auto"""	)	)






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







						"""simple docstring"""







						__UpperCamelCase					:	int             =      AutoConfig.from_pretrained("""bert-base-uncased"""	)
						self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)






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







						"""simple docstring"""







						__UpperCamelCase					:	Dict             =      AutoConfig.from_pretrained(lowerCAmelCase	)
						self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)






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







						"""simple docstring"""







						__UpperCamelCase					:	Optional[int]             =      AutoConfig.from_pretrained(lowerCAmelCase	)
						self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						Tuple	)	->       Any:







						"""simple docstring"""







						__UpperCamelCase					:	List[Any]             =      AutoConfig.for_model("""roberta"""	)
						self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)






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







						"""simple docstring"""







						with tempfile.TemporaryDirectory() as tmp_dir:
											# This model name contains bert and roberta, but roberta ends up being picked.
											__UpperCamelCase					:	Union[str, Any]             =      os.path.join(lowerCAmelCase					,     """fake-roberta"""	)
											os.makedirs(lowerCAmelCase					,     exist_ok=lowerCAmelCase	)
											with open(os.path.join(lowerCAmelCase					,     """config.json"""	)					,     """w"""	) as f:
																f.write(json.dumps({}	)	)
											__UpperCamelCase					:	Dict             =      AutoConfig.from_pretrained(lowerCAmelCase	)
											self.assertEqual(type(lowerCAmelCase	)					,     lowerCAmelCase	)






	def 					lowerCamelCase__     (					self		:						int	)	->       str:







						"""simple docstring"""







						try:
											AutoConfig.register("""custom"""					,     lowerCAmelCase	)
											# Wrong model type will raise an error
											with self.assertRaises(lowerCAmelCase	):
																AutoConfig.register("""model"""					,     lowerCAmelCase	)
											# Trying to register something existing in the Transformers library will raise an error
											with self.assertRaises(lowerCAmelCase	):
																AutoConfig.register("""bert"""					,     lowerCAmelCase	)

											# Now that the config is registered, it can be used as any other config with the auto-API
											__UpperCamelCase					:	Optional[Any]             =      CustomConfig()
											with tempfile.TemporaryDirectory() as tmp_dir:
																config.save_pretrained(lowerCAmelCase	)
																__UpperCamelCase					:	List[str]             =      AutoConfig.from_pretrained(lowerCAmelCase	)
																self.assertIsInstance(lowerCAmelCase					,     lowerCAmelCase	)

						finally:
											if "custom" in CONFIG_MAPPING._extra_content:
																del CONFIG_MAPPING._extra_content["custom"]






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







						"""simple docstring"""







						with self.assertRaisesRegex(
						    lowerCAmelCase					,     """bert-base is not a local folder and is not a valid model identifier"""	):
											__UpperCamelCase					:	Tuple             =      AutoConfig.from_pretrained("""bert-base"""	)






	def 					lowerCamelCase__     (					self		:						Dict	)	->       Any:







						"""simple docstring"""







						with self.assertRaisesRegex(
						    lowerCAmelCase					,     R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""	):
											__UpperCamelCase					:	Optional[Any]             =      AutoConfig.from_pretrained(lowerCAmelCase					,     revision="""aaaaaa"""	)






	def 					lowerCamelCase__     (					self		:						str	)	->       List[str]:







						"""simple docstring"""







						with self.assertRaisesRegex(
						    lowerCAmelCase					,     """hf-internal-testing/no-config-test-repo does not appear to have a file named config.json."""					,     ):
											__UpperCamelCase					:	Union[str, Any]             =      AutoConfig.from_pretrained("""hf-internal-testing/no-config-test-repo"""	)






	def 					lowerCamelCase__     (					self		:						Optional[int]	)	->       Optional[int]:







						"""simple docstring"""







						with self.assertRaises(lowerCAmelCase	):
											__UpperCamelCase					:	str             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""	)
						# If remote code is disabled, we can't load this config.
						with self.assertRaises(lowerCAmelCase	):
											__UpperCamelCase					:	Tuple             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""					,     trust_remote_code=lowerCAmelCase	)

						__UpperCamelCase					:	List[str]             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""					,     trust_remote_code=lowerCAmelCase	)
						self.assertEqual(config.__class__.__name__					,     """NewModelConfig"""	)

						# Test config can be reloaded.
						with tempfile.TemporaryDirectory() as tmp_dir:
											config.save_pretrained(lowerCAmelCase	)
											__UpperCamelCase					:	str             =      AutoConfig.from_pretrained(lowerCAmelCase					,     trust_remote_code=lowerCAmelCase	)
						self.assertEqual(reloaded_config.__class__.__name__					,     """NewModelConfig"""	)






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







						"""simple docstring"""







						class    SCREAMING_SNAKE_CASE_				(       _UpperCamelCase   ):



							"""simple docstring"""


							__magic_name__   :       int      	=		'new-model'

						try:
											AutoConfig.register("""new-model"""					,     lowerCAmelCase	)
											# If remote code is not set, the default is to use local
											__UpperCamelCase					:	Tuple             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""	)
											self.assertEqual(config.__class__.__name__					,     """NewModelConfigLocal"""	)

											# If remote code is disabled, we load the local one.
											__UpperCamelCase					:	Any             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""					,     trust_remote_code=lowerCAmelCase	)
											self.assertEqual(config.__class__.__name__					,     """NewModelConfigLocal"""	)

											# If remote is enabled, we load from the Hub
											__UpperCamelCase					:	List[Any]             =      AutoConfig.from_pretrained("""hf-internal-testing/test_dynamic_model"""					,     trust_remote_code=lowerCAmelCase	)
											self.assertEqual(config.__class__.__name__					,     """NewModelConfig"""	)

						finally:
											if "new-model" in CONFIG_MAPPING._extra_content:
																del CONFIG_MAPPING._extra_content["new-model"]






