

from __future__ import annotations

import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import is_tf_available, is_vision_available

from ...test_modeling_tf_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_tf_bert import TFBertModelTester
from ..clip.test_modeling_tf_clip import TFCLIPVisionModelTester
from ..deit.test_modeling_tf_deit import TFDeiTModelTester
from ..roberta.test_modeling_tf_roberta import TFRobertaModelTester
from ..vit.test_modeling_tf_vit import TFViTModelTester


if is_tf_available():
	from transformers import (
	    TFBertModel,
	    TFCLIPVisionModel,
	    TFDeiTModel,
	    TFRobertaModel,
	    TFVisionTextDualEncoderModel,
	    TFViTModel,
	    VisionTextDualEncoderConfig,
	)

if is_vision_available():
	from PIL import Image

	from transformers import VisionTextDualEncoderProcessor
def 							lowerCAmelCase(     __lowerCamelCase						):





					if isinstance(lowerCAmelCase__				,    collections.abc.Iterable						):
										return x
					return (x, x)


@require_tf
class 						a__							:



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



			def __SCREAMING_SNAKE_CASE							(						self   )			->   List[str]:
								pass



			def __SCREAMING_SNAKE_CASE							(						self   )			->   Any:
								pass



			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase=None				,						**UpperCAmelCase   )			->   Union[str, Any]:
								__a		 =       VisionTextDualEncoderConfig.from_vision_text_configs(__UpperCAmelCase				,						__UpperCAmelCase   )

								__a		 =       TFVisionTextDualEncoderModel(__UpperCAmelCase   )

								__a		 =       model(input_ids=__UpperCAmelCase				,						pixel_values=__UpperCAmelCase				,						attention_mask=__UpperCAmelCase   )

								self.assertEqual(output['text_embeds'].shape				,						(input_ids.shape[0], config.projection_dim)   )
								self.assertEqual(output['image_embeds'].shape				,						(pixel_values.shape[0], config.projection_dim)   )



			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase=None				,						**UpperCAmelCase   )			->   int:
								__a		 =       self.get_vision_text_model(__UpperCAmelCase				,						__UpperCAmelCase   )
								__a		 =       TFVisionTextDualEncoderModel(vision_model=__UpperCAmelCase				,						text_model=__UpperCAmelCase   )

								__a		 =       model(input_ids=__UpperCAmelCase				,						pixel_values=__UpperCAmelCase				,						attention_mask=__UpperCAmelCase   )

								self.assertEqual(output['text_embeds'].shape				,						(input_ids.shape[0], model.config.projection_dim)   )
								self.assertEqual(output['image_embeds'].shape				,						(pixel_values.shape[0], model.config.projection_dim)   )



			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase=None				,						**UpperCAmelCase   )			->   List[Any]:
								__a		 =       self.get_vision_text_model(__UpperCAmelCase				,						__UpperCAmelCase   )
								__a		 =       {"""vision_model""": vision_model, """text_model""": text_model}
								__a		 =       TFVisionTextDualEncoderModel.from_vision_text_pretrained(**__UpperCAmelCase   )

								__a		 =       model(input_ids=__UpperCAmelCase				,						pixel_values=__UpperCAmelCase				,						attention_mask=__UpperCAmelCase   )

								self.assertEqual(output['text_embeds'].shape				,						(input_ids.shape[0], model.config.projection_dim)   )
								self.assertEqual(output['image_embeds'].shape				,						(pixel_values.shape[0], model.config.projection_dim)   )



			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase=None				,						**UpperCAmelCase   )			->   Optional[int]:
								__a		 =       self.get_vision_text_model(__UpperCAmelCase				,						__UpperCAmelCase   )
								__a		 =       TFVisionTextDualEncoderModel(vision_model=__UpperCAmelCase				,						text_model=__UpperCAmelCase   )

								__a		 =       model(input_ids=__UpperCAmelCase				,						pixel_values=__UpperCAmelCase				,						attention_mask=__UpperCAmelCase   )
								__a		 =       output[0].numpy()

								with tempfile.TemporaryDirectory() as tmpdirname:
													model.save_pretrained(__UpperCAmelCase   )
													__a		 =       TFVisionTextDualEncoderModel.from_pretrained(__UpperCAmelCase   )

													__a		 =       model(input_ids=__UpperCAmelCase				,						pixel_values=__UpperCAmelCase				,						attention_mask=__UpperCAmelCase   )
													__a		 =       after_output[0].numpy()
													__a		 =       np.amax(np.abs(out_a - out_a   )   )
													self.assertLessEqual(__UpperCAmelCase				,						1e-5   )



			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase=None				,						**UpperCAmelCase   )			->   Optional[int]:
								__a		 =       self.get_vision_text_model(__UpperCAmelCase				,						__UpperCAmelCase   )
								__a		 =       TFVisionTextDualEncoderModel(vision_model=__UpperCAmelCase				,						text_model=__UpperCAmelCase   )

								__a		 =       model(
								    input_ids=__UpperCAmelCase				,						pixel_values=__UpperCAmelCase				,						attention_mask=__UpperCAmelCase				,						output_attentions=__UpperCAmelCase   )

								__a		 =       output.vision_model_output.attentions
								self.assertEqual(len(__UpperCAmelCase   )				,						vision_config.num_hidden_layers   )

								# in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
								__a		 =       to_atuple(vision_model.config.image_size   )
								__a		 =       to_atuple(vision_model.config.patch_size   )
								__a		 =       (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
								__a		 =       num_patches + 1
								self.assertEqual(vision_attentions[0].shape[-3:]				,						(vision_config.num_attention_heads, seq_len, seq_len)   )

								__a		 =       output.text_model_output.attentions
								self.assertEqual(len(__UpperCAmelCase   )				,						text_config.num_hidden_layers   )

								self.assertEqual(
								    text_attentions[0].shape[-3:]				,						(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])				,						)



			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   )			->   Union[str, Any]:
								__a		 =       np.abs((a - b)   ).max()
								self.assertLessEqual(__UpperCAmelCase				,						__UpperCAmelCase				,						f'''Difference between torch and flax is {diff} (>= {tol}).'''   )



			def __SCREAMING_SNAKE_CASE							(						self   )			->   List[Any]:
								__a		 =       self.prepare_config_and_inputs()
								self.check_vision_text_dual_encoder_model(**__UpperCAmelCase   )



			def __SCREAMING_SNAKE_CASE							(						self   )			->   Optional[Any]:
								__a		 =       self.prepare_config_and_inputs()
								self.check_model_from_pretrained_configs(**__UpperCAmelCase   )



			def __SCREAMING_SNAKE_CASE							(						self   )			->   Optional[Any]:
								__a		 =       self.prepare_config_and_inputs()
								self.check_vision_text_dual_encoder_from_pretrained(**__UpperCAmelCase   )



			def __SCREAMING_SNAKE_CASE							(						self   )			->   Dict:
								__a		 =       self.prepare_config_and_inputs()
								self.check_save_load(**__UpperCAmelCase   )



			def __SCREAMING_SNAKE_CASE							(						self   )			->   Union[str, Any]:
								__a		 =       self.prepare_config_and_inputs()
								self.check_vision_text_output_attention(**__UpperCAmelCase   )







			@slow
			def __SCREAMING_SNAKE_CASE							(						self   )			->   Tuple:
								__a		 =       self.get_pretrained_model_and_inputs()

								__a		 =       model_a(**__UpperCAmelCase   )
								__a		 =       outputs[0].numpy()

								with tempfile.TemporaryDirectory() as tmp_dirname:
													model_a.save_pretrained(__UpperCAmelCase   )
													__a		 =       TFVisionTextDualEncoderModel.from_pretrained(__UpperCAmelCase   )

													__a		 =       model_a(**__UpperCAmelCase   )
													__a		 =       after_outputs[0].numpy()
													__a		 =       np.amax(np.abs(out_a - out_a   )   )
													self.assertLessEqual(__UpperCAmelCase				,						1e-5   )






