



import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    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 ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
			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	:       List[Any]			,     lowercase	:       Union[str, Any]			,     lowercase	:       Tuple=13			,     lowercase	:       Union[str, Any]=30			,     lowercase	:       List[str]=2			,     lowercase	:       Optional[Any]=3			,     lowercase	:       List[Any]=True			,     lowercase	:       Dict=True			,     lowercase	:       List[Any]=32			,     lowercase	:       int=5			,     lowercase	:       Tuple=4			,     lowercase	:       List[str]=37			,     lowercase	:       Any="gelu"			,     lowercase	:       Any=0.1			,     lowercase	:       Dict=0.1			,     lowercase	:       Optional[Any]=10			,     lowercase	:       Tuple=0.02			,     lowercase	:       str=None			,     lowercase	:       int=2			,     )						->					Tuple:
									"""simple docstring"""







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

									# in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
									__lowercase					=				(image_size // patch_size) ** 2
									__lowercase					=				num_patches + 1





			def 		snake_case__			(				self	:       str   )						->					str:
									"""simple docstring"""







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

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

									__lowercase					=				self.get_config()

									return config, pixel_values, labels





			def 		snake_case__			(				self	:       Optional[Any]   )						->					Optional[int]:
									"""simple docstring"""







									return ViTConfig(
									    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=lowercase			,     initializer_range=self.initializer_range			,     encoder_stride=self.encoder_stride			,     )





			def 		snake_case__			(				self	:       Optional[int]			,     lowercase	:       Tuple			,     lowercase	:       Tuple			,     lowercase	:       Optional[int]   )						->					int:
									"""simple docstring"""







									__lowercase					=				ViTModel(config=lowercase   )
									model.to(lowercase   )
									model.eval()
									__lowercase					=				model(lowercase   )
									self.parent.assertEqual(result.last_hidden_state.shape			,     (self.batch_size, self.seq_length, self.hidden_size)   )





			def 		snake_case__			(				self	:       Tuple			,     lowercase	:       str			,     lowercase	:       Optional[int]			,     lowercase	:       Dict   )						->					Any:
									"""simple docstring"""







									__lowercase					=				ViTForMaskedImageModeling(config=lowercase   )
									model.to(lowercase   )
									model.eval()
									__lowercase					=				model(lowercase   )
									self.parent.assertEqual(
									    result.reconstruction.shape			,     (self.batch_size, self.num_channels, self.image_size, self.image_size)   )

									# test greyscale images
									__lowercase					=				1
									__lowercase					=				ViTForMaskedImageModeling(lowercase   )
									model.to(lowercase   )
									model.eval()

									__lowercase					=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
									__lowercase					=				model(lowercase   )
									self.parent.assertEqual(result.reconstruction.shape			,     (self.batch_size, 1, self.image_size, self.image_size)   )





			def 		snake_case__			(				self	:       Optional[Any]			,     lowercase	:       str			,     lowercase	:       int			,     lowercase	:       List[str]   )						->					Any:
									"""simple docstring"""







									__lowercase					=				self.type_sequence_label_size
									__lowercase					=				ViTForImageClassification(lowercase   )
									model.to(lowercase   )
									model.eval()
									__lowercase					=				model(lowercase			,     labels=lowercase   )
									self.parent.assertEqual(result.logits.shape			,     (self.batch_size, self.type_sequence_label_size)   )

									# test greyscale images
									__lowercase					=				1
									__lowercase					=				ViTForImageClassification(lowercase   )
									model.to(lowercase   )
									model.eval()

									__lowercase					=				floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
									__lowercase					=				model(lowercase   )
									self.parent.assertEqual(result.logits.shape			,     (self.batch_size, self.type_sequence_label_size)   )







			def 		snake_case__			(				self	:       Tuple   )						->					Tuple:
									"""simple docstring"""







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





@require_torch
class 			_lowerCAmelCase   ( _UpperCAmelCase    ,					_UpperCAmelCase    ,					unittest.TestCase			):






			"""simple docstring"""





			lowercase__      :  str      =			(
			    (
			        ViTModel,
			        ViTForImageClassification,
			        ViTForMaskedImageModeling,
			    )
			    if is_torch_available()
			    else ()
			)
			lowercase__      :  Any      =			(
			    {"""feature-extraction""": ViTModel, """image-classification""": ViTForImageClassification}
			    if is_torch_available()
			    else {}
			)
			lowercase__      :  Tuple      =			True

			lowercase__      :  int      =			False
			lowercase__      :  Any      =			False
			lowercase__      :  List[Any]      =			False





			def 		snake_case__			(				self	:       List[str]   )						->					int:
									"""simple docstring"""







									__lowercase					=				ViTModelTester(self   )
									__lowercase					=				ConfigTester(self			,     config_class=lowercase			,     has_text_modality=lowercase			,     hidden_size=37   )





			def 		snake_case__			(				self	:       Union[str, Any]   )						->					int:
									"""simple docstring"""







									self.config_tester.run_common_tests()





			@unittest.skip(reason="""ViT does not use inputs_embeds"""   )
			def 		snake_case__			(				self	:       Dict   )						->					List[Any]:
									"""simple docstring"""







									pass





			def 		snake_case__			(				self	:       Union[str, Any]   )						->					Union[str, Any]:
									"""simple docstring"""







									__lowercase	,			__lowercase					=				self.model_tester.prepare_config_and_inputs_for_common()

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





			def 		snake_case__			(				self	:       int   )						->					Dict:
									"""simple docstring"""







									__lowercase	,			__lowercase					=				self.model_tester.prepare_config_and_inputs_for_common()

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

															__lowercase					=				["""pixel_values"""]
															self.assertListEqual(arg_names[:1]			,     lowercase   )





			def 		snake_case__			(				self	:       List[Any]   )						->					Optional[int]:
									"""simple docstring"""







									__lowercase					=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowercase   )





			def 		snake_case__			(				self	:       Optional[int]   )						->					List[Any]:
									"""simple docstring"""







									__lowercase					=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_image_modeling(*lowercase   )





			def 		snake_case__			(				self	:       Union[str, Any]   )						->					List[Any]:
									"""simple docstring"""







									__lowercase					=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*lowercase   )







			@slow
			def 		snake_case__			(				self	:       List[Any]   )						->					Dict:
									"""simple docstring"""







									for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															__lowercase					=				ViTModel.from_pretrained(lowercase   )
															self.assertIsNotNone(lowercase   )






