



import inspect
import unittest

from transformers import MobileViTVaConfig
from transformers.testing_utils import require_torch, require_torch_multi_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 transformers import MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation, MobileViTVaModel
     from transformers.models.mobilevitva.modeling_mobilevitva import (
         MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST,
         make_divisible,
     )


if is_vision_available():
     from PIL import Image

     from transformers import MobileViTImageProcessor





class    lowerCamelCase__  (						UpperCAmelCase_):
     """simple docstring"""



     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.config_class(**self.inputs_dict    )
           self.parent.assertTrue(hasattr(__a				,			"width_multiplier"    )    )





class    lowerCamelCase__  :
     """simple docstring"""



     def __init__(self				,			__a				,			__a=13				,			__a=64				,			__a=2				,			__a=3				,			__a="swish"				,			__a=3				,			__a=32				,			__a=0.1				,			__a=0.02				,			__a=True				,			__a=True				,			__a=10				,			__a=None				,			__a=0.25				,			__a=0.0				,			__a=0.0				,			):






           '''simple docstring'''



           lowerCamelCase								=		parent
           lowerCamelCase								=		batch_size
           lowerCamelCase								=		image_size
           lowerCamelCase								=		patch_size
           lowerCamelCase								=		num_channels
           lowerCamelCase								=		make_divisible(5_12 * width_multiplier				,			divisor=8    )
           lowerCamelCase								=		hidden_act
           lowerCamelCase								=		conv_kernel_size
           lowerCamelCase								=		output_stride
           lowerCamelCase								=		classifier_dropout_prob
           lowerCamelCase								=		use_labels
           lowerCamelCase								=		is_training
           lowerCamelCase								=		num_labels
           lowerCamelCase								=		initializer_range
           lowerCamelCase								=		scope
           lowerCamelCase								=		width_multiplier
           lowerCamelCase								=		ffn_dropout
           lowerCamelCase								=		attn_dropout







     def 						_a      (self    ):






           '''simple docstring'''



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

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

           lowerCamelCase								=		self.get_config()

           return config, pixel_values, labels, pixel_labels







     def 						_a      (self    ):






           '''simple docstring'''



           return MobileViTVaConfig(
               image_size=self.image_size				,			patch_size=self.patch_size				,			num_channels=self.num_channels				,			hidden_act=self.hidden_act				,			conv_kernel_size=self.conv_kernel_size				,			output_stride=self.output_stride				,			classifier_dropout_prob=self.classifier_dropout_prob				,			initializer_range=self.initializer_range				,			width_multiplier=self.width_multiplier				,			ffn_dropout=self.ffn_dropout_prob				,			attn_dropout=self.attn_dropout_prob				,			)







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






           '''simple docstring'''



           lowerCamelCase								=		MobileViTVaModel(config=__a    )
           model.to(__a    )
           model.eval()
           lowerCamelCase								=		model(__a    )
           self.parent.assertEqual(
               result.last_hidden_state.shape				,			(
                   self.batch_size,
                   self.last_hidden_size,
                   self.image_size // self.output_stride,
                   self.image_size // self.output_stride,
               )				,			)







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






           '''simple docstring'''



           lowerCamelCase								=		self.num_labels
           lowerCamelCase								=		MobileViTVaForImageClassification(__a    )
           model.to(__a    )
           model.eval()
           lowerCamelCase								=		model(__a				,			labels=__a    )
           self.parent.assertEqual(result.logits.shape				,			(self.batch_size, self.num_labels)    )







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






           '''simple docstring'''



           lowerCamelCase								=		self.num_labels
           lowerCamelCase								=		MobileViTVaForSemanticSegmentation(__a    )
           model.to(__a    )
           model.eval()
           lowerCamelCase								=		model(__a    )
           self.parent.assertEqual(
               result.logits.shape				,			(
                   self.batch_size,
                   self.num_labels,
                   self.image_size // self.output_stride,
                   self.image_size // self.output_stride,
               )				,			)
           lowerCamelCase								=		model(__a				,			labels=__a    )
           self.parent.assertEqual(
               result.logits.shape				,			(
                   self.batch_size,
                   self.num_labels,
                   self.image_size // self.output_stride,
                   self.image_size // self.output_stride,
               )				,			)







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.prepare_config_and_inputs()
           lowerCamelCase			,							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       =			(
         (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
         if is_torch_available()
         else ()
     )

     _A       =			(
         {
             'feature-extraction': MobileViTVaModel,
             'image-classification': MobileViTVaForImageClassification,
             'image-segmentation': MobileViTVaForSemanticSegmentation,
         }
         if is_torch_available()
         else {}
     )

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







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		MobileViTVaModelTester(self    )
           lowerCamelCase								=		MobileViTVaConfigTester(self				,			config_class=__a				,			has_text_modality=__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           self.config_tester.run_common_tests()







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






           '''simple docstring'''



           pass







     @unittest.skip(reason="MobileViTV2 does not support input and output embeddings"    )
     def 						_a      (self    ):






           '''simple docstring'''



           pass







     @unittest.skip(reason="MobileViTV2 does not output attentions"    )
     def 						_a      (self    ):






           '''simple docstring'''



           pass







     @require_torch_multi_gpu
     @unittest.skip(reason="Got `CUDA error: misaligned address` for tests after this one being run."    )
     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







     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'''



           def check_hidden_states_output(__a				,			__a				,			__a    ):
                 lowerCamelCase								=		model_class(__a    )
                 model.to(__a    )
                 model.eval()

                 with torch.no_grad():
                       lowerCamelCase								=		model(**self._prepare_for_class(__a				,			__a    )    )

                 lowerCamelCase								=		outputs.hidden_states

                 lowerCamelCase								=		5
                 self.assertEqual(len(__a    )				,			__a    )

                 # MobileViTV2's feature maps are of shape (batch_size, num_channels, height, width)
                 # with the width and height being successively divided by 2.
                 lowerCamelCase								=		2
                 for i in range(len(__a    )    ):
                       self.assertListEqual(
                           list(hidden_states[i].shape[-2:]    )				,			[self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]				,			)
                       divisor *= 2

                 self.assertEqual(self.model_tester.output_stride				,			divisor // 2    )

           lowerCamelCase			,							lowerCamelCase								=		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 lowerCamelCase								=		True
                 check_hidden_states_output(__a				,			__a				,			__a    )

                 # check that output_hidden_states also work using config
                 del inputs_dict["output_hidden_states"]
                 lowerCamelCase								=		True

                 check_hidden_states_output(__a				,			__a				,			__a    )







     def 						_a      (self    ):






           '''simple docstring'''



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







     def 						_a      (self    ):






           '''simple docstring'''



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







     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase								=		MobileViTVaModel.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 (
               MobileViTImageProcessor.from_pretrained("apple/mobilevitv2-1.0-imagenet1k-256"    )
               if is_vision_available()
               else None
           )







     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		MobileViTVaForImageClassification.from_pretrained("apple/mobilevitv2-1.0-imagenet1k-256"    ).to(
               __a    )

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

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

           # verify the logits
           lowerCamelCase								=		torch.Size((1, 10_00)    )
           self.assertEqual(outputs.logits.shape				,			__a    )

           lowerCamelCase								=		torch.tensor([-1.6_336E00, -7.3_204E-02, -5.1_883E-01]    ).to(__a    )

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







     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		MobileViTVaForSemanticSegmentation.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3"    )
           lowerCamelCase								=		model.to(__a    )

           lowerCamelCase								=		MobileViTImageProcessor.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3"    )

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

           # forward pass
           with torch.no_grad():
                 lowerCamelCase								=		model(**__a    )
           lowerCamelCase								=		outputs.logits

           # verify the logits
           lowerCamelCase								=		torch.Size((1, 21, 32, 32)    )
           self.assertEqual(logits.shape				,			__a    )

           lowerCamelCase								=		torch.tensor(
               [
                   [[7.0863, 7.1525, 6.8201], [6.6931, 6.8770, 6.8933], [6.2978, 7.0366, 6.9636]],
                   [[-3.7134, -3.6712, -3.6675], [-3.5825, -3.3549, -3.4777], [-3.3435, -3.3979, -3.2857]],
                   [[-2.9329, -2.8003, -2.7369], [-3.0564, -2.4780, -2.0207], [-2.6889, -1.9298, -1.7640]],
               ]				,			device=__a				,			)

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







     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		MobileViTVaForSemanticSegmentation.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3"    )
           lowerCamelCase								=		model.to(__a    )

           lowerCamelCase								=		MobileViTImageProcessor.from_pretrained("shehan97/mobilevitv2-1.0-voc-deeplabv3"    )

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

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

           lowerCamelCase								=		outputs.logits.detach().cpu()

           lowerCamelCase								=		image_processor.post_process_semantic_segmentation(outputs=__a				,			target_sizes=[(50, 60)]    )
           lowerCamelCase								=		torch.Size((50, 60)    )
           self.assertEqual(segmentation[0].shape				,			__a    )

           lowerCamelCase								=		image_processor.post_process_semantic_segmentation(outputs=__a    )
           lowerCamelCase								=		torch.Size((32, 32)    )
           self.assertEqual(segmentation[0].shape				,			__a    )




from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch





@require_torch
@require_tf
class    lowerCamelCase__  (						UpperCAmelCase_):
     """simple docstring"""



     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		SMALL_MODEL_IDENTIFIER
           lowerCamelCase								=		"pt"
           lowerCamelCase								=		"tf"







     def 						_a      (self				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		AutoModel.from_pretrained(self.test_model    )
           model_pt.save_pretrained(__a    )







     def 						_a      (self				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		TFAutoModel.from_pretrained(self.test_model				,			from_pt=__a    )
           model_tf.save_pretrained(__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		"mock_framework"

           # Framework provided - return whatever the user provides
           lowerCamelCase								=		FeaturesManager.determine_framework(self.test_model				,			__a    )
           self.assertEqual(__a				,			__a    )

           # Local checkpoint and framework provided - return provided framework
           # PyTorch checkpoint
           with TemporaryDirectory() as local_pt_ckpt:
                 self._setup_pt_ckpt(__a    )
                 lowerCamelCase								=		FeaturesManager.determine_framework(__a				,			__a    )
                 self.assertEqual(__a				,			__a    )

           # TensorFlow checkpoint
           with TemporaryDirectory() as local_tf_ckpt:
                 self._setup_tf_ckpt(__a    )
                 lowerCamelCase								=		FeaturesManager.determine_framework(__a				,			__a    )
                 self.assertEqual(__a				,			__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           with TemporaryDirectory() as local_pt_ckpt:
                 self._setup_pt_ckpt(__a    )
                 lowerCamelCase								=		FeaturesManager.determine_framework(__a    )
                 self.assertEqual(__a				,			self.framework_pt    )

           # TensorFlow checkpoint
           with TemporaryDirectory() as local_tf_ckpt:
                 self._setup_tf_ckpt(__a    )
                 lowerCamelCase								=		FeaturesManager.determine_framework(__a    )
                 self.assertEqual(__a				,			self.framework_tf    )

           # Invalid local checkpoint
           with TemporaryDirectory() as local_invalid_ckpt:
                 with self.assertRaises(__a    ):
                       lowerCamelCase								=		FeaturesManager.determine_framework(__a    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		MagicMock(return_value=__a    )
           with patch("transformers.onnx.features.is_tf_available"				,			__a    ):
                 lowerCamelCase								=		FeaturesManager.determine_framework(self.test_model    )
                 self.assertEqual(__a				,			self.framework_pt    )

           # PyTorch not in environment -> use TensorFlow
           lowerCamelCase								=		MagicMock(return_value=__a    )
           with patch("transformers.onnx.features.is_torch_available"				,			__a    ):
                 lowerCamelCase								=		FeaturesManager.determine_framework(self.test_model    )
                 self.assertEqual(__a				,			self.framework_tf    )

           # Both in environment -> use PyTorch
           lowerCamelCase								=		MagicMock(return_value=__a    )
           lowerCamelCase								=		MagicMock(return_value=__a    )
           with patch("transformers.onnx.features.is_tf_available"				,			__a    ), patch(
               "transformers.onnx.features.is_torch_available"				,			__a    ):
                 lowerCamelCase								=		FeaturesManager.determine_framework(self.test_model    )
                 self.assertEqual(__a				,			self.framework_pt    )

           # Both not in environment -> raise error
           lowerCamelCase								=		MagicMock(return_value=__a    )
           lowerCamelCase								=		MagicMock(return_value=__a    )
           with patch("transformers.onnx.features.is_tf_available"				,			__a    ), patch(
               "transformers.onnx.features.is_torch_available"				,			__a    ):
                 with self.assertRaises(__a    ):
                       lowerCamelCase								=		FeaturesManager.determine_framework(self.test_model    )