






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__     (			_UpperCAmelCase,					unittest.TestCase	):
	a_         =KandinskyVaaPriorPipeline
	a_         =["""prompt"""]
	a_         =["""prompt""", """negative_prompt"""]
	a_         =[
	    """num_images_per_prompt""",
	    """generator""",
	    """num_inference_steps""",
	    """latents""",
	    """negative_prompt""",
	    """guidance_scale""",
	    """output_type""",
	    """return_dict""",
	]
	a_         =False

	@property
	def 		UpperCAmelCase  (							self       )->					Tuple:







							'''simple docstring'''





							return 32

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







							'''simple docstring'''





							return 32

	@property
	def 		UpperCAmelCase  (							self       )->					List[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       )->					int:







							'''simple docstring'''





							return 100

	@property
	def 		UpperCAmelCase  (							self       )->					List[str]:







							'''simple docstring'''





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

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







							'''simple docstring'''





							torch.manual_seed(0       )
							lowerCAmelCase__       =						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-0_5		,    num_attention_heads=4		,    num_hidden_layers=5		,    pad_token_id=1		,    vocab_size=1000		,    )
							return CLIPTextModelWithProjection(__UpperCAmelCase       )

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







							'''simple docstring'''





							torch.manual_seed(0       )

							lowerCAmelCase__       =						{
							    "num_attention_heads": 2,
							    "attention_head_dim": 12,
							    "embedding_dim": self.text_embedder_hidden_size,
							    "num_layers": 1,
							}

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

	@property
	def 		UpperCAmelCase  (							self       )->					str:







							'''simple docstring'''





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

							lowerCAmelCase__       =						CLIPVisionModelWithProjection(__UpperCAmelCase       )
							return model

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







							'''simple docstring'''





							lowerCAmelCase__       =						CLIPImageProcessor(
							    crop_size=224		,    do_center_crop=__UpperCAmelCase		,    do_normalize=__UpperCAmelCase		,    do_resize=__UpperCAmelCase		,    image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]		,    image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]		,    resample=3		,    size=224		,    )

							return image_processor

	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





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

							lowerCAmelCase__       =						UnCLIPScheduler(
							    variance_type="fixed_small_log"		,    prediction_type="sample"		,    num_train_timesteps=1000		,    clip_sample=__UpperCAmelCase		,    clip_sample_range=10.0		,    )

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

							return components

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase=0       )->					List[str]:







							'''simple docstring'''





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

	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						"cpu"

							lowerCAmelCase__       =						self.get_dummy_components()

							lowerCAmelCase__       =						self.pipeline_class(**__UpperCAmelCase       )
							lowerCAmelCase__       =						pipe.to(__UpperCAmelCase       )

							pipe.set_progress_bar_config(disable=__UpperCAmelCase       )

							lowerCAmelCase__       =						pipe(**self.get_dummy_inputs(__UpperCAmelCase       )       )
							lowerCAmelCase__       =						output.image_embeds

							lowerCAmelCase__       =						pipe(
							    **self.get_dummy_inputs(__UpperCAmelCase       )		,    return_dict=__UpperCAmelCase		,    )[0]

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

							assert image.shape == (1, 32)

							lowerCAmelCase__       =						np.array(
							    [-0.0_532, 1.7_120, 0.3_656, -1.0_852, -0.8_946, -1.1_756, 0.4_348, 0.2_482, 0.5_146, -0.1_156]       )

							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       )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						torch_device == "cpu"
							lowerCAmelCase__       =						True
							lowerCAmelCase__       =						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       )->					List[str]:







							'''simple docstring'''





							lowerCAmelCase__       =						torch_device == "cpu"
							lowerCAmelCase__       =						False

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














import unittest

from transformers import AutoTokenizer, NystromformerConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    NystromformerForMaskedLM,
	    NystromformerForMultipleChoice,
	    NystromformerForQuestionAnswering,
	    NystromformerForSequenceClassification,
	    NystromformerForTokenClassification,
	    NystromformerModel,
	)
	from transformers.models.nystromformer.modeling_nystromformer import NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


