






"""simple docstring"""






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       snake_case_       :



	"""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=None  ,					__a=2  ,					):


					"""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__							 =       scope
					A__							 =       encoder_stride

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


	def       _UpperCAmelCase						(      self						):


					"""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       _UpperCAmelCase						(      self						):


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


	def       _UpperCAmelCase						(      self  ,					__a  ,					__a  ,					__a						):


					"""simple docstring"""

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


	def       _UpperCAmelCase						(      self  ,					__a  ,					__a  ,					__a						):


					"""simple docstring"""

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

					# test greyscale images
					A__							 =       1
					A__							 =       ViTForMaskedImageModeling(__a						)
					model.to(__a						)
					model.eval()

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


	def       _UpperCAmelCase						(      self  ,					__a  ,					__a  ,					__a						):


					"""simple docstring"""

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

					# test greyscale images
					A__							 =       1
					A__							 =       ViTForImageClassification(__a						)
					model.to(__a						)
					model.eval()

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






	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

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





@require_torch
class       snake_case_       (    _lowerCamelCase    ,							_lowerCamelCase    ,							unittest.TestCase       ):



	"""simple docstring"""

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

	SCREAMING_SNAKE_CASE_:							Tuple   			=					False
	SCREAMING_SNAKE_CASE_:							List[Any]   			=					False
	SCREAMING_SNAKE_CASE_:							Dict   			=					False


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       ViTModelTester(self						)
					A__							 =       ConfigTester(self  ,					config_class=__a  ,					has_text_modality=__a  ,					hidden_size=37						)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					self.config_tester.run_common_tests()


	@unittest.skip(reason='ViT does not use inputs_embeds'						)
	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					pass


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__ ,			A__							 =       self.model_tester.prepare_config_and_inputs_for_common()

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


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__ ,			A__							 =       self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
									A__							 =       model_class(__a						)
									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]  ,					__a						)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

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


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_image_modeling(*__a						)


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_image_classification(*__a						)






	@slow
	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

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





def  __lowerCamelCase					(					):
				A__							 =       Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
				return image







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



	"""simple docstring"""

	@cached_property
	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

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


	@slow
	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       ViTForImageClassification.from_pretrained('google/vit-base-patch16-224'						).to(__a						)

					A__							 =       self.default_image_processor
					A__							 =       prepare_img()
					A__							 =       image_processor(images=__a  ,					return_tensors='pt'						).to(__a						)

					# forward pass
					with torch.no_grad():
									A__							 =       model(**__a						)

					# verify the logits
					A__							 =       torch.Size((1, 1000)						)
					self.assertEqual(outputs.logits.shape  ,					__a						)

					A__							 =       torch.tensor([-0.2744, 0.8215, -0.0836]						).to(__a						)

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


	@slow
	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       ViTModel.from_pretrained('facebook/dino-vits8'						).to(__a						)

					A__							 =       ViTImageProcessor.from_pretrained('facebook/dino-vits8'  ,					size=480						)
					A__							 =       prepare_img()
					A__							 =       image_processor(images=__a  ,					return_tensors='pt'						)
					A__							 =       inputs.pixel_values.to(__a						)

					# forward pass
					with torch.no_grad():
									A__							 =       model(__a  ,					interpolate_pos_encoding=__a						)

					# verify the logits
					A__							 =       torch.Size((1, 3601, 384)						)
					self.assertEqual(outputs.last_hidden_state.shape  ,					__a						)

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

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






	@slow
	@require_accelerate
	@require_torch_gpu
	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

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

					A__							 =       prepare_img()
					A__							 =       image_processor(images=__a  ,					return_tensors='pt'						)
					A__							 =       inputs.pixel_values.to(__a						)

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










"""simple docstring"""






from typing import Union

import fire
import torch
from tqdm import tqdm



def  __lowerCamelCase					(					lowerCAmelCase__       ,lowerCAmelCase__ = "cpu"       ,lowerCAmelCase__ = None				):
				A__							 =       torch.load(lowerCAmelCase__       ,map_location=lowerCAmelCase__				)
				for k, v in tqdm(state_dict.items()				):
								if not isinstance(lowerCAmelCase__       ,torch.Tensor				):
												raise TypeError('FP16 conversion only works on paths that are saved state dicts, like pytorch_model.bin'				)
								A__							 =       v.half()
				if save_path is None:  # overwrite src_path
								A__							 =       src_path
				torch.save(lowerCAmelCase__       ,lowerCAmelCase__				)


if __name__ == "__main__":
					fire.Fire(convert)


