


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    (  UpperCamelCase_       ,				unittest.TestCase	):
      UpperCamelCase__				:    Union[str, Any] =KandinskyVaaPriorPipeline
      UpperCamelCase__				:    Optional[Any] =['prompt']
      UpperCamelCase__				:    Union[str, Any] =['prompt', 'negative_prompt']
      UpperCamelCase__				:    List[Any] =[
          'num_images_per_prompt',
          'generator',
          'num_inference_steps',
          'latents',
          'negative_prompt',
          'guidance_scale',
          'output_type',
          'return_dict',
      ]
      UpperCamelCase__				:    Optional[Any] =False






      @property
      def 			lowerCamelCase			(   self		:	Dict					)   ->		Optional[int]:
             """simple docstring"""




             return 32






      @property
      def 			lowerCamelCase			(   self		:	int					)   ->		List[str]:
             """simple docstring"""




             return 32






      @property
      def 			lowerCamelCase			(   self		:	Optional[Any]					)   ->		int:
             """simple docstring"""




             return self.time_input_dim






      @property
      def 			lowerCamelCase			(   self		:	Dict					)   ->		Tuple:
             """simple docstring"""




             return self.time_input_dim * 4






      @property
      def 			lowerCamelCase			(   self		:	int					)   ->		Dict:
             """simple docstring"""




             return 100






      @property
      def 			lowerCamelCase			(   self		:	Tuple					)   ->		Union[str, Any]:
             """simple docstring"""




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






      @property
      def 			lowerCamelCase			(   self		:	List[str]					)   ->		List[str]:
             """simple docstring"""




             torch.manual_seed(0					)
             _lowerCamelCase			:						int            =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(lowercase_					)






      @property
      def 			lowerCamelCase			(   self		:	Any					)   ->		Optional[Any]:
             """simple docstring"""




             torch.manual_seed(0					)

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

             _lowerCamelCase			:						Dict            =PriorTransformer(**lowercase_					)
             # 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
             _lowerCamelCase			:						Optional[Any]            =nn.Parameter(torch.ones(model.clip_std.shape					)					)
             return model






      @property
      def 			lowerCamelCase			(   self		:	Tuple					)   ->		List[Any]:
             """simple docstring"""




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

             _lowerCamelCase			:						int            =CLIPVisionModelWithProjection(lowercase_					)
             return model






      @property
      def 			lowerCamelCase			(   self		:	Optional[int]					)   ->		Dict:
             """simple docstring"""




             _lowerCamelCase			:						Dict            =CLIPImageProcessor(
                 crop_size=224     ,  do_center_crop=lowercase_     ,  do_normalize=lowercase_     ,  do_resize=lowercase_     ,  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		:	str					)   ->		List[str]:
             """simple docstring"""




             _lowerCamelCase			:						List[str]            =self.dummy_prior
             _lowerCamelCase			:						List[str]            =self.dummy_image_encoder
             _lowerCamelCase			:						List[Any]            =self.dummy_text_encoder
             _lowerCamelCase			:						List[str]            =self.dummy_tokenizer
             _lowerCamelCase			:						List[str]            =self.dummy_image_processor

             _lowerCamelCase			:						Dict            =UnCLIPScheduler(
                 variance_type='fixed_small_log'     ,  prediction_type='sample'     ,  num_train_timesteps=1000     ,  clip_sample=lowercase_     ,  clip_sample_range=10.0     ,  )

             _lowerCamelCase			:						Optional[Any]            ={
                 'prior': prior,
                 'image_encoder': image_encoder,
                 'text_encoder': text_encoder,
                 'tokenizer': tokenizer,
                 'scheduler': scheduler,
                 'image_processor': image_processor,
             }

             return components






      def 			lowerCamelCase			(   self		:	Union[str, Any]     ,  lowercase_		:	Optional[int]     ,  lowercase_		:	List[str]=0					)   ->		Union[str, Any]:
             """simple docstring"""




             if str(lowercase_					).startswith('mps'					):
                    _lowerCamelCase			:						Any            =torch.manual_seed(lowercase_					)
             else:
                    _lowerCamelCase			:						List[str]            =torch.Generator(device=lowercase_					).manual_seed(lowercase_					)
             _lowerCamelCase			:						Dict            ={
                 'prompt': 'horse',
                 'generator': generator,
                 'guidance_scale': 4.0,
                 'num_inference_steps': 2,
                 'output_type': 'np',
             }
             return inputs






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




             _lowerCamelCase			:						List[Any]            ='cpu'

             _lowerCamelCase			:						Any            =self.get_dummy_components()

             _lowerCamelCase			:						str            =self.pipeline_class(**lowercase_					)
             _lowerCamelCase			:						int            =pipe.to(lowercase_					)

             pipe.set_progress_bar_config(disable=lowercase_					)

             _lowerCamelCase			:						Dict            =pipe(**self.get_dummy_inputs(lowercase_					)					)
             _lowerCamelCase			:						Dict            =output.image_embeds

             _lowerCamelCase			:						List[str]            =pipe(
                 **self.get_dummy_inputs(lowercase_					)     ,  return_dict=lowercase_     ,  )[0]

             _lowerCamelCase			:						Tuple            =image[0, -10:]
             _lowerCamelCase			:						Union[str, Any]            =image_from_tuple[0, -10:]

             assert image.shape == (1, 32)

             _lowerCamelCase			:						int            =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					)   ->		Union[str, Any]:
             """simple docstring"""




             _lowerCamelCase			:						Any            =torch_device == 'cpu'
             _lowerCamelCase			:						Dict            =True
             _lowerCamelCase			:						Dict            =False

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






      @skip_mps
      def 			lowerCamelCase			(   self		:	List[Any]					)   ->		int:
             """simple docstring"""




             _lowerCamelCase			:						Optional[int]            =torch_device == 'cpu'
             _lowerCamelCase			:						int            =False

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






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