@require_tf
class 						a__							(       __SCREAMING_SNAKE_CASE				,						unittest.TestCase		):



			def __SCREAMING_SNAKE_CASE							(						self   )			->   Any:
								__a		 =       TFVisionTextDualEncoderModel.from_vision_text_pretrained(
								    'hf-internal-testing/tiny-random-vit'				,						'hf-internal-testing/tiny-random-bert'   )
								__a		 =       1_3
								__a		 =       floats_tensor(
								    [
								        batch_size,
								        model.vision_model.config.num_channels,
								        model.vision_model.config.image_size,
								        model.vision_model.config.image_size,
								    ]   )
								__a		 =       ids_tensor([batch_size, 4]				,						model.text_model.config.vocab_size   )
								__a		 =       random_attention_mask([batch_size, 4]   )
								__a		 =       {"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

								return model, inputs



			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase				,						UpperCAmelCase   )			->   Optional[int]:
								__a		 =       TFViTModel(__UpperCAmelCase				,						name='vision_model'   )
								__a		 =       TFBertModel(__UpperCAmelCase				,						name='text_model'   )
								return vision_model, text_model







			def __SCREAMING_SNAKE_CASE							(						self   )			->   int:
								__a		 =       TFViTModelTester(self   )
								__a		 =       TFBertModelTester(self   )
								__a		 =       vit_model_tester.prepare_config_and_inputs()
								__a		 =       bert_model_tester.prepare_config_and_inputs()

								__a		 =       vision_config_and_inputs

								(
								    __a					
								)		 =       text_config_and_inputs

								return {
								    "text_config": text_config,
								    "vision_config": vision_config,
								    "pixel_values": pixel_values,
								    "attention_mask": input_mask,
								    "input_ids": input_ids,
								    "text_token_type_ids": token_type_ids,
								    "text_sequence_labels": sequence_labels,
								    "text_token_labels": token_labels,
								    "text_choice_labels": choice_labels,
								}






