




from __future__ import annotations

import copy
import inspect
import json
import math
import os
import tempfile
import unittest
from importlib import import_module

import numpy as np

from transformers import ViTMAEConfig
from transformers.file_utils import cached_property, is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, require_vision, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
				import tensorflow as tf

				from transformers import TFViTMAEForPreTraining, TFViTMAEModel


if is_vision_available():
				from PIL import Image

				from transformers import ViTImageProcessor




class    _A					:






		def __init__(     self    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE=13    ,		_SCREAMING_SNAKE_CASE=30    ,		_SCREAMING_SNAKE_CASE=2    ,		_SCREAMING_SNAKE_CASE=3    ,		_SCREAMING_SNAKE_CASE=True    ,		_SCREAMING_SNAKE_CASE=True    ,		_SCREAMING_SNAKE_CASE=32    ,		_SCREAMING_SNAKE_CASE=2    ,		_SCREAMING_SNAKE_CASE=4    ,		_SCREAMING_SNAKE_CASE=37    ,		_SCREAMING_SNAKE_CASE="gelu"    ,		_SCREAMING_SNAKE_CASE=0.1    ,		_SCREAMING_SNAKE_CASE=0.1    ,		_SCREAMING_SNAKE_CASE=10    ,		_SCREAMING_SNAKE_CASE=0.02    ,		_SCREAMING_SNAKE_CASE=3    ,		_SCREAMING_SNAKE_CASE=0.6    ,		_SCREAMING_SNAKE_CASE=None    ,		):
								_UpperCAmelCase	     =      parent
								_UpperCAmelCase	     =      batch_size
								_UpperCAmelCase	     =      image_size
								_UpperCAmelCase	     =      patch_size
								_UpperCAmelCase	     =      num_channels
								_UpperCAmelCase	     =      is_training
								_UpperCAmelCase	     =      use_labels
								_UpperCAmelCase	     =      hidden_size
								_UpperCAmelCase	     =      num_hidden_layers
								_UpperCAmelCase	     =      num_attention_heads
								_UpperCAmelCase	     =      intermediate_size
								_UpperCAmelCase	     =      hidden_act
								_UpperCAmelCase	     =      hidden_dropout_prob
								_UpperCAmelCase	     =      attention_probs_dropout_prob
								_UpperCAmelCase	     =      type_sequence_label_size
								_UpperCAmelCase	     =      initializer_range
								_UpperCAmelCase	     =      mask_ratio
								_UpperCAmelCase	     =      scope

								# in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
								# (we add 1 for the [CLS] token)
								_UpperCAmelCase	     =      (image_size // patch_size) ** 2
								_UpperCAmelCase	     =      int(math.ceil((1 - mask_ratio) * (num_patches + 1)      )      )






		def 					UpperCAmelCase			(     self      ):
								_UpperCAmelCase	     =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

								_UpperCAmelCase	     =      None
								if self.use_labels:
														_UpperCAmelCase	     =      ids_tensor([self.batch_size]    ,		self.type_sequence_label_size      )

								_UpperCAmelCase	     =      self.get_config()

								return config, pixel_values, labels






		def 					UpperCAmelCase			(     self      ):
								return ViTMAEConfig(
								    image_size=self.image_size    ,		patch_size=self.patch_size    ,		num_channels=self.num_channels    ,		hidden_size=self.hidden_size    ,		num_hidden_layers=self.num_hidden_layers    ,		num_attention_heads=self.num_attention_heads    ,		intermediate_size=self.intermediate_size    ,		decoder_hidden_size=self.hidden_size    ,		decoder_num_hidden_layers=self.num_hidden_layers    ,		decoder_num_attention_heads=self.num_attention_heads    ,		decoder_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    ,		is_decoder=__A    ,		initializer_range=self.initializer_range    ,		mask_ratio=self.mask_ratio    ,		)






		def 					UpperCAmelCase			(     self    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE      ):
								_UpperCAmelCase	     =      TFViTMAEModel(config=__A      )
								_UpperCAmelCase	     =      model(__A    ,		training=__A      )
								self.parent.assertEqual(result.last_hidden_state.shape    ,		(self.batch_size, self.seq_length, self.hidden_size)      )






		def 					UpperCAmelCase			(     self    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE      ):
								_UpperCAmelCase	     =      TFViTMAEForPreTraining(__A      )
								_UpperCAmelCase	     =      model(__A    ,		training=__A      )
								# expected sequence length = num_patches
								_UpperCAmelCase	     =      (self.image_size // self.patch_size) ** 2
								_UpperCAmelCase	     =      self.patch_size**2 * self.num_channels
								self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, num_patches, expected_num_channels)      )

								# test greyscale images
								_UpperCAmelCase	     =      1
								_UpperCAmelCase	     =      TFViTMAEForPreTraining(__A      )

								_UpperCAmelCase	     =      floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
								_UpperCAmelCase	     =      model(__A    ,		training=__A      )
								_UpperCAmelCase	     =      self.patch_size**2
								self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, num_patches, expected_num_channels)      )






		def 					UpperCAmelCase			(     self      ):
								_UpperCAmelCase	     =      self.prepare_config_and_inputs()
								(_UpperCAmelCase)	     =      config_and_inputs
								_UpperCAmelCase	     =      {"""pixel_values""": pixel_values}
								return config, inputs_dict




