






import random
import unittest

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

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

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class 					__lowercase					(		_UpperCAmelCase	,					unittest.TestCase     ):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE        = KandinskyVaaPriorEmbaEmbPipeline
     SCREAMING_SNAKE_CASE        = ["prompt", "image"]
     SCREAMING_SNAKE_CASE        = ["prompt", "image"]
     SCREAMING_SNAKE_CASE        = [
         "num_images_per_prompt",
         "strength",
         "generator",
         "num_inference_steps",
         "latents",
         "negative_prompt",
         "guidance_scale",
         "output_type",
         "return_dict",
     ]
     SCREAMING_SNAKE_CASE        = False




     @property
     def 						lowerCAmelCase_   ( self				:	List[str] ):


         """simple docstring"""



         return 32




     @property
     def 						lowerCAmelCase_   ( self				:	int ):


         """simple docstring"""



         return 32




     @property
     def 						lowerCAmelCase_   ( self				:	Tuple ):


         """simple docstring"""



         return self.time_input_dim




     @property
     def 						lowerCAmelCase_   ( self				:	Optional[Any] ):


         """simple docstring"""



         return self.time_input_dim * 4




     @property
     def 						lowerCAmelCase_   ( self				:	Any ):


         """simple docstring"""



         return 100




     @property
     def 						lowerCAmelCase_   ( self				:	Optional[int] ):


         """simple docstring"""



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




     @property
     def 						lowerCAmelCase_   ( self				:	List[Any] ):


         """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-05       ,		num_attention_heads=4       ,		num_hidden_layers=5       ,		pad_token_id=1       ,		vocab_size=1_000       ,		)
         return CLIPTextModelWithProjection(lowercase__ )




     @property
     def 						lowerCAmelCase_   ( self				:	Tuple ):


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




     @property
     def 						lowerCAmelCase_   ( self				:	List[Any] ):


         """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(lowercase__ )
         return model




     @property
     def 						lowerCAmelCase_   ( self				:	Any ):


         """simple docstring"""



         __A   			=		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				:	List[str] ):


         """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=lowercase__       ,		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				:	Dict       ,		UpperCamelCase_				:	Optional[int]       ,		UpperCamelCase_				:	List[str]=0 ):


         """simple docstring"""



         if str(lowercase__ ).startswith("""mps""" ):
             __A   			=		torch.manual_seed(lowercase__ )
         else:
             __A   			=		torch.Generator(device=lowercase__ ).manual_seed(lowercase__ )

         __A   			=		floats_tensor((1, 3, 64, 64)       ,		rng=random.Random(lowercase__ ) ).to(lowercase__ )
         __A   			=		image.cpu().permute(0       ,		2       ,		3       ,		1 )[0]
         __A   			=		Image.fromarray(np.uinta(lowercase__ ) ).convert("""RGB""" ).resize((256, 256) )

         __A   			=		{
             "prompt": "horse",
             "image": init_image,
             "strength": 0.5,
             "generator": generator,
             "guidance_scale": 4.0,
             "num_inference_steps": 2,
             "output_type": "np",
         }
         return inputs




     def 						lowerCAmelCase_   ( self				:	str ):


         """simple docstring"""



         __A   			=		"cpu"

         __A   			=		self.get_dummy_components()

         __A   			=		self.pipeline_class(**lowercase__ )
         __A   			=		pipe.to(lowercase__ )

         pipe.set_progress_bar_config(disable=lowercase__ )

         __A   			=		pipe(**self.get_dummy_inputs(lowercase__ ) )
         __A   			=		output.image_embeds

         __A   			=		pipe(
             **self.get_dummy_inputs(lowercase__ )       ,		return_dict=lowercase__       ,		)[0]

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

         assert image.shape == (1, 32)

         __A   			=		np.array(
             [
                 0.1071284,
                 1.3330271,
                 0.61260223,
                 -0.6691065,
                 -0.3846852,
                 -1.0303661,
                 0.22716111,
                 0.03348901,
                 0.30040675,
                 -0.24805029,
             ] )

         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				:	List[Any] ):


         """simple docstring"""



         __A   			=		torch_device == "cpu"
         __A   			=		True
         __A   			=		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				:	Union[str, Any] ):


         """simple docstring"""



         __A   			=		torch_device == "cpu"
         __A   			=		False

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







'''simple docstring'''






from ..utils import DummyObject, requires_backends
class 	SCREAMING_SNAKE_CASE__			(				metaclass=_UpperCAmelCase      ):
    _A			    =							["torch", "transformers", "invisible_watermark"]




    def __init__(					self				,     *lowercase__				,     **lowercase__			):






          """simple docstring"""






          requires_backends(self				,     ["torch", "transformers", "invisible_watermark"]			)




    @classmethod
    def   __lowerCamelCase      (					cls				,     *lowercase__				,     **lowercase__			):






          """simple docstring"""






          requires_backends(cls				,     ["torch", "transformers", "invisible_watermark"]			)







    @classmethod
    def   __lowerCamelCase      (					cls				,     *lowercase__				,     **lowercase__			):






          """simple docstring"""






          requires_backends(cls				,     ["torch", "transformers", "invisible_watermark"]			)






class 	SCREAMING_SNAKE_CASE__			(				metaclass=_UpperCAmelCase      ):
    _A			    =							["torch", "transformers", "invisible_watermark"]




    def __init__(					self				,     *lowercase__				,     **lowercase__			):






          """simple docstring"""






          requires_backends(self				,     ["torch", "transformers", "invisible_watermark"]			)




    @classmethod
    def   __lowerCamelCase      (					cls				,     *lowercase__				,     **lowercase__			):






          """simple docstring"""






          requires_backends(cls				,     ["torch", "transformers", "invisible_watermark"]			)







    @classmethod
    def   __lowerCamelCase      (					cls				,     *lowercase__				,     **lowercase__			):






          """simple docstring"""






          requires_backends(cls				,     ["torch", "transformers", "invisible_watermark"]			)

