





import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class 					A__   (       __snake_case       ,   unittest.TestCase			):
      _UpperCAmelCase	:int						    =   KandinskyVaaPriorPipeline
      _UpperCAmelCase	:int						    =   ['prompt']
      _UpperCAmelCase	:str						    =   ['prompt', 'negative_prompt']
      _UpperCAmelCase	:Tuple						    =   [
          'num_images_per_prompt',
          'generator',
          'num_inference_steps',
          'latents',
          'negative_prompt',
          'guidance_scale',
          'output_type',
          'return_dict',
      ]
      _UpperCAmelCase	:Tuple						    =   False





      @property
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       return 32





      @property
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       return 32





      @property
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       return self.time_input_dim





      @property
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       return self.time_input_dim * 4





      @property
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       return 100





      @property
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )
       return tokenizer





      @property
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       torch.manual_seed(0   )
       UpperCamelCase	:					List[str]       	=   CLIPTextConfig(
           bos_token_id=0       ,							eos_token_id=2       ,							hidden_size=self.text_embedder_hidden_size       ,							projection_dim=self.text_embedder_hidden_size       ,							intermediate_size=37       ,							layer_norm_eps=1e-05       ,							num_attention_heads=4       ,							num_hidden_layers=5       ,							pad_token_id=1       ,							vocab_size=1000       ,							)
       return CLIPTextModelWithProjection(A_   )





      @property
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       torch.manual_seed(0   )

       UpperCamelCase	:					Optional[Any]       	=   {
           "num_attention_heads": 2,
           "attention_head_dim": 12,
           "embedding_dim": self.text_embedder_hidden_size,
           "num_layers": 1,
       }

       UpperCamelCase	:					List[str]       	=   PriorTransformer(**A_   )
       # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
       UpperCamelCase	:					Union[str, Any]       	=   nn.Parameter(torch.ones(model.clip_std.shape   )   )
       return model





      @property
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       torch.manual_seed(0   )
       UpperCamelCase	:					int       	=   CLIPVisionConfig(
           hidden_size=self.text_embedder_hidden_size       ,							image_size=224       ,							projection_dim=self.text_embedder_hidden_size       ,							intermediate_size=37       ,							num_attention_heads=4       ,							num_channels=3       ,							num_hidden_layers=5       ,							patch_size=14       ,							)

       UpperCamelCase	:					Dict       	=   CLIPVisionModelWithProjection(A_   )
       return model





      @property
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Dict       	=   CLIPImageProcessor(
           crop_size=224       ,							do_center_crop=A_       ,							do_normalize=A_       ,							do_resize=A_       ,							image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]       ,							image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]       ,							resample=3       ,							size=224       ,							)

       return image_processor





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Tuple       	=   self.dummy_prior
       UpperCamelCase	:					Any       	=   self.dummy_image_encoder
       UpperCamelCase	:					List[str]       	=   self.dummy_text_encoder
       UpperCamelCase	:					Dict       	=   self.dummy_tokenizer
       UpperCamelCase	:					Optional[Any]       	=   self.dummy_image_processor

       UpperCamelCase	:					List[str]       	=   UnCLIPScheduler(
           variance_type="fixed_small_log"       ,							prediction_type="sample"       ,							num_train_timesteps=1000       ,							clip_sample=A_       ,							clip_sample_range=10.0       ,							)

       UpperCamelCase	:					Optional[int]       	=   {
           "prior": prior,
           "image_encoder": image_encoder,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "scheduler": scheduler,
           "image_processor": image_processor,
       }

       return components





      def        __UpperCamelCase( self       ,							A_       ,							A_=0   ):
       '''simple docstring'''




       if str(A_   ).startswith("mps"   ):
        UpperCamelCase	:					List[Any]       	=   torch.manual_seed(A_   )
       else:
        UpperCamelCase	:					int       	=   torch.Generator(device=A_   ).manual_seed(A_   )
       UpperCamelCase	:					Union[str, Any]       	=   {
           "prompt": "horse",
           "generator": generator,
           "guidance_scale": 4.0,
           "num_inference_steps": 2,
           "output_type": "np",
       }
       return inputs





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   "cpu"

       UpperCamelCase	:					Dict       	=   self.get_dummy_components()

       UpperCamelCase	:					Union[str, Any]       	=   self.pipeline_class(**A_   )
       UpperCamelCase	:					str       	=   pipe.to(A_   )

       pipe.set_progress_bar_config(disable=A_   )

       UpperCamelCase	:					str       	=   pipe(**self.get_dummy_inputs(A_   )   )
       UpperCamelCase	:					Any       	=   output.image_embeds

       UpperCamelCase	:					int       	=   pipe(
           **self.get_dummy_inputs(A_   )       ,							return_dict=A_       ,							)[0]

       UpperCamelCase	:					List[str]       	=   image[0, -10:]
       UpperCamelCase	:					Dict       	=   image_from_tuple[0, -10:]

       assert image.shape == (1, 32)

       UpperCamelCase	:					Any       	=   np.array(
           [-0.05_32, 1.71_20, 0.36_56, -1.08_52, -0.89_46, -1.17_56, 0.43_48, 0.24_82, 0.51_46, -0.11_56]   )

       assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-2
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice   ).max() < 1e-2





      @skip_mps
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   torch_device == "cpu"
       UpperCamelCase	:					Dict       	=   True
       UpperCamelCase	:					str       	=   False

       self._test_inference_batch_single_identical(
           test_max_difference=A_       ,							relax_max_difference=A_       ,							test_mean_pixel_difference=A_       ,							)





      @skip_mps
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[int]       	=   torch_device == "cpu"
       UpperCamelCase	:					Tuple       	=   False

       self._test_attention_slicing_forward_pass(
           test_max_difference=A_       ,							test_mean_pixel_difference=A_       ,							)







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__lowerCamelCase       :     Any   =     {
    """configuration_electra""": ["""ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP""", """ElectraConfig""", """ElectraOnnxConfig"""],
    """tokenization_electra""": ["""ElectraTokenizer"""],
}

