



import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
     import torch
     from torch import nn

     from transformers import ViTMAEForPreTraining, ViTMAEModel
     from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import ViTImageProcessor





class    lowerCamelCase__  :
     """simple docstring"""



     def __init__(self				,			__a				,			__a=13				,			__a=30				,			__a=2				,			__a=3				,			__a=True				,			__a=True				,			__a=32				,			__a=5				,			__a=4				,			__a=37				,			__a="gelu"				,			__a=0.1				,			__a=0.1				,			__a=10				,			__a=0.02				,			__a=3				,			__a=0.6				,			__a=None				,			):






           '''simple docstring'''



           lowerCamelCase								=		parent
           lowerCamelCase								=		batch_size
           lowerCamelCase								=		image_size
           lowerCamelCase								=		patch_size
           lowerCamelCase								=		num_channels
           lowerCamelCase								=		is_training
           lowerCamelCase								=		use_labels
           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								=		type_sequence_label_size
           lowerCamelCase								=		initializer_range
           lowerCamelCase								=		mask_ratio
           lowerCamelCase								=		scope

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







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

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

           lowerCamelCase								=		self.get_config()

           return config, pixel_values, labels







     def 						_a      (self    ):






           '''simple docstring'''



           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				,			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 						_a      (self				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		ViTMAEModel(config=__a    )
           model.to(__a    )
           model.eval()
           lowerCamelCase								=		model(__a    )
           self.parent.assertEqual(result.last_hidden_state.shape				,			(self.batch_size, self.seq_length, self.hidden_size)    )







     def 						_a      (self				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		ViTMAEForPreTraining(__a    )
           model.to(__a    )
           model.eval()
           lowerCamelCase								=		model(__a    )
           lowerCamelCase								=		(self.image_size // self.patch_size) ** 2
           lowerCamelCase								=		self.patch_size**2 * self.num_channels
           self.parent.assertEqual(result.logits.shape				,			(self.batch_size, num_patches, expected_num_channels)    )

           # test greyscale images
           lowerCamelCase								=		1
           lowerCamelCase								=		ViTMAEForPreTraining(__a    )
           model.to(__a    )
           model.eval()
           lowerCamelCase								=		floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
           lowerCamelCase								=		model(__a    )
           lowerCamelCase								=		self.patch_size**2
           self.parent.assertEqual(result.logits.shape				,			(self.batch_size, num_patches, expected_num_channels)    )







     def 						_a      (self    ):






           '''simple docstring'''



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





@require_torch
class    lowerCamelCase__  (						UpperCAmelCase_       ,    UpperCAmelCase_       ,    unittest.TestCase):
     """simple docstring"""



     _A       =			(ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
     _A       =			{'feature-extraction': ViTMAEModel} if is_torch_available() else {}

     _A       =			False
     _A       =			False
     _A       =			False
     _A       =			False







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		ViTMAEModelTester(self    )
           lowerCamelCase								=		ConfigTester(self				,			config_class=__a				,			has_text_modality=__a				,			hidden_size=37    )







     def 						_a      (self    ):






           '''simple docstring'''



           self.config_tester.run_common_tests()







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






           '''simple docstring'''



           pass







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase			,							lowerCamelCase								=		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 lowerCamelCase								=		model_class(__a    )
                 self.assertIsInstance(model.get_input_embeddings()				,			(nn.Module)    )
                 lowerCamelCase								=		model.get_output_embeddings()
                 self.assertTrue(x is None or isinstance(__a				,			nn.Linear    )    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase			,							lowerCamelCase								=		self.model_tester.prepare_config_and_inputs_for_common()

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

                 lowerCamelCase								=		["pixel_values"]
                 self.assertListEqual(arg_names[:1]				,			__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_pretraining(*__a    )







     def 						_a      (self				,			__a				,			__a				,			__a    ):






           '''simple docstring'''



           np.random.seed(2    )

           lowerCamelCase								=		int((pt_model.config.image_size // pt_model.config.patch_size) ** 2    )
           lowerCamelCase								=		np.random.uniform(size=(self.model_tester.batch_size, num_patches)    )
           lowerCamelCase								=		torch.from_numpy(__a    )

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

           super().check_pt_tf_models(__a				,			__a				,			__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase			,							lowerCamelCase								=		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 lowerCamelCase								=		model_class(__a    )
                 model.to(__a    )
                 model.eval()
                 # make random mask reproducible
                 torch.manual_seed(2    )
                 with torch.no_grad():
                       lowerCamelCase								=		model(**self._prepare_for_class(__a				,			__a    )    )

                 lowerCamelCase								=		outputs[0].cpu().numpy()
                 lowerCamelCase								=		0

                 with tempfile.TemporaryDirectory() as tmpdirname:
                       model.save_pretrained(__a    )
                       lowerCamelCase								=		model_class.from_pretrained(__a    )
                       model.to(__a    )
                       # make random mask reproducible
                       torch.manual_seed(2    )
                       with torch.no_grad():
                             lowerCamelCase								=		model(**self._prepare_for_class(__a				,			__a    )    )

                       # Make sure we don't have nans
                       lowerCamelCase								=		after_outputs[0].cpu().numpy()
                       lowerCamelCase								=		0
                       lowerCamelCase								=		np.amax(np.abs(out_a - out_a    )    )
                       self.assertLessEqual(__a				,			1E-5    )







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






           '''simple docstring'''



           pass







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






           '''simple docstring'''



           pass







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






           '''simple docstring'''



           pass







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






           '''simple docstring'''



           pass







     @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."    )
     def 						_a      (self    ):






           '''simple docstring'''



           pass







     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		ViTMAEModel.from_pretrained(__a    )
                 self.assertIsNotNone(__a    )



def     __lowercase(	):


      """simple docstring"""



      lowerCamelCase								=		Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png" )
      return image






@require_torch
@require_vision
class    lowerCamelCase__  (						unittest.TestCase):
     """simple docstring"""



     @cached_property
     def 						_a      (self    ):






           '''simple docstring'''



           return ViTImageProcessor.from_pretrained("facebook/vit-mae-base"    ) if is_vision_available() else None







     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           np.random.seed(2    )

           lowerCamelCase								=		ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"    ).to(__a    )

           lowerCamelCase								=		self.default_image_processor
           lowerCamelCase								=		prepare_img()
           lowerCamelCase								=		image_processor(images=__a				,			return_tensors="pt"    ).to(__a    )

           # 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)
           lowerCamelCase								=		ViTMAEConfig()
           lowerCamelCase								=		int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2    )
           lowerCamelCase								=		np.random.uniform(size=(1, num_patches)    )

           # forward pass
           with torch.no_grad():
                 lowerCamelCase								=		model(**__a				,			noise=torch.from_numpy(__a    ).to(device=__a    )    )

           # verify the logits
           lowerCamelCase								=		torch.Size((1, 1_96, 7_68)    )
           self.assertEqual(outputs.logits.shape				,			__a    )

           lowerCamelCase								=		torch.tensor(
               [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]]    )

           self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]				,			expected_slice.to(__a    )				,			atol=1E-4    )    )




import gc
import unittest

from diffusers import FlaxStableDiffusionInpaintPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
     import jax
     import jax.numpy as jnp
     from flax.jax_utils import replicate
     from flax.training.common_utils import shard





@slow
@require_flax
class    lowerCamelCase__  (						unittest.TestCase):
     """simple docstring"""



     def 						_a      (self    ):






           '''simple docstring'''



           super().tearDown()
           gc.collect()







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		load_image(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
               "/sd2-inpaint/init_image.png"    )
           lowerCamelCase								=		load_image(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"    )

           lowerCamelCase								=		"xvjiarui/stable-diffusion-2-inpainting"
           lowerCamelCase			,							lowerCamelCase								=		FlaxStableDiffusionInpaintPipeline.from_pretrained(__a				,			safety_checker=__a    )

           lowerCamelCase								=		"Face of a yellow cat, high resolution, sitting on a park bench"

           lowerCamelCase								=		jax.random.PRNGKey(0    )
           lowerCamelCase								=		50

           lowerCamelCase								=		jax.device_count()
           lowerCamelCase								=		num_samples * [prompt]
           lowerCamelCase								=		num_samples * [init_image]
           lowerCamelCase								=		num_samples * [mask_image]
           lowerCamelCase			,							lowerCamelCase			,							lowerCamelCase								=		pipeline.prepare_inputs(__a				,			__a				,			__a    )

           # shard inputs and rng
           lowerCamelCase								=		replicate(__a    )
           lowerCamelCase								=		jax.random.split(__a				,			jax.device_count()    )
           lowerCamelCase								=		shard(__a    )
           lowerCamelCase								=		shard(__a    )
           lowerCamelCase								=		shard(__a    )

           lowerCamelCase								=		pipeline(
               __a				,			__a				,			__a				,			__a				,			__a				,			__a				,			jit=__a    )

           lowerCamelCase								=		output.images.reshape(__a				,			5_12				,			5_12				,			3    )

           lowerCamelCase								=		images[0, 2_53:2_56, 2_53:2_56, -1]

           lowerCamelCase								=		jnp.asarray(jax.device_get(image_slice.flatten()    )    )
           lowerCamelCase								=		jnp.array(
               [0.361_1307, 0.3764_9736, 0.375_7408, 0.3821_3953, 0.3929_5167, 0.384_1631, 0.4155_4978, 0.413_7475, 0.421_7084]    )
           print(F"""output_slice: {output_slice}"""    )
           assert jnp.abs(output_slice - expected_slice    ).max() < 1E-2