def  UpperCAmelCase__    (    ) ->      str:
						__lowercase					=				Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""   )
						return image







@require_torch
@require_vision
class 			_lowerCAmelCase   ( unittest.TestCase			):






			"""simple docstring"""





			@cached_property
			def 		snake_case__			(				self	:       Optional[Any]   )						->					Union[str, Any]:
									"""simple docstring"""







									return ViTImageProcessor.from_pretrained("""google/vit-base-patch16-224"""   ) if is_vision_available() else None





			@slow
			def 		snake_case__			(				self	:       int   )						->					List[str]:
									"""simple docstring"""







									__lowercase					=				ViTForImageClassification.from_pretrained("""google/vit-base-patch16-224"""   ).to(lowercase   )

									__lowercase					=				self.default_image_processor
									__lowercase					=				prepare_img()
									__lowercase					=				image_processor(images=lowercase			,     return_tensors="""pt"""   ).to(lowercase   )

									# forward pass
									with torch.no_grad():
															__lowercase					=				model(**lowercase   )

									# verify the logits
									__lowercase					=				torch.Size((1, 1_000)   )
									self.assertEqual(outputs.logits.shape			,     lowercase   )

									__lowercase					=				torch.tensor([-0.2744, 0.8215, -0.0836]   ).to(lowercase   )

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





			@slow
			def 		snake_case__			(				self	:       Union[str, Any]   )						->					Any:
									"""simple docstring"""







									__lowercase					=				ViTModel.from_pretrained("""facebook/dino-vits8"""   ).to(lowercase   )

									__lowercase					=				ViTImageProcessor.from_pretrained("""facebook/dino-vits8"""			,     size=480   )
									__lowercase					=				prepare_img()
									__lowercase					=				image_processor(images=lowercase			,     return_tensors="""pt"""   )
									__lowercase					=				inputs.pixel_values.to(lowercase   )

									# forward pass
									with torch.no_grad():
															__lowercase					=				model(lowercase			,     interpolate_pos_encoding=lowercase   )

									# verify the logits
									__lowercase					=				torch.Size((1, 3_601, 384)   )
									self.assertEqual(outputs.last_hidden_state.shape			,     lowercase   )

									__lowercase					=				torch.tensor(
									    [[4.2340, 4.3906, -6.6692], [4.5463, 1.8928, -6.7257], [4.4429, 0.8496, -5.8585]]   ).to(lowercase   )

									self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]			,     lowercase			,     atol=1E-4   )   )







			@slow
			@require_accelerate
			@require_torch_gpu
			def 		snake_case__			(				self	:       Dict   )						->					Optional[int]:
									"""simple docstring"""







									__lowercase					=				ViTModel.from_pretrained("""facebook/dino-vits8"""			,     torch_dtype=torch.floataa			,     device_map="""auto"""   )
									__lowercase					=				self.default_image_processor

									__lowercase					=				prepare_img()
									__lowercase					=				image_processor(images=lowercase			,     return_tensors="""pt"""   )
									__lowercase					=				inputs.pixel_values.to(lowercase   )

									# forward pass to make sure inference works in fp16
									with torch.no_grad():
															__lowercase					=				model(lowercase   )










def  UpperCAmelCase__    (    lowercase__ = 100   ) ->      int:
						__lowercase					=				n * (n + 1) * (2 * n + 1) / 6
						__lowercase					=				(n * (n + 1) / 2) ** 2
						return int(square_of_sum - sum_of_squares   )


if __name__ == "__main__":
			print(F"""{solution() = }""")





