
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 				UpperCamelCase__		(	__SCREAMING_SNAKE_CASE       ,    unittest.TestCase	):






       """simple docstring"""


       UpperCAmelCase_				       =KandinskyPriorPipeline
       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     _UpperCamelCase	(							self	)				->				Union[str, Any]:
          return 32



       @property
       def     _UpperCamelCase	(							self	)				->				List[str]:
          return 32



       @property
       def     _UpperCamelCase	(							self	)				->				Tuple:
          return self.time_input_dim



       @property
       def     _UpperCamelCase	(							self	)				->				str:
          return self.time_input_dim * 4



       @property
       def     _UpperCamelCase	(							self	)				->				str:
          return 100



       @property
       def     _UpperCamelCase	(							self	)				->				List[str]:
          SCREAMING_SNAKE_CASE_	=      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''	)
          return tokenizer



       @property
       def     _UpperCamelCase	(							self	)				->				Any:
          torch.manual_seed(0	)
          SCREAMING_SNAKE_CASE_	=      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     _UpperCamelCase	(							self	)				->				List[Any]:
          torch.manual_seed(0	)

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

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



       @property
       def     _UpperCamelCase	(							self	)				->				Union[str, Any]:
          torch.manual_seed(0	)
          SCREAMING_SNAKE_CASE_	=      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 ,     )

          SCREAMING_SNAKE_CASE_	=      CLIPVisionModelWithProjection(_A	)
          return model



       @property
       def     _UpperCamelCase	(							self	)				->				List[str]:
          SCREAMING_SNAKE_CASE_	=      CLIPImageProcessor(
              crop_size=224 ,     do_center_crop=_A ,     do_normalize=_A ,     do_resize=_A ,     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=224 ,     )

          return image_processor



       def     _UpperCamelCase	(							self	)				->				Optional[Any]:
          SCREAMING_SNAKE_CASE_	=      self.dummy_prior
          SCREAMING_SNAKE_CASE_	=      self.dummy_image_encoder
          SCREAMING_SNAKE_CASE_	=      self.dummy_text_encoder
          SCREAMING_SNAKE_CASE_	=      self.dummy_tokenizer
          SCREAMING_SNAKE_CASE_	=      self.dummy_image_processor

          SCREAMING_SNAKE_CASE_	=      UnCLIPScheduler(
              variance_type='''fixed_small_log''' ,     prediction_type='''sample''' ,     num_train_timesteps=1000 ,     clip_sample=_A ,     clip_sample_range=10.0 ,     )

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

          return components



       def     _UpperCamelCase	(							self ,     _A ,     _A=0	)				->				Any:
          if str(_A	).startswith('''mps'''	):
             SCREAMING_SNAKE_CASE_	=      torch.manual_seed(_A	)
          else:
             SCREAMING_SNAKE_CASE_	=      torch.Generator(device=_A	).manual_seed(_A	)
          SCREAMING_SNAKE_CASE_	=      {
              '''prompt''': '''horse''',
              '''generator''': generator,
              '''guidance_scale''': 4.0,
              '''num_inference_steps''': 2,
              '''output_type''': '''np''',
          }
          return inputs



       def     _UpperCamelCase	(							self	)				->				Dict:
          SCREAMING_SNAKE_CASE_	=      '''cpu'''

          SCREAMING_SNAKE_CASE_	=      self.get_dummy_components()

          SCREAMING_SNAKE_CASE_	=      self.pipeline_class(**_A	)
          SCREAMING_SNAKE_CASE_	=      pipe.to(_A	)

          pipe.set_progress_bar_config(disable=_A	)

          SCREAMING_SNAKE_CASE_	=      pipe(**self.get_dummy_inputs(_A	)	)
          SCREAMING_SNAKE_CASE_	=      output.image_embeds

          SCREAMING_SNAKE_CASE_	=      pipe(
              **self.get_dummy_inputs(_A	) ,     return_dict=_A ,     )[0]

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

          assert image.shape == (1, 32)

          SCREAMING_SNAKE_CASE_	=      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     _UpperCamelCase	(							self	)				->				List[str]:
          SCREAMING_SNAKE_CASE_	=      torch_device == '''cpu'''
          SCREAMING_SNAKE_CASE_	=      True
          SCREAMING_SNAKE_CASE_	=      False

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




       @skip_mps
       def     _UpperCamelCase	(							self	)				->				Any:
          SCREAMING_SNAKE_CASE_	=      torch_device == '''cpu'''
          SCREAMING_SNAKE_CASE_	=      False

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




import torch

from transformers import PreTrainedModel

from .custom_configuration import CustomConfig, NoSuperInitConfig





class 				UpperCamelCase__		(	__SCREAMING_SNAKE_CASE	):






       """simple docstring"""


       UpperCAmelCase_				       =CustomConfig



       def __init__(							self ,     _A	)				->				int:
          super().__init__(_A	)
          SCREAMING_SNAKE_CASE_	=      torch.nn.Linear(config.hidden_size ,     config.hidden_size	)



       def     _UpperCamelCase	(							self ,     _A	)				->				Optional[int]:
          return self.linear(_A	)




       def     _UpperCamelCase	(							self ,     _A	)				->				int:
          pass







class 				UpperCamelCase__		(	__SCREAMING_SNAKE_CASE	):






       """simple docstring"""


       UpperCAmelCase_				       =NoSuperInitConfig



       def __init__(							self ,     _A	)				->				Optional[int]:
          super().__init__(_A	)
          SCREAMING_SNAKE_CASE_	=      torch.nn.Linear(config.attribute ,     config.attribute	)



       def     _UpperCamelCase	(							self ,     _A	)				->				Any:
          return self.linear(_A	)




       def     _UpperCamelCase	(							self ,     _A	)				->				Tuple:
          pass


