






"""simple docstring"""

import math






def _lowerCAmelCase	(  __lowerCamelCase:int   ):

						'''simple docstring'''






						assert isinstance(__snake_case ,						__snake_case   ) and (
						    number >= 0
						), "'number' must been an int and positive"

						if 1 < number < 4:
												# 2 and 3 are primes
												return True
						elif number < 2 or not number % 2:
												# Negatives, 0, 1 and all even numbers are not primes
												return False

						__magic_name__            =     range(3 ,						int(math.sqrt(__snake_case   ) + 1   ) ,						2   )
						return not any(not number % i for i in odd_numbers   )






def _lowerCAmelCase	(  __lowerCamelCase:str ,						__lowerCamelCase:Tuple=1 ,						**__lowerCamelCase:Optional[Any]   ):

						'''simple docstring'''






						__magic_name__            =     factor * value
						__magic_name__            =     value

						while not is_prime(__snake_case   ):
												value += 1 if not ("desc" in kwargs and kwargs["desc"] is True) else -1

						if value == first_value_val:
												return next_prime(value + 1 ,						**__snake_case   )
						return value











"""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 KandinskyVaaPriorPipeline, 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_							(      snake_case_					,							unittest.TestCase  ):
				UpperCAmelCase__        =				KandinskyVaaPriorPipeline
				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     _snake_case					(							self							:						Dict	)				->					Dict:
										return 3_2







				@property
				def     _snake_case					(							self							:						Tuple	)				->					Tuple:
										return 3_2







				@property
				def     _snake_case					(							self							:						Dict	)				->					Optional[int]:
										return self.time_input_dim







				@property
				def     _snake_case					(							self							:						List[Any]	)				->					str:
										return self.time_input_dim * 4







				@property
				def     _snake_case					(							self							:						List[str]	)				->					List[Any]:
										return 1_0_0







				@property
				def     _snake_case					(							self							:						List[str]	)				->					Tuple:
										__magic_name__            =     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"	)
										return tokenizer







				@property
				def     _snake_case					(							self							:						Optional[int]	)				->					Optional[int]:
										torch.manual_seed(0	)
										__magic_name__            =     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(__lowerCamelCase	)







				@property
				def     _snake_case					(							self							:						List[str]	)				->					Union[str, Any]:
										torch.manual_seed(0	)

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

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







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

										__magic_name__            =     CLIPVisionModelWithProjection(__lowerCamelCase	)
										return model







				@property
				def     _snake_case					(							self							:						Union[str, Any]	)				->					int:
										__magic_name__            =     CLIPImageProcessor(
										    crop_size=2_2_4						,    do_center_crop=__lowerCamelCase						,    do_normalize=__lowerCamelCase						,    do_resize=__lowerCamelCase						,    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     _snake_case					(							self							:						List[Any]	)				->					Optional[Any]:
										__magic_name__            =     self.dummy_prior
										__magic_name__            =     self.dummy_image_encoder
										__magic_name__            =     self.dummy_text_encoder
										__magic_name__            =     self.dummy_tokenizer
										__magic_name__            =     self.dummy_image_processor

										__magic_name__            =     UnCLIPScheduler(
										    variance_type="fixed_small_log"						,    prediction_type="sample"						,    num_train_timesteps=1_0_0_0						,    clip_sample=__lowerCamelCase						,    clip_sample_range=10.0						,    )

										__magic_name__            =     {
										    "prior": prior,
										    "image_encoder": image_encoder,
										    "text_encoder": text_encoder,
										    "tokenizer": tokenizer,
										    "scheduler": scheduler,
										    "image_processor": image_processor,
										}

										return components







				def     _snake_case					(							self							:						int						,    __lowerCamelCase							:						List[Any]						,    __lowerCamelCase							:						str=0	)				->					Optional[int]:
										if str(__lowerCamelCase	).startswith("mps"	):
																__magic_name__            =     torch.manual_seed(__lowerCamelCase	)
										else:
																__magic_name__            =     torch.Generator(device=__lowerCamelCase	).manual_seed(__lowerCamelCase	)
										__magic_name__            =     {
										    "prompt": "horse",
										    "generator": generator,
										    "guidance_scale": 4.0,
										    "num_inference_steps": 2,
										    "output_type": "np",
										}
										return inputs







				def     _snake_case					(							self							:						Optional[Any]	)				->					Optional[Any]:
										__magic_name__            =     "cpu"

										__magic_name__            =     self.get_dummy_components()

										__magic_name__            =     self.pipeline_class(**__lowerCamelCase	)
										__magic_name__            =     pipe.to(__lowerCamelCase	)

										pipe.set_progress_bar_config(disable=__lowerCamelCase	)

										__magic_name__            =     pipe(**self.get_dummy_inputs(__lowerCamelCase	)	)
										__magic_name__            =     output.image_embeds

										__magic_name__            =     pipe(
										    **self.get_dummy_inputs(__lowerCamelCase	)						,    return_dict=__lowerCamelCase						,    )[0]

										__magic_name__            =     image[0, -1_0:]
										__magic_name__            =     image_from_tuple[0, -1_0:]

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

										__magic_name__            =     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     _snake_case					(							self							:						Dict	)				->					List[str]:
										__magic_name__            =     torch_device == "cpu"
										__magic_name__            =     True
										__magic_name__            =     False

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







				@skip_mps
				def     _snake_case					(							self							:						Tuple	)				->					Tuple:
										__magic_name__            =     torch_device == "cpu"
										__magic_name__            =     False

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