@require_tf
class 						a__							(       __SCREAMING_SNAKE_CASE				,						unittest.TestCase		):



			def __SCREAMING_SNAKE_CASE							(						self   )			->   str:
								__a		 =       TFVisionTextDualEncoderModel.from_vision_text_pretrained(
								    'Rocketknight1/tiny-random-deit-tf'				,						'hf-internal-testing/tiny-random-roberta'   )
								__a		 =       1_3
								__a		 =       floats_tensor(
								    [
								        batch_size,
								        model.vision_model.config.num_channels,
								        model.vision_model.config.image_size,
								        model.vision_model.config.image_size,
								    ]   )
								__a		 =       ids_tensor([batch_size, 4]				,						model.text_model.config.vocab_size   )
								__a		 =       random_attention_mask([batch_size, 4]   )
								__a		 =       {"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

								return model, inputs



			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase=None				,						**UpperCAmelCase   )			->   str:
								__a		 =       self.get_vision_text_model(__UpperCAmelCase				,						__UpperCAmelCase   )
								__a		 =       TFVisionTextDualEncoderModel(vision_model=__UpperCAmelCase				,						text_model=__UpperCAmelCase   )

								__a		 =       model(
								    input_ids=__UpperCAmelCase				,						pixel_values=__UpperCAmelCase				,						attention_mask=__UpperCAmelCase				,						output_attentions=__UpperCAmelCase   )

								__a		 =       output.vision_model_output.attentions
								self.assertEqual(len(__UpperCAmelCase   )				,						vision_config.num_hidden_layers   )

								# in DEiT, the seq_len equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
								__a		 =       to_atuple(vision_model.config.image_size   )
								__a		 =       to_atuple(vision_model.config.patch_size   )
								__a		 =       (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
								__a		 =       num_patches + 2
								self.assertEqual(vision_attentions[0].shape[-3:]				,						(vision_config.num_attention_heads, seq_len, seq_len)   )

								__a		 =       output.text_model_output.attentions
								self.assertEqual(len(__UpperCAmelCase   )				,						text_config.num_hidden_layers   )

								self.assertEqual(
								    text_attentions[0].shape[-3:]				,						(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])				,						)



			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase				,						UpperCAmelCase   )			->   List[Any]:
								__a		 =       TFDeiTModel(__UpperCAmelCase				,						name='vision_model'   )
								__a		 =       TFRobertaModel(__UpperCAmelCase				,						name='text_model'   )
								return vision_model, text_model







			def __SCREAMING_SNAKE_CASE							(						self   )			->   Union[str, Any]:
								__a		 =       TFDeiTModelTester(self   )
								__a		 =       TFRobertaModelTester(self   )
								__a		 =       vit_model_tester.prepare_config_and_inputs()
								__a		 =       bert_model_tester.prepare_config_and_inputs()

								__a		 =       vision_config_and_inputs

								(
								    __a					
								)		 =       text_config_and_inputs

								return {
								    "text_config": text_config,
								    "vision_config": vision_config,
								    "pixel_values": pixel_values,
								    "attention_mask": input_mask,
								    "input_ids": input_ids,
								    "text_token_type_ids": token_type_ids,
								    "text_sequence_labels": sequence_labels,
								    "text_token_labels": token_labels,
								    "text_choice_labels": choice_labels,
								}






