


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__(    snake_case__						):






    '''simple docstring'''





    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      Any						=   self.config_class(**self.inputs_dict	)
       self.parent.assertTrue(hasattr(__snake_case       ,   '''width_multiplier'''	)	)





class  lowerCAmelCase__:






    '''simple docstring'''





    def __init__(	self				:   str       ,   __snake_case				:   Union[str, Any]       ,   __snake_case				:   int=13       ,   __snake_case				:   Tuple=64       ,   __snake_case				:   str=2       ,   __snake_case				:   Dict=3       ,   __snake_case				:   List[Any]="swish"       ,   __snake_case				:   Optional[int]=3       ,   __snake_case				:   Any=32       ,   __snake_case				:   Union[str, Any]=0.1       ,   __snake_case				:   Optional[int]=0.02       ,   __snake_case				:   str=True       ,   __snake_case				:   Dict=True       ,   __snake_case				:   Optional[Any]=10       ,   __snake_case				:   str=None       ,   __snake_case				:   Tuple=0.25       ,   __snake_case				:   Dict=0.0       ,   __snake_case				:   Any=0.0       ,   ):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   parent
       UpperCAmelCase_							:      int						=   batch_size
       UpperCAmelCase_							:      List[str]						=   image_size
       UpperCAmelCase_							:      str						=   patch_size
       UpperCAmelCase_							:      int						=   num_channels
       UpperCAmelCase_							:      List[Any]						=   make_divisible(512 * width_multiplier       ,   divisor=8	)
       UpperCAmelCase_							:      Dict						=   hidden_act
       UpperCAmelCase_							:      Union[str, Any]						=   conv_kernel_size
       UpperCAmelCase_							:      Optional[int]						=   output_stride
       UpperCAmelCase_							:      Optional[int]						=   classifier_dropout_prob
       UpperCAmelCase_							:      List[Any]						=   use_labels
       UpperCAmelCase_							:      Optional[Any]						=   is_training
       UpperCAmelCase_							:      Tuple						=   num_labels
       UpperCAmelCase_							:      Tuple						=   initializer_range
       UpperCAmelCase_							:      Dict						=   scope
       UpperCAmelCase_							:      Optional[int]						=   width_multiplier
       UpperCAmelCase_							:      int						=   ffn_dropout
       UpperCAmelCase_							:      Dict						=   attn_dropout





    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       UpperCAmelCase_							:      Tuple						=   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

       UpperCAmelCase_							:      int						=   None
       UpperCAmelCase_							:      List[str]						=   None
       if self.use_labels:
          UpperCAmelCase_							:      Dict						=   ids_tensor([self.batch_size]       ,   self.num_labels	)
          UpperCAmelCase_							:      List[Any]						=   ids_tensor([self.batch_size, self.image_size, self.image_size]       ,   self.num_labels	)

       UpperCAmelCase_							:      List[str]						=   self.get_config()

       return config, pixel_values, labels, pixel_labels





    def 					_lowerCamelCase (	self				:   List[str]	):





       '''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 					_lowerCamelCase (	self				:   Any       ,   __snake_case				:   List[str]       ,   __snake_case				:   Union[str, Any]       ,   __snake_case				:   int       ,   __snake_case				:   List[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   MobileViTVaModel(config=__snake_case	)
       model.to(__snake_case	)
       model.eval()
       UpperCAmelCase_							:      Union[str, Any]						=   model(__snake_case	)
       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 					_lowerCamelCase (	self				:   Optional[int]       ,   __snake_case				:   Tuple       ,   __snake_case				:   Union[str, Any]       ,   __snake_case				:   List[str]       ,   __snake_case				:   Any	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   self.num_labels
       UpperCAmelCase_							:      str						=   MobileViTVaForImageClassification(__snake_case	)
       model.to(__snake_case	)
       model.eval()
       UpperCAmelCase_							:      Optional[int]						=   model(__snake_case       ,   labels=__snake_case	)
       self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.num_labels)	)





    def 					_lowerCamelCase (	self				:   str       ,   __snake_case				:   List[str]       ,   __snake_case				:   List[Any]       ,   __snake_case				:   List[str]       ,   __snake_case				:   str	):





       '''simple docstring'''
       UpperCAmelCase_							:      Dict						=   self.num_labels
       UpperCAmelCase_							:      int						=   MobileViTVaForSemanticSegmentation(__snake_case	)
       model.to(__snake_case	)
       model.eval()
       UpperCAmelCase_							:      Union[str, Any]						=   model(__snake_case	)
       self.parent.assertEqual(
           result.logits.shape       ,   (
               self.batch_size,
               self.num_labels,
               self.image_size // self.output_stride,
               self.image_size // self.output_stride,
           )       ,   )
       UpperCAmelCase_							:      Optional[Any]						=   model(__snake_case       ,   labels=__snake_case	)
       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 					_lowerCamelCase (	self				:   Tuple	):





       '''simple docstring'''
       UpperCAmelCase_							:      Any						=   self.prepare_config_and_inputs()
       UpperCAmelCase_		,      UpperCAmelCase_		,      UpperCAmelCase_		,      UpperCAmelCase_							:      List[Any]						=   config_and_inputs
       UpperCAmelCase_							:      Tuple						=   {'''pixel_values''': pixel_values}
       return config, inputs_dict





@require_torch
class  lowerCAmelCase__(    snake_case__    ,    snake_case__    ,    unittest.TestCase						):






    '''simple docstring'''





    A_  :     List[str]			= (
        (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
        if is_torch_available()
        else ()
    )

    A_  :     Union[str, Any]			= (
        {
            'feature-extraction': MobileViTVaModel,
            'image-classification': MobileViTVaForImageClassification,
            'image-segmentation': MobileViTVaForSemanticSegmentation,
        }
        if is_torch_available()
        else {}
    )

    A_  :     Union[str, Any]			= False
    A_  :     List[str]			= False
    A_  :     List[str]			= False
    A_  :     List[Any]			= False





    def 					_lowerCamelCase (	self				:   Union[str, Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      str						=   MobileViTVaModelTester(self	)
       UpperCAmelCase_							:      int						=   MobileViTVaConfigTester(self       ,   config_class=__snake_case       ,   has_text_modality=__snake_case	)





    def 					_lowerCamelCase (	self				:   Optional[int]	):





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





    @unittest.skip(reason='''MobileViTV2 does not use inputs_embeds'''	)
    def 					_lowerCamelCase (	self				:   int	):





       '''simple docstring'''
       pass





    @unittest.skip(reason='''MobileViTV2 does not support input and output embeddings'''	)
    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       pass





    @unittest.skip(reason='''MobileViTV2 does not output attentions'''	)
    def 					_lowerCamelCase (	self				:   Optional[Any]	):





       '''simple docstring'''
       pass





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





       '''simple docstring'''
       pass





    def 					_lowerCamelCase (	self				:   Union[str, Any]	):





       '''simple docstring'''
       UpperCAmelCase_		,      UpperCAmelCase_							:      List[str]						=   self.model_tester.prepare_config_and_inputs_for_common()

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

          UpperCAmelCase_							:      int						=   ['''pixel_values''']
          self.assertListEqual(arg_names[:1]       ,   __snake_case	)





    def 					_lowerCamelCase (	self				:   Optional[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[Any]						=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*__snake_case	)





    def 					_lowerCamelCase (	self				:   Tuple	):





       '''simple docstring'''
       def check_hidden_states_output(__snake_case				:   Optional[int]       ,   __snake_case				:   Any       ,   __snake_case				:   Union[str, Any]	):
          UpperCAmelCase_							:      Optional[int]						=   model_class(__snake_case	)
          model.to(__snake_case	)
          model.eval()

          with torch.no_grad():
             UpperCAmelCase_							:      int						=   model(**self._prepare_for_class(__snake_case       ,   __snake_case	)	)

          UpperCAmelCase_							:      str						=   outputs.hidden_states

          UpperCAmelCase_							:      Union[str, Any]						=   5
          self.assertEqual(len(__snake_case	)       ,   __snake_case	)

          # MobileViTV2's feature maps are of shape (batch_size, num_channels, height, width)
          # with the width and height being successively divided by 2.
          UpperCAmelCase_							:      List[str]						=   2
          for i in range(len(__snake_case	)	):
             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	)

       UpperCAmelCase_		,      UpperCAmelCase_							:      Any						=   self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
          UpperCAmelCase_							:      Optional[Any]						=   True
          check_hidden_states_output(__snake_case       ,   __snake_case       ,   __snake_case	)

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          UpperCAmelCase_							:      Optional[int]						=   True

          check_hidden_states_output(__snake_case       ,   __snake_case       ,   __snake_case	)





    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[str]						=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*__snake_case	)





    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[Any]						=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_semantic_segmentation(*__snake_case	)





    @slow
    def 					_lowerCamelCase (	self				:   List[str]	):





       '''simple docstring'''
       for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          UpperCAmelCase_							:      Optional[int]						=   MobileViTVaModel.from_pretrained(__snake_case	)
          self.assertIsNotNone(__snake_case	)





def      snake_case_						( ):
   UpperCAmelCase_							:      Optional[int]						=   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[int]	):





       '''simple docstring'''
       return (
           MobileViTImageProcessor.from_pretrained('''apple/mobilevitv2-1.0-imagenet1k-256'''	)
           if is_vision_available()
           else None
       )





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





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[int]						=   MobileViTVaForImageClassification.from_pretrained('''apple/mobilevitv2-1.0-imagenet1k-256'''	).to(
           __snake_case	)

       UpperCAmelCase_							:      Union[str, Any]						=   self.default_image_processor
       UpperCAmelCase_							:      Tuple						=   prepare_img()
       UpperCAmelCase_							:      Optional[int]						=   image_processor(images=__snake_case       ,   return_tensors='''pt'''	).to(__snake_case	)

       # forward pass
       with torch.no_grad():
          UpperCAmelCase_							:      Union[str, Any]						=   model(**__snake_case	)

       # verify the logits
       UpperCAmelCase_							:      List[str]						=   torch.Size((1, 1_000)	)
       self.assertEqual(outputs.logits.shape       ,   __snake_case	)

       UpperCAmelCase_							:      Tuple						=   torch.tensor([-1.6336E00, -7.3204E-02, -5.1883E-01]	).to(__snake_case	)

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





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





       '''simple docstring'''
       UpperCAmelCase_							:      int						=   MobileViTVaForSemanticSegmentation.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''	)
       UpperCAmelCase_							:      Union[str, Any]						=   model.to(__snake_case	)

       UpperCAmelCase_							:      List[str]						=   MobileViTImageProcessor.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''	)

       UpperCAmelCase_							:      List[Any]						=   prepare_img()
       UpperCAmelCase_							:      Any						=   image_processor(images=__snake_case       ,   return_tensors='''pt'''	).to(__snake_case	)

       # forward pass
       with torch.no_grad():
          UpperCAmelCase_							:      Optional[int]						=   model(**__snake_case	)
       UpperCAmelCase_							:      Dict						=   outputs.logits

       # verify the logits
       UpperCAmelCase_							:      Optional[Any]						=   torch.Size((1, 21, 32, 32)	)
       self.assertEqual(logits.shape       ,   __snake_case	)

       UpperCAmelCase_							:      Optional[int]						=   torch.tensor(
           [
               [[7.0_863, 7.1_525, 6.8_201], [6.6_931, 6.8_770, 6.8_933], [6.2_978, 7.0_366, 6.9_636]],
               [[-3.7_134, -3.6_712, -3.6_675], [-3.5_825, -3.3_549, -3.4_777], [-3.3_435, -3.3_979, -3.2_857]],
               [[-2.9_329, -2.8_003, -2.7_369], [-3.0_564, -2.4_780, -2.0_207], [-2.6_889, -1.9_298, -1.7_640]],
           ]       ,   device=__snake_case       ,   )

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





    @slow
    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[str]						=   MobileViTVaForSemanticSegmentation.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''	)
       UpperCAmelCase_							:      Optional[int]						=   model.to(__snake_case	)

       UpperCAmelCase_							:      str						=   MobileViTImageProcessor.from_pretrained('''shehan97/mobilevitv2-1.0-voc-deeplabv3'''	)

       UpperCAmelCase_							:      Optional[int]						=   prepare_img()
       UpperCAmelCase_							:      Tuple						=   image_processor(images=__snake_case       ,   return_tensors='''pt'''	).to(__snake_case	)

       # forward pass
       with torch.no_grad():
          UpperCAmelCase_							:      str						=   model(**__snake_case	)

       UpperCAmelCase_							:      Any						=   outputs.logits.detach().cpu()

       UpperCAmelCase_							:      Optional[Any]						=   image_processor.post_process_semantic_segmentation(outputs=__snake_case       ,   target_sizes=[(50, 60)]	)
       UpperCAmelCase_							:      int						=   torch.Size((50, 60)	)
       self.assertEqual(segmentation[0].shape       ,   __snake_case	)

       UpperCAmelCase_							:      Optional[Any]						=   image_processor.post_process_semantic_segmentation(outputs=__snake_case	)
       UpperCAmelCase_							:      str						=   torch.Size((32, 32)	)
       self.assertEqual(segmentation[0].shape       ,   __snake_case	)



import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCamelCase				:  Tuple      					=							logging.get_logger(__name__)





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     Union[str, Any]			= 'encoder-decoder'
    A_  :     Optional[int]			= True





    def __init__(	self				:   Dict       ,   **__snake_case				:   Union[str, Any]	):





       '''simple docstring'''
       super().__init__(**__snake_case	)
       assert (
           "encoder" in kwargs and "decoder" in kwargs
       ), "Config has to be initialized with encoder and decoder config"
       UpperCAmelCase_							:      int						=   kwargs.pop('''encoder'''	)
       UpperCAmelCase_							:      List[Any]						=   encoder_config.pop('''model_type'''	)
       UpperCAmelCase_							:      int						=   kwargs.pop('''decoder'''	)
       UpperCAmelCase_							:      Union[str, Any]						=   decoder_config.pop('''model_type'''	)

       from ..auto.configuration_auto import AutoConfig

       UpperCAmelCase_							:      Optional[int]						=   AutoConfig.for_model(__snake_case       ,   **__snake_case	)
       UpperCAmelCase_							:      Any						=   AutoConfig.for_model(__snake_case       ,   **__snake_case	)
       UpperCAmelCase_							:      List[Any]						=   True





    @classmethod
    def 					_lowerCamelCase (	cls				:   Optional[Any]       ,   __snake_case				:   PretrainedConfig       ,   __snake_case				:   PretrainedConfig       ,   **__snake_case				:   Union[str, Any]	):





       '''simple docstring'''
       logger.info('''Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config'''	)
       UpperCAmelCase_							:      Dict						=   True
       UpperCAmelCase_							:      str						=   True

       return cls(encoder=encoder_config.to_dict()       ,   decoder=decoder_config.to_dict()       ,   **__snake_case	)





    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      int						=   copy.deepcopy(self.__dict__	)
       UpperCAmelCase_							:      Tuple						=   self.encoder.to_dict()
       UpperCAmelCase_							:      Tuple						=   self.decoder.to_dict()
       UpperCAmelCase_							:      Tuple						=   self.__class__.model_type
       return output