





import inspect
import unittest

from transformers import RegNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
     import jax
     import jax.numpy as jnp

     from transformers.models.regnet.modeling_flax_regnet import FlaxRegNetForImageClassification, FlaxRegNetModel

if is_vision_available():
     from PIL import Image

     from transformers import AutoImageProcessor




class 					_lowerCamelCase			(     unittest.TestCase     ):



      """simple docstring"""






      def __init__(			self					,    UpperCAmelCase					,    UpperCAmelCase=3					,    UpperCAmelCase=32					,    UpperCAmelCase=3					,    UpperCAmelCase=10					,    UpperCAmelCase=[10, 20, 30, 40]					,    UpperCAmelCase=[1, 1, 2, 1]					,    UpperCAmelCase=True					,    UpperCAmelCase=True					,    UpperCAmelCase="relu"					,    UpperCAmelCase=3					,    UpperCAmelCase=None					,    )							-> Tuple:






        '''simple docstring'''






        __snake_case :					List[str]          =							parent
        __snake_case :					List[Any]          =							batch_size
        __snake_case :					List[str]          =							image_size
        __snake_case :					Any          =							num_channels
        __snake_case :					Optional[int]          =							embeddings_size
        __snake_case :					Any          =							hidden_sizes
        __snake_case :					Optional[int]          =							depths
        __snake_case :					int          =							is_training
        __snake_case :					List[str]          =							use_labels
        __snake_case :					str          =							hidden_act
        __snake_case :					Dict          =							num_labels
        __snake_case :					str          =							scope
        __snake_case :					List[str]          =							len(__SCREAMING_SNAKE_CASE )







      def     UpperCAmelCase		(			self )							-> Tuple:






        '''simple docstring'''






        __snake_case :					Optional[int]          =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

        __snake_case :					int          =							self.get_config()

        return config, pixel_values







      def     UpperCAmelCase		(			self )							-> Optional[Any]:






        '''simple docstring'''






        return RegNetConfig(
            num_channels=self.num_channels					,    embeddings_size=self.embeddings_size					,    hidden_sizes=self.hidden_sizes					,    depths=self.depths					,    hidden_act=self.hidden_act					,    num_labels=self.num_labels					,    image_size=self.image_size					,    )







      def     UpperCAmelCase		(			self					,    UpperCAmelCase					,    UpperCAmelCase )							-> Optional[Any]:






        '''simple docstring'''






        __snake_case :					List[Any]          =							FlaxRegNetModel(config=__SCREAMING_SNAKE_CASE )
        __snake_case :					Tuple          =							model(__SCREAMING_SNAKE_CASE )

        # Output shape (b, c, h, w)
        self.parent.assertEqual(
            result.last_hidden_state.shape					,    (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)					,    )







      def     UpperCAmelCase		(			self					,    UpperCAmelCase					,    UpperCAmelCase )							-> int:






        '''simple docstring'''






        __snake_case :					Dict          =							self.num_labels
        __snake_case :					Dict          =							FlaxRegNetForImageClassification(config=__SCREAMING_SNAKE_CASE )
        __snake_case :					Dict          =							model(__SCREAMING_SNAKE_CASE )
        self.parent.assertEqual(result.logits.shape					,    (self.batch_size, self.num_labels) )







      def     UpperCAmelCase		(			self )							-> str:






        '''simple docstring'''






        __snake_case :					List[str]          =							self.prepare_config_and_inputs()
        __snake_case :					Dict          =							config_and_inputs
        __snake_case :					List[Any]          =							{'''pixel_values''': pixel_values}
        return config, inputs_dict




