





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 SCREAMING_SNAKE_CASE_    (       snake_case_							,					unittest.TestCase       ):
						__magic_name__:      List[Any]			   =							KandinskyPriorPipeline
						__magic_name__:      Optional[Any]			   =							["prompt"]
						__magic_name__:      Any			   =							["prompt", "negative_prompt"]
						__magic_name__:      str			   =							[
						    "num_images_per_prompt",
						    "generator",
						    "num_inference_steps",
						    "latents",
						    "negative_prompt",
						    "guidance_scale",
						    "output_type",
						    "return_dict",
						]
						__magic_name__:      Union[str, Any]			   =							False






						@property
						def UpperCAmelCase_					(   self   :   int						)				->     List[Any]:







							"""simple docstring"""


							return 32






						@property
						def UpperCAmelCase_					(   self   :   Union[str, Any]						)				->     Any:







							"""simple docstring"""


							return 32






						@property
						def UpperCAmelCase_					(   self   :   int						)				->     Optional[Any]:







							"""simple docstring"""


							return self.time_input_dim






						@property
						def UpperCAmelCase_					(   self   :   List[Any]						)				->     Tuple:







							"""simple docstring"""


							return self.time_input_dim * 4






						@property
						def UpperCAmelCase_					(   self   :   Optional[Any]						)				->     Dict:







							"""simple docstring"""


							return 100






						@property
						def UpperCAmelCase_					(   self   :   Optional[int]						)				->     Any:







							"""simple docstring"""


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






						@property
						def UpperCAmelCase_					(   self   :   Union[str, Any]						)				->     int:







							"""simple docstring"""


							torch.manual_seed(0						)
							snake_case_		:      List[Any]           =						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]						)				->     Any:







							"""simple docstring"""


							torch.manual_seed(0						)

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

							snake_case_		:      Optional[Any]           =						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
							snake_case_		:      Union[str, Any]           =						nn.Parameter(torch.ones(model.clip_std.shape						)						)
							return model






						@property
						def UpperCAmelCase_					(   self   :   str						)				->     Union[str, Any]:







							"""simple docstring"""


							torch.manual_seed(0						)
							snake_case_		:      Optional[Any]           =						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		, )

							snake_case_		:      Tuple           =						CLIPVisionModelWithProjection(_A						)
							return model






						@property
						def UpperCAmelCase_					(   self   :   Union[str, Any]						)				->     List[str]:







							"""simple docstring"""


							snake_case_		:      Optional[Any]           =						CLIPImageProcessor(
							    crop_size=224		, do_center_crop=_A		, do_normalize=_A		, do_resize=_A		, 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=224		, )

							return image_processor






						def UpperCAmelCase_					(   self   :   Any						)				->     int:







							"""simple docstring"""


							snake_case_		:      Tuple           =						self.dummy_prior
							snake_case_		:      Optional[int]           =						self.dummy_image_encoder
							snake_case_		:      List[str]           =						self.dummy_text_encoder
							snake_case_		:      List[str]           =						self.dummy_tokenizer
							snake_case_		:      List[str]           =						self.dummy_image_processor

							snake_case_		:      int           =						UnCLIPScheduler(
							    variance_type='fixed_small_log'		, prediction_type='sample'		, num_train_timesteps=1000		, clip_sample=_A		, clip_sample_range=1_0.0		, )

							snake_case_		:      List[str]           =						{
							    'prior': prior,
							    'image_encoder': image_encoder,
							    'text_encoder': text_encoder,
							    'tokenizer': tokenizer,
							    'scheduler': scheduler,
							    'image_processor': image_processor,
							}

							return components






						def UpperCAmelCase_					(   self   :   List[Any]		, _A   :   Optional[int]		, _A   :   Optional[Any]=0						)				->     List[str]:







							"""simple docstring"""


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






						def UpperCAmelCase_					(   self   :   Optional[Any]						)				->     str:







							"""simple docstring"""


							snake_case_		:      Optional[Any]           =						'cpu'

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

							snake_case_		:      Union[str, Any]           =						self.pipeline_class(**_A						)
							snake_case_		:      str           =						pipe.to(_A						)

							pipe.set_progress_bar_config(disable=_A						)

							snake_case_		:      List[Any]           =						pipe(**self.get_dummy_inputs(_A						)						)
							snake_case_		:      List[str]           =						output.image_embeds

							snake_case_		:      Optional[int]           =						pipe(
							    **self.get_dummy_inputs(_A						)		, return_dict=_A		, )[0]

							snake_case_		:      Tuple           =						image[0, -10:]
							snake_case_		:      Union[str, Any]           =						image_from_tuple[0, -10:]

							assert image.shape == (1, 32)

							snake_case_		:      Optional[Any]           =						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 UpperCAmelCase_					(   self   :   List[str]						)				->     Optional[int]:







							"""simple docstring"""


							snake_case_		:      Union[str, Any]           =						torch_device == 'cpu'
							snake_case_		:      Tuple           =						True
							snake_case_		:      str           =						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   :   List[str]						)				->     Dict:







							"""simple docstring"""


							snake_case_		:      List[str]           =						torch_device == 'cpu'
							snake_case_		:      Tuple           =						False

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









