


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):


       """simple docstring"""




       A__											=			DanceDiffusionPipeline
       A__											=			UNCONDITIONAL_AUDIO_GENERATION_PARAMS
       A__											=			PipelineTesterMixin.required_optional_params - {
           "callback",
           "latents",
           "callback_steps",
           "output_type",
           "num_images_per_prompt",
       }
       A__											=			UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
       A__											=			False
       A__											=			False





       def 						lowerCAmelCase     (   self							:			Tuple       ):


             '''simple docstring'''



             torch.manual_seed(0       )
             lowerCamelCase__		:				Tuple    =  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=_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")					,      )
             lowerCamelCase__		:				Tuple    =  IPNDMScheduler()

             lowerCamelCase__		:				int    =  {
                 "unet": unet,
                 "scheduler": scheduler,
             }
             return components





       def 						lowerCAmelCase     (   self							:			str					,      __lowerCamelCase							:			Any					,      __lowerCamelCase							:			Optional[int]=0       ):


             '''simple docstring'''



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





       def 						lowerCAmelCase     (   self							:			str       ):


             '''simple docstring'''



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

             lowerCamelCase__		:				List[Any]    =  self.get_dummy_inputs(_UpperCamelCase       )
             lowerCamelCase__		:				int    =  pipe(**_UpperCamelCase       )
             lowerCamelCase__		:				Optional[Any]    =  output.audios

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

             assert audio.shape == (1, 2, components["unet"].sample_size)
             lowerCamelCase__		:				Optional[int]    =  np.array([-0.7_2_6_5, 1.0_0_0_0, -0.8_3_8_8, 0.1_1_7_5, 0.9_4_9_8, -1.0_0_0_0]       )
             assert np.abs(audio_slice.flatten() - expected_slice       ).max() < 1E-2





       @skip_mps
       def 						lowerCAmelCase     (   self							:			str       ):


             '''simple docstring'''



             return super().test_save_load_local()





       @skip_mps
       def 						lowerCAmelCase     (   self							:			str       ):


             '''simple docstring'''



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





       @skip_mps
       def 						lowerCAmelCase     (   self							:			Tuple       ):


             '''simple docstring'''



             return super().test_save_load_optional_components()





       @skip_mps
       def 						lowerCAmelCase     (   self							:			Dict       ):


             '''simple docstring'''



             return super().test_attention_slicing_forward_pass()





       def 						lowerCAmelCase     (   self							:			Optional[Any]       ):


             '''simple docstring'''



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







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


       """simple docstring"""




       def 						lowerCAmelCase     (   self							:			Optional[int]       ):


             '''simple docstring'''



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





       def 						lowerCAmelCase     (   self							:			List[Any]       ):


             '''simple docstring'''



             lowerCamelCase__		:				List[str]    =  torch_device

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

             lowerCamelCase__		:				Optional[int]    =  torch.manual_seed(0       )
             lowerCamelCase__		:				Union[str, Any]    =  pipe(generator=_UpperCamelCase					,      num_inference_steps=100					,      audio_length_in_s=4.0_9_6       )
             lowerCamelCase__		:				Tuple    =  output.audios

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

             assert audio.shape == (1, 2, pipe.unet.sample_size)
             lowerCamelCase__		:				List[Any]    =  np.array([-0.0_1_9_2, -0.0_2_3_1, -0.0_3_1_8, -0.0_0_5_9, 0.0_0_0_2, -0.0_0_2_0]       )

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





       def 						lowerCAmelCase     (   self							:			Optional[Any]       ):


             '''simple docstring'''



             lowerCamelCase__		:				List[Any]    =  torch_device

             lowerCamelCase__		:				Any    =  DanceDiffusionPipeline.from_pretrained("harmonai/maestro-150k"					,      torch_dtype=torch.floataa       )
             lowerCamelCase__		:				int    =  pipe.to(_UpperCamelCase       )
             pipe.set_progress_bar_config(disable=_UpperCamelCase       )

             lowerCamelCase__		:				List[Any]    =  torch.manual_seed(0       )
             lowerCamelCase__		:				Any    =  pipe(generator=_UpperCamelCase					,      num_inference_steps=100					,      audio_length_in_s=4.0_9_6       )
             lowerCamelCase__		:				str    =  output.audios

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

             assert audio.shape == (1, 2, pipe.unet.sample_size)
             lowerCamelCase__		:				Union[str, Any]    =  np.array([-0.0_3_6_7, -0.0_4_8_8, -0.0_7_7_1, -0.0_5_2_5, -0.0_4_4_4, -0.0_3_4_1]       )

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




import inspect
import unittest

from transformers import ConvNextVaConfig
from transformers.models.auto import get_values
from transformers.models.auto.modeling_auto import MODEL_FOR_BACKBONE_MAPPING_NAMES, MODEL_MAPPING_NAMES
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch

			from transformers import ConvNextVaBackbone, ConvNextVaForImageClassification, ConvNextVaModel
			from transformers.models.convnextva.modeling_convnextva import CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
			from PIL import Image

			from transformers import AutoImageProcessor





class 	_lowerCAmelCase					:




				def __init__(					self				,				_UpperCamelCase				,				_UpperCamelCase=13				,				_UpperCamelCase=32				,				_UpperCamelCase=3				,				_UpperCamelCase=4				,				_UpperCamelCase=[10, 20, 30, 40]				,				_UpperCamelCase=[2, 2, 3, 2]				,				_UpperCamelCase=True				,				_UpperCamelCase=True				,				_UpperCamelCase=37				,				_UpperCamelCase="gelu"				,				_UpperCamelCase=10				,				_UpperCamelCase=0.02				,				_UpperCamelCase=["stage2", "stage3", "stage4"]				,				_UpperCamelCase=[2, 3, 4]				,				_UpperCamelCase=None				,				)				->      Union[str, Any]:
							lowerCAmelCase_						=	parent
							lowerCAmelCase_						=	batch_size
							lowerCAmelCase_						=	image_size
							lowerCAmelCase_						=	num_channels
							lowerCAmelCase_						=	num_stages
							lowerCAmelCase_						=	hidden_sizes
							lowerCAmelCase_						=	depths
							lowerCAmelCase_						=	is_training
							lowerCAmelCase_						=	use_labels
							lowerCAmelCase_						=	intermediate_size
							lowerCAmelCase_						=	hidden_act
							lowerCAmelCase_						=	num_labels
							lowerCAmelCase_						=	initializer_range
							lowerCAmelCase_						=	out_features
							lowerCAmelCase_						=	out_indices
							lowerCAmelCase_						=	scope




				def      __a						(					self     )				->      Optional[int]:
							lowerCAmelCase_						=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

							lowerCAmelCase_						=	None
							if self.use_labels:
										lowerCAmelCase_						=	ids_tensor([self.batch_size]				,				self.num_labels     )

							lowerCAmelCase_						=	self.get_config()

							return config, pixel_values, labels




				def      __a						(					self     )				->      List[str]:
							return ConvNextVaConfig(
							    num_channels=self.num_channels				,				hidden_sizes=self.hidden_sizes				,				depths=self.depths				,				num_stages=self.num_stages				,				hidden_act=self.hidden_act				,				is_decoder=_UpperCamelCase				,				initializer_range=self.initializer_range				,				out_features=self.out_features				,				out_indices=self.out_indices				,				num_labels=self.num_labels				,				)




				def      __a						(					self				,				_UpperCamelCase				,				_UpperCamelCase				,				_UpperCamelCase     )				->      str:
							lowerCAmelCase_						=	ConvNextVaModel(config=_UpperCamelCase     )
							model.to(_UpperCamelCase     )
							model.eval()
							lowerCAmelCase_						=	model(_UpperCamelCase     )
							# expected last hidden states: B, C, H // 32, W // 32
							self.parent.assertEqual(
							    result.last_hidden_state.shape				,				(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)				,				)




				def      __a						(					self				,				_UpperCamelCase				,				_UpperCamelCase				,				_UpperCamelCase     )				->      Optional[int]:
							lowerCAmelCase_						=	ConvNextVaForImageClassification(_UpperCamelCase     )
							model.to(_UpperCamelCase     )
							model.eval()
							lowerCAmelCase_						=	model(_UpperCamelCase				,				labels=_UpperCamelCase     )
							self.parent.assertEqual(result.logits.shape				,				(self.batch_size, self.num_labels)     )




				def      __a						(					self				,				_UpperCamelCase				,				_UpperCamelCase				,				_UpperCamelCase     )				->      Tuple:
							lowerCAmelCase_						=	ConvNextVaBackbone(config=_UpperCamelCase     )
							model.to(_UpperCamelCase     )
							model.eval()
							lowerCAmelCase_						=	model(_UpperCamelCase     )

							# verify hidden states
							self.parent.assertEqual(len(result.feature_maps     )				,				len(config.out_features     )     )
							self.parent.assertListEqual(list(result.feature_maps[0].shape     )				,				[self.batch_size, self.hidden_sizes[1], 4, 4]     )

							# verify channels
							self.parent.assertEqual(len(model.channels     )				,				len(config.out_features     )     )
							self.parent.assertListEqual(model.channels				,				config.hidden_sizes[1:]     )

							# verify backbone works with out_features=None
							lowerCAmelCase_						=	None
							lowerCAmelCase_						=	ConvNextVaBackbone(config=_UpperCamelCase     )
							model.to(_UpperCamelCase     )
							model.eval()
							lowerCAmelCase_						=	model(_UpperCamelCase     )

							# verify feature maps
							self.parent.assertEqual(len(result.feature_maps     )				,				1     )
							self.parent.assertListEqual(list(result.feature_maps[0].shape     )				,				[self.batch_size, self.hidden_sizes[-1], 1, 1]     )

							# verify channels
							self.parent.assertEqual(len(model.channels     )				,				1     )
							self.parent.assertListEqual(model.channels				,				[config.hidden_sizes[-1]]     )




				def      __a						(					self     )				->      Any:
							lowerCAmelCase_						=	self.prepare_config_and_inputs()
							lowerCAmelCase_	,  lowerCAmelCase_	,  lowerCAmelCase_						=	config_and_inputs
							lowerCAmelCase_						=	{"pixel_values": pixel_values}
							return config, inputs_dict




				def      __a						(					self     )				->      str:
							lowerCAmelCase_						=	self.prepare_config_and_inputs()
							lowerCAmelCase_	,  lowerCAmelCase_	,  lowerCAmelCase_						=	config_and_inputs
							lowerCAmelCase_						=	{"pixel_values": pixel_values, "labels": labels}
							return config, inputs_dict





