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__						(   UpperCamelCase_	,							unittest.TestCase):
						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        __A       (       self	:				Optional[int]						):





												'''simple docstring'''





												torch.manual_seed(0						)
												SCREAMING_SNAKE_CASE							:    Any							=							UNetaDModel(
												    block_out_channels=(32, 32, 64)    ,		extra_in_channels=16    ,		sample_size=512    ,		sample_rate=1_6000    ,		in_channels=2    ,		out_channels=2    ,		flip_sin_to_cos=UpperCamelCase__    ,		use_timestep_embedding=UpperCamelCase__    ,		time_embedding_type='''fourier'''    ,		mid_block_type='''UNetMidBlock1D'''    ,		down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''')    ,		up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''')    ,		)
												SCREAMING_SNAKE_CASE							:    Optional[int]							=							IPNDMScheduler()

												SCREAMING_SNAKE_CASE							:    Any							=							{
												    '''unet''': unet,
												    '''scheduler''': scheduler,
												}
												return components


						def        __A       (       self	:				List[Any]    ,		UpperCamelCase__	:				Optional[int]    ,		UpperCamelCase__	:				List[Any]=0						):





												'''simple docstring'''





												if str(UpperCamelCase__						).startswith('''mps'''						):
																		SCREAMING_SNAKE_CASE							:    Optional[int]							=							torch.manual_seed(UpperCamelCase__						)
												else:
																		SCREAMING_SNAKE_CASE							:    List[str]							=							torch.Generator(device=UpperCamelCase__						).manual_seed(UpperCamelCase__						)
												SCREAMING_SNAKE_CASE							:    List[str]							=							{
												    '''batch_size''': 1,
												    '''generator''': generator,
												    '''num_inference_steps''': 4,
												}
												return inputs


						def        __A       (       self	:				str						):





												'''simple docstring'''





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

												SCREAMING_SNAKE_CASE							:    List[str]							=							self.get_dummy_inputs(UpperCamelCase__						)
												SCREAMING_SNAKE_CASE							:    int							=							pipe(**UpperCamelCase__						)
												SCREAMING_SNAKE_CASE							:    int							=							output.audios

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

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





												'''simple docstring'''





												return super().test_save_load_local()


						@skip_mps
						def        __A       (       self	:				Tuple						):





												'''simple docstring'''





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


						@skip_mps
						def        __A       (       self	:				Optional[int]						):





												'''simple docstring'''





												return super().test_save_load_optional_components()


						@skip_mps
						def        __A       (       self	:				Optional[int]						):





												'''simple docstring'''





												return super().test_attention_slicing_forward_pass()



						def        __A       (       self	:				Union[str, 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	:				Tuple						):





												'''simple docstring'''





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


						def        __A       (       self	:				str						):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:    Optional[Any]							=							torch_device

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

												SCREAMING_SNAKE_CASE							:    str							=							torch.manual_seed(0						)
												SCREAMING_SNAKE_CASE							:    List[str]							=							pipe(generator=UpperCamelCase__    ,		num_inference_steps=100    ,		audio_length_in_s=4.096						)
												SCREAMING_SNAKE_CASE							:    Tuple							=							output.audios

												SCREAMING_SNAKE_CASE							:    Union[str, Any]							=							audio[0, -3:, -3:]

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





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:    Union[str, Any]							=							torch_device

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

												SCREAMING_SNAKE_CASE							:    Optional[Any]							=							torch.manual_seed(0						)
												SCREAMING_SNAKE_CASE							:    str							=							pipe(generator=UpperCamelCase__    ,		num_inference_steps=100    ,		audio_length_in_s=4.096						)
												SCREAMING_SNAKE_CASE							:    List[Any]							=							output.audios

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

												assert audio.shape == (1, 2, pipe.unet.sample_size)
												SCREAMING_SNAKE_CASE							:    Union[str, Any]							=							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 json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    CLIPConfig,
    CLIPImageProcessor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER


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

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_image_processing import CustomImageProcessor  # noqa E402





class 					lowercase__						(   unittest.TestCase):


						def        __A       (       self	:				List[Any]						):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:    Dict							=							0


						def        __A       (       self	:				Dict						):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:    str							=							AutoImageProcessor.from_pretrained('''openai/clip-vit-base-patch32'''						)
												self.assertIsInstance(UpperCamelCase__    ,		UpperCamelCase__						)


						def        __A       (       self	:				Dict						):





												'''simple docstring'''





												with tempfile.TemporaryDirectory() as tmpdirname:
																		SCREAMING_SNAKE_CASE							:    str							=							Path(UpperCamelCase__						) / '''preprocessor_config.json'''
																		SCREAMING_SNAKE_CASE							:    Dict							=							Path(UpperCamelCase__						) / '''config.json'''
																		json.dump(
																		    {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}    ,		open(UpperCamelCase__    ,		'''w'''						)    ,		)
																		json.dump({'''model_type''': '''clip'''}    ,		open(UpperCamelCase__    ,		'''w'''						)						)

																		SCREAMING_SNAKE_CASE							:    Any							=							AutoImageProcessor.from_pretrained(UpperCamelCase__						)
																		self.assertIsInstance(UpperCamelCase__    ,		UpperCamelCase__						)


						def        __A       (       self	:				int						):





												'''simple docstring'''





												with tempfile.TemporaryDirectory() as tmpdirname:
																		SCREAMING_SNAKE_CASE							:    Tuple							=							Path(UpperCamelCase__						) / '''preprocessor_config.json'''
																		SCREAMING_SNAKE_CASE							:    Dict							=							Path(UpperCamelCase__						) / '''config.json'''
																		json.dump(
																		    {'''feature_extractor_type''': '''CLIPFeatureExtractor''', '''processor_class''': '''CLIPProcessor'''}    ,		open(UpperCamelCase__    ,		'''w'''						)    ,		)
																		json.dump({'''model_type''': '''clip'''}    ,		open(UpperCamelCase__    ,		'''w'''						)						)

																		SCREAMING_SNAKE_CASE							:    Optional[Any]							=							AutoImageProcessor.from_pretrained(UpperCamelCase__						)
																		self.assertIsInstance(UpperCamelCase__    ,		UpperCamelCase__						)


						def        __A       (       self	:				List[str]						):





												'''simple docstring'''





												with tempfile.TemporaryDirectory() as tmpdirname:
																		SCREAMING_SNAKE_CASE							:    Dict							=							CLIPConfig()

																		# Create a dummy config file with image_proceesor_type
																		SCREAMING_SNAKE_CASE							:    Dict							=							Path(UpperCamelCase__						) / '''preprocessor_config.json'''
																		SCREAMING_SNAKE_CASE							:    List[str]							=							Path(UpperCamelCase__						) / '''config.json'''
																		json.dump(
																		    {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}    ,		open(UpperCamelCase__    ,		'''w'''						)    ,		)
																		json.dump({'''model_type''': '''clip'''}    ,		open(UpperCamelCase__    ,		'''w'''						)						)

																		# remove image_processor_type to make sure config.json alone is enough to load image processor locally
																		SCREAMING_SNAKE_CASE							:    List[Any]							=							AutoImageProcessor.from_pretrained(UpperCamelCase__						).to_dict()

																		config_dict.pop('''image_processor_type'''						)
																		SCREAMING_SNAKE_CASE							:    str							=							CLIPImageProcessor(**UpperCamelCase__						)

																		# save in new folder
																		model_config.save_pretrained(UpperCamelCase__						)
																		config.save_pretrained(UpperCamelCase__						)

																		SCREAMING_SNAKE_CASE							:    Union[str, Any]							=							AutoImageProcessor.from_pretrained(UpperCamelCase__						)

																		# make sure private variable is not incorrectly saved
																		SCREAMING_SNAKE_CASE							:    List[Any]							=							json.loads(config.to_json_string()						)
																		self.assertTrue('''_processor_class''' not in dict_as_saved						)

												self.assertIsInstance(UpperCamelCase__    ,		UpperCamelCase__						)


						def        __A       (       self	:				Tuple						):





												'''simple docstring'''





												with tempfile.TemporaryDirectory() as tmpdirname:
																		SCREAMING_SNAKE_CASE							:    Optional[int]							=							Path(UpperCamelCase__						) / '''preprocessor_config.json'''
																		json.dump(
																		    {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}    ,		open(UpperCamelCase__    ,		'''w'''						)    ,		)

																		SCREAMING_SNAKE_CASE							:    Dict							=							AutoImageProcessor.from_pretrained(UpperCamelCase__						)
																		self.assertIsInstance(UpperCamelCase__    ,		UpperCamelCase__						)


						def        __A       (       self	:				int						):





												'''simple docstring'''





												with self.assertRaisesRegex(
												    UpperCamelCase__    ,		'''clip-base is not a local folder and is not a valid model identifier'''						):
																		SCREAMING_SNAKE_CASE							:    List[str]							=							AutoImageProcessor.from_pretrained('''clip-base'''						)


						def        __A       (       self	:				List[str]						):





												'''simple docstring'''





												with self.assertRaisesRegex(
												    UpperCamelCase__    ,		r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''						):
																		SCREAMING_SNAKE_CASE							:    Tuple							=							AutoImageProcessor.from_pretrained(UpperCamelCase__    ,		revision='''aaaaaa'''						)


						def        __A       (       self	:				Dict						):





												'''simple docstring'''





												with self.assertRaisesRegex(
												    UpperCamelCase__    ,		'''hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json.'''    ,		):
																		SCREAMING_SNAKE_CASE							:    Dict							=							AutoImageProcessor.from_pretrained('''hf-internal-testing/config-no-model'''						)


						def        __A       (       self	:				List[Any]						):





												'''simple docstring'''





												with self.assertRaises(UpperCamelCase__						):
																		SCREAMING_SNAKE_CASE							:    str							=							AutoImageProcessor.from_pretrained('''hf-internal-testing/test_dynamic_image_processor'''						)
												# If remote code is disabled, we can't load this config.
												with self.assertRaises(UpperCamelCase__						):
																		SCREAMING_SNAKE_CASE							:    List[Any]							=							AutoImageProcessor.from_pretrained(
																		    '''hf-internal-testing/test_dynamic_image_processor'''    ,		trust_remote_code=UpperCamelCase__						)

												SCREAMING_SNAKE_CASE							:    List[Any]							=							AutoImageProcessor.from_pretrained(
												    '''hf-internal-testing/test_dynamic_image_processor'''    ,		trust_remote_code=UpperCamelCase__						)
												self.assertEqual(image_processor.__class__.__name__    ,		'''NewImageProcessor'''						)

												# Test image processor can be reloaded.
												with tempfile.TemporaryDirectory() as tmp_dir:
																		image_processor.save_pretrained(UpperCamelCase__						)
																		SCREAMING_SNAKE_CASE							:    Optional[int]							=							AutoImageProcessor.from_pretrained(UpperCamelCase__    ,		trust_remote_code=UpperCamelCase__						)
												self.assertEqual(reloaded_image_processor.__class__.__name__    ,		'''NewImageProcessor'''						)


						def        __A       (       self	:				Optional[Any]						):





												'''simple docstring'''





												try:
																		AutoConfig.register('''custom'''    ,		UpperCamelCase__						)
																		AutoImageProcessor.register(UpperCamelCase__    ,		UpperCamelCase__						)
																		# Trying to register something existing in the Transformers library will raise an error
																		with self.assertRaises(UpperCamelCase__						):
																								AutoImageProcessor.register(UpperCamelCase__    ,		UpperCamelCase__						)

																		with tempfile.TemporaryDirectory() as tmpdirname:
																								SCREAMING_SNAKE_CASE							:    Union[str, Any]							=							Path(UpperCamelCase__						) / '''preprocessor_config.json'''
																								SCREAMING_SNAKE_CASE							:    Any							=							Path(UpperCamelCase__						) / '''config.json'''
																								json.dump(
																								    {'''feature_extractor_type''': '''CLIPFeatureExtractor''', '''processor_class''': '''CLIPProcessor'''}    ,		open(UpperCamelCase__    ,		'''w'''						)    ,		)
																								json.dump({'''model_type''': '''clip'''}    ,		open(UpperCamelCase__    ,		'''w'''						)						)

																								SCREAMING_SNAKE_CASE							:    Tuple							=							CustomImageProcessor.from_pretrained(UpperCamelCase__						)

																		# Now that the config is registered, it can be used as any other config with the auto-API
																		with tempfile.TemporaryDirectory() as tmp_dir:
																								image_processor.save_pretrained(UpperCamelCase__						)
																								SCREAMING_SNAKE_CASE							:    Any							=							AutoImageProcessor.from_pretrained(UpperCamelCase__						)
																								self.assertIsInstance(UpperCamelCase__    ,		UpperCamelCase__						)

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



						def        __A       (       self	:				Any						):





												'''simple docstring'''





												class 					lowercase__						(   UpperCamelCase_):
																		UpperCamelCase_     =		True

												try:
																		AutoConfig.register('''custom'''    ,		UpperCamelCase__						)
																		AutoImageProcessor.register(UpperCamelCase__    ,		UpperCamelCase__						)
																		# If remote code is not set, the default is to use local
																		SCREAMING_SNAKE_CASE							:    str							=							AutoImageProcessor.from_pretrained('''hf-internal-testing/test_dynamic_image_processor'''						)
																		self.assertEqual(image_processor.__class__.__name__    ,		'''NewImageProcessor'''						)
																		self.assertTrue(image_processor.is_local						)

																		# If remote code is disabled, we load the local one.
																		SCREAMING_SNAKE_CASE							:    Union[str, Any]							=							AutoImageProcessor.from_pretrained(
																		    '''hf-internal-testing/test_dynamic_image_processor'''    ,		trust_remote_code=UpperCamelCase__						)
																		self.assertEqual(image_processor.__class__.__name__    ,		'''NewImageProcessor'''						)
																		self.assertTrue(image_processor.is_local						)

																		# If remote is enabled, we load from the Hub
																		SCREAMING_SNAKE_CASE							:    Optional[Any]							=							AutoImageProcessor.from_pretrained(
																		    '''hf-internal-testing/test_dynamic_image_processor'''    ,		trust_remote_code=UpperCamelCase__						)
																		self.assertEqual(image_processor.__class__.__name__    ,		'''NewImageProcessor'''						)
																		self.assertTrue(not hasattr(UpperCamelCase__    ,		'''is_local'''						)						)

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





