





import unittest

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

from diffusers import KandinskyPriorPipeline, 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						(					A__ ,     unittest.TestCase ):
 _lowerCAmelCase : Union[str, Any] 				=				KandinskyPriorPipeline
 _lowerCAmelCase : Any 				=				['''prompt''']
 _lowerCAmelCase : Union[str, Any] 				=				['''prompt''', '''negative_prompt''']
 _lowerCAmelCase : Optional[int] 				=				[
     '''num_images_per_prompt''',
     '''generator''',
     '''num_inference_steps''',
     '''latents''',
     '''negative_prompt''',
     '''guidance_scale''',
     '''output_type''',
     '''return_dict''',
 ]
 _lowerCAmelCase : List[Any] 				=				False







 @property
 def       __lowercase      (	self      :     int		):


      '''simple docstring'''






      return 32







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


      '''simple docstring'''






      return 32







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


      '''simple docstring'''






      return self.time_input_dim







 @property
 def       __lowercase      (	self      :     Dict		):


      '''simple docstring'''






      return self.time_input_dim * 4







 @property
 def       __lowercase      (	self      :     Any		):


      '''simple docstring'''






      return 1_00







 @property
 def       __lowercase      (	self      :     int		):


      '''simple docstring'''






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







 @property
 def       __lowercase      (	self      :     str		):


      '''simple docstring'''






      torch.manual_seed(0		)
      UpperCamelCase__    :     Optional[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-0_5							,     num_attention_heads=4							,     num_hidden_layers=5							,     pad_token_id=1							,     vocab_size=10_00							,     )
      return CLIPTextModelWithProjection(SCREAMING_SNAKE_CASE		)







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


      '''simple docstring'''






      torch.manual_seed(0		)

      UpperCamelCase__    :     Optional[int]        =						{
          "num_attention_heads": 2,
          "attention_head_dim": 12,
          "embedding_dim": self.text_embedder_hidden_size,
          "num_layers": 1,
      }

      UpperCamelCase__    :     Tuple        =						PriorTransformer(**SCREAMING_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__    :     Union[str, Any]        =						nn.Parameter(torch.ones(model.clip_std.shape		)		)
      return model







 @property
 def       __lowercase      (	self      :     int		):


      '''simple docstring'''






      torch.manual_seed(0		)
      UpperCamelCase__    :     Tuple        =						CLIPVisionConfig(
          hidden_size=self.text_embedder_hidden_size							,     image_size=2_24							,     projection_dim=self.text_embedder_hidden_size							,     intermediate_size=37							,     num_attention_heads=4							,     num_channels=3							,     num_hidden_layers=5							,     patch_size=14							,     )

      UpperCamelCase__    :     int        =						CLIPVisionModelWithProjection(SCREAMING_SNAKE_CASE		)
      return model







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


      '''simple docstring'''






      UpperCamelCase__    :     str        =						CLIPImageProcessor(
          crop_size=2_24							,     do_center_crop=SCREAMING_SNAKE_CASE							,     do_normalize=SCREAMING_SNAKE_CASE							,     do_resize=SCREAMING_SNAKE_CASE							,     image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3]							,     image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1]							,     resample=3							,     size=2_24							,     )

      return image_processor







 def       __lowercase      (	self      :     Optional[Any]		):


      '''simple docstring'''






      UpperCamelCase__    :     Union[str, Any]        =						self.dummy_prior
      UpperCamelCase__    :     Optional[Any]        =						self.dummy_image_encoder
      UpperCamelCase__    :     List[Any]        =						self.dummy_text_encoder
      UpperCamelCase__    :     Tuple        =						self.dummy_tokenizer
      UpperCamelCase__    :     Optional[int]        =						self.dummy_image_processor

      UpperCamelCase__    :     List[str]        =						UnCLIPScheduler(
          variance_type="fixed_small_log"							,     prediction_type="sample"							,     num_train_timesteps=10_00							,     clip_sample=SCREAMING_SNAKE_CASE							,     clip_sample_range=10.0							,     )

      UpperCamelCase__    :     Union[str, Any]        =						{
          "prior": prior,
          "image_encoder": image_encoder,
          "text_encoder": text_encoder,
          "tokenizer": tokenizer,
          "scheduler": scheduler,
          "image_processor": image_processor,
      }

      return components







 def       __lowercase      (	self      :     Optional[int]							,     SCREAMING_SNAKE_CASE      :     int							,     SCREAMING_SNAKE_CASE      :     str=0		):


      '''simple docstring'''






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







 def       __lowercase      (	self      :     Optional[int]		):


      '''simple docstring'''






      UpperCamelCase__    :     int        =						"cpu"

      UpperCamelCase__    :     Optional[Any]        =						self.get_dummy_components()

      UpperCamelCase__    :     Optional[int]        =						self.pipeline_class(**SCREAMING_SNAKE_CASE		)
      UpperCamelCase__    :     Dict        =						pipe.to(SCREAMING_SNAKE_CASE		)

      pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE		)

      UpperCamelCase__    :     Tuple        =						pipe(**self.get_dummy_inputs(SCREAMING_SNAKE_CASE		)		)
      UpperCamelCase__    :     str        =						output.image_embeds

      UpperCamelCase__    :     str        =						pipe(
          **self.get_dummy_inputs(SCREAMING_SNAKE_CASE		)							,     return_dict=SCREAMING_SNAKE_CASE							,     )[0]

      UpperCamelCase__    :     int        =						image[0, -10:]
      UpperCamelCase__    :     str        =						image_from_tuple[0, -10:]

      assert image.shape == (1, 32)

      UpperCamelCase__    :     Dict        =						np.array(
          [-0.0_5_3_2, 1.7_1_2_0, 0.3_6_5_6, -1.0_8_5_2, -0.8_9_4_6, -1.1_7_5_6, 0.4_3_4_8, 0.2_4_8_2, 0.5_1_4_6, -0.1_1_5_6]		)

      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       __lowercase      (	self      :     Optional[int]		):


      '''simple docstring'''






      UpperCamelCase__    :     Optional[int]        =						torch_device == "cpu"
      UpperCamelCase__    :     List[str]        =						True
      UpperCamelCase__    :     Tuple        =						False

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







 @skip_mps
 def       __lowercase      (	self      :     Union[str, Any]		):


      '''simple docstring'''






      UpperCamelCase__    :     int        =						torch_device == "cpu"
      UpperCamelCase__    :     List[str]        =						False

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







from __future__ import annotations







def 	SCREAMING_SNAKE_CASE							(      __lowerCAmelCase    )				->			list[int]:
     assert (
         isinstance(__lowerCAmelCase      ,					__lowerCAmelCase    ) and mask > 0
     ), f'mask needs to be positive integer, your input {mask}'
     UpperCamelCase__    :     Any        =						[]
     UpperCamelCase__    :     int        =						mask

     while submask:
          all_submasks.append(__lowerCAmelCase    )
          UpperCamelCase__    :     List[str]        =						(submask - 1) & mask

     return all_submasks


if __name__ == "__main__":
 import doctest

 doctest.testmod()