




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      :
							def __init__(		self		:int						, __magic_name__		:List[Any]						, __magic_name__		:List[str]=13						, __magic_name__		:Dict=30						, __magic_name__		:Optional[int]=2						, __magic_name__		:Optional[Any]=3						, __magic_name__		:Optional[Any]=True						, __magic_name__		:int=True						, __magic_name__		:Union[str, Any]=32						, __magic_name__		:List[Any]=5						, __magic_name__		:Union[str, Any]=4						, __magic_name__		:str=37						, __magic_name__		:Any="gelu"						, __magic_name__		:str=0.1						, __magic_name__		:Optional[Any]=0.1						, __magic_name__		:List[Any]=10						, __magic_name__		:Union[str, Any]=0.02						, __magic_name__		:Tuple=3						, __magic_name__		:str=0.6						, __magic_name__		:Dict=None						, ):







								'''simple docstring'''







								a									=				parent
								a									=				batch_size
								a									=				image_size
								a									=				patch_size
								a									=				num_channels
								a									=				is_training
								a									=				use_labels
								a									=				hidden_size
								a									=				num_hidden_layers
								a									=				num_attention_heads
								a									=				intermediate_size
								a									=				hidden_act
								a									=				hidden_dropout_prob
								a									=				attention_probs_dropout_prob
								a									=				type_sequence_label_size
								a									=				initializer_range
								a									=				mask_ratio
								a									=				scope

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







								'''simple docstring'''







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

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

								a									=				self.get_config()

								return config, pixel_values, labels
							def  lowerCamelCase__				(		self		:List[Any]							):







								'''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=_SCREAMING_SNAKE_CASE						, initializer_range=self.initializer_range						, mask_ratio=self.mask_ratio						, )
							def  lowerCamelCase__				(		self		:Optional[Any]						, __magic_name__		:Optional[int]						, __magic_name__		:Tuple						, __magic_name__		:Union[str, Any]							):







								'''simple docstring'''







								a									=				ViTMAEModel(config=_SCREAMING_SNAKE_CASE							)
								model.to(_SCREAMING_SNAKE_CASE							)
								model.eval()
								a									=				model(_SCREAMING_SNAKE_CASE							)
								self.parent.assertEqual(result.last_hidden_state.shape						, (self.batch_size, self.seq_length, self.hidden_size)							)
							def  lowerCamelCase__				(		self		:Tuple						, __magic_name__		:Optional[int]						, __magic_name__		:Optional[int]						, __magic_name__		:List[str]							):







								'''simple docstring'''







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

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





							def  lowerCamelCase__				(		self		:List[Any]							):







								'''simple docstring'''







								a									=				self.prepare_config_and_inputs()
								a									=				config_and_inputs
								a									=				{'pixel_values': pixel_values}
								return config, inputs_dict







@require_torch
class 						__lowerCAmelCase      (      _lowerCamelCase ,			_lowerCamelCase ,			unittest.TestCase					):
							UpperCamelCase__   	=    (ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
							UpperCamelCase__   	=    {'''feature-extraction''': ViTMAEModel} if is_torch_available() else {}

							UpperCamelCase__   	=    False
							UpperCamelCase__   	=    False
							UpperCamelCase__   	=    False
							UpperCamelCase__   	=    False
							def  lowerCamelCase__				(		self		:Optional[Any]							):







								'''simple docstring'''







								a									=				ViTMAEModelTester(self							)
								a									=				ConfigTester(self						, config_class=_SCREAMING_SNAKE_CASE						, has_text_modality=_SCREAMING_SNAKE_CASE						, hidden_size=37							)
							def  lowerCamelCase__				(		self		:Any							):







								'''simple docstring'''







								self.config_tester.run_common_tests()
							@unittest.skip(reason="""ViTMAE does not use inputs_embeds"""							)
							def  lowerCamelCase__				(		self		:Optional[int]							):







								'''simple docstring'''







								pass
							def  lowerCamelCase__				(		self		:Dict							):







								'''simple docstring'''







								a									=				self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
									a									=				model_class(_SCREAMING_SNAKE_CASE							)
									self.assertIsInstance(model.get_input_embeddings()						, (nn.Module)							)
									a									=				model.get_output_embeddings()
									self.assertTrue(x is None or isinstance(_SCREAMING_SNAKE_CASE						, nn.Linear							)							)
							def  lowerCamelCase__				(		self		:Optional[Any]							):







								'''simple docstring'''







								a									=				self.model_tester.prepare_config_and_inputs_for_common()

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

									a									=				['pixel_values']
									self.assertListEqual(arg_names[:1]						, _SCREAMING_SNAKE_CASE							)
							def  lowerCamelCase__				(		self		:Any							):







								'''simple docstring'''







								a									=				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE							)
							def  lowerCamelCase__				(		self		:Optional[int]							):







								'''simple docstring'''







								a									=				self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*_SCREAMING_SNAKE_CASE							)
							def  lowerCamelCase__				(		self		:int						, __magic_name__		:Any						, __magic_name__		:Optional[int]						, __magic_name__		:Dict							):







								'''simple docstring'''







								np.random.seed(2							)

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

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

								super().check_pt_tf_models(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE							)
							def  lowerCamelCase__				(		self		:Optional[int]							):







								'''simple docstring'''







								a									=				self.model_tester.prepare_config_and_inputs_for_common()

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

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

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

										# Make sure we don't have nans
										a									=				after_outputs[0].cpu().numpy()
										a									=				0
										a									=				np.amax(np.abs(out_a - out_a							)							)
										self.assertLessEqual(_SCREAMING_SNAKE_CASE						, 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  lowerCamelCase__				(		self		:int							):







								'''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  lowerCamelCase__				(		self		:Any							):







								'''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  lowerCamelCase__				(		self		:Union[str, Any]							):







								'''simple docstring'''







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







								'''simple docstring'''







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







								'''simple docstring'''







								pass





							@slow
							def  lowerCamelCase__				(		self		:Union[str, Any]							):







								'''simple docstring'''







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







def 						__A						(      ) ->			Dict:
	a									=				Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""       )
	return image



@require_torch
@require_vision
class 						__lowerCAmelCase      (      unittest.TestCase					):
							@cached_property
							def  lowerCamelCase__				(		self		:List[Any]							):







								'''simple docstring'''







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





							@slow
							def  lowerCamelCase__				(		self		:Optional[int]							):







								'''simple docstring'''







								np.random.seed(2							)

								a									=				ViTMAEForPreTraining.from_pretrained("""facebook/vit-mae-base"""							).to(_SCREAMING_SNAKE_CASE							)

								a									=				self.default_image_processor
								a									=				prepare_img()
								a									=				image_processor(images=_SCREAMING_SNAKE_CASE						, return_tensors="""pt"""							).to(_SCREAMING_SNAKE_CASE							)

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

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

								# verify the logits
								a									=				torch.Size((1, 196, 768)							)
								self.assertEqual(outputs.logits.shape						, _SCREAMING_SNAKE_CASE							)

								a									=				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(_SCREAMING_SNAKE_CASE							)						, atol=1E-4							)							)






"""simple docstring"""

import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()
@slow
@require_torch_gpu
class     lowerCAmelCase__						(	unittest.TestCase ):


	'''simple docstring'''







	def 			_lowerCAmelCase     (      self				:       List[Any]    )						->	str:





					"""simple docstring"""


					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()





	def 			_lowerCAmelCase     (      self				:       Any    )						->	str:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							int							=     StableDiffusionKDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'    )
					SCREAMING_SNAKE_CASE						:							Dict							=     sd_pipe.to(_SCREAMING_SNAKE_CASE    )
					sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					sd_pipe.set_scheduler('sample_euler'    )

					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     'A painting of a squirrel eating a burger'
					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     torch.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     sd_pipe([prompt]     ,       generator=_SCREAMING_SNAKE_CASE     ,       guidance_scale=9.0     ,       num_inference_steps=20     ,       output_type='np'    )

					SCREAMING_SNAKE_CASE						:							str							=     output.images

					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     image[0, -3:, -3:, -1]

					assert image.shape == (1, 512, 512, 3)
					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     np.array([0.0_4_4_7, 0.0_4_9_2, 0.0_4_6_8, 0.0_4_0_8, 0.0_3_8_3, 0.0_4_0_8, 0.0_3_5_4, 0.0_3_8_0, 0.0_3_3_9]    )

					assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2





	def 			_lowerCAmelCase     (      self				:       Dict    )						->	Optional[int]:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     StableDiffusionKDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'    )
					SCREAMING_SNAKE_CASE						:							List[Any]							=     sd_pipe.to(_SCREAMING_SNAKE_CASE    )
					sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					sd_pipe.set_scheduler('sample_euler'    )

					SCREAMING_SNAKE_CASE						:							List[str]							=     'A painting of a squirrel eating a burger'
					SCREAMING_SNAKE_CASE						:							Optional[int]							=     torch.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							Dict							=     sd_pipe([prompt]     ,       generator=_SCREAMING_SNAKE_CASE     ,       guidance_scale=9.0     ,       num_inference_steps=20     ,       output_type='np'    )

					SCREAMING_SNAKE_CASE						:							Optional[int]							=     output.images

					SCREAMING_SNAKE_CASE						:							List[str]							=     image[0, -3:, -3:, -1]

					assert image.shape == (1, 512, 512, 3)
					SCREAMING_SNAKE_CASE						:							List[Any]							=     np.array([0.1_2_3_7, 0.1_3_2_0, 0.1_4_3_8, 0.1_3_5_9, 0.1_3_9_0, 0.1_1_3_2, 0.1_2_7_7, 0.1_1_7_5, 0.1_1_1_2]    )

					assert np.abs(image_slice.flatten() - expected_slice    ).max() < 5E-1







	def 			_lowerCAmelCase     (      self				:       Dict    )						->	Optional[Any]:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Optional[int]							=     StableDiffusionKDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'    )
					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     sd_pipe.to(_SCREAMING_SNAKE_CASE    )
					sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					sd_pipe.set_scheduler('sample_dpmpp_2m'    )

					SCREAMING_SNAKE_CASE						:							Optional[int]							=     'A painting of a squirrel eating a burger'
					SCREAMING_SNAKE_CASE						:							str							=     torch.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     sd_pipe(
					    [prompt]     ,       generator=_SCREAMING_SNAKE_CASE     ,       guidance_scale=7.5     ,       num_inference_steps=15     ,       output_type='np'     ,       use_karras_sigmas=_SCREAMING_SNAKE_CASE     ,       )

					SCREAMING_SNAKE_CASE						:							int							=     output.images

					SCREAMING_SNAKE_CASE						:							List[Any]							=     image[0, -3:, -3:, -1]

					assert image.shape == (1, 512, 512, 3)
					SCREAMING_SNAKE_CASE						:							int							=     np.array(
					    [0.1_1_3_8_1_6_8_9, 0.1_2_1_1_2_9_2_1, 0.1_3_8_9_4_5_7, 0.1_2_5_4_9_6_0_6, 0.1_2_4_4_9_6_4, 0.1_0_8_3_1_5_1_7, 0.1_1_5_6_2_8_6_6, 0.1_0_8_6_7_8_1_6, 0.1_0_4_9_9_0_4_8]    )

					assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2