@require_tf
class    _A					(		_lowercase							,  _lowercase							,  unittest.TestCase  ):
		__a            =							(TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
		__a            =							{'''feature-extraction''': TFViTMAEModel} if is_tf_available() else {}

		__a            =							False
		__a            =							False
		__a            =							False
		__a            =							False






		def 					UpperCAmelCase			(     self      ):
								_UpperCAmelCase	     =      TFViTMAEModelTester(self      )
								_UpperCAmelCase	     =      ConfigTester(self    ,		config_class=__A    ,		has_text_modality=__A    ,		hidden_size=37      )






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






		@unittest.skip(reason="""ViTMAE does not use inputs_embeds"""      )
		def 					UpperCAmelCase			(     self      ):
								pass






		def 					UpperCAmelCase			(     self      ):
								_UpperCAmelCase	     =      self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
														_UpperCAmelCase	     =      model_class(__A      )
														self.assertIsInstance(model.get_input_embeddings()    ,		(tf.keras.layers.Layer)      )
														_UpperCAmelCase	     =      model.get_output_embeddings()
														self.assertTrue(x is None or isinstance(__A    ,		tf.keras.layers.Layer      )      )






		def 					UpperCAmelCase			(     self      ):
								_UpperCAmelCase	     =      self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
														_UpperCAmelCase	     =      model_class(__A      )
														_UpperCAmelCase	     =      inspect.signature(model.call      )
														# signature.parameters is an OrderedDict => so arg_names order is deterministic
														_UpperCAmelCase	     =      [*signature.parameters.keys()]

														_UpperCAmelCase	     =      ["""pixel_values"""]
														self.assertListEqual(arg_names[:1]    ,		__A      )






		def 					UpperCAmelCase			(     self      ):
								_UpperCAmelCase	     =      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__A      )






		def 					UpperCAmelCase			(     self      ):
								_UpperCAmelCase	     =      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*__A      )






		def 					UpperCAmelCase			(     self      ):
								# make the mask reproducible
								np.random.seed(2      )

								_UpperCAmelCase	     =      self.model_tester.prepare_config_and_inputs_for_common()

								_UpperCAmelCase	     =      int((config.image_size // config.patch_size) ** 2      )
								_UpperCAmelCase	     =      np.random.uniform(size=(self.model_tester.batch_size, num_patches)      )

								for model_class in self.all_model_classes:
														_UpperCAmelCase	     =      model_class(__A      )
														_UpperCAmelCase	     =      self._prepare_for_class(__A    ,		__A      )

														_UpperCAmelCase	     =      model(__A    ,		noise=__A      )

														_UpperCAmelCase	     =      copy.deepcopy(self._prepare_for_class(__A    ,		__A      )      )
														_UpperCAmelCase	     =      model(**__A    ,		noise=__A      )
														_UpperCAmelCase	     =      outputs_dict[0].numpy()
														_UpperCAmelCase	     =      outputs_keywords[0].numpy()

														self.assertLess(np.sum(np.abs(output_dict - output_keywords      )      )    ,		1e-6      )






		def 					UpperCAmelCase			(     self      ):
								# make the mask reproducible
								np.random.seed(2      )

								_UpperCAmelCase	     =      self.model_tester.prepare_config_and_inputs_for_common()

								_UpperCAmelCase	     =      int((config.image_size // config.patch_size) ** 2      )
								_UpperCAmelCase	     =      np.random.uniform(size=(self.model_tester.batch_size, num_patches)      )

								def prepare_numpy_arrays(_SCREAMING_SNAKE_CASE      ):
														_UpperCAmelCase	     =      {}
														for k, v in inputs_dict.items():
																				if tf.is_tensor(__A      ):
																										_UpperCAmelCase	     =      v.numpy()
																				else:
																										_UpperCAmelCase	     =      np.array(__A      )

														return inputs_np_dict

								for model_class in self.all_model_classes:
														_UpperCAmelCase	     =      model_class(__A      )

														_UpperCAmelCase	     =      self._prepare_for_class(__A    ,		__A      )
														_UpperCAmelCase	     =      prepare_numpy_arrays(__A      )

														_UpperCAmelCase	     =      model(__A    ,		noise=__A      )
														_UpperCAmelCase	     =      model(**__A    ,		noise=__A      )
														self.assert_outputs_same(__A    ,		__A      )






		def 					UpperCAmelCase			(     self    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE      ):
								# make masks reproducible
								np.random.seed(2      )

								_UpperCAmelCase	     =      int((tf_model.config.image_size // tf_model.config.patch_size) ** 2      )
								_UpperCAmelCase	     =      np.random.uniform(size=(self.model_tester.batch_size, num_patches)      )
								_UpperCAmelCase	     =      tf.constant(__A      )

								# Add `noise` argument.
								# PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
								_UpperCAmelCase	     =      tf_noise

								super().check_pt_tf_models(__A    ,		__A    ,		__A      )






		def 					UpperCAmelCase			(     self      ):
								# make mask reproducible
								np.random.seed(2      )

								_UpperCAmelCase	     =      self.model_tester.prepare_config_and_inputs_for_common()

								_UpperCAmelCase	     =      {
								    module_member
								    for model_class in self.all_model_classes
								    for module in (import_module(model_class.__module__      ),)
								    for module_member_name in dir(__A      )
								    if module_member_name.endswith("""MainLayer"""      )
								    # This condition is required, since `modeling_tf_clip.py` has 3 classes whose names end with `MainLayer`.
								    and module_member_name[: -len("""MainLayer"""      )] == model_class.__name__[: -len("""Model"""      )]
								    for module_member in (getattr(__A    ,		__A      ),)
								    if isinstance(__A    ,		__A      )
								    and tf.keras.layers.Layer in module_member.__bases__
								    and getattr(__A    ,		"""_keras_serializable"""    ,		__A      )
								}

								_UpperCAmelCase	     =      int((config.image_size // config.patch_size) ** 2      )
								_UpperCAmelCase	     =      np.random.uniform(size=(self.model_tester.batch_size, num_patches)      )
								_UpperCAmelCase	     =      tf.convert_to_tensor(__A      )
								inputs_dict.update({"""noise""": noise}      )

								for main_layer_class in tf_main_layer_classes:
														_UpperCAmelCase	     =      main_layer_class(__A      )

														_UpperCAmelCase	     =      {
														    name: tf.keras.Input(tensor.shape[1:]    ,		dtype=tensor.dtype      ) for name, tensor in inputs_dict.items()
														}

														_UpperCAmelCase	     =      tf.keras.Model(__A    ,		outputs=main_layer(__A      )      )
														_UpperCAmelCase	     =      model(__A      )

														with tempfile.TemporaryDirectory() as tmpdirname:
																				_UpperCAmelCase	     =      os.path.join(__A    ,		"""keras_model.h5"""      )
																				model.save(__A      )
																				_UpperCAmelCase	     =      tf.keras.models.load_model(
																				    __A    ,		custom_objects={main_layer_class.__name__: main_layer_class}      )
																				assert isinstance(__A    ,		tf.keras.Model      )
																				_UpperCAmelCase	     =      model(__A      )
																				self.assert_outputs_same(__A    ,		__A      )






		@slow
		def 					UpperCAmelCase			(     self      ):
								# make mask reproducible
								np.random.seed(2      )

								_UpperCAmelCase	     =      self.model_tester.prepare_config_and_inputs_for_common()

								_UpperCAmelCase	     =      int((config.image_size // config.patch_size) ** 2      )
								_UpperCAmelCase	     =      np.random.uniform(size=(self.model_tester.batch_size, num_patches)      )

								for model_class in self.all_model_classes:
														_UpperCAmelCase	     =      model_class(__A      )
														_UpperCAmelCase	     =      self._prepare_for_class(__A    ,		__A      )
														_UpperCAmelCase	     =      model(__A    ,		noise=__A      )

														if model_class.__name__ == "TFViTMAEModel":
																				_UpperCAmelCase	     =      outputs.last_hidden_state.numpy()
																				_UpperCAmelCase	     =      0
														else:
																				_UpperCAmelCase	     =      outputs.logits.numpy()
																				_UpperCAmelCase	     =      0

														with tempfile.TemporaryDirectory() as tmpdirname:
																				model.save_pretrained(__A    ,		saved_model=__A      )
																				_UpperCAmelCase	     =      model_class.from_pretrained(__A      )
																				_UpperCAmelCase	     =      model(__A    ,		noise=__A      )

																				if model_class.__name__ == "TFViTMAEModel":
																										_UpperCAmelCase	     =      after_outputs["""last_hidden_state"""].numpy()
																										_UpperCAmelCase	     =      0
																				else:
																										_UpperCAmelCase	     =      after_outputs["""logits"""].numpy()
																										_UpperCAmelCase	     =      0

																				_UpperCAmelCase	     =      np.amax(np.abs(out_a - out_a      )      )
																				self.assertLessEqual(__A    ,		1e-5      )






		def 					UpperCAmelCase			(     self      ):
								# make mask reproducible
								np.random.seed(2      )

								_UpperCAmelCase	     =      self.model_tester.prepare_config_and_inputs_for_common()

								_UpperCAmelCase	     =      int((config.image_size // config.patch_size) ** 2      )
								_UpperCAmelCase	     =      np.random.uniform(size=(self.model_tester.batch_size, num_patches)      )

								for model_class in self.all_model_classes:
														_UpperCAmelCase	     =      model_class(__A      )
														_UpperCAmelCase	     =      self._prepare_for_class(__A    ,		__A      )

														_UpperCAmelCase	     =      model(__A    ,		noise=__A      )
														_UpperCAmelCase	     =      model.get_config()
														# make sure that returned config is jsonifiable, which is required by keras
														json.dumps(__A      )
														_UpperCAmelCase	     =      model_class.from_config(model.get_config()      )
														# make sure it also accepts a normal config
														_UpperCAmelCase	     =      model_class.from_config(model.config      )
														_UpperCAmelCase	     =      new_model(__A      )  # Build model
														new_model.set_weights(model.get_weights()      )
														_UpperCAmelCase	     =      new_model(__A    ,		noise=__A      )

														self.assert_outputs_same(__A    ,		__A      )






		@unittest.skip(
		    reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load
    to get deterministic results."""      )
		def 					UpperCAmelCase			(     self      ):
								pass






		@unittest.skip(reason="""ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"""      )
		def 					UpperCAmelCase			(     self      ):
								pass






		@slow
		def 					UpperCAmelCase			(     self      ):
								_UpperCAmelCase	     =      TFViTMAEModel.from_pretrained("""google/vit-base-patch16-224"""      )
								self.assertIsNotNone(__A      )




def 							_SCREAMING_SNAKE_CASE   (		)    ->					Optional[Any]:
						_UpperCAmelCase	     =      Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""							)
						return image




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






		@cached_property
		def 					UpperCAmelCase			(     self      ):
								return ViTImageProcessor.from_pretrained("""facebook/vit-mae-base"""      ) if is_vision_available() else None






		@slow
		def 					UpperCAmelCase			(     self      ):
								# make random mask reproducible across the PT and TF model
								np.random.seed(2      )

								_UpperCAmelCase	     =      TFViTMAEForPreTraining.from_pretrained("""facebook/vit-mae-base"""      )

								_UpperCAmelCase	     =      self.default_image_processor
								_UpperCAmelCase	     =      prepare_img()
								_UpperCAmelCase	     =      image_processor(images=__A    ,		return_tensors="""tf"""      )

								# prepare a noise vector that will be also used for testing the TF model
								# (this way we can ensure that the PT and TF models operate on the same inputs)
								_UpperCAmelCase	     =      ViTMAEConfig()
								_UpperCAmelCase	     =      int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2      )
								_UpperCAmelCase	     =      np.random.uniform(size=(1, num_patches)      )

								# forward pass
								_UpperCAmelCase	     =      model(**__A    ,		noise=__A      )

								# verify the logits
								_UpperCAmelCase	     =      tf.convert_to_tensor([1, 196, 768]      )
								self.assertEqual(outputs.logits.shape    ,		__A      )

								_UpperCAmelCase	     =      tf.convert_to_tensor(
								    [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]]      )

								tf.debugging.assert_near(outputs.logits[0, :3, :3]    ,		__A    ,		atol=1e-4      )


import warnings

from diffusers import StableDiffusionInpaintPipeline as StableDiffusionInpaintPipeline  # noqa F401


warnings.warn(
    '''The `inpainting.py` script is outdated. Please use directly `from diffusers import'''
    ''' StableDiffusionInpaintPipeline` instead.'''
)