@require_tf
class 						a__							(       __SCREAMING_SNAKE_CASE				,						unittest.TestCase		):



			def __SCREAMING_SNAKE_CASE							(						self   )			->   int:
								__a		 =       TFVisionTextDualEncoderModel.from_vision_text_pretrained(
								    'Rocketknight1/tiny-random-clip-tf'				,						'hf-internal-testing/tiny-random-bert'   )
								__a		 =       1_3
								__a		 =       floats_tensor(
								    [
								        batch_size,
								        model.vision_model.config.num_channels,
								        model.vision_model.config.image_size,
								        model.vision_model.config.image_size,
								    ]   )
								__a		 =       ids_tensor([batch_size, 4]				,						model.text_model.config.vocab_size   )
								__a		 =       random_attention_mask([batch_size, 4]   )
								__a		 =       {"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

								return model, inputs



			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase				,						UpperCAmelCase   )			->   Any:
								__a		 =       TFCLIPVisionModel(__UpperCAmelCase				,						name='vision_model'   )
								__a		 =       TFBertModel(__UpperCAmelCase				,						name='text_model'   )
								return vision_model, text_model







			def __SCREAMING_SNAKE_CASE							(						self   )			->   Union[str, Any]:
								__a		 =       TFCLIPVisionModelTester(self   )
								__a		 =       TFBertModelTester(self   )
								__a		 =       clip_model_tester.prepare_config_and_inputs()
								__a		 =       bert_model_tester.prepare_config_and_inputs()

								__a		 =       vision_config_and_inputs

								(
								    __a					
								)		 =       text_config_and_inputs

								return {
								    "text_config": text_config,
								    "vision_config": vision_config,
								    "pixel_values": pixel_values,
								    "attention_mask": input_mask,
								    "input_ids": input_ids,
								    "text_token_type_ids": token_type_ids,
								    "text_sequence_labels": sequence_labels,
								    "text_token_labels": token_labels,
								    "text_choice_labels": choice_labels,
								}






