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


				"""simple docstring"""



				__UpperCamelCase   :		Dict       =     DanceDiffusionPipeline
				__UpperCamelCase   :		str       =     UNCONDITIONAL_AUDIO_GENERATION_PARAMS
				__UpperCamelCase   :		Any       =     PipelineTesterMixin.required_optional_params - {
				    '''callback''',
				    '''latents''',
				    '''callback_steps''',
				    '''output_type''',
				    '''num_images_per_prompt''',
				}
				__UpperCamelCase   :		str       =     UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
				__UpperCamelCase   :		int       =     False
				__UpperCamelCase   :		str       =     False



				def 			__magic_name__					(self					)      ->	Any:



											"""simple docstring"""




											torch.manual_seed(0					)
											SCREAMING_SNAKE_CASE__		: Dict     =     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=SCREAMING_SNAKE_CASE__     , use_timestep_embedding=SCREAMING_SNAKE_CASE__     , time_embedding_type="""fourier"""     , mid_block_type="""UNetMidBlock1D"""     , down_block_types=("""DownBlock1DNoSkip""", """DownBlock1D""", """AttnDownBlock1D""")     , up_block_types=("""AttnUpBlock1D""", """UpBlock1D""", """UpBlock1DNoSkip""")     , )
											SCREAMING_SNAKE_CASE__		: Tuple     =     IPNDMScheduler()

											SCREAMING_SNAKE_CASE__		: str     =     {
											    """unet""": unet,
											    """scheduler""": scheduler,
											}
											return components



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__=0					)      ->	Tuple:



											"""simple docstring"""




											if str(SCREAMING_SNAKE_CASE__					).startswith("""mps"""					):
																		SCREAMING_SNAKE_CASE__		: List[str]     =     torch.manual_seed(SCREAMING_SNAKE_CASE__					)
											else:
																		SCREAMING_SNAKE_CASE__		: Dict     =     torch.Generator(device=SCREAMING_SNAKE_CASE__					).manual_seed(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: List[str]     =     {
											    """batch_size""": 1,
											    """generator""": generator,
											    """num_inference_steps""": 4,
											}
											return inputs



				def 			__magic_name__					(self					)      ->	Dict:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: str     =     """cpu"""  # ensure determinism for the device-dependent torch.Generator
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     self.get_dummy_components()
											SCREAMING_SNAKE_CASE__		: Dict     =     DanceDiffusionPipeline(**SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: int     =     pipe.to(SCREAMING_SNAKE_CASE__					)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)

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

											SCREAMING_SNAKE_CASE__		: Tuple     =     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 			__magic_name__					(self					)      ->	Dict:



											"""simple docstring"""




											return super().test_save_load_local()



				@skip_mps
				def 			__magic_name__					(self					)      ->	Dict:



											"""simple docstring"""




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



				@skip_mps
				def 			__magic_name__					(self					)      ->	Any:



											"""simple docstring"""




											return super().test_save_load_optional_components()



				@skip_mps
				def 			__magic_name__					(self					)      ->	Any:



											"""simple docstring"""




											return super().test_attention_slicing_forward_pass()



				def 			__magic_name__					(self					)      ->	str:



											"""simple docstring"""




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






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


				"""simple docstring"""



				def 			__magic_name__					(self					)      ->	str:



											"""simple docstring"""




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



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



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: List[Any]     =     torch_device

											SCREAMING_SNAKE_CASE__		: Tuple     =     DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k"""					)
											SCREAMING_SNAKE_CASE__		: Dict     =     pipe.to(SCREAMING_SNAKE_CASE__					)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: List[str]     =     torch.manual_seed(0					)
											SCREAMING_SNAKE_CASE__		: int     =     pipe(generator=SCREAMING_SNAKE_CASE__     , num_inference_steps=1_00     , audio_length_in_s=4.096					)
											SCREAMING_SNAKE_CASE__		: int     =     output.audios

											SCREAMING_SNAKE_CASE__		: List[str]     =     audio[0, -3:, -3:]

											assert audio.shape == (1, 2, pipe.unet.sample_size)
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     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 			__magic_name__					(self					)      ->	Optional[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Optional[int]     =     torch_device

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

											SCREAMING_SNAKE_CASE__		: Dict     =     torch.manual_seed(0					)
											SCREAMING_SNAKE_CASE__		: Any     =     pipe(generator=SCREAMING_SNAKE_CASE__     , num_inference_steps=1_00     , audio_length_in_s=4.096					)
											SCREAMING_SNAKE_CASE__		: int     =     output.audios

											SCREAMING_SNAKE_CASE__		: Tuple     =     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






"""simple docstring"""




import copy
from collections import OrderedDict
from typing import Dict, Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging
from ..auto import CONFIG_MAPPING


UpperCAmelCase__   :					Optional[Any]												=				logging.get_logger(__name__)

UpperCAmelCase__   :					str												=				{
    'facebook/detr-resnet-50': 'https://huggingface.co/facebook/detr-resnet-50/resolve/main/config.json',
    # See all DETR models at https://huggingface.co/models?filter=detr
}



class 			lowerCAmelCase_		(a__						):


				"""simple docstring"""



				__UpperCamelCase   :		List[Any]       =     '''detr'''
				__UpperCamelCase   :		List[Any]       =     ['''past_key_values''']
				__UpperCamelCase   :		List[Any]       =     {
				    '''hidden_size''': '''d_model''',
				    '''num_attention_heads''': '''encoder_attention_heads''',
				}



				def __init__(self     , SCREAMING_SNAKE_CASE__=True     , SCREAMING_SNAKE_CASE__=None     , SCREAMING_SNAKE_CASE__=3     , SCREAMING_SNAKE_CASE__=1_00     , SCREAMING_SNAKE_CASE__=6     , SCREAMING_SNAKE_CASE__=20_48     , SCREAMING_SNAKE_CASE__=8     , SCREAMING_SNAKE_CASE__=6     , SCREAMING_SNAKE_CASE__=20_48     , SCREAMING_SNAKE_CASE__=8     , SCREAMING_SNAKE_CASE__=0.0     , SCREAMING_SNAKE_CASE__=0.0     , SCREAMING_SNAKE_CASE__=True     , SCREAMING_SNAKE_CASE__="relu"     , SCREAMING_SNAKE_CASE__=2_56     , SCREAMING_SNAKE_CASE__=0.1     , SCREAMING_SNAKE_CASE__=0.0     , SCREAMING_SNAKE_CASE__=0.0     , SCREAMING_SNAKE_CASE__=0.02     , SCREAMING_SNAKE_CASE__=1.0     , SCREAMING_SNAKE_CASE__=False     , SCREAMING_SNAKE_CASE__="sine"     , SCREAMING_SNAKE_CASE__="resnet50"     , SCREAMING_SNAKE_CASE__=True     , SCREAMING_SNAKE_CASE__=False     , SCREAMING_SNAKE_CASE__=1     , SCREAMING_SNAKE_CASE__=5     , SCREAMING_SNAKE_CASE__=2     , SCREAMING_SNAKE_CASE__=1     , SCREAMING_SNAKE_CASE__=1     , SCREAMING_SNAKE_CASE__=5     , SCREAMING_SNAKE_CASE__=2     , SCREAMING_SNAKE_CASE__=0.1     , **SCREAMING_SNAKE_CASE__     , )      ->	Optional[Any]:



											"""simple docstring"""




											if backbone_config is not None and use_timm_backbone:
																		raise ValueError("""You can't specify both `backbone_config` and `use_timm_backbone`."""					)

											if not use_timm_backbone:
																		if backbone_config is None:
																									logger.info("""`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone."""					)
																									SCREAMING_SNAKE_CASE__		: Any     =     CONFIG_MAPPING["""resnet"""](out_features=["""stage4"""]					)
																		elif isinstance(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					):
																									SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     backbone_config.get("""model_type"""					)
																									SCREAMING_SNAKE_CASE__		: int     =     CONFIG_MAPPING[backbone_model_type]
																									SCREAMING_SNAKE_CASE__		: str     =     config_class.from_dict(SCREAMING_SNAKE_CASE__					)
																		# set timm attributes to None
																		SCREAMING_SNAKE_CASE__   ,   SCREAMING_SNAKE_CASE__   ,   SCREAMING_SNAKE_CASE__		: List[str]     =     None, None, None

											SCREAMING_SNAKE_CASE__		: Optional[int]     =     use_timm_backbone
											SCREAMING_SNAKE_CASE__		: Tuple     =     backbone_config
											SCREAMING_SNAKE_CASE__		: List[Any]     =     num_channels
											SCREAMING_SNAKE_CASE__		: Tuple     =     num_queries
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     d_model
											SCREAMING_SNAKE_CASE__		: str     =     encoder_ffn_dim
											SCREAMING_SNAKE_CASE__		: str     =     encoder_layers
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     encoder_attention_heads
											SCREAMING_SNAKE_CASE__		: Any     =     decoder_ffn_dim
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     decoder_layers
											SCREAMING_SNAKE_CASE__		: Dict     =     decoder_attention_heads
											SCREAMING_SNAKE_CASE__		: List[str]     =     dropout
											SCREAMING_SNAKE_CASE__		: Tuple     =     attention_dropout
											SCREAMING_SNAKE_CASE__		: Tuple     =     activation_dropout
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     activation_function
											SCREAMING_SNAKE_CASE__		: Any     =     init_std
											SCREAMING_SNAKE_CASE__		: Dict     =     init_xavier_std
											SCREAMING_SNAKE_CASE__		: Any     =     encoder_layerdrop
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     decoder_layerdrop
											SCREAMING_SNAKE_CASE__		: Dict     =     encoder_layers
											SCREAMING_SNAKE_CASE__		: List[str]     =     auxiliary_loss
											SCREAMING_SNAKE_CASE__		: List[Any]     =     position_embedding_type
											SCREAMING_SNAKE_CASE__		: List[str]     =     backbone
											SCREAMING_SNAKE_CASE__		: Dict     =     use_pretrained_backbone
											SCREAMING_SNAKE_CASE__		: Any     =     dilation
											# Hungarian matcher
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     class_cost
											SCREAMING_SNAKE_CASE__		: Tuple     =     bbox_cost
											SCREAMING_SNAKE_CASE__		: List[Any]     =     giou_cost
											# Loss coefficients
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     mask_loss_coefficient
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     dice_loss_coefficient
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     bbox_loss_coefficient
											SCREAMING_SNAKE_CASE__		: Any     =     giou_loss_coefficient
											SCREAMING_SNAKE_CASE__		: List[str]     =     eos_coefficient
											super().__init__(is_encoder_decoder=SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)



				@property
				def 			__magic_name__					(self					)      ->	int:



											"""simple docstring"""




											return self.encoder_attention_heads



				@property
				def 			__magic_name__					(self					)      ->	int:



											"""simple docstring"""




											return self.d_model



				@classmethod
				def 			__magic_name__					(cls     , SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)      ->	Tuple:



											"""simple docstring"""




											return cls(backbone_config=SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)



				def 			__magic_name__					(self					)      ->	Dict[str, any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Dict     =     copy.deepcopy(self.__dict__					)
											if output["backbone_config"] is not None:
																		SCREAMING_SNAKE_CASE__		: Any     =     self.backbone_config.to_dict()
											SCREAMING_SNAKE_CASE__		: Dict     =     self.__class__.model_type
											return output






class 			lowerCAmelCase_		(a__						):


				"""simple docstring"""



				__UpperCamelCase   :		Any       =     version.parse('''1.11'''						)



				@property
				def 			__magic_name__					(self					)      ->	Mapping[str, Mapping[int, str]]:



											"""simple docstring"""




											return OrderedDict(
											    [
											        ("""pixel_values""", {0: """batch""", 1: """num_channels""", 2: """height""", 3: """width"""}),
											        ("""pixel_mask""", {0: """batch"""}),
											    ]					)



				@property
				def 			__magic_name__					(self					)      ->	float:



											"""simple docstring"""




											return 1E-5



				@property
				def 			__magic_name__					(self					)      ->	int:



											"""simple docstring"""




											return 12




