



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 						lowercase    (							SCREAMING_SNAKE_CASE_      ,    unittest.TestCase):
						'''simple docstring'''


						UpperCAmelCase			:		Dict		=		KandinskyVaaPriorPipeline
						UpperCAmelCase			:		Optional[int]		=		['prompt']
						UpperCAmelCase			:		str		=		['prompt', 'negative_prompt']
						UpperCAmelCase			:		List[str]		=		[
						    'num_images_per_prompt',
						    'generator',
						    'num_inference_steps',
						    'latents',
						    'negative_prompt',
						    'guidance_scale',
						    'output_type',
						    'return_dict',
						]
						UpperCAmelCase			:		Any		=		False







						@property
						def       lowerCamelCase_				(  self    :				str	):



											'''simple docstring'''


											return 32







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



											'''simple docstring'''


											return 32







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



											'''simple docstring'''


											return self.time_input_dim







						@property
						def       lowerCamelCase_				(  self    :				Any	):



											'''simple docstring'''


											return self.time_input_dim * 4







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



											'''simple docstring'''


											return 100







						@property
						def       lowerCamelCase_				(  self    :				str	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)
											return tokenizer







						@property
						def       lowerCamelCase_				(  self    :				Dict	):



											'''simple docstring'''


											torch.manual_seed(0	)
											SCREAMING_SNAKE_CASE							:     Dict						=				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(snake_case	)







						@property
						def       lowerCamelCase_				(  self    :				Dict	):



											'''simple docstring'''


											torch.manual_seed(0	)

											SCREAMING_SNAKE_CASE							:     List[str]						=				{
											    'num_attention_heads': 2,
											    'attention_head_dim': 12,
											    'embedding_dim': self.text_embedder_hidden_size,
											    'num_layers': 1,
											}

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







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



											'''simple docstring'''


											torch.manual_seed(0	)
											SCREAMING_SNAKE_CASE							:     Union[str, 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       ,	)

											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				CLIPVisionModelWithProjection(snake_case	)
											return model







						@property
						def       lowerCamelCase_				(  self    :				str	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				CLIPImageProcessor(
											    crop_size=224       ,	do_center_crop=snake_case       ,	do_normalize=snake_case       ,	do_resize=snake_case       ,	image_mean=[0.48145466, 0.4578275, 0.40821073]       ,	image_std=[0.26862954, 0.26130258, 0.27577711]       ,	resample=3       ,	size=224       ,	)

											return image_processor







						def       lowerCamelCase_				(  self    :				Tuple	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     str						=				self.dummy_prior
											SCREAMING_SNAKE_CASE							:     Tuple						=				self.dummy_image_encoder
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				self.dummy_text_encoder
											SCREAMING_SNAKE_CASE							:     int						=				self.dummy_tokenizer
											SCREAMING_SNAKE_CASE							:     Any						=				self.dummy_image_processor

											SCREAMING_SNAKE_CASE							:     Optional[int]						=				UnCLIPScheduler(
											    variance_type='fixed_small_log'       ,	prediction_type='sample'       ,	num_train_timesteps=1000       ,	clip_sample=snake_case       ,	clip_sample_range=10.0       ,	)

											SCREAMING_SNAKE_CASE							:     Optional[int]						=				{
											    'prior': prior,
											    'image_encoder': image_encoder,
											    'text_encoder': text_encoder,
											    'tokenizer': tokenizer,
											    'scheduler': scheduler,
											    'image_processor': image_processor,
											}

											return components







						def       lowerCamelCase_				(  self    :				Any       ,	snake_case    :				Dict       ,	snake_case    :				Optional[Any]=0	):



											'''simple docstring'''


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







						def       lowerCamelCase_				(  self    :				List[str]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     List[str]						=				'cpu'

											SCREAMING_SNAKE_CASE							:     Optional[int]						=				self.get_dummy_components()

											SCREAMING_SNAKE_CASE							:     Optional[int]						=				self.pipeline_class(**snake_case	)
											SCREAMING_SNAKE_CASE							:     List[str]						=				pipe.to(snake_case	)

											pipe.set_progress_bar_config(disable=snake_case	)

											SCREAMING_SNAKE_CASE							:     str						=				pipe(**self.get_dummy_inputs(snake_case	)	)
											SCREAMING_SNAKE_CASE							:     Any						=				output.image_embeds

											SCREAMING_SNAKE_CASE							:     int						=				pipe(
											    **self.get_dummy_inputs(snake_case	)       ,	return_dict=snake_case       ,	)[0]

											SCREAMING_SNAKE_CASE							:     Any						=				image[0, -10:]
											SCREAMING_SNAKE_CASE							:     int						=				image_from_tuple[0, -10:]

											assert image.shape == (1, 32)

											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				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       lowerCamelCase_				(  self    :				Optional[Any]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Any						=				torch_device == 'cpu'
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				True
											SCREAMING_SNAKE_CASE							:     str						=				False

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







						@skip_mps
						def       lowerCamelCase_				(  self    :				Optional[Any]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				torch_device == 'cpu'
											SCREAMING_SNAKE_CASE							:     List[Any]						=				False

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




from typing import Callable, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowerCamelCase    :		int		      =    logging.get_logger(__name__)

_lowerCamelCase    :		Union[str, Any]		      =    {
    """microsoft/xprophetnet-large-wiki100-cased""": (
        """https://huggingface.co/microsoft/xprophetnet-large-wiki100-cased/resolve/main/config.json"""
    ),
}


class 						lowercase    (							SCREAMING_SNAKE_CASE_):
						'''simple docstring'''


						UpperCAmelCase			:		Tuple		=		'xlm-prophetnet'
						UpperCAmelCase			:		int		=		['past_key_values']
						UpperCAmelCase			:		List[Any]		=		{
						    'num_attention_heads': 'num_encoder_attention_heads',
						}







						def __init__(  self    :				List[Any]       ,	snake_case    :				Optional[float] = 0.1       ,	snake_case    :				Optional[Union[str, Callable]] = "gelu"       ,	snake_case    :				Optional[int] = 30522       ,	snake_case    :				Optional[int] = 1024       ,	snake_case    :				Optional[int] = 4096       ,	snake_case    :				Optional[int] = 12       ,	snake_case    :				Optional[int] = 16       ,	snake_case    :				Optional[int] = 4096       ,	snake_case    :				Optional[int] = 12       ,	snake_case    :				Optional[int] = 16       ,	snake_case    :				Optional[float] = 0.1       ,	snake_case    :				Optional[float] = 0.1       ,	snake_case    :				Optional[int] = 512       ,	snake_case    :				Optional[float] = 0.02       ,	snake_case    :				Optional[bool] = True       ,	snake_case    :				Optional[bool] = True       ,	snake_case    :				Optional[int] = 0       ,	snake_case    :				Optional[int] = 2       ,	snake_case    :				Optional[int] = 32       ,	snake_case    :				Optional[int] = 128       ,	snake_case    :				Optional[bool] = False       ,	snake_case    :				Optional[float] = 0.0       ,	snake_case    :				Optional[bool] = True       ,	snake_case    :				Optional[int] = 0       ,	snake_case    :				Optional[int] = 1       ,	snake_case    :				Optional[int] = 2       ,	**snake_case    :				Optional[int]       ,	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				vocab_size
											SCREAMING_SNAKE_CASE							:     int						=				hidden_size
											SCREAMING_SNAKE_CASE							:     Dict						=				encoder_ffn_dim
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				num_encoder_layers
											SCREAMING_SNAKE_CASE							:     str						=				num_encoder_attention_heads
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				decoder_ffn_dim
											SCREAMING_SNAKE_CASE							:     Any						=				num_decoder_layers
											SCREAMING_SNAKE_CASE							:     List[str]						=				num_decoder_attention_heads
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				max_position_embeddings
											SCREAMING_SNAKE_CASE							:     Any						=				init_std  # Normal(0, this parameter)
											SCREAMING_SNAKE_CASE							:     List[str]						=				activation_function

											# parameters for xlmprophetnet
											SCREAMING_SNAKE_CASE							:     Tuple						=				ngram
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				num_buckets
											SCREAMING_SNAKE_CASE							:     Any						=				relative_max_distance
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				disable_ngram_loss
											SCREAMING_SNAKE_CASE							:     List[str]						=				eps

											# 3 Types of Dropout
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				attention_dropout
											SCREAMING_SNAKE_CASE							:     Any						=				activation_dropout
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				dropout

											SCREAMING_SNAKE_CASE							:     Tuple						=				use_cache

											super().__init__(
											    pad_token_id=snake_case       ,	bos_token_id=snake_case       ,	eos_token_id=snake_case       ,	is_encoder_decoder=snake_case       ,	add_cross_attention=snake_case       ,	decoder_start_token_id=snake_case       ,	**snake_case       ,	)







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



											'''simple docstring'''


											return self.num_encoder_layers + self.num_decoder_layers







						@num_hidden_layers.setter
						def       lowerCamelCase_				(  self    :				List[Any]       ,	snake_case    :				Optional[int]	):



											'''simple docstring'''


											raise NotImplementedError(
											    'This model does not support the setting of `num_hidden_layers`. Please set `num_encoder_layers` and'
											    ' `num_decoder_layers`.'	)