@require_vision
@require_tf
class 						a__							(       unittest.TestCase		):



			@slow
			def __SCREAMING_SNAKE_CASE							(						self   )			->   Any:
								__a		 =       TFVisionTextDualEncoderModel.from_pretrained(
								    'clip-italian/clip-italian'				,						logit_scale_init_value=1.0				,						from_pt=__UpperCAmelCase   )
								__a		 =       VisionTextDualEncoderProcessor.from_pretrained('clip-italian/clip-italian'   )

								__a		 =       Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
								__a		 =       processor(
								    text=['una foto di un gatto', 'una foto di un cane']				,						images=__UpperCAmelCase				,						padding=__UpperCAmelCase				,						return_tensors='np'   )

								__a		 =       model(**__UpperCAmelCase   )

								# verify the logits
								self.assertEqual(outputs.logits_per_image.shape				,						(inputs.pixel_values.shape[0], inputs.input_ids.shape[0])   )
								self.assertEqual(
								    outputs.logits_per_text.shape				,						(inputs.input_ids.shape[0], inputs.pixel_values.shape[0])				,						)

								__a		 =       np.array([[1.2_284_727, 0.3_104_122]]   )

								self.assertTrue(np.allclose(outputs.logits_per_image.numpy()				,						__UpperCAmelCase				,						atol=1e-3   )   )





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 						a__							:



			def __SCREAMING_SNAKE_CASE							(						self   )			->   Optional[Any]:
								torch.manual_seed(0   )
								__a		 =       TaEncoderModel.from_pretrained('hf-internal-testing/tiny-random-t5'   )

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

								torch.manual_seed(0   )
								__a		 =       UNetaDConditionModel(
								    sample_size=3_2				,						layers_per_block=1				,						block_out_channels=[3_2, 6_4]				,						down_block_types=[
								        'ResnetDownsampleBlock2D',
								        'SimpleCrossAttnDownBlock2D',
								    ]				,						mid_block_type='UNetMidBlock2DSimpleCrossAttn'				,						up_block_types=['SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D']				,						in_channels=3				,						out_channels=6				,						cross_attention_dim=3_2				,						encoder_hid_dim=3_2				,						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   )
								__a		 =       DDPMScheduler(
								    num_train_timesteps=1_0_0_0				,						beta_schedule='squaredcos_cap_v2'				,						beta_start=0.0_001				,						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   )
								__a		 =       IFWatermarker()

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



			def __SCREAMING_SNAKE_CASE							(						self   )			->   Optional[int]:
								torch.manual_seed(0   )
								__a		 =       TaEncoderModel.from_pretrained('hf-internal-testing/tiny-random-t5'   )

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

								torch.manual_seed(0   )
								__a		 =       UNetaDConditionModel(
								    sample_size=3_2				,						layers_per_block=[1, 2]				,						block_out_channels=[3_2, 6_4]				,						down_block_types=[
								        'ResnetDownsampleBlock2D',
								        'SimpleCrossAttnDownBlock2D',
								    ]				,						mid_block_type='UNetMidBlock2DSimpleCrossAttn'				,						up_block_types=['SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D']				,						in_channels=6				,						out_channels=6				,						cross_attention_dim=3_2				,						encoder_hid_dim=3_2				,						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.414				,						time_embedding_act_fn='gelu'				,						time_embedding_dim=3_2				,						)
								unet.set_attn_processor(AttnAddedKVProcessor()   )  # For reproducibility tests

								torch.manual_seed(0   )
								__a		 =       DDPMScheduler(
								    num_train_timesteps=1_0_0_0				,						beta_schedule='squaredcos_cap_v2'				,						beta_start=0.0_001				,						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   )
								__a		 =       DDPMScheduler(
								    num_train_timesteps=1_0_0_0				,						beta_schedule='squaredcos_cap_v2'				,						beta_start=0.0_001				,						beta_end=0.02				,						)

								torch.manual_seed(0   )
								__a		 =       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 __SCREAMING_SNAKE_CASE							(						self   )			->   str:
								__a		 =       self.get_dummy_components()
								__a		 =       self.pipeline_class(**UpperCAmelCase   )
								pipe.to(UpperCAmelCase   )
								pipe.set_progress_bar_config(disable=UpperCAmelCase   )

								__a		 =       self.get_dummy_inputs(UpperCAmelCase   )

								__a		 =       inputs['prompt']
								__a		 =       inputs['generator']
								__a		 =       inputs['num_inference_steps']
								__a		 =       inputs['output_type']

								if "image" in inputs:
													__a		 =       inputs['image']
								else:
													__a		 =       None

								if "mask_image" in inputs:
													__a		 =       inputs['mask_image']
								else:
													__a		 =       None

								if "original_image" in inputs:
													__a		 =       inputs['original_image']
								else:
													__a		 =       None

								__a ,				__a		 =       pipe.encode_prompt(UpperCAmelCase   )

								# inputs with prompt converted to embeddings
								__a		 =       {
								    '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:
													__a		 =       image

								if mask_image is not None:
													__a		 =       mask_image

								if original_image is not None:
													__a		 =       original_image

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

								__a		 =       pipe(**UpperCAmelCase   )[0]

								with tempfile.TemporaryDirectory() as tmpdir:
													pipe.save_pretrained(UpperCAmelCase   )
													__a		 =       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.'''				,						)

								__a		 =       self.get_dummy_inputs(UpperCAmelCase   )

								__a		 =       inputs['generator']
								__a		 =       inputs['num_inference_steps']
								__a		 =       inputs['output_type']

								# inputs with prompt converted to embeddings
								__a		 =       {
								    '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:
													__a		 =       image

								if mask_image is not None:
													__a		 =       mask_image

								if original_image is not None:
													__a		 =       original_image

								__a		 =       pipe_loaded(**UpperCAmelCase   )[0]

								__a		 =       np.abs(to_np(UpperCAmelCase   ) - to_np(UpperCAmelCase   )   ).max()
								self.assertLess(UpperCAmelCase				,						1e-4   )







			def __SCREAMING_SNAKE_CASE							(						self   )			->   int:
								__a		 =       self.get_dummy_components()
								__a		 =       self.pipeline_class(**UpperCAmelCase   )
								pipe.to(UpperCAmelCase   )
								pipe.set_progress_bar_config(disable=UpperCAmelCase   )

								__a		 =       self.get_dummy_inputs(UpperCAmelCase   )
								__a		 =       pipe(**UpperCAmelCase   )[0]

								with tempfile.TemporaryDirectory() as tmpdir:
													pipe.save_pretrained(UpperCAmelCase   )
													__a		 =       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

								__a		 =       self.get_dummy_inputs(UpperCAmelCase   )
								__a		 =       pipe_loaded(**UpperCAmelCase   )[0]

								__a		 =       np.abs(to_np(UpperCAmelCase   ) - to_np(UpperCAmelCase   )   ).max()
								self.assertLess(UpperCAmelCase				,						1e-4   )




