






"""simple docstring"""







import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()







class    lowerCamelCase__   (			_UpperCAmelCase       ,	unittest.TestCase					):
    SCREAMING_SNAKE_CASE						=     PriorTransformer
    SCREAMING_SNAKE_CASE						=     '''hidden_states'''






    @property
    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			4
          UpperCAmelCase							     =			8
          UpperCAmelCase							     =			7

          UpperCAmelCase							     =			floats_tensor((batch_size, embedding_dim)   ).to(lowerCAmelCase__   )

          UpperCAmelCase							     =			floats_tensor((batch_size, embedding_dim)   ).to(lowerCAmelCase__   )
          UpperCAmelCase							     =			floats_tensor((batch_size, num_embeddings, embedding_dim)   ).to(lowerCAmelCase__   )

          return {
              "hidden_states": hidden_states,
              "timestep": 2,
              "proj_embedding": proj_embedding,
              "encoder_hidden_states": encoder_hidden_states,
          }






    def 							_UpperCamelCase				(					self      ,A=0   ):
          torch.manual_seed(lowerCAmelCase__   )
          UpperCAmelCase							     =			4
          UpperCAmelCase							     =			8
          UpperCAmelCase							     =			7

          UpperCAmelCase							     =			torch.randn((batch_size, embedding_dim)   ).to(lowerCAmelCase__   )

          UpperCAmelCase							     =			torch.randn((batch_size, embedding_dim)   ).to(lowerCAmelCase__   )
          UpperCAmelCase							     =			torch.randn((batch_size, num_embeddings, embedding_dim)   ).to(lowerCAmelCase__   )

          return {
              "hidden_states": hidden_states,
              "timestep": 2,
              "proj_embedding": proj_embedding,
              "encoder_hidden_states": encoder_hidden_states,
          }






    @property
    def 							_UpperCamelCase				(					self   ):
          return (4, 8)






    @property
    def 							_UpperCamelCase				(					self   ):
          return (4, 8)






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			{
              """num_attention_heads""": 2,
              """attention_head_dim""": 4,
              """num_layers""": 2,
              """embedding_dim""": 8,
              """num_embeddings""": 7,
              """additional_embeddings""": 4,
          }
          UpperCAmelCase							     =			self.dummy_input
          return init_dict, inputs_dict






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase     ,   UpperCAmelCase							     =			PriorTransformer.from_pretrained(
              """hf-internal-testing/prior-dummy"""      ,output_loading_info=lowerCAmelCase__   )
          self.assertIsNotNone(lowerCAmelCase__   )
          self.assertEqual(len(loading_info["""missing_keys"""]   )      ,0   )

          model.to(lowerCAmelCase__   )
          UpperCAmelCase							     =			model(**self.dummy_input   )[0]

          assert hidden_states is not None, "Make sure output is not None"






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase     ,   UpperCAmelCase							     =			self.prepare_init_args_and_inputs_for_common()

          UpperCAmelCase							     =			self.model_class(**lowerCAmelCase__   )
          UpperCAmelCase							     =			inspect.signature(model.forward   )
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          UpperCAmelCase							     =			[*signature.parameters.keys()]

          UpperCAmelCase							     =			["""hidden_states""", """timestep"""]
          self.assertListEqual(arg_names[:2]      ,lowerCAmelCase__   )






    def 							_UpperCamelCase				(					self   ):
          UpperCAmelCase							     =			PriorTransformer.from_pretrained("""hf-internal-testing/prior-dummy"""   )
          UpperCAmelCase							     =			model.to(lowerCAmelCase__   )

          if hasattr(lowerCAmelCase__      ,"""set_default_attn_processor"""   ):
                model.set_default_attn_processor()

          UpperCAmelCase							     =			self.get_dummy_seed_input()

          with torch.no_grad():
                UpperCAmelCase							     =			model(**lowerCAmelCase__   )[0]

          UpperCAmelCase							     =			output[0, :5].flatten().cpu()
          print(lowerCAmelCase__   )

          # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
          # the expected output slices are not the same for CPU and GPU.
          UpperCAmelCase							     =			torch.tensor([-1.3436, -0.2870, 0.7538, 0.4368, -0.0239]   )
          self.assertTrue(torch_all_close(lowerCAmelCase__      ,lowerCAmelCase__      ,rtol=1e-2   )   )







