





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      A       (	lowerCAmelCase__		,      unittest.TestCase						):
      _snake_case							     =PriorTransformer
      _snake_case							     ="hidden_states"



      @property
      def 	lowerCAmelCase__    (	self:						Dict       )						->    int:
         '''simple docstring'''

         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 	lowerCAmelCase__    (	self:						List[Any]		,  _lowerCAmelCase:						Union[str, Any]=0       )						->    Optional[int]:
         '''simple docstring'''

         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 	lowerCAmelCase__    (	self:						Dict       )						->    Union[str, Any]:
         '''simple docstring'''

         return (4, 8)



      @property
      def 	lowerCAmelCase__    (	self:						Union[str, Any]       )						->    int:
         '''simple docstring'''

         return (4, 8)



      def 	lowerCAmelCase__    (	self:						Tuple       )						->    Union[str, Any]:
         '''simple docstring'''

         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 	lowerCAmelCase__    (	self:						Any       )						->    Optional[int]:
         '''simple docstring'''

         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 	lowerCAmelCase__    (	self:						int       )						->    str:
         '''simple docstring'''

         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 	lowerCAmelCase__    (	self:						List[str]       )						->    int:
         '''simple docstring'''

         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.34_36, -0.28_70, 0.75_38, 0.43_68, -0.02_39]       )
         self.assertTrue(torch_all_close(_lowerCamelCase		,  _lowerCamelCase		,  rtol=1e-2       )       )

@slow
class      A       (	unittest.TestCase						):



      def 	lowerCAmelCase__    (	self:						Union[str, Any]		,  _lowerCAmelCase:						Union[str, Any]=1		,  _lowerCAmelCase:						Optional[int]=768		,  _lowerCAmelCase:						Optional[Any]=77		,  _lowerCAmelCase:						Tuple=0       )						->    Optional[Any]:
         '''simple docstring'''

         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 	lowerCAmelCase__    (	self:						List[str]       )						->    Optional[Any]:
         '''simple docstring'''

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






      @parameterized.expand(
          [
              # fmt: off
              [13, [-0.58_61,  0.12_83, -0.09_31,  0.08_82,  0.44_76,  0.13_29, -0.04_98,  0.06_40]],
              [37, [-0.49_13,  0.01_10, -0.04_83,  0.05_41,  0.49_54, -0.01_70,  0.03_54,  0.16_51]],
              # fmt: on
          ]       )
      def 	lowerCAmelCase__    (	self:						Tuple		,  _lowerCAmelCase:						Dict		,  _lowerCAmelCase:						str       )						->    Optional[int]:
         '''simple docstring'''

         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


__lowercase  :   Union[str, Any]						      =logging.get_logger(__name__)

__lowercase  :   List[Any]						      ={
    """tanreinama/GPTSAN-2.8B-spout_is_uniform""": (
        """https://huggingface.co/tanreinama/GPTSAN-2.8B-spout_is_uniform/resolve/main/config.json"""
    ),
}

class      A       (	__lowercase						):
      _snake_case							     ='''gptsan-japanese'''
      _snake_case							     =[
          '''past_key_values''',
      ]
      _snake_case							     ={
          '''hidden_size''': '''d_model''',
          '''num_attention_heads''': '''num_heads''',
          '''num_hidden_layers''': '''num_layers''',
      }



      def __init__(	self:						Optional[Any]		,  _lowerCAmelCase:						List[Any]=3_6000		,  _lowerCAmelCase:						List[Any]=1280		,  _lowerCAmelCase:						str=1024		,  _lowerCAmelCase:						Any=8192		,  _lowerCAmelCase:						str=4096		,  _lowerCAmelCase:						int=128		,  _lowerCAmelCase:						int=10		,  _lowerCAmelCase:						Dict=0		,  _lowerCAmelCase:						Any=16		,  _lowerCAmelCase:						Optional[int]=16		,  _lowerCAmelCase:						List[Any]=128		,  _lowerCAmelCase:						Tuple=0.0		,  _lowerCAmelCase:						Optional[Any]=1e-5		,  _lowerCAmelCase:						int=False		,  _lowerCAmelCase:						Optional[Any]=0.0		,  _lowerCAmelCase:						str="float32"		,  _lowerCAmelCase:						Dict=False		,  _lowerCAmelCase:						Any=False		,  _lowerCAmelCase:						int=False		,  _lowerCAmelCase:						Union[str, Any]=0.0_02		,  _lowerCAmelCase:						Optional[int]=False		,  _lowerCAmelCase:						int=True		,  _lowerCAmelCase:						List[str]=3_5998		,  _lowerCAmelCase:						Optional[int]=3_5995		,  _lowerCAmelCase:						Dict=3_5999		,  **_lowerCAmelCase:						Optional[Any]		,  )						->    List[Any]:
         '''simple docstring'''

         UpperCAmelCase_											=vocab_size
         UpperCAmelCase_											=max_position_embeddings
         UpperCAmelCase_											=d_model
         UpperCAmelCase_											=d_ff
         UpperCAmelCase_											=d_ext
         UpperCAmelCase_											=d_spout
         UpperCAmelCase_											=num_switch_layers
         UpperCAmelCase_											=num_ext_layers
         UpperCAmelCase_											=num_switch_layers + num_ext_layers
         UpperCAmelCase_											=num_heads
         UpperCAmelCase_											=num_experts
         UpperCAmelCase_											=expert_capacity
         UpperCAmelCase_											=dropout_rate
         UpperCAmelCase_											=layer_norm_epsilon
         UpperCAmelCase_											=router_bias
         UpperCAmelCase_											=router_jitter_noise
         UpperCAmelCase_											=router_dtype
         UpperCAmelCase_											=router_ignore_padding_tokens
         UpperCAmelCase_											=output_hidden_states
         UpperCAmelCase_											=output_attentions
         UpperCAmelCase_											=initializer_factor
         UpperCAmelCase_											=output_router_logits
         UpperCAmelCase_											=use_cache

         super().__init__(
             separator_token_id=_lowerCAmelCase		,  pad_token_id=_lowerCAmelCase		,  eos_token_id=_lowerCAmelCase		,  **_lowerCAmelCase		,  )




