


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 							_UpperCamelCase (						lowerCamelCase__						):

   '''simple docstring'''







   def   _snake_case    (  self  : Optional[int]							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[int]       						=						self.config_class(**self.inputs_dict							)
         self.parent.assertTrue(hasattr(__lowerCamelCase       ,  """width_multiplier"""							)							)



class 							_UpperCamelCase :

   '''simple docstring'''







   def __init__(  self  : Optional[Any]       ,  _lowerCamelCase  : Any       ,  _lowerCamelCase  : Union[str, Any]=1_3       ,  _lowerCamelCase  : Optional[int]=6_4       ,  _lowerCamelCase  : Union[str, Any]=2       ,  _lowerCamelCase  : List[Any]=3       ,  _lowerCamelCase  : Tuple="swish"       ,  _lowerCamelCase  : List[Any]=3       ,  _lowerCamelCase  : Union[str, Any]=3_2       ,  _lowerCamelCase  : Tuple=0.1       ,  _lowerCamelCase  : Dict=0.02       ,  _lowerCamelCase  : List[str]=True       ,  _lowerCamelCase  : int=True       ,  _lowerCamelCase  : List[Any]=1_0       ,  _lowerCamelCase  : Optional[int]=None       ,  _lowerCamelCase  : Union[str, Any]=0.25       ,  _lowerCamelCase  : int=0.0       ,  _lowerCamelCase  : Optional[int]=0.0       ,  ):






         '''simple docstring'''



         __lowerCamelCase     :      Dict       						=						parent
         __lowerCamelCase     :      List[str]       						=						batch_size
         __lowerCamelCase     :      Dict       						=						image_size
         __lowerCamelCase     :      Tuple       						=						patch_size
         __lowerCamelCase     :      List[Any]       						=						num_channels
         __lowerCamelCase     :      Tuple       						=						make_divisible(5_1_2 * width_multiplier       ,  divisor=8							)
         __lowerCamelCase     :      Dict       						=						hidden_act
         __lowerCamelCase     :      Optional[int]       						=						conv_kernel_size
         __lowerCamelCase     :      int       						=						output_stride
         __lowerCamelCase     :      List[Any]       						=						classifier_dropout_prob
         __lowerCamelCase     :      Dict       						=						use_labels
         __lowerCamelCase     :      Tuple       						=						is_training
         __lowerCamelCase     :      str       						=						num_labels
         __lowerCamelCase     :      Optional[Any]       						=						initializer_range
         __lowerCamelCase     :      Tuple       						=						scope
         __lowerCamelCase     :      Tuple       						=						width_multiplier
         __lowerCamelCase     :      Optional[int]       						=						ffn_dropout
         __lowerCamelCase     :      Dict       						=						attn_dropout

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






         '''simple docstring'''



         __lowerCamelCase     :      str       						=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

         __lowerCamelCase     :      Union[str, Any]       						=						None
         __lowerCamelCase     :      Dict       						=						None
         if self.use_labels:
               __lowerCamelCase     :      Optional[Any]       						=						ids_tensor([self.batch_size]       ,  self.num_labels							)
               __lowerCamelCase     :      Optional[Any]       						=						ids_tensor([self.batch_size, self.image_size, self.image_size]       ,  self.num_labels							)

         __lowerCamelCase     :      Tuple       						=						self.get_config()

         return config, pixel_values, labels, pixel_labels

   def   _snake_case    (  self  : int							):






         '''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   _snake_case    (  self  : List[str]       ,  _lowerCamelCase  : Dict       ,  _lowerCamelCase  : List[Any]       ,  _lowerCamelCase  : List[Any]       ,  _lowerCamelCase  : Tuple							):






         '''simple docstring'''



         __lowerCamelCase     :      int       						=						MobileViTVaModel(config=__lowerCamelCase							)
         model.to(__lowerCamelCase							)
         model.eval()
         __lowerCamelCase     :      Dict       						=						model(__lowerCamelCase							)
         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   _snake_case    (  self  : Tuple       ,  _lowerCamelCase  : Optional[int]       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : Tuple       ,  _lowerCamelCase  : Any							):






         '''simple docstring'''



         __lowerCamelCase     :      List[str]       						=						self.num_labels
         __lowerCamelCase     :      Union[str, Any]       						=						MobileViTVaForImageClassification(__lowerCamelCase							)
         model.to(__lowerCamelCase							)
         model.eval()
         __lowerCamelCase     :      Dict       						=						model(__lowerCamelCase       ,  labels=__lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.num_labels)							)

   def   _snake_case    (  self  : Dict       ,  _lowerCamelCase  : Tuple       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : List[str]							):






         '''simple docstring'''



         __lowerCamelCase     :      Any       						=						self.num_labels
         __lowerCamelCase     :      Optional[int]       						=						MobileViTVaForSemanticSegmentation(__lowerCamelCase							)
         model.to(__lowerCamelCase							)
         model.eval()
         __lowerCamelCase     :      Optional[Any]       						=						model(__lowerCamelCase							)
         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     :      List[str]       						=						model(__lowerCamelCase       ,  labels=__lowerCamelCase							)
         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   _snake_case    (  self  : List[Any]							):






         '''simple docstring'''



         __lowerCamelCase     :      Union[str, Any]       						=						self.prepare_config_and_inputs()
         __lowerCamelCase     :      Optional[int]       						=						config_and_inputs
         __lowerCamelCase     :      Dict       						=						{'''pixel_values''': pixel_values}
         return config, inputs_dict



@require_torch
class 							_UpperCamelCase (						lowerCamelCase__,lowerCamelCase__,unittest.TestCase						):

   '''simple docstring'''







   a_				:							Any    = (
       (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
       if is_torch_available()
       else ()
   )

   a_				:							int    = (
       {
           "feature-extraction": MobileViTVaModel,
           "image-classification": MobileViTVaForImageClassification,
           "image-segmentation": MobileViTVaForSemanticSegmentation,
       }
       if is_torch_available()
       else {}
   )

   a_				:							Tuple    = False
   a_				:							str    = False
   a_				:							List[str]    = False
   a_				:							Optional[Any]    = False

   def   _snake_case    (  self  : List[Any]							):






         '''simple docstring'''



         __lowerCamelCase     :      Dict       						=						MobileViTVaModelTester(self							)
         __lowerCamelCase     :      Optional[int]       						=						MobileViTVaConfigTester(self       ,  config_class=__lowerCamelCase       ,  has_text_modality=__lowerCamelCase							)

   def   _snake_case    (  self  : List[str]							):






         '''simple docstring'''



         self.config_tester.run_common_tests()

   @unittest.skip(reason="""MobileViTV2 does not use inputs_embeds"""							)
   def   _snake_case    (  self  : Optional[int]							):






         '''simple docstring'''



         pass

   @unittest.skip(reason="""MobileViTV2 does not support input and output embeddings"""							)
   def   _snake_case    (  self  : Union[str, Any]							):






         '''simple docstring'''



         pass

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






         '''simple docstring'''



         pass

   @require_torch_multi_gpu
   @unittest.skip(reason="""Got `CUDA error: misaligned address` for tests after this one being run."""							)
   def   _snake_case    (  self  : Optional[Any]							):






         '''simple docstring'''



         pass

   @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""							)
   def   _snake_case    (  self  : Union[str, Any]							):






         '''simple docstring'''



         pass

   def   _snake_case    (  self  : int							):






         '''simple docstring'''



         __lowerCamelCase     :      List[Any]       						=						self.model_tester.prepare_config_and_inputs_for_common()

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

               __lowerCamelCase     :      Union[str, Any]       						=						['''pixel_values''']
               self.assertListEqual(arg_names[:1]       ,  __lowerCamelCase							)

   def   _snake_case    (  self  : List[str]							):






         '''simple docstring'''



         __lowerCamelCase     :      List[Any]       						=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__lowerCamelCase							)

   def   _snake_case    (  self  : List[Any]							):






         '''simple docstring'''



         def check_hidden_states_output(_lowerCamelCase  : Dict       ,  _lowerCamelCase  : Dict       ,  _lowerCamelCase  : Any							):
               __lowerCamelCase     :      Union[str, Any]       						=						model_class(__lowerCamelCase							)
               model.to(__lowerCamelCase							)
               model.eval()

               with torch.no_grad():
                     __lowerCamelCase     :      Optional[Any]       						=						model(**self._prepare_for_class(__lowerCamelCase       ,  __lowerCamelCase							)							)

               __lowerCamelCase     :      Dict       						=						outputs.hidden_states

               __lowerCamelCase     :      Any       						=						5
               self.assertEqual(len(__lowerCamelCase							)       ,  __lowerCamelCase							)

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

         for model_class in self.all_model_classes:
               __lowerCamelCase     :      Any       						=						True
               check_hidden_states_output(__lowerCamelCase       ,  __lowerCamelCase       ,  __lowerCamelCase							)

               # check that output_hidden_states also work using config
               del inputs_dict["output_hidden_states"]
               __lowerCamelCase     :      Tuple       						=						True

               check_hidden_states_output(__lowerCamelCase       ,  __lowerCamelCase       ,  __lowerCamelCase							)

   def   _snake_case    (  self  : int							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[Any]       						=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*__lowerCamelCase							)

   def   _snake_case    (  self  : int							):






         '''simple docstring'''



         __lowerCamelCase     :      str       						=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_semantic_segmentation(*__lowerCamelCase							)


   @slow
   def   _snake_case    (  self  : Tuple							):






         '''simple docstring'''



         for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __lowerCamelCase     :      Optional[int]       						=						MobileViTVaModel.from_pretrained(__lowerCamelCase							)
               self.assertIsNotNone(__lowerCamelCase							)





def   _UpperCAmelCase       (     ):
      """simple docstring"""




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




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

   '''simple docstring'''







   @cached_property
   def   _snake_case    (  self  : int							):






         '''simple docstring'''



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

   @slow
   def   _snake_case    (  self  : List[Any]							):






         '''simple docstring'''



         __lowerCamelCase     :      str       						=						MobileViTVaForImageClassification.from_pretrained("""apple/mobilevitv2-1.0-imagenet1k-256"""							).to(
             __lowerCamelCase							)

         __lowerCamelCase     :      str       						=						self.default_image_processor
         __lowerCamelCase     :      Union[str, Any]       						=						prepare_img()
         __lowerCamelCase     :      int       						=						image_processor(images=__lowerCamelCase       ,  return_tensors="""pt"""							).to(__lowerCamelCase							)

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

         # verify the logits
         __lowerCamelCase     :      str       						=						torch.Size((1, 1_0_0_0)							)
         self.assertEqual(outputs.logits.shape       ,  __lowerCamelCase							)

         __lowerCamelCase     :      Dict       						=						torch.tensor([-1.63_36E00, -7.32_04E-02, -5.18_83E-01]							).to(__lowerCamelCase							)

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

   @slow
   def   _snake_case    (  self  : Any							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[int]       						=						MobileViTVaForSemanticSegmentation.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""							)
         __lowerCamelCase     :      Optional[Any]       						=						model.to(__lowerCamelCase							)

         __lowerCamelCase     :      List[Any]       						=						MobileViTImageProcessor.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""							)

         __lowerCamelCase     :      Optional[Any]       						=						prepare_img()
         __lowerCamelCase     :      Union[str, Any]       						=						image_processor(images=__lowerCamelCase       ,  return_tensors="""pt"""							).to(__lowerCamelCase							)

         # forward pass
         with torch.no_grad():
               __lowerCamelCase     :      List[Any]       						=						model(**__lowerCamelCase							)
         __lowerCamelCase     :      Tuple       						=						outputs.logits

         # verify the logits
         __lowerCamelCase     :      Optional[int]       						=						torch.Size((1, 2_1, 3_2, 3_2)							)
         self.assertEqual(logits.shape       ,  __lowerCamelCase							)

         __lowerCamelCase     :      List[Any]       						=						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=__lowerCamelCase       ,  )

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


   @slow
   def   _snake_case    (  self  : str							):






         '''simple docstring'''



         __lowerCamelCase     :      Dict       						=						MobileViTVaForSemanticSegmentation.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""							)
         __lowerCamelCase     :      Dict       						=						model.to(__lowerCamelCase							)

         __lowerCamelCase     :      Tuple       						=						MobileViTImageProcessor.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""							)

         __lowerCamelCase     :      int       						=						prepare_img()
         __lowerCamelCase     :      Any       						=						image_processor(images=__lowerCamelCase       ,  return_tensors="""pt"""							).to(__lowerCamelCase							)

         # forward pass
         with torch.no_grad():
               __lowerCamelCase     :      List[Any]       						=						model(**__lowerCamelCase							)

         __lowerCamelCase     :      Optional[Any]       						=						outputs.logits.detach().cpu()

         __lowerCamelCase     :      Optional[Any]       						=						image_processor.post_process_semantic_segmentation(outputs=__lowerCamelCase       ,  target_sizes=[(5_0, 6_0)]							)
         __lowerCamelCase     :      Tuple       						=						torch.Size((5_0, 6_0)							)
         self.assertEqual(segmentation[0].shape       ,  __lowerCamelCase							)

         __lowerCamelCase     :      Optional[Any]       						=						image_processor.post_process_semantic_segmentation(outputs=__lowerCamelCase							)
         __lowerCamelCase     :      Optional[Any]       						=						torch.Size((3_2, 3_2)							)
         self.assertEqual(segmentation[0].shape       ,  __lowerCamelCase							)







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__UpperCamelCase    :						Any     =							{
    'configuration_timesformer': ['TIMESFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TimesformerConfig'],
}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __UpperCamelCase    :						int     =							[
        'TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
        'TimesformerModel',
        'TimesformerForVideoClassification',
        'TimesformerPreTrainedModel',
    ]

if TYPE_CHECKING:
    from .configuration_timesformer import TIMESFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, TimesformerConfig

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_timesformer import (
            TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
            TimesformerForVideoClassification,
            TimesformerModel,
            TimesformerPreTrainedModel,
        )

else:
    import sys

    __UpperCamelCase    :						List[Any]     =							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



