


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							:							str     ,					_lowerCAmelCase							:							List[str]     ,					_lowerCAmelCase							:							Any=1_3     ,					_lowerCAmelCase							:							int=3_0     ,					_lowerCAmelCase							:							str=2     ,					_lowerCAmelCase							:							str=3     ,					_lowerCAmelCase							:							Dict=True     ,					_lowerCAmelCase							:							int=True     ,					_lowerCAmelCase							:							List[str]=3_2     ,					_lowerCAmelCase							:							Tuple=5     ,					_lowerCAmelCase							:							Optional[Any]=4     ,					_lowerCAmelCase							:							List[Any]=3_7     ,					_lowerCAmelCase							:							Dict="gelu"     ,					_lowerCAmelCase							:							int=0.1     ,					_lowerCAmelCase							:							Tuple=0.1     ,					_lowerCAmelCase							:							int=1_0     ,					_lowerCAmelCase							:							str=0.02     ,					_lowerCAmelCase							:							List[Any]=None     ,					_lowerCAmelCase							:							int=2     ,					) ->			List[Any]:




     """simple docstring"""
     snake_case_            =       parent
     snake_case_            =       batch_size
     snake_case_            =       image_size
     snake_case_            =       patch_size
     snake_case_            =       num_channels
     snake_case_            =       is_training
     snake_case_            =       use_labels
     snake_case_            =       hidden_size
     snake_case_            =       num_hidden_layers
     snake_case_            =       num_attention_heads
     snake_case_            =       intermediate_size
     snake_case_            =       hidden_act
     snake_case_            =       hidden_dropout_prob
     snake_case_            =       attention_probs_dropout_prob
     snake_case_            =       type_sequence_label_size
     snake_case_            =       initializer_range
     snake_case_            =       scope
     snake_case_            =       encoder_stride

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



  def        lowerCAmelCase__	(  self							:							Union[str, Any]	) ->			int:




     """simple docstring"""
     snake_case_            =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

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

     snake_case_            =       self.get_config()

     return config, pixel_values, labels



  def        lowerCAmelCase__	(  self							:							int	) ->			Dict:




     """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=__UpperCamelCase     ,					initializer_range=self.initializer_range     ,					encoder_stride=self.encoder_stride     ,					)



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




     """simple docstring"""
     snake_case_            =       ViTModel(config=__UpperCamelCase	)
     model.to(__UpperCamelCase	)
     model.eval()
     snake_case_            =       model(__UpperCamelCase	)
     self.parent.assertEqual(result.last_hidden_state.shape     ,					(self.batch_size, self.seq_length, self.hidden_size)	)



  def        lowerCAmelCase__	(  self							:							Dict     ,					_lowerCAmelCase							:							List[str]     ,					_lowerCAmelCase							:							List[Any]     ,					_lowerCAmelCase							:							str	) ->			Optional[Any]:




     """simple docstring"""
     snake_case_            =       ViTForMaskedImageModeling(config=__UpperCamelCase	)
     model.to(__UpperCamelCase	)
     model.eval()
     snake_case_            =       model(__UpperCamelCase	)
     self.parent.assertEqual(
         result.reconstruction.shape     ,					(self.batch_size, self.num_channels, self.image_size, self.image_size)	)

     # test greyscale images
     snake_case_            =       1
     snake_case_            =       ViTForMaskedImageModeling(__UpperCamelCase	)
     model.to(__UpperCamelCase	)
     model.eval()

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



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




     """simple docstring"""
     snake_case_            =       self.type_sequence_label_size
     snake_case_            =       ViTForImageClassification(__UpperCamelCase	)
     model.to(__UpperCamelCase	)
     model.eval()
     snake_case_            =       model(__UpperCamelCase     ,					labels=__UpperCamelCase	)
     self.parent.assertEqual(result.logits.shape     ,					(self.batch_size, self.type_sequence_label_size)	)

     # test greyscale images
     snake_case_            =       1
     snake_case_            =       ViTForImageClassification(__UpperCamelCase	)
     model.to(__UpperCamelCase	)
     model.eval()

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



  def        lowerCAmelCase__	(  self							:							Union[str, Any]	) ->			int:




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






@require_torch
class      __lowerCAmelCase  (			__snake_case			,	__snake_case			,	unittest.TestCase       ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE		     =							(
      (
          ViTModel,
          ViTForImageClassification,
          ViTForMaskedImageModeling,
      )
      if is_torch_available()
      else ()
  )
  _SCREAMING_SNAKE_CASE		     =							(
      {'feature-extraction': ViTModel, 'image-classification': ViTForImageClassification}
      if is_torch_available()
      else {}
  )
  _SCREAMING_SNAKE_CASE		     =							True

  _SCREAMING_SNAKE_CASE		     =							False
  _SCREAMING_SNAKE_CASE		     =							False
  _SCREAMING_SNAKE_CASE		     =							False



  def        lowerCAmelCase__	(  self							:							int	) ->			Dict:




     """simple docstring"""
     snake_case_            =       ViTModelTester(self	)
     snake_case_            =       ConfigTester(self     ,					config_class=__UpperCamelCase     ,					has_text_modality=__UpperCamelCase     ,					hidden_size=3_7	)



  def        lowerCAmelCase__	(  self							:							Optional[Any]	) ->			Optional[Any]:




     """simple docstring"""
     self.config_tester.run_common_tests()



  @unittest.skip(reason="ViT does not use inputs_embeds"	)
  def        lowerCAmelCase__	(  self							:							Union[str, Any]	) ->			Tuple:




     """simple docstring"""
     pass



  def        lowerCAmelCase__	(  self							:							List[Any]	) ->			Optional[Any]:




     """simple docstring"""
     snake_case_     ,     snake_case_            =       self.model_tester.prepare_config_and_inputs_for_common()

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



  def        lowerCAmelCase__	(  self							:							List[str]	) ->			Tuple:




     """simple docstring"""
     snake_case_     ,     snake_case_            =       self.model_tester.prepare_config_and_inputs_for_common()

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

        snake_case_            =       ["pixel_values"]
        self.assertListEqual(arg_names[:1]     ,					__UpperCamelCase	)



  def        lowerCAmelCase__	(  self							:							List[str]	) ->			str:




     """simple docstring"""
     snake_case_            =       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__UpperCamelCase	)



  def        lowerCAmelCase__	(  self							:							int	) ->			List[str]:




     """simple docstring"""
     snake_case_            =       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_image_modeling(*__UpperCamelCase	)



  def        lowerCAmelCase__	(  self							:							Any	) ->			List[Any]:




     """simple docstring"""
     snake_case_            =       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*__UpperCamelCase	)



  @slow
  def        lowerCAmelCase__	(  self							:							Dict	) ->			List[str]:




     """simple docstring"""


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






def       _lowerCAmelCase    (				)->Union[str, Any]:
   '''simple docstring'''






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



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






  """simple docstring"""





  @cached_property
  def        lowerCAmelCase__	(  self							:							Optional[Any]	) ->			Optional[int]:




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



  @slow
  def        lowerCAmelCase__	(  self							:							List[str]	) ->			Optional[int]:




     """simple docstring"""
     snake_case_            =       ViTForImageClassification.from_pretrained("google/vit-base-patch16-224"	).to(__UpperCamelCase	)

     snake_case_            =       self.default_image_processor
     snake_case_            =       prepare_img()
     snake_case_            =       image_processor(images=__UpperCamelCase     ,					return_tensors="pt"	).to(__UpperCamelCase	)

     # forward pass
     with torch.no_grad():
        snake_case_            =       model(**__UpperCamelCase	)

     # verify the logits
     snake_case_            =       torch.Size((1, 1_0_0_0)	)
     self.assertEqual(outputs.logits.shape     ,					__UpperCamelCase	)

     snake_case_            =       torch.tensor([-0.2_744, 0.8_215, -0.0_836]	).to(__UpperCamelCase	)

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



  @slow
  def        lowerCAmelCase__	(  self							:							Optional[int]	) ->			List[str]:




     """simple docstring"""
     snake_case_            =       ViTModel.from_pretrained("facebook/dino-vits8"	).to(__UpperCamelCase	)

     snake_case_            =       ViTImageProcessor.from_pretrained("facebook/dino-vits8"     ,					size=4_8_0	)
     snake_case_            =       prepare_img()
     snake_case_            =       image_processor(images=__UpperCamelCase     ,					return_tensors="pt"	)
     snake_case_            =       inputs.pixel_values.to(__UpperCamelCase	)

     # forward pass
     with torch.no_grad():
        snake_case_            =       model(__UpperCamelCase     ,					interpolate_pos_encoding=__UpperCamelCase	)

     # verify the logits
     snake_case_            =       torch.Size((1, 3_6_0_1, 3_8_4)	)
     self.assertEqual(outputs.last_hidden_state.shape     ,					__UpperCamelCase	)

     snake_case_            =       torch.tensor(
         [[4.2_340, 4.3_906, -6.6_692], [4.5_463, 1.8_928, -6.7_257], [4.4_429, 0.8_496, -5.8_585]]	).to(__UpperCamelCase	)

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



  @slow
  @require_accelerate
  @require_torch_gpu
  def        lowerCAmelCase__	(  self							:							Any	) ->			Optional[Any]:




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

     snake_case_            =       prepare_img()
     snake_case_            =       image_processor(images=__UpperCamelCase     ,					return_tensors="pt"	)
     snake_case_            =       inputs.pixel_values.to(__UpperCamelCase	)

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







import unittest

from accelerate import debug_launcher
from accelerate.test_utils import require_cpu, test_ops, test_script


@require_cpu
class SCREAMING_SNAKE_CASE	(		unittest.TestCase				):






			"""simple docstring"""




			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				debug_launcher(test_script.main				)







			def       __lowerCAmelCase			(     self				):

				"""simple docstring"""


				debug_launcher(test_ops.main				)