class      lowercase__     :

	def __init__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase=13		,    __UpperCAmelCase=7		,    __UpperCAmelCase=True		,    __UpperCAmelCase=True		,    __UpperCAmelCase=True		,    __UpperCAmelCase=True		,    __UpperCAmelCase=99		,    __UpperCAmelCase=32		,    __UpperCAmelCase=5		,    __UpperCAmelCase=4		,    __UpperCAmelCase=37		,    __UpperCAmelCase="gelu"		,    __UpperCAmelCase=0.1		,    __UpperCAmelCase=0.1		,    __UpperCAmelCase=512		,    __UpperCAmelCase=16		,    __UpperCAmelCase=2		,    __UpperCAmelCase=0.02		,    __UpperCAmelCase=3		,    __UpperCAmelCase=4		,    __UpperCAmelCase=None		,    )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						parent
							lowerCAmelCase__       =						batch_size
							lowerCAmelCase__       =						seq_length
							lowerCAmelCase__       =						is_training
							lowerCAmelCase__       =						use_input_mask
							lowerCAmelCase__       =						use_token_type_ids
							lowerCAmelCase__       =						use_labels
							lowerCAmelCase__       =						vocab_size
							lowerCAmelCase__       =						hidden_size
							lowerCAmelCase__       =						num_hidden_layers
							lowerCAmelCase__       =						num_attention_heads
							lowerCAmelCase__       =						intermediate_size
							lowerCAmelCase__       =						hidden_act
							lowerCAmelCase__       =						hidden_dropout_prob
							lowerCAmelCase__       =						attention_probs_dropout_prob
							lowerCAmelCase__       =						max_position_embeddings
							lowerCAmelCase__       =						type_vocab_size
							lowerCAmelCase__       =						type_sequence_label_size
							lowerCAmelCase__       =						initializer_range
							lowerCAmelCase__       =						num_labels
							lowerCAmelCase__       =						num_choices
							lowerCAmelCase__       =						scope

	def 		UpperCAmelCase  (							self       )->					Optional[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						ids_tensor([self.batch_size, self.seq_length]		,    self.vocab_size       )

							lowerCAmelCase__       =						None
							if self.use_input_mask:
													lowerCAmelCase__       =						random_attention_mask([self.batch_size, self.seq_length]       )

							lowerCAmelCase__       =						None
							if self.use_token_type_ids:
													lowerCAmelCase__       =						ids_tensor([self.batch_size, self.seq_length]		,    self.type_vocab_size       )

							lowerCAmelCase__       =						None
							lowerCAmelCase__       =						None
							lowerCAmelCase__       =						None
							if self.use_labels:
													lowerCAmelCase__       =						ids_tensor([self.batch_size]		,    self.type_sequence_label_size       )
													lowerCAmelCase__       =						ids_tensor([self.batch_size, self.seq_length]		,    self.num_labels       )
													lowerCAmelCase__       =						ids_tensor([self.batch_size]		,    self.num_choices       )

							lowerCAmelCase__       =						self.get_config()

							return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

	def 		UpperCAmelCase  (							self       )->					Optional[Any]:







							'''simple docstring'''





							return NystromformerConfig(
							    vocab_size=self.vocab_size		,    hidden_size=self.hidden_size		,    num_hidden_layers=self.num_hidden_layers		,    num_attention_heads=self.num_attention_heads		,    intermediate_size=self.intermediate_size		,    hidden_act=self.hidden_act		,    hidden_dropout_prob=self.hidden_dropout_prob		,    attention_probs_dropout_prob=self.attention_probs_dropout_prob		,    max_position_embeddings=self.max_position_embeddings		,    type_vocab_size=self.type_vocab_size		,    is_decoder=__UpperCAmelCase		,    initializer_range=self.initializer_range		,    )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					int:







							'''simple docstring'''





							lowerCAmelCase__       =						NystromformerModel(config=__UpperCAmelCase       )
							model.to(__UpperCAmelCase       )
							model.eval()
							lowerCAmelCase__       =						model(__UpperCAmelCase		,    attention_mask=__UpperCAmelCase		,    token_type_ids=__UpperCAmelCase       )
							lowerCAmelCase__       =						model(__UpperCAmelCase		,    token_type_ids=__UpperCAmelCase       )
							lowerCAmelCase__       =						model(__UpperCAmelCase       )
							self.parent.assertEqual(result.last_hidden_state.shape		,    (self.batch_size, self.seq_length, self.hidden_size)       )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					str:







							'''simple docstring'''





							lowerCAmelCase__       =						NystromformerForMaskedLM(config=__UpperCAmelCase       )
							model.to(__UpperCAmelCase       )
							model.eval()
							lowerCAmelCase__       =						model(__UpperCAmelCase		,    attention_mask=__UpperCAmelCase		,    token_type_ids=__UpperCAmelCase		,    labels=__UpperCAmelCase       )
							self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.seq_length, self.vocab_size)       )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						NystromformerForQuestionAnswering(config=__UpperCAmelCase       )
							model.to(__UpperCAmelCase       )
							model.eval()
							lowerCAmelCase__       =						model(
							    __UpperCAmelCase		,    attention_mask=__UpperCAmelCase		,    token_type_ids=__UpperCAmelCase		,    start_positions=__UpperCAmelCase		,    end_positions=__UpperCAmelCase		,    )
							self.parent.assertEqual(result.start_logits.shape		,    (self.batch_size, self.seq_length)       )
							self.parent.assertEqual(result.end_logits.shape		,    (self.batch_size, self.seq_length)       )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					Dict:







							'''simple docstring'''





							lowerCAmelCase__       =						self.num_labels
							lowerCAmelCase__       =						NystromformerForSequenceClassification(__UpperCAmelCase       )
							model.to(__UpperCAmelCase       )
							model.eval()
							lowerCAmelCase__       =						model(__UpperCAmelCase		,    attention_mask=__UpperCAmelCase		,    token_type_ids=__UpperCAmelCase		,    labels=__UpperCAmelCase       )
							self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.num_labels)       )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					str:







							'''simple docstring'''





							lowerCAmelCase__       =						self.num_labels
							lowerCAmelCase__       =						NystromformerForTokenClassification(config=__UpperCAmelCase       )
							model.to(__UpperCAmelCase       )
							model.eval()
							lowerCAmelCase__       =						model(__UpperCAmelCase		,    attention_mask=__UpperCAmelCase		,    token_type_ids=__UpperCAmelCase		,    labels=__UpperCAmelCase       )
							self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.seq_length, self.num_labels)       )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					Optional[int]:







							'''simple docstring'''





							lowerCAmelCase__       =						self.num_choices
							lowerCAmelCase__       =						NystromformerForMultipleChoice(config=__UpperCAmelCase       )
							model.to(__UpperCAmelCase       )
							model.eval()
							lowerCAmelCase__       =						input_ids.unsqueeze(1       ).expand(-1		,    self.num_choices		,    -1       ).contiguous()
							lowerCAmelCase__       =						token_type_ids.unsqueeze(1       ).expand(-1		,    self.num_choices		,    -1       ).contiguous()
							lowerCAmelCase__       =						input_mask.unsqueeze(1       ).expand(-1		,    self.num_choices		,    -1       ).contiguous()
							lowerCAmelCase__       =						model(
							    __UpperCAmelCase		,    attention_mask=__UpperCAmelCase		,    token_type_ids=__UpperCAmelCase		,    labels=__UpperCAmelCase		,    )
							self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.num_choices)       )


	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						self.prepare_config_and_inputs()
							(
							    (
							    lowerCAmelCase__
							)      ,							(
							    lowerCAmelCase__
							)      ,							(
							    lowerCAmelCase__
							)      ,							(
							    lowerCAmelCase__
							)      ,							(
							    lowerCAmelCase__
							)      ,							(
							    lowerCAmelCase__
							)      ,							(
							    lowerCAmelCase__
							)      ,							
							)       =						config_and_inputs
							lowerCAmelCase__       =						{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
							return config, inputs_dict





@require_torch
class      lowercase__     (			_UpperCAmelCase,					_UpperCAmelCase,					unittest.TestCase	):
	a_         =(
	    (
	        NystromformerModel,
	        NystromformerForMaskedLM,
	        NystromformerForMultipleChoice,
	        NystromformerForQuestionAnswering,
	        NystromformerForSequenceClassification,
	        NystromformerForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	a_         =(
	    {
	        """feature-extraction""": NystromformerModel,
	        """fill-mask""": NystromformerForMaskedLM,
	        """question-answering""": NystromformerForQuestionAnswering,
	        """text-classification""": NystromformerForSequenceClassification,
	        """token-classification""": NystromformerForTokenClassification,
	        """zero-shot""": NystromformerForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	a_         =False
	a_         =False

	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							lowerCAmelCase__       =						NystromformerModelTester(self       )
							lowerCAmelCase__       =						ConfigTester(self		,    config_class=__UpperCAmelCase		,    hidden_size=37       )

	def 		UpperCAmelCase  (							self       )->					Dict:







							'''simple docstring'''





							self.config_tester.run_common_tests()

	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__UpperCAmelCase       )

	def 		UpperCAmelCase  (							self       )->					Union[str, Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
													lowerCAmelCase__       =						type
													self.model_tester.create_and_check_model(*__UpperCAmelCase       )

	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase       )

	def 		UpperCAmelCase  (							self       )->					Optional[int]:







							'''simple docstring'''





							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*__UpperCAmelCase       )

	def 		UpperCAmelCase  (							self       )->					Union[str, Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*__UpperCAmelCase       )

	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*__UpperCAmelCase       )

	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*__UpperCAmelCase       )


	@slow
	def 		UpperCAmelCase  (							self       )->					Optional[int]:







							'''simple docstring'''





							for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													lowerCAmelCase__       =						NystromformerModel.from_pretrained(__UpperCAmelCase       )
													self.assertIsNotNone(__UpperCAmelCase       )





@require_torch
class      lowercase__     (			unittest.TestCase	):

	@slow
	def 		UpperCAmelCase  (							self       )->					Dict:







							'''simple docstring'''





							lowerCAmelCase__       =						NystromformerModel.from_pretrained("uw-madison/nystromformer-512"       )
							lowerCAmelCase__       =						torch.tensor([[0, 1, 2, 3, 4, 5]]       )

							with torch.no_grad():
													lowerCAmelCase__       =						model(__UpperCAmelCase       )[0]

							lowerCAmelCase__       =						torch.Size((1, 6, 768)       )
							self.assertEqual(output.shape		,    __UpperCAmelCase       )

							lowerCAmelCase__       =						torch.tensor(
							    [[[-0.4_532, -0.0_936, 0.5_137], [-0.2_676, 0.0_628, 0.6_186], [-0.3_629, -0.1_726, 0.4_716]]]       )

							self.assertTrue(torch.allclose(output[:, :3, :3]		,    __UpperCAmelCase		,    atol=1E-4       )       )


	@slow
	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						"the [MASK] of Belgium is Brussels"

							lowerCAmelCase__       =						AutoTokenizer.from_pretrained("uw-madison/nystromformer-512"       )
							lowerCAmelCase__       =						NystromformerForMaskedLM.from_pretrained("uw-madison/nystromformer-512"       )

							lowerCAmelCase__       =						tokenizer(__UpperCAmelCase		,    return_tensors="pt"       )

							with torch.no_grad():
													lowerCAmelCase__       =						model(encoding.input_ids       ).logits

							lowerCAmelCase__       =						token_logits[:, 2, :].argmax(-1       )[0]

							self.assertEqual(tokenizer.decode(__UpperCAmelCase       )		,    "capital"       )