@require_flax
class 					_lowerCamelCase			(     UpperCamelCase_			, unittest.TestCase     ):



      """simple docstring"""






      UpperCAmelCase_    :   int      =(FlaxRegNetModel, FlaxRegNetForImageClassification) if is_flax_available() else ()

      UpperCAmelCase_    :   Union[str, Any]      =False
      UpperCAmelCase_    :   Union[str, Any]      =False
      UpperCAmelCase_    :   Optional[Any]      =False







      def     UpperCAmelCase		(			self )							-> None:






        '''simple docstring'''






        __snake_case :					List[Any]          =							FlaxRegNetModelTester(self )
        __snake_case :					int          =							ConfigTester(self					,    config_class=__SCREAMING_SNAKE_CASE					,    has_text_modality=__SCREAMING_SNAKE_CASE )







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






        '''simple docstring'''






        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     UpperCAmelCase		(			self )							-> int:






        '''simple docstring'''






        return







      def     UpperCAmelCase		(			self )							-> Any:






        '''simple docstring'''






        __snake_case :					Optional[Any]          =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE )







      def     UpperCAmelCase		(			self )							-> Tuple:






        '''simple docstring'''






        __snake_case :					List[str]          =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*__SCREAMING_SNAKE_CASE )







      @unittest.skip(reason="RegNet does not use inputs_embeds" )
      def     UpperCAmelCase		(			self )							-> Dict:






        '''simple docstring'''






        pass







      @unittest.skip(reason="RegNet does not support input and output embeddings" )
      def     UpperCAmelCase		(			self )							-> Any:






        '''simple docstring'''






        pass







      def     UpperCAmelCase		(			self )							-> Dict:






        '''simple docstring'''






        __snake_case :					List[Any]          =							self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          __snake_case :					List[Any]          =							model_class(__SCREAMING_SNAKE_CASE )
          __snake_case :					Optional[Any]          =							inspect.signature(model.__call__ )
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          __snake_case :					int          =							[*signature.parameters.keys()]

          __snake_case :					Optional[Any]          =							['''pixel_values''']
          self.assertListEqual(arg_names[:1]					,    __SCREAMING_SNAKE_CASE )







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






        '''simple docstring'''






        def check_hidden_states_output(UpperCAmelCase					,    UpperCAmelCase					,    UpperCAmelCase ):
          __snake_case :					List[Any]          =							model_class(__SCREAMING_SNAKE_CASE )

          __snake_case :					Optional[int]          =							model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE ) )

          __snake_case :					Tuple          =							outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

          __snake_case :					str          =							self.model_tester.num_stages
          self.assertEqual(len(__SCREAMING_SNAKE_CASE )					,    expected_num_stages + 1 )

        __snake_case :					Any          =							self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          __snake_case :					Optional[Any]          =							True
          check_hidden_states_output(__SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE )

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          __snake_case :					Optional[Any]          =							True

          check_hidden_states_output(__SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE )







      def     UpperCAmelCase		(			self )							-> Dict:






        '''simple docstring'''






        __snake_case :					Any          =							self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          with self.subTest(model_class.__name__ ):
            __snake_case :					int          =							self._prepare_for_class(__SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE )
            __snake_case :					Union[str, Any]          =							model_class(__SCREAMING_SNAKE_CASE )

            @jax.jit
            def model_jitted(UpperCAmelCase					,    **UpperCAmelCase ):
              return model(pixel_values=__SCREAMING_SNAKE_CASE					,    **__SCREAMING_SNAKE_CASE )

            with self.subTest("JIT Enabled" ):
              __snake_case :					Any          =							model_jitted(**__SCREAMING_SNAKE_CASE ).to_tuple()

            with self.subTest("JIT Disabled" ):
              with jax.disable_jit():
                __snake_case :					Optional[int]          =							model_jitted(**__SCREAMING_SNAKE_CASE ).to_tuple()

            self.assertEqual(len(__SCREAMING_SNAKE_CASE )					,    len(__SCREAMING_SNAKE_CASE ) )
            for jitted_output, output in zip(__SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE ):
              self.assertEqual(jitted_output.shape					,    output.shape )







def    lowerCAmelCase__(						)			->     Union[str, Any]:
  __snake_case :					int          =							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"					)
  return image