@slow
class    lowerCamelCase__   (			unittest.TestCase					):






    def 							_UpperCamelCase				(					self      ,A=1      ,A=768      ,A=77      ,A=0   ):
          torch.manual_seed(lowerCAmelCase__   )
          UpperCAmelCase							     =			batch_size
          UpperCAmelCase							     =			embedding_dim
          UpperCAmelCase							     =			num_embeddings

          UpperCAmelCase							     =			torch.randn((batch_size, embedding_dim)   ).to(lowerCAmelCase__   )

          UpperCAmelCase							     =			torch.randn((batch_size, embedding_dim)   ).to(lowerCAmelCase__   )
          UpperCAmelCase							     =			torch.randn((batch_size, num_embeddings, embedding_dim)   ).to(lowerCAmelCase__   )

          return {
              "hidden_states": hidden_states,
              "timestep": 2,
              "proj_embedding": proj_embedding,
              "encoder_hidden_states": encoder_hidden_states,
          }






    def 							_UpperCamelCase				(					self   ):
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()






    @parameterized.expand(
        [
            # fmt: off
            [13, [-0.5861,  0.1283, -0.0931,  0.0882,  0.4476,  0.1329, -0.0498,  0.0640]],
            [37, [-0.4913,  0.0110, -0.0483,  0.0541,  0.4954, -0.0170,  0.0354,  0.1651]],
            # fmt: on
        ]   )
    def 							_UpperCamelCase				(					self      ,A      ,A   ):
          UpperCAmelCase							     =			PriorTransformer.from_pretrained("""kandinsky-community/kandinsky-2-1-prior"""      ,subfolder="""prior"""   )
          model.to(lowerCAmelCase__   )
          UpperCAmelCase							     =			self.get_dummy_seed_input(seed=lowerCAmelCase__   )

          with torch.no_grad():
                UpperCAmelCase							     =			model(**lowerCAmelCase__   )[0]

          assert list(sample.shape   ) == [1, 768]

          UpperCAmelCase							     =			sample[0, :8].flatten().cpu()
          print(lowerCAmelCase__   )
          UpperCAmelCase							     =			torch.tensor(lowerCAmelCase__   )

          assert torch_all_close(lowerCAmelCase__      ,lowerCAmelCase__      ,atol=1e-3   )




from ...configuration_utils import PretrainedConfig
from ...utils import logging


__a    :	List[Any]         =    logging.get_logger(__name__)

__a    :	List[Any]         =    {
    """MIT/ast-finetuned-audioset-10-10-0.4593""": (
        """https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593/resolve/main/config.json"""
    ),
}




class   _UpperCamelCase (		_UpperCAmelCase					):


			"""simple docstring"""






			__a    :					Optional[Any]     					=		'''audio-spectrogram-transformer'''





			def __init__(  self ,  lowerCAmelCase__=7_68 ,  lowerCAmelCase__=12 ,  lowerCAmelCase__=12 ,  lowerCAmelCase__=30_72 ,  lowerCAmelCase__="gelu" ,  lowerCAmelCase__=0.0 ,  lowerCAmelCase__=0.0 ,  lowerCAmelCase__=0.02 ,  lowerCAmelCase__=1E-12 ,  lowerCAmelCase__=16 ,  lowerCAmelCase__=True ,  lowerCAmelCase__=10 ,  lowerCAmelCase__=10 ,  lowerCAmelCase__=10_24 ,  lowerCAmelCase__=1_28 ,  **lowerCAmelCase__ ,  )	->   Optional[Any]:

							'''simple docstring'''


							super().__init__(**lowerCAmelCase__		)

							__lowercase					     =				hidden_size
							__lowercase					     =				num_hidden_layers
							__lowercase					     =				num_attention_heads
							__lowercase					     =				intermediate_size
							__lowercase					     =				hidden_act
							__lowercase					     =				hidden_dropout_prob
							__lowercase					     =				attention_probs_dropout_prob
							__lowercase					     =				initializer_range
							__lowercase					     =				layer_norm_eps
							__lowercase					     =				patch_size
							__lowercase					     =				qkv_bias
							__lowercase					     =				frequency_stride
							__lowercase					     =				time_stride
							__lowercase					     =				max_length
							__lowercase					     =				num_mel_bins