lowerCamelCase		   =			logging.get_logger(__name__)

lowerCamelCase		   =			{
    'sayakpaul/vit-msn-base': 'https://huggingface.co/sayakpaul/vit-msn-base/resolve/main/config.json',
    # See all ViT MSN models at https://huggingface.co/models?filter=vit_msn
}




class A    (  UpperCamelCase_	):
      UpperCamelCase__				:    Union[str, Any] ='vit_msn'






      def __init__(   self		:	Union[str, Any]     ,  lowercase_		:	List[str]=768     ,  lowercase_		:	Optional[int]=12     ,  lowercase_		:	List[str]=12     ,  lowercase_		:	List[Any]=3072     ,  lowercase_		:	int="gelu"     ,  lowercase_		:	List[Any]=0.0     ,  lowercase_		:	Optional[int]=0.0     ,  lowercase_		:	Dict=0.02     ,  lowercase_		:	Any=1E-06     ,  lowercase_		:	Union[str, Any]=224     ,  lowercase_		:	Optional[int]=16     ,  lowercase_		:	List[Any]=3     ,  lowercase_		:	Any=True     ,  **lowercase_		:	Tuple     ,  )   ->		Dict:
             """simple docstring"""




             super().__init__(**lowercase_					)

             _lowerCamelCase			:						Dict            =hidden_size
             _lowerCamelCase			:						Any            =num_hidden_layers
             _lowerCamelCase			:						Union[str, Any]            =num_attention_heads
             _lowerCamelCase			:						List[Any]            =intermediate_size
             _lowerCamelCase			:						Tuple            =hidden_act
             _lowerCamelCase			:						Any            =hidden_dropout_prob
             _lowerCamelCase			:						Dict            =attention_probs_dropout_prob
             _lowerCamelCase			:						Any            =initializer_range
             _lowerCamelCase			:						List[Any]            =layer_norm_eps
             _lowerCamelCase			:						Optional[int]            =image_size
             _lowerCamelCase			:						Optional[int]            =patch_size
             _lowerCamelCase			:						List[Any]            =num_channels
             _lowerCamelCase			:						Optional[Any]            =qkv_bias