@require_flax
class 					_lowerCamelCase			(     unittest.TestCase     ):



      """simple docstring"""






      @cached_property
      def     UpperCAmelCase		(			self )							-> List[str]:






        '''simple docstring'''






        return AutoImageProcessor.from_pretrained("facebook/regnet-y-040" ) if is_vision_available() else None







      @slow
      def     UpperCAmelCase		(			self )							-> Dict:






        '''simple docstring'''






        __snake_case :					Optional[int]          =							FlaxRegNetForImageClassification.from_pretrained("facebook/regnet-y-040" )

        __snake_case :					Any          =							self.default_image_processor
        __snake_case :					Tuple          =							prepare_img()
        __snake_case :					Tuple          =							image_processor(images=__SCREAMING_SNAKE_CASE					,    return_tensors="np" )

        __snake_case :					Tuple          =							model(**__SCREAMING_SNAKE_CASE )

        # verify the logits
        __snake_case :					Union[str, Any]          =							(1, 1000)
        self.assertEqual(outputs.logits.shape					,    __SCREAMING_SNAKE_CASE )

        __snake_case :					List[str]          =							jnp.array([-0.4_180, -1.5_051, -3.4_836] )

        self.assertTrue(jnp.allclose(outputs.logits[0, :3]					,    __SCREAMING_SNAKE_CASE					,    atol=1E-4 ) )



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





     """simple docstring"""







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






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Any:




           """simple docstring"""







           torch.manual_seed(0	)
           UpperCamelCase__     : List[str]											=					UNetaDModel(
               block_out_channels=(3_2, 3_2, 6_4)    ,    extra_in_channels=1_6    ,    sample_size=5_1_2    ,    sample_rate=1_6_0_0_0    ,    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''')    ,    )
           UpperCamelCase__     : Union[str, Any]											=					IPNDMScheduler()

           UpperCamelCase__     : List[str]											=					{
               '''unet''': unet,
               '''scheduler''': scheduler,
           }
           return components






     def  __SCREAMING_SNAKE_CASE      (				self    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE=0	)	-> Any:




           """simple docstring"""







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






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




           """simple docstring"""







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

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

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

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




           """simple docstring"""







           return super().test_save_load_local()






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Tuple:




           """simple docstring"""







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






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> str:




           """simple docstring"""







           return super().test_save_load_optional_components()






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




           """simple docstring"""







           return super().test_attention_slicing_forward_pass()






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Tuple:




           """simple docstring"""







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




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





     """simple docstring"""







     def  __SCREAMING_SNAKE_CASE      (				self	)	-> int:




           """simple docstring"""







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






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




           """simple docstring"""







           UpperCamelCase__     : Any											=					torch_device

           UpperCamelCase__     : Any											=					DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''	)
           UpperCamelCase__     : int											=					pipe.to(__SCREAMING_SNAKE_CASE	)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : int											=					torch.manual_seed(0	)
           UpperCamelCase__     : Optional[Any]											=					pipe(generator=__SCREAMING_SNAKE_CASE    ,    num_inference_steps=1_0_0    ,    audio_length_in_s=4.096	)
           UpperCamelCase__     : str											=					output.audios

           UpperCamelCase__     : List[Any]											=					audio[0, -3:, -3:]

           assert audio.shape == (1, 2, pipe.unet.sample_size)
           UpperCamelCase__     : Tuple											=					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  __SCREAMING_SNAKE_CASE      (				self	)	-> Union[str, Any]:




           """simple docstring"""







           UpperCamelCase__     : Any											=					torch_device

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

           UpperCamelCase__     : int											=					torch.manual_seed(0	)
           UpperCamelCase__     : Union[str, Any]											=					pipe(generator=__SCREAMING_SNAKE_CASE    ,    num_inference_steps=1_0_0    ,    audio_length_in_s=4.096	)
           UpperCamelCase__     : List[Any]											=					output.audios

           UpperCamelCase__     : List[Any]											=					audio[0, -3:, -3:]

           assert audio.shape == (1, 2, pipe.unet.sample_size)
           UpperCamelCase__     : Optional[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


