






"""simple docstring"""



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        lowerCAmelCase__       (  __magic_name__				,		unittest.TestCase					):
   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 	__a				(						self	:     Tuple				):


        '''simple docstring'''
        return 3_2
   @property
   def 	__a				(						self	:     Tuple				):


        '''simple docstring'''
        return 3_2
   @property
   def 	__a				(						self	:     Any				):


        '''simple docstring'''
        return self.time_input_dim
   @property
   def 	__a				(						self	:     Dict				):


        '''simple docstring'''
        return self.time_input_dim * 4
   @property
   def 	__a				(						self	:     Any				):


        '''simple docstring'''
        return 1_0_0
   @property
   def 	__a				(						self	:     int				):


        '''simple docstring'''
        UpperCAmelCase__		:  Tuple				=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				)
        return tokenizer
   @property
   def 	__a				(						self	:     List[str]				):


        '''simple docstring'''
        torch.manual_seed(0				)
        UpperCAmelCase__		:  Dict				=				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=3_7				,       layer_norm_eps=1e-05				,       num_attention_heads=4				,       num_hidden_layers=5				,       pad_token_id=1				,       vocab_size=1_0_0_0				,       )
        return CLIPTextModelWithProjection(snake_case__				)
   @property
   def 	__a				(						self	:     Dict				):


        '''simple docstring'''
        torch.manual_seed(0				)

        UpperCAmelCase__		:  int				=				{
            "num_attention_heads": 2,
            "attention_head_dim": 1_2,
            "embedding_dim": self.text_embedder_hidden_size,
            "num_layers": 1,
        }

        UpperCAmelCase__		:  Optional[int]				=				PriorTransformer(**snake_case__				)
        # 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__		:  List[str]				=				nn.Parameter(torch.ones(model.clip_std.shape				)				)
        return model
   @property
   def 	__a				(						self	:     List[str]				):


        '''simple docstring'''
        torch.manual_seed(0				)
        UpperCAmelCase__		:  Optional[int]				=				CLIPVisionConfig(
            hidden_size=self.text_embedder_hidden_size				,       image_size=2_2_4				,       projection_dim=self.text_embedder_hidden_size				,       intermediate_size=3_7				,       num_attention_heads=4				,       num_channels=3				,       num_hidden_layers=5				,       patch_size=1_4				,       )

        UpperCAmelCase__		:  Any				=				CLIPVisionModelWithProjection(snake_case__				)
        return model
   @property
   def 	__a				(						self	:     Tuple				):


        '''simple docstring'''
        UpperCAmelCase__		:  List[str]				=				CLIPImageProcessor(
            crop_size=2_2_4				,       do_center_crop=snake_case__				,       do_normalize=snake_case__				,       do_resize=snake_case__				,       image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]				,       image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]				,       resample=3				,       size=2_2_4				,       )

        return image_processor
   def 	__a				(						self	:     Optional[Any]				):


        '''simple docstring'''
        UpperCAmelCase__		:  str				=				self.dummy_prior
        UpperCAmelCase__		:  List[Any]				=				self.dummy_image_encoder
        UpperCAmelCase__		:  List[str]				=				self.dummy_text_encoder
        UpperCAmelCase__		:  int				=				self.dummy_tokenizer
        UpperCAmelCase__		:  Union[str, Any]				=				self.dummy_image_processor

        UpperCAmelCase__		:  List[Any]				=				UnCLIPScheduler(
            variance_type="fixed_small_log"				,       prediction_type="sample"				,       num_train_timesteps=1_0_0_0				,       clip_sample=snake_case__				,       clip_sample_range=10.0				,       )

        UpperCAmelCase__		:  List[str]				=				{
            "prior": prior,
            "image_encoder": image_encoder,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "scheduler": scheduler,
            "image_processor": image_processor,
        }

        return components
   def 	__a				(						self	:     str				,       snake_case__	:     List[str]				,       snake_case__	:     Tuple=0				):


        '''simple docstring'''
        if str(snake_case__				).startswith("mps"				):
             UpperCAmelCase__		:  List[Any]				=				torch.manual_seed(snake_case__				)
        else:
             UpperCAmelCase__		:  int				=				torch.Generator(device=snake_case__				).manual_seed(snake_case__				)

        UpperCAmelCase__		:  Any				=				floats_tensor((1, 3, 6_4, 6_4)				,       rng=random.Random(snake_case__				)				).to(snake_case__				)
        UpperCAmelCase__		:  Any				=				image.cpu().permute(0				,       2				,       3				,       1				)[0]
        UpperCAmelCase__		:  Any				=				Image.fromarray(np.uinta(snake_case__				)				).convert("RGB"				).resize((2_5_6, 2_5_6)				)

        UpperCAmelCase__		:  Optional[int]				=				{
            "prompt": "horse",
            "image": init_image,
            "strength": 0.5,
            "generator": generator,
            "guidance_scale": 4.0,
            "num_inference_steps": 2,
            "output_type": "np",
        }
        return inputs
   def 	__a				(						self	:     Optional[Any]				):


        '''simple docstring'''
        UpperCAmelCase__		:  str				=				"cpu"

        UpperCAmelCase__		:  Union[str, Any]				=				self.get_dummy_components()

        UpperCAmelCase__		:  List[str]				=				self.pipeline_class(**snake_case__				)
        UpperCAmelCase__		:  int				=				pipe.to(snake_case__				)

        pipe.set_progress_bar_config(disable=snake_case__				)

        UpperCAmelCase__		:  Optional[int]				=				pipe(**self.get_dummy_inputs(snake_case__				)				)
        UpperCAmelCase__		:  Union[str, Any]				=				output.image_embeds

        UpperCAmelCase__		:  str				=				pipe(
            **self.get_dummy_inputs(snake_case__				)				,       return_dict=snake_case__				,       )[0]

        UpperCAmelCase__		:  Tuple				=				image[0, -1_0:]
        UpperCAmelCase__		:  Any				=				image_from_tuple[0, -1_0:]

        assert image.shape == (1, 3_2)

        UpperCAmelCase__		:  Optional[int]				=				np.array(
            [
                0.107_1284,
                1.333_0271,
                0.6126_0223,
                -0.669_1065,
                -0.384_6852,
                -1.030_3661,
                0.2271_6111,
                0.0334_8901,
                0.3004_0675,
                -0.2480_5029,
            ]				)

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


        '''simple docstring'''
        UpperCAmelCase__		:  List[str]				=				torch_device == "cpu"
        UpperCAmelCase__		:  int				=				True
        UpperCAmelCase__		:  Dict				=				False

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




   @skip_mps
   def 	__a				(						self	:     Any				):


        '''simple docstring'''
        UpperCAmelCase__		:  Any				=				torch_device == "cpu"
        UpperCAmelCase__		:  Union[str, Any]				=				False

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





'''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"]			)

