






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






      '''simple docstring'''




      __snake_case	    =    KandinskyVaaPriorPipeline
      __snake_case	    =    ['''prompt''']
      __snake_case	    =    ['''prompt''', '''negative_prompt''']
      __snake_case	    =    [
          '''num_images_per_prompt''',
          '''generator''',
          '''num_inference_steps''',
          '''latents''',
          '''negative_prompt''',
          '''guidance_scale''',
          '''output_type''',
          '''return_dict''',
      ]
      __snake_case	    =    False







      @property
      def  __lowerCAmelCase   ( self :		Any    )       ->Optional[int]:




       """simple docstring"""


       return 32







      @property
      def  __lowerCAmelCase   ( self :		str    )       ->Optional[int]:




       """simple docstring"""


       return 32







      @property
      def  __lowerCAmelCase   ( self :		int    )       ->List[str]:




       """simple docstring"""


       return self.time_input_dim







      @property
      def  __lowerCAmelCase   ( self :		List[Any]    )       ->Any:




       """simple docstring"""


       return self.time_input_dim * 4







      @property
      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Optional[int]:




       """simple docstring"""


       return 100







      @property
      def  __lowerCAmelCase   ( self :		List[Any]    )       ->Optional[int]:




       """simple docstring"""


       a						    = CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )
       return tokenizer







      @property
      def  __lowerCAmelCase   ( self :		Union[str, Any]    )       ->str:




       """simple docstring"""


       torch.manual_seed(0    )
       a						    = 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-0_5					,   num_attention_heads=4					,   num_hidden_layers=5					,   pad_token_id=1					,   vocab_size=1_000					,   )
       return CLIPTextModelWithProjection(__UpperCAmelCase    )







      @property
      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Optional[Any]:




       """simple docstring"""


       torch.manual_seed(0    )

       a						    = {
           '''num_attention_heads''': 2,
           '''attention_head_dim''': 12,
           '''embedding_dim''': self.text_embedder_hidden_size,
           '''num_layers''': 1,
       }

       a						    = PriorTransformer(**__UpperCAmelCase    )
       # 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
       a						    = nn.Parameter(torch.ones(model.clip_std.shape    )    )
       return model







      @property
      def  __lowerCAmelCase   ( self :		List[str]    )       ->str:




       """simple docstring"""


       torch.manual_seed(0    )
       a						    = 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					,   )

       a						    = CLIPVisionModelWithProjection(__UpperCAmelCase    )
       return model







      @property
      def  __lowerCAmelCase   ( self :		int    )       ->Optional[int]:




       """simple docstring"""


       a						    = CLIPImageProcessor(
           crop_size=224					,   do_center_crop=__UpperCAmelCase					,   do_normalize=__UpperCAmelCase					,   do_resize=__UpperCAmelCase					,   image_mean=[0.48145466, 0.4578275, 0.40821073]					,   image_std=[0.26862954, 0.26130258, 0.27577711]					,   resample=3					,   size=224					,   )

       return image_processor







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




       """simple docstring"""


       a						    = self.dummy_prior
       a						    = self.dummy_image_encoder
       a						    = self.dummy_text_encoder
       a						    = self.dummy_tokenizer
       a						    = self.dummy_image_processor

       a						    = UnCLIPScheduler(
           variance_type='''fixed_small_log'''					,   prediction_type='''sample'''					,   num_train_timesteps=1_000					,   clip_sample=__UpperCAmelCase					,   clip_sample_range=10.0					,   )

       a						    = {
           '''prior''': prior,
           '''image_encoder''': image_encoder,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''scheduler''': scheduler,
           '''image_processor''': image_processor,
       }

       return components







      def  __lowerCAmelCase   ( self :		List[Any]					,   __UpperCAmelCase :		Optional[Any]					,   __UpperCAmelCase :		Optional[int]=0    )       ->str:




       """simple docstring"""


       if str(__UpperCAmelCase    ).startswith('''mps'''    ):
        a						    = torch.manual_seed(__UpperCAmelCase    )
       else:
        a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(__UpperCAmelCase    )
       a						    = {
           '''prompt''': '''horse''',
           '''generator''': generator,
           '''guidance_scale''': 4.0,
           '''num_inference_steps''': 2,
           '''output_type''': '''np''',
       }
       return inputs







      def  __lowerCAmelCase   ( self :		Any    )       ->Tuple:




       """simple docstring"""


       a						    = '''cpu'''

       a						    = self.get_dummy_components()

       a						    = self.pipeline_class(**__UpperCAmelCase    )
       a						    = pipe.to(__UpperCAmelCase    )

       pipe.set_progress_bar_config(disable=__UpperCAmelCase    )

       a						    = pipe(**self.get_dummy_inputs(__UpperCAmelCase    )    )
       a						    = output.image_embeds

       a						    = pipe(
           **self.get_dummy_inputs(__UpperCAmelCase    )					,   return_dict=__UpperCAmelCase					,   )[0]

       a						    = image[0, -10:]
       a						    = image_from_tuple[0, -10:]

       assert image.shape == (1, 32)

       a						    = np.array(
           [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]    )

       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  __lowerCAmelCase   ( self :		Dict    )       ->Dict:




       """simple docstring"""


       a						    = torch_device == '''cpu'''
       a						    = True
       a						    = False

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







      @skip_mps
      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->List[Any]:




       """simple docstring"""


       a						    = torch_device == '''cpu'''
       a						    = False

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










from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


UpperCAmelCase__							  =						{
    "configuration_rag": ["RagConfig"],
    "retrieval_rag": ["RagRetriever"],
    "tokenization_rag": ["RagTokenizer"],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCAmelCase__							  =						[
         "RagModel",
         "RagPreTrainedModel",
         "RagSequenceForGeneration",
         "RagTokenForGeneration",
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCAmelCase__							  =						[
         "TFRagModel",
         "TFRagPreTrainedModel",
         "TFRagSequenceForGeneration",
         "TFRagTokenForGeneration",
     ]


if TYPE_CHECKING:
     from .configuration_rag import RagConfig
     from .retrieval_rag import RagRetriever
     from .tokenization_rag import RagTokenizer

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_rag import RagModel, RagPreTrainedModel, RagSequenceForGeneration, RagTokenForGeneration

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_rag import (
              TFRagModel,
              TFRagPreTrainedModel,
              TFRagSequenceForGeneration,
              TFRagTokenForGeneration,
          )

else:
     import sys

     UpperCAmelCase__							  =						_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


