
'''simple docstring'''







import unittest

from transformers import MraConfig, 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, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
						import torch

						from transformers import (
						    MraForMaskedLM,
						    MraForMultipleChoice,
						    MraForQuestionAnswering,
						    MraForSequenceClassification,
						    MraForTokenClassification,
						    MraModel,
						)
						from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST


class 			lowerCAmelCase :




		def __init__(				self					,						UpperCamelCase					,						UpperCamelCase=2					,						UpperCamelCase=8					,						UpperCamelCase=True					,						UpperCamelCase=True					,						UpperCamelCase=True					,						UpperCamelCase=True					,						UpperCamelCase=99					,						UpperCamelCase=16					,						UpperCamelCase=5					,						UpperCamelCase=2					,						UpperCamelCase=36					,						UpperCamelCase="gelu"					,						UpperCamelCase=0.0					,						UpperCamelCase=0.0					,						UpperCamelCase=512					,						UpperCamelCase=16					,						UpperCamelCase=2					,						UpperCamelCase=0.02					,						UpperCamelCase=3					,						UpperCamelCase=4					,						UpperCamelCase=None					,						):
								_SCREAMING_SNAKE_CASE				= parent
								_SCREAMING_SNAKE_CASE				= batch_size
								_SCREAMING_SNAKE_CASE				= seq_length
								_SCREAMING_SNAKE_CASE				= is_training
								_SCREAMING_SNAKE_CASE				= use_input_mask
								_SCREAMING_SNAKE_CASE				= use_token_type_ids
								_SCREAMING_SNAKE_CASE				= use_labels
								_SCREAMING_SNAKE_CASE				= vocab_size
								_SCREAMING_SNAKE_CASE				= hidden_size
								_SCREAMING_SNAKE_CASE				= num_hidden_layers
								_SCREAMING_SNAKE_CASE				= num_attention_heads
								_SCREAMING_SNAKE_CASE				= intermediate_size
								_SCREAMING_SNAKE_CASE				= hidden_act
								_SCREAMING_SNAKE_CASE				= hidden_dropout_prob
								_SCREAMING_SNAKE_CASE				= attention_probs_dropout_prob
								_SCREAMING_SNAKE_CASE				= max_position_embeddings
								_SCREAMING_SNAKE_CASE				= type_vocab_size
								_SCREAMING_SNAKE_CASE				= type_sequence_label_size
								_SCREAMING_SNAKE_CASE				= initializer_range
								_SCREAMING_SNAKE_CASE				= num_labels
								_SCREAMING_SNAKE_CASE				= num_choices
								_SCREAMING_SNAKE_CASE				= scope




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= ids_tensor([self.batch_size, self.seq_length]					,						self.vocab_size      )

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

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

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

								_SCREAMING_SNAKE_CASE				= self.get_config()

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




		def     lowercase       (				self      ):
								return MraConfig(
								    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     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= self.get_config()
								_SCREAMING_SNAKE_CASE				= 300
								return config




		def     lowercase       (				self      ):
								(
								    (
								    _SCREAMING_SNAKE_CASE							
								),							(
								    _SCREAMING_SNAKE_CASE							
								),							(
								    _SCREAMING_SNAKE_CASE							
								),							(
								    _SCREAMING_SNAKE_CASE							
								),							(
								    _SCREAMING_SNAKE_CASE							
								),							(
								    _SCREAMING_SNAKE_CASE							
								),							(
								    _SCREAMING_SNAKE_CASE							
								),														
								)				= self.prepare_config_and_inputs()

								_SCREAMING_SNAKE_CASE				= True
								_SCREAMING_SNAKE_CASE				= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
								_SCREAMING_SNAKE_CASE				= ids_tensor([self.batch_size, self.seq_length]					,						vocab_size=2      )

								return (
								    config,
								    input_ids,
								    token_type_ids,
								    input_mask,
								    sequence_labels,
								    token_labels,
								    choice_labels,
								    encoder_hidden_states,
								    encoder_attention_mask,
								)




		def     lowercase       (				self					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase      ):
								_SCREAMING_SNAKE_CASE				= MraModel(config=UpperCamelCase      )
								model.to(UpperCamelCase      )
								model.eval()
								_SCREAMING_SNAKE_CASE				= model(UpperCamelCase					,						attention_mask=UpperCamelCase					,						token_type_ids=UpperCamelCase      )
								_SCREAMING_SNAKE_CASE				= model(UpperCamelCase					,						token_type_ids=UpperCamelCase      )
								_SCREAMING_SNAKE_CASE				= model(UpperCamelCase      )
								self.parent.assertEqual(result.last_hidden_state.shape					,						(self.batch_size, self.seq_length, self.hidden_size)      )




		def     lowercase       (				self					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						):
								_SCREAMING_SNAKE_CASE				= True
								_SCREAMING_SNAKE_CASE				= MraModel(UpperCamelCase      )
								model.to(UpperCamelCase      )
								model.eval()
								_SCREAMING_SNAKE_CASE				= model(
								    UpperCamelCase					,						attention_mask=UpperCamelCase					,						token_type_ids=UpperCamelCase					,						encoder_hidden_states=UpperCamelCase					,						encoder_attention_mask=UpperCamelCase					,						)
								_SCREAMING_SNAKE_CASE				= model(
								    UpperCamelCase					,						attention_mask=UpperCamelCase					,						token_type_ids=UpperCamelCase					,						encoder_hidden_states=UpperCamelCase					,						)
								_SCREAMING_SNAKE_CASE				= model(UpperCamelCase					,						attention_mask=UpperCamelCase					,						token_type_ids=UpperCamelCase      )
								self.parent.assertEqual(result.last_hidden_state.shape					,						(self.batch_size, self.seq_length, self.hidden_size)      )




		def     lowercase       (				self					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase      ):
								_SCREAMING_SNAKE_CASE				= MraForMaskedLM(config=UpperCamelCase      )
								model.to(UpperCamelCase      )
								model.eval()
								_SCREAMING_SNAKE_CASE				= 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     lowercase       (				self					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase      ):
								_SCREAMING_SNAKE_CASE				= MraForQuestionAnswering(config=UpperCamelCase      )
								model.to(UpperCamelCase      )
								model.eval()
								_SCREAMING_SNAKE_CASE				= 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     lowercase       (				self					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase      ):
								_SCREAMING_SNAKE_CASE				= self.num_labels
								_SCREAMING_SNAKE_CASE				= MraForSequenceClassification(UpperCamelCase      )
								model.to(UpperCamelCase      )
								model.eval()
								_SCREAMING_SNAKE_CASE				= model(UpperCamelCase					,						attention_mask=UpperCamelCase					,						token_type_ids=UpperCamelCase					,						labels=UpperCamelCase      )
								self.parent.assertEqual(result.logits.shape					,						(self.batch_size, self.num_labels)      )




		def     lowercase       (				self					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase      ):
								_SCREAMING_SNAKE_CASE				= self.num_labels
								_SCREAMING_SNAKE_CASE				= MraForTokenClassification(config=UpperCamelCase      )
								model.to(UpperCamelCase      )
								model.eval()
								_SCREAMING_SNAKE_CASE				= 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     lowercase       (				self					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase      ):
								_SCREAMING_SNAKE_CASE				= self.num_choices
								_SCREAMING_SNAKE_CASE				= MraForMultipleChoice(config=UpperCamelCase      )
								model.to(UpperCamelCase      )
								model.eval()
								_SCREAMING_SNAKE_CASE				= input_ids.unsqueeze(1      ).expand(-1					,						self.num_choices					,						-1      ).contiguous()
								_SCREAMING_SNAKE_CASE				= token_type_ids.unsqueeze(1      ).expand(-1					,						self.num_choices					,						-1      ).contiguous()
								_SCREAMING_SNAKE_CASE				= input_mask.unsqueeze(1      ).expand(-1					,						self.num_choices					,						-1      ).contiguous()
								_SCREAMING_SNAKE_CASE				= model(
								    UpperCamelCase					,						attention_mask=UpperCamelCase					,						token_type_ids=UpperCamelCase					,						labels=UpperCamelCase					,						)
								self.parent.assertEqual(result.logits.shape					,						(self.batch_size, self.num_choices)      )




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


@require_torch
class 			lowerCAmelCase (				__UpperCAmelCase		,     unittest.TestCase			):
		a							:					Dict       =					(
		    (
		        MraModel,
		        MraForMaskedLM,
		        MraForMultipleChoice,
		        MraForQuestionAnswering,
		        MraForSequenceClassification,
		        MraForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		a							:					Dict       =					False
		a							:					Optional[Any]       =					False
		a							:					Optional[Any]       =					False
		a							:					List[Any]       =					False

		a							:					Any       =					()




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= MraModelTester(self      )
								_SCREAMING_SNAKE_CASE				= ConfigTester(self					,						config_class=UpperCamelCase					,						hidden_size=37      )




		def     lowercase       (				self      ):
								self.config_tester.run_common_tests()




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= self.model_tester.prepare_config_and_inputs()
								for type in ["absolute", "relative_key", "relative_key_query"]:
														_SCREAMING_SNAKE_CASE				= type
														self.model_tester.create_and_check_model(*UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*UpperCamelCase      )




		@slow
		def     lowercase       (				self      ):
								for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_SCREAMING_SNAKE_CASE				= MraModel.from_pretrained(UpperCamelCase      )
														self.assertIsNotNone(UpperCamelCase      )




		@unittest.skip(reason="MRA does not output attentions"      )
		def     lowercase       (				self      ):
								return


@require_torch
class 			lowerCAmelCase (				unittest.TestCase			):




		@slow
		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= MraModel.from_pretrained("uw-madison/mra-base-512-4"      )
								_SCREAMING_SNAKE_CASE				= torch.arange(256      ).unsqueeze(0      )

								with torch.no_grad():
														_SCREAMING_SNAKE_CASE				= model(UpperCamelCase      )[0]

								_SCREAMING_SNAKE_CASE				= torch.Size((1, 256, 768)      )
								self.assertEqual(output.shape					,						UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= torch.tensor(
								    [[[-0.01_40, 0.08_30, -0.03_81], [0.15_46, 0.14_02, 0.02_20], [0.11_62, 0.08_51, 0.01_65]]]      )

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




		@slow
		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= MraForMaskedLM.from_pretrained("uw-madison/mra-base-512-4"      )
								_SCREAMING_SNAKE_CASE				= torch.arange(256      ).unsqueeze(0      )

								with torch.no_grad():
														_SCREAMING_SNAKE_CASE				= model(UpperCamelCase      )[0]

								_SCREAMING_SNAKE_CASE				= 50_265

								_SCREAMING_SNAKE_CASE				= torch.Size((1, 256, vocab_size)      )
								self.assertEqual(output.shape					,						UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= torch.tensor(
								    [[[9.25_95, -3.60_38, 11.88_19], [9.38_69, -3.26_93, 11.09_56], [11.85_24, -3.49_38, 13.12_10]]]      )

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




		@slow
		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= MraForMaskedLM.from_pretrained("uw-madison/mra-base-4096-8-d3"      )
								_SCREAMING_SNAKE_CASE				= torch.arange(4_096      ).unsqueeze(0      )

								with torch.no_grad():
														_SCREAMING_SNAKE_CASE				= model(UpperCamelCase      )[0]

								_SCREAMING_SNAKE_CASE				= 50_265

								_SCREAMING_SNAKE_CASE				= torch.Size((1, 4_096, vocab_size)      )
								self.assertEqual(output.shape					,						UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= torch.tensor(
								    [[[5.47_89, -2.35_64, 7.50_64], [7.90_67, -1.33_69, 9.96_68], [9.07_12, -1.81_06, 7.03_80]]]      )

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

'''simple docstring'''







import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np


class 			lowerCAmelCase :




		def     lowercase       (				self      ):
								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"      )

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"      )

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= UNetaDConditionModel(
								    sample_size=32					,						layers_per_block=1					,						block_out_channels=[32, 64]					,						down_block_types=[
								        "ResnetDownsampleBlock2D",
								        "SimpleCrossAttnDownBlock2D",
								    ]					,						mid_block_type="UNetMidBlock2DSimpleCrossAttn"					,						up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]					,						in_channels=3					,						out_channels=6					,						cross_attention_dim=32					,						encoder_hid_dim=32					,						attention_head_dim=8					,						addition_embed_type="text"					,						addition_embed_type_num_heads=2					,						cross_attention_norm="group_norm"					,						resnet_time_scale_shift="scale_shift"					,						act_fn="gelu"					,						)
								unet.set_attn_processor(AttnAddedKVProcessor()      )  # For reproducibility tests

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= DDPMScheduler(
								    num_train_timesteps=1_000					,						beta_schedule="squaredcos_cap_v2"					,						beta_start=0.00_01					,						beta_end=0.02					,						thresholding=UpperCamelCase					,						dynamic_thresholding_ratio=0.95					,						sample_max_value=1.0					,						prediction_type="epsilon"					,						variance_type="learned_range"					,						)

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= IFWatermarker()

								return {
								    "text_encoder": text_encoder,
								    "tokenizer": tokenizer,
								    "unet": unet,
								    "scheduler": scheduler,
								    "watermarker": watermarker,
								    "safety_checker": None,
								    "feature_extractor": None,
								}




		def     lowercase       (				self      ):
								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"      )

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"      )

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= UNetaDConditionModel(
								    sample_size=32					,						layers_per_block=[1, 2]					,						block_out_channels=[32, 64]					,						down_block_types=[
								        "ResnetDownsampleBlock2D",
								        "SimpleCrossAttnDownBlock2D",
								    ]					,						mid_block_type="UNetMidBlock2DSimpleCrossAttn"					,						up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]					,						in_channels=6					,						out_channels=6					,						cross_attention_dim=32					,						encoder_hid_dim=32					,						attention_head_dim=8					,						addition_embed_type="text"					,						addition_embed_type_num_heads=2					,						cross_attention_norm="group_norm"					,						resnet_time_scale_shift="scale_shift"					,						act_fn="gelu"					,						class_embed_type="timestep"					,						mid_block_scale_factor=1.4_14					,						time_embedding_act_fn="gelu"					,						time_embedding_dim=32					,						)
								unet.set_attn_processor(AttnAddedKVProcessor()      )  # For reproducibility tests

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= DDPMScheduler(
								    num_train_timesteps=1_000					,						beta_schedule="squaredcos_cap_v2"					,						beta_start=0.00_01					,						beta_end=0.02					,						thresholding=UpperCamelCase					,						dynamic_thresholding_ratio=0.95					,						sample_max_value=1.0					,						prediction_type="epsilon"					,						variance_type="learned_range"					,						)

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= DDPMScheduler(
								    num_train_timesteps=1_000					,						beta_schedule="squaredcos_cap_v2"					,						beta_start=0.00_01					,						beta_end=0.02					,						)

								torch.manual_seed(0      )
								_SCREAMING_SNAKE_CASE				= IFWatermarker()

								return {
								    "text_encoder": text_encoder,
								    "tokenizer": tokenizer,
								    "unet": unet,
								    "scheduler": scheduler,
								    "image_noising_scheduler": image_noising_scheduler,
								    "watermarker": watermarker,
								    "safety_checker": None,
								    "feature_extractor": None,
								}




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= self.get_dummy_components()
								_SCREAMING_SNAKE_CASE				= self.pipeline_class(**UpperCamelCase      )
								pipe.to(UpperCamelCase      )
								pipe.set_progress_bar_config(disable=UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= self.get_dummy_inputs(UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= inputs["prompt"]
								_SCREAMING_SNAKE_CASE				= inputs["generator"]
								_SCREAMING_SNAKE_CASE				= inputs["num_inference_steps"]
								_SCREAMING_SNAKE_CASE				= inputs["output_type"]

								if "image" in inputs:
														_SCREAMING_SNAKE_CASE				= inputs["image"]
								else:
														_SCREAMING_SNAKE_CASE				= None

								if "mask_image" in inputs:
														_SCREAMING_SNAKE_CASE				= inputs["mask_image"]
								else:
														_SCREAMING_SNAKE_CASE				= None

								if "original_image" in inputs:
														_SCREAMING_SNAKE_CASE				= inputs["original_image"]
								else:
														_SCREAMING_SNAKE_CASE				= None

								_SCREAMING_SNAKE_CASE,							_SCREAMING_SNAKE_CASE				= pipe.encode_prompt(UpperCamelCase      )

								# inputs with prompt converted to embeddings
								_SCREAMING_SNAKE_CASE				= {
								    "prompt_embeds": prompt_embeds,
								    "negative_prompt_embeds": negative_prompt_embeds,
								    "generator": generator,
								    "num_inference_steps": num_inference_steps,
								    "output_type": output_type,
								}

								if image is not None:
														_SCREAMING_SNAKE_CASE				= image

								if mask_image is not None:
														_SCREAMING_SNAKE_CASE				= mask_image

								if original_image is not None:
														_SCREAMING_SNAKE_CASE				= original_image

								# set all optional components to None
								for optional_component in pipe._optional_components:
														setattr(UpperCamelCase					,						UpperCamelCase					,						UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= pipe(**UpperCamelCase      )[0]

								with tempfile.TemporaryDirectory() as tmpdir:
														pipe.save_pretrained(UpperCamelCase      )
														_SCREAMING_SNAKE_CASE				= self.pipeline_class.from_pretrained(UpperCamelCase      )
														pipe_loaded.to(UpperCamelCase      )
														pipe_loaded.set_progress_bar_config(disable=UpperCamelCase      )

								pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()      )  # For reproducibility tests

								for optional_component in pipe._optional_components:
														self.assertTrue(
														    getattr(UpperCamelCase					,						UpperCamelCase      ) is None					,						F'`{optional_component}` did not stay set to None after loading.'					,						)

								_SCREAMING_SNAKE_CASE				= self.get_dummy_inputs(UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= inputs["generator"]
								_SCREAMING_SNAKE_CASE				= inputs["num_inference_steps"]
								_SCREAMING_SNAKE_CASE				= inputs["output_type"]

								# inputs with prompt converted to embeddings
								_SCREAMING_SNAKE_CASE				= {
								    "prompt_embeds": prompt_embeds,
								    "negative_prompt_embeds": negative_prompt_embeds,
								    "generator": generator,
								    "num_inference_steps": num_inference_steps,
								    "output_type": output_type,
								}

								if image is not None:
														_SCREAMING_SNAKE_CASE				= image

								if mask_image is not None:
														_SCREAMING_SNAKE_CASE				= mask_image

								if original_image is not None:
														_SCREAMING_SNAKE_CASE				= original_image

								_SCREAMING_SNAKE_CASE				= pipe_loaded(**UpperCamelCase      )[0]

								_SCREAMING_SNAKE_CASE				= np.abs(to_np(UpperCamelCase      ) - to_np(UpperCamelCase      )      ).max()
								self.assertLess(UpperCamelCase					,						1e-4      )




		def     lowercase       (				self      ):
								_SCREAMING_SNAKE_CASE				= self.get_dummy_components()
								_SCREAMING_SNAKE_CASE				= self.pipeline_class(**UpperCamelCase      )
								pipe.to(UpperCamelCase      )
								pipe.set_progress_bar_config(disable=UpperCamelCase      )

								_SCREAMING_SNAKE_CASE				= self.get_dummy_inputs(UpperCamelCase      )
								_SCREAMING_SNAKE_CASE				= pipe(**UpperCamelCase      )[0]

								with tempfile.TemporaryDirectory() as tmpdir:
														pipe.save_pretrained(UpperCamelCase      )
														_SCREAMING_SNAKE_CASE				= self.pipeline_class.from_pretrained(UpperCamelCase      )
														pipe_loaded.to(UpperCamelCase      )
														pipe_loaded.set_progress_bar_config(disable=UpperCamelCase      )

								pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()      )  # For reproducibility tests

								_SCREAMING_SNAKE_CASE				= self.get_dummy_inputs(UpperCamelCase      )
								_SCREAMING_SNAKE_CASE				= pipe_loaded(**UpperCamelCase      )[0]

								_SCREAMING_SNAKE_CASE				= np.abs(to_np(UpperCamelCase      ) - to_np(UpperCamelCase      )      ).max()
								self.assertLess(UpperCamelCase					,						1e-4      )