




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  ):



 '''simple docstring'''


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

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

       """simple docstring"""






       return 32

 @property
 def 							a__					(							self   :    str    )  ->      Any:

       """simple docstring"""






       return 32

 @property
 def 							a__					(							self   :    Tuple    )  ->      int:

       """simple docstring"""






       return self.time_input_dim

 @property
 def 							a__					(							self   :    str    )  ->      str:

       """simple docstring"""






       return self.time_input_dim * 4

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

       """simple docstring"""






       return 100

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

       """simple docstring"""






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

 @property
 def 							a__					(							self   :    Any    )  ->      str:

       """simple docstring"""






       torch.manual_seed(0    )
       lowerCamelCase_							=       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 							a__					(							self   :    Optional[int]    )  ->      Any:

       """simple docstring"""






       torch.manual_seed(0    )

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

       lowerCamelCase_							=       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
       lowerCamelCase_							=       nn.Parameter(torch.ones(model.clip_std.shape    )    )
       return model

 @property
 def 							a__					(							self   :    int    )  ->      Tuple:

       """simple docstring"""






       torch.manual_seed(0    )
       lowerCamelCase_							=       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_							=       CLIPVisionModelWithProjection(A_    )
       return model

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

       """simple docstring"""






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

       return image_processor

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

       """simple docstring"""






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

       lowerCamelCase_							=       UnCLIPScheduler(
           variance_type='fixed_small_log'		,				prediction_type='sample'		,				num_train_timesteps=1000		,				clip_sample=A_		,				clip_sample_range=10.0		,				)

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

       return components

 def 							a__					(							self   :    Dict		,				A_   :    Optional[int]		,				A_   :    Union[str, Any]=0    )  ->      List[Any]:

       """simple docstring"""






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

 def 							a__					(							self   :    Dict    )  ->      Tuple:

       """simple docstring"""






       lowerCamelCase_							=       'cpu'

       lowerCamelCase_							=       self.get_dummy_components()

       lowerCamelCase_							=       self.pipeline_class(**A_    )
       lowerCamelCase_							=       pipe.to(A_    )

       pipe.set_progress_bar_config(disable=A_    )

       lowerCamelCase_							=       pipe(**self.get_dummy_inputs(A_    )    )
       lowerCamelCase_							=       output.image_embeds

       lowerCamelCase_							=       pipe(
           **self.get_dummy_inputs(A_    )		,				return_dict=A_		,				)[0]

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

       assert image.shape == (1, 32)

       lowerCamelCase_							=       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 							a__					(							self   :    Dict    )  ->      Any:

       """simple docstring"""






       lowerCamelCase_							=       torch_device == 'cpu'
       lowerCamelCase_							=       True
       lowerCamelCase_							=       False

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





 @skip_mps
 def 							a__					(							self   :    List[str]    )  ->      Optional[int]:

       """simple docstring"""






       lowerCamelCase_							=       torch_device == 'cpu'
       lowerCamelCase_							=       False

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








import numpy as np







def  lowerCAmelCase_	(lowerCAmelCase__:		np.ndarray		,						lowerCAmelCase__:		float  ):





			"""simple docstring"""


			return np.where(vector > 0		,						lowerCAmelCase__		,						(alpha * (np.exp(lowerCAmelCase__  ) - 1))  )


if __name__ == "__main__":
	import doctest

	doctest.testmod()