@require_torch
class 	_lowerCAmelCase					(	__a					,      __a					,      unittest.TestCase       ):
				_lowercase					=(
				    (
				        ConvNextVaModel,
				        ConvNextVaForImageClassification,
				        ConvNextVaBackbone,
				    )
				    if is_torch_available()
				    else ()
				)
				_lowercase					=(
				    {'''feature-extraction''': ConvNextVaModel, '''image-classification''': ConvNextVaForImageClassification}
				    if is_torch_available()
				    else {}
				)

				_lowercase					=False
				_lowercase					=False
				_lowercase					=False
				_lowercase					=False
				_lowercase					=False




				def      __a						(					self     )				->      Tuple:
							lowerCAmelCase_						=	ConvNextVaModelTester(self     )
							lowerCAmelCase_						=	ConfigTester(self				,				config_class=_UpperCamelCase				,				has_text_modality=_UpperCamelCase				,				hidden_size=37     )




				def      __a						(					self     )				->      List[Any]:
							self.create_and_test_config_common_properties()
							self.config_tester.create_and_test_config_to_json_string()
							self.config_tester.create_and_test_config_to_json_file()
							self.config_tester.create_and_test_config_from_and_save_pretrained()
							self.config_tester.create_and_test_config_with_num_labels()
							self.config_tester.check_config_can_be_init_without_params()
							self.config_tester.check_config_arguments_init()




				def      __a						(					self     )				->      Any:
							return




				@unittest.skip(reason="ConvNextV2 does not use inputs_embeds"     )
				def      __a						(					self     )				->      List[Any]:
							pass




				@unittest.skip(reason="ConvNextV2 does not support input and output embeddings"     )
				def      __a						(					self     )				->      int:
							pass




				@unittest.skip(reason="ConvNextV2 does not use feedforward chunking"     )
				def      __a						(					self     )				->      str:
							pass




				def      __a						(					self     )				->      Union[str, Any]:
							if not self.model_tester.is_training:
										return

							for model_class in self.all_model_classes:
										lowerCAmelCase_	,  lowerCAmelCase_						=	self.model_tester.prepare_config_and_inputs_with_labels()
										lowerCAmelCase_						=	True

										if model_class.__name__ in [
										    *get_values(_UpperCamelCase     ),
										    *get_values(_UpperCamelCase     ),
										]:
													continue

										lowerCAmelCase_						=	model_class(_UpperCamelCase     )
										model.to(_UpperCamelCase     )
										model.train()
										lowerCAmelCase_						=	self._prepare_for_class(_UpperCamelCase				,				_UpperCamelCase				,				return_labels=_UpperCamelCase     )
										lowerCAmelCase_						=	model(**_UpperCamelCase     ).loss
										loss.backward()




				def      __a						(					self     )				->      Optional[int]:
							if not self.model_tester.is_training:
										return

							for model_class in self.all_model_classes:
										lowerCAmelCase_	,  lowerCAmelCase_						=	self.model_tester.prepare_config_and_inputs_with_labels()
										lowerCAmelCase_						=	False
										lowerCAmelCase_						=	True

										if (
										    model_class.__name__
										    in [*get_values(_UpperCamelCase     ), *get_values(_UpperCamelCase     )]
										    or not model_class.supports_gradient_checkpointing
										):
													continue

										lowerCAmelCase_						=	model_class(_UpperCamelCase     )
										model.to(_UpperCamelCase     )
										model.gradient_checkpointing_enable()
										model.train()
										lowerCAmelCase_						=	self._prepare_for_class(_UpperCamelCase				,				_UpperCamelCase				,				return_labels=_UpperCamelCase     )
										lowerCAmelCase_						=	model(**_UpperCamelCase     ).loss
										loss.backward()




				def      __a						(					self     )				->      Any:
							lowerCAmelCase_	,  lowerCAmelCase_						=	self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										lowerCAmelCase_						=	model_class(_UpperCamelCase     )
										lowerCAmelCase_						=	inspect.signature(model.forward     )
										# signature.parameters is an OrderedDict => so arg_names order is deterministic
										lowerCAmelCase_						=	[*signature.parameters.keys()]

										lowerCAmelCase_						=	["pixel_values"]
										self.assertListEqual(arg_names[:1]				,				_UpperCamelCase     )




				def      __a						(					self     )				->      Tuple:
							lowerCAmelCase_						=	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_UpperCamelCase     )




				def      __a						(					self     )				->      List[Any]:
							def check_hidden_states_output(_UpperCamelCase				,				_UpperCamelCase				,				_UpperCamelCase     ):
										lowerCAmelCase_						=	model_class(_UpperCamelCase     )
										model.to(_UpperCamelCase     )
										model.eval()

										with torch.no_grad():
													lowerCAmelCase_						=	model(**self._prepare_for_class(_UpperCamelCase				,				_UpperCamelCase     )     )

										lowerCAmelCase_						=	outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

										lowerCAmelCase_						=	self.model_tester.num_stages
										self.assertEqual(len(_UpperCamelCase     )				,				expected_num_stages + 1     )

										# ConvNextV2's feature maps are of shape (batch_size, num_channels, height, width)
										self.assertListEqual(
										    list(hidden_states[0].shape[-2:]     )				,				[self.model_tester.image_size // 4, self.model_tester.image_size // 4]				,				)

							lowerCAmelCase_	,  lowerCAmelCase_						=	self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										lowerCAmelCase_						=	True
										check_hidden_states_output(_UpperCamelCase				,				_UpperCamelCase				,				_UpperCamelCase     )

										# check that output_hidden_states also work using config
										del inputs_dict["output_hidden_states"]
										lowerCAmelCase_						=	True

										check_hidden_states_output(_UpperCamelCase				,				_UpperCamelCase				,				_UpperCamelCase     )




				def      __a						(					self     )				->      Optional[int]:
							lowerCAmelCase_						=	self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*_UpperCamelCase     )




				@slow
				def      __a						(					self     )				->      Optional[Any]:
							for model_name in CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCAmelCase_						=	ConvNextVaModel.from_pretrained(_UpperCamelCase     )
										self.assertIsNotNone(_UpperCamelCase     )






def    lowerCamelCase__				(				):




			"""simple docstring"""


			lowerCAmelCase_						=	Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"				)
			return image





@require_torch
@require_vision
class 	_lowerCAmelCase					(	unittest.TestCase       ):




				@cached_property
				def      __a						(					self     )				->      int:
							return AutoImageProcessor.from_pretrained("facebook/convnextv2-tiny-1k-224"     ) if is_vision_available() else None




				@slow
				def      __a						(					self     )				->      List[str]:
							lowerCAmelCase_						=	ConvNextVaForImageClassification.from_pretrained("facebook/convnextv2-tiny-1k-224"     ).to(_UpperCamelCase     )

							lowerCAmelCase_						=	self.default_image_processor
							lowerCAmelCase_						=	prepare_img()
							lowerCAmelCase_						=	preprocessor(images=_UpperCamelCase				,				return_tensors="pt"     ).to(_UpperCamelCase     )

							# forward pass
							with torch.no_grad():
										lowerCAmelCase_						=	model(**_UpperCamelCase     )

							# verify the logits
							lowerCAmelCase_						=	torch.Size((1, 1_000)     )
							self.assertEqual(outputs.logits.shape				,				_UpperCamelCase     )

							lowerCAmelCase_						=	torch.tensor([0.9996, 0.1966, -0.4386]     ).to(_UpperCamelCase     )
							self.assertTrue(torch.allclose(outputs.logits[0, :3]				,				_UpperCamelCase				,				atol=1e-4     )     )