'''simple docstring'''







import torch

from diffusers import KDPMaAncestralDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest






class        a__     (      lowerCamelCase    ):
     '''simple docstring'''







     _snake_case							   =	(KDPMaAncestralDiscreteScheduler,)
     _snake_case							   =	10






     def      _UpperCamelCase (self     :						int					,					**lowercase__     :						Optional[int]       )	->							Dict:
           __A		 =   {
               "num_train_timesteps": 1100,
               "beta_start": 0.00_01,
               "beta_end": 0.02,
               "beta_schedule": "linear",
           }

           config.update(**lowercase__       )
           return config






     def      _UpperCamelCase (self     :						str       )	->							Optional[Any]:
           for timesteps in [10, 50, 100, 1000]:
                 self.check_over_configs(num_train_timesteps=lowercase__       )






     def      _UpperCamelCase (self     :						List[Any]       )	->							int:
           for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01]					,					[0.00_02, 0.0_02, 0.02]       ):
                 self.check_over_configs(beta_start=lowercase__					,					beta_end=lowercase__       )






     def      _UpperCamelCase (self     :						int       )	->							Tuple:
           for schedule in ["linear", "scaled_linear"]:
                 self.check_over_configs(beta_schedule=lowercase__       )






     def      _UpperCamelCase (self     :						int       )	->							Any:
           if torch_device == "mps":
                 return
           __A		 =   self.scheduler_classes[0]
           __A		 =   self.get_scheduler_config()
           __A		 =   scheduler_class(**lowercase__       )

           scheduler.set_timesteps(self.num_inference_steps       )

           __A		 =   torch.manual_seed(0       )

           __A		 =   self.dummy_model()
           __A		 =   self.dummy_sample_deter * scheduler.init_noise_sigma
           __A		 =   sample.to(lowercase__       )

           for i, t in enumerate(scheduler.timesteps       ):
                 __A		 =   scheduler.scale_model_input(lowercase__					,					lowercase__       )

                 __A		 =   model(lowercase__					,					lowercase__       )

                 __A		 =   scheduler.step(lowercase__					,					lowercase__					,					lowercase__					,					generator=lowercase__       )
                 __A		 =   output.prev_sample

           __A		 =   torch.sum(torch.abs(lowercase__       )       )
           __A		 =   torch.mean(torch.abs(lowercase__       )       )

           assert abs(result_sum.item() - 1_38_49.38_77       ) < 1e-2
           assert abs(result_mean.item() - 18.03_31       ) < 5e-3






     def      _UpperCamelCase (self     :						Optional[Any]       )	->							str:
           for prediction_type in ["epsilon", "v_prediction"]:
                 self.check_over_configs(prediction_type=lowercase__       )






     def      _UpperCamelCase (self     :						List[Any]       )	->							List[Any]:
           if torch_device == "mps":
                 return
           __A		 =   self.scheduler_classes[0]
           __A		 =   self.get_scheduler_config(prediction_type="v_prediction"       )
           __A		 =   scheduler_class(**lowercase__       )

           scheduler.set_timesteps(self.num_inference_steps       )

           __A		 =   self.dummy_model()
           __A		 =   self.dummy_sample_deter * scheduler.init_noise_sigma
           __A		 =   sample.to(lowercase__       )

           __A		 =   torch.manual_seed(0       )

           for i, t in enumerate(scheduler.timesteps       ):
                 __A		 =   scheduler.scale_model_input(lowercase__					,					lowercase__       )

                 __A		 =   model(lowercase__					,					lowercase__       )

                 __A		 =   scheduler.step(lowercase__					,					lowercase__					,					lowercase__					,					generator=lowercase__       )
                 __A		 =   output.prev_sample

           __A		 =   torch.sum(torch.abs(lowercase__       )       )
           __A		 =   torch.mean(torch.abs(lowercase__       )       )

           assert abs(result_sum.item() - 3_28.99_70       ) < 1e-2
           assert abs(result_mean.item() - 0.42_84       ) < 1e-3






     def      _UpperCamelCase (self     :						Optional[Any]       )	->							int:
           if torch_device == "mps":
                 return
           __A		 =   self.scheduler_classes[0]
           __A		 =   self.get_scheduler_config()
           __A		 =   scheduler_class(**lowercase__       )

           scheduler.set_timesteps(self.num_inference_steps					,					device=lowercase__       )
           __A		 =   torch.manual_seed(0       )

           __A		 =   self.dummy_model()
           __A		 =   self.dummy_sample_deter.to(lowercase__       ) * scheduler.init_noise_sigma

           for t in scheduler.timesteps:
                 __A		 =   scheduler.scale_model_input(lowercase__					,					lowercase__       )

                 __A		 =   model(lowercase__					,					lowercase__       )

                 __A		 =   scheduler.step(lowercase__					,					lowercase__					,					lowercase__					,					generator=lowercase__       )
                 __A		 =   output.prev_sample

           __A		 =   torch.sum(torch.abs(lowercase__       )       )
           __A		 =   torch.mean(torch.abs(lowercase__       )       )

           assert abs(result_sum.item() - 1_38_49.38_18       ) < 1e-1
           assert abs(result_mean.item() - 18.03_31       ) < 1e-3