try:
    if not is_tokenizers_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCamelCase       :     Dict   =     ["""ElectraTokenizerFast"""]

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCamelCase       :     Tuple   =     [
        """ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """ElectraForCausalLM""",
        """ElectraForMaskedLM""",
        """ElectraForMultipleChoice""",
        """ElectraForPreTraining""",
        """ElectraForQuestionAnswering""",
        """ElectraForSequenceClassification""",
        """ElectraForTokenClassification""",
        """ElectraModel""",
        """ElectraPreTrainedModel""",
        """load_tf_weights_in_electra""",
    ]

try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCamelCase       :     List[Any]   =     [
        """TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """TFElectraForMaskedLM""",
        """TFElectraForMultipleChoice""",
        """TFElectraForPreTraining""",
        """TFElectraForQuestionAnswering""",
        """TFElectraForSequenceClassification""",
        """TFElectraForTokenClassification""",
        """TFElectraModel""",
        """TFElectraPreTrainedModel""",
    ]

try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCamelCase       :     List[str]   =     [
        """FlaxElectraForCausalLM""",
        """FlaxElectraForMaskedLM""",
        """FlaxElectraForMultipleChoice""",
        """FlaxElectraForPreTraining""",
        """FlaxElectraForQuestionAnswering""",
        """FlaxElectraForSequenceClassification""",
        """FlaxElectraForTokenClassification""",
        """FlaxElectraModel""",
        """FlaxElectraPreTrainedModel""",
    ]


if TYPE_CHECKING:
    from .configuration_electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig, ElectraOnnxConfig
    from .tokenization_electra import ElectraTokenizer

    try:
        if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_electra_fast import ElectraTokenizerFast

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_electra import (
            ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
            ElectraForCausalLM,
            ElectraForMaskedLM,
            ElectraForMultipleChoice,
            ElectraForPreTraining,
            ElectraForQuestionAnswering,
            ElectraForSequenceClassification,
            ElectraForTokenClassification,
            ElectraModel,
            ElectraPreTrainedModel,
            load_tf_weights_in_electra,
        )

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_electra import (
            TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
            TFElectraForMaskedLM,
            TFElectraForMultipleChoice,
            TFElectraForPreTraining,
            TFElectraForQuestionAnswering,
            TFElectraForSequenceClassification,
            TFElectraForTokenClassification,
            TFElectraModel,
            TFElectraPreTrainedModel,
        )

    try:
        if not is_flax_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_flax_electra import (
            FlaxElectraForCausalLM,
            FlaxElectraForMaskedLM,
            FlaxElectraForMultipleChoice,
            FlaxElectraForPreTraining,
            FlaxElectraForQuestionAnswering,
            FlaxElectraForSequenceClassification,
            FlaxElectraForTokenClassification,
            FlaxElectraModel,
            FlaxElectraPreTrainedModel,
        )

else:
    import sys

    __lowerCamelCase       :     Any   =     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
