


'''simple docstring'''

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   __lowerCAmelCase						(lowercase_  ,					unittest.TestCase					):





						'''simple docstring'''






						lowerCAmelCase__ :					Any   =							KandinskyPriorPipeline
						lowerCAmelCase__ :					str   =							["""prompt"""]
						lowerCAmelCase__ :					Optional[int]   =							["""prompt""", """negative_prompt"""]
						lowerCAmelCase__ :					Optional[int]   =							[
						    """num_images_per_prompt""",
						    """generator""",
						    """num_inference_steps""",
						    """latents""",
						    """negative_prompt""",
						    """guidance_scale""",
						    """output_type""",
						    """return_dict""",
						]
						lowerCAmelCase__ :					Tuple   =							False



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

													'''simple docstring'''







													return 32



						@property
						def 				UpperCamelCase__		(self   :   int			):

													'''simple docstring'''







													return 32



						@property
						def 				UpperCamelCase__		(self   :   str			):

													'''simple docstring'''







													return self.time_input_dim



						@property
						def 				UpperCamelCase__		(self   :   Any			):

													'''simple docstring'''







													return self.time_input_dim * 4



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

													'''simple docstring'''







													return 100



						@property
						def 				UpperCamelCase__		(self   :   int			):

													'''simple docstring'''







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



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

													'''simple docstring'''







													torch.manual_seed(0			)
													lowercase__				      =			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(UpperCamelCase			)



						@property
						def 				UpperCamelCase__		(self   :   Any			):

													'''simple docstring'''







													torch.manual_seed(0			)

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

													lowercase__				      =			PriorTransformer(**UpperCamelCase			)
													# 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
													lowercase__				      =			nn.Parameter(torch.ones(model.clip_std.shape			)			)
													return model



						@property
						def 				UpperCamelCase__		(self   :   Union[str, Any]			):

													'''simple docstring'''







													torch.manual_seed(0			)
													lowercase__				      =			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       ,						)

													lowercase__				      =			CLIPVisionModelWithProjection(UpperCamelCase			)
													return model



						@property
						def 				UpperCamelCase__		(self   :   Union[str, Any]			):

													'''simple docstring'''







													lowercase__				      =			CLIPImageProcessor(
													    crop_size=224       ,						do_center_crop=UpperCamelCase       ,						do_normalize=UpperCamelCase       ,						do_resize=UpperCamelCase       ,						image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]       ,						image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]       ,						resample=3       ,						size=224       ,						)

													return image_processor



						def 				UpperCamelCase__		(self   :   str			):

													'''simple docstring'''







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

													lowercase__				      =			UnCLIPScheduler(
													    variance_type='''fixed_small_log'''       ,						prediction_type='''sample'''       ,						num_train_timesteps=1000       ,						clip_sample=UpperCamelCase       ,						clip_sample_range=10.0       ,						)

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

													return components



						def 				UpperCamelCase__		(self   :   Union[str, Any]       ,						UpperCamelCase   :   Tuple       ,						UpperCamelCase   :   Optional[int]=0			):

													'''simple docstring'''







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



						def 				UpperCamelCase__		(self   :   Optional[int]			):

													'''simple docstring'''







													lowercase__				      =			'''cpu'''

													lowercase__				      =			self.get_dummy_components()

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

													pipe.set_progress_bar_config(disable=UpperCamelCase			)

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

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

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

													assert image.shape == (1, 32)

													lowercase__				      =			np.array(
													    [-0.05_32, 1.71_20, 0.36_56, -1.08_52, -0.89_46, -1.17_56, 0.43_48, 0.24_82, 0.51_46, -0.11_56]			)

													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   :   Union[str, Any]			):

													'''simple docstring'''







													lowercase__				      =			torch_device == '''cpu'''
													lowercase__				      =			True
													lowercase__				      =			False

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






						@skip_mps
						def 				UpperCamelCase__		(self   :   Dict			):

													'''simple docstring'''







													lowercase__				      =			torch_device == '''cpu'''
													lowercase__				      =			False

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








'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCamelCase							:       str  						=     {'configuration_nat': ['NAT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'NatConfig']}


try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase							:       Tuple  						=     [
     'NAT_PRETRAINED_MODEL_ARCHIVE_LIST',
     'NatForImageClassification',
     'NatModel',
     'NatPreTrainedModel',
     'NatBackbone',
 ]

if TYPE_CHECKING:
 from .configuration_nat import NAT_PRETRAINED_CONFIG_ARCHIVE_MAP, NatConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_nat import (
      NAT_PRETRAINED_MODEL_ARCHIVE_LIST,
      NatBackbone,
      NatForImageClassification,
      NatModel,
      NatPreTrainedModel,
  )

else:
 import sys

 lowerCamelCase							:       Optional[Any]  						=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




