






'''simple docstring'''






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 a_    (    __UpperCAmelCase							):







   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[int]:
      """simple docstring"""






      __snake_case				:      Optional[Any]     							= self.config_class(**self.inputs_dict)
      self.parent.assertTrue(hasattr(UpperCAmelCase_					,   'width_multiplier'))




class a_    :







   def __init__(self					,   __a					,   __a=1_3					,   __a=6_4					,   __a=2					,   __a=3					,   __a="swish"					,   __a=3					,   __a=3_2					,   __a=0.1					,   __a=0.02					,   __a=True					,   __a=True					,   __a=1_0					,   __a=None					,   __a=0.25					,   __a=0.0					,   __a=0.0					,   )			->      Tuple:
      """simple docstring"""






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







   def     SCREAMING_SNAKE_CASE__					(self)			->      Tuple:
      """simple docstring"""






      __snake_case				:      Optional[int]     							= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

      __snake_case				:      Any     							= self.get_config()

      return config, pixel_values, labels, pixel_labels







   def     SCREAMING_SNAKE_CASE__					(self)			->      Tuple:
      """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     SCREAMING_SNAKE_CASE__					(self					,   __a					,   __a					,   __a					,   __a)			->      str:
      """simple docstring"""






      __snake_case				:      int     							= MobileViTVaModel(config=UpperCAmelCase_)
      model.to(UpperCAmelCase_)
      model.eval()
      __snake_case				:      List[str]     							= model(UpperCAmelCase_)
      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     SCREAMING_SNAKE_CASE__					(self					,   __a					,   __a					,   __a					,   __a)			->      str:
      """simple docstring"""






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







   def     SCREAMING_SNAKE_CASE__					(self					,   __a					,   __a					,   __a					,   __a)			->      Optional[Any]:
      """simple docstring"""






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






      __snake_case				:      Tuple     							= self.prepare_config_and_inputs()
      __snake_case							,__snake_case							,__snake_case							,__snake_case				:      Optional[Any]     							= config_and_inputs
      __snake_case				:      str     							= {'pixel_values': pixel_values}
      return config, inputs_dict




@require_torch
class a_    (    __UpperCAmelCase		,      __UpperCAmelCase		,      unittest.TestCase							):
   _snake_case        =							(
       (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
       if is_torch_available()
       else ()
   )

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

   _snake_case        =							False
   _snake_case        =							False
   _snake_case        =							False
   _snake_case        =							False







   def     SCREAMING_SNAKE_CASE__					(self)			->      Dict:
      """simple docstring"""






      __snake_case				:      Any     							= MobileViTVaModelTester(self)
      __snake_case				:      List[Any]     							= MobileViTVaConfigTester(self					,   config_class=UpperCAmelCase_					,   has_text_modality=UpperCAmelCase_)







   def     SCREAMING_SNAKE_CASE__					(self)			->      Union[str, Any]:
      """simple docstring"""






      self.config_tester.run_common_tests()







   @unittest.skip(reason='MobileViTV2 does not use inputs_embeds')
   def     SCREAMING_SNAKE_CASE__					(self)			->      int:
      """simple docstring"""






      pass







   @unittest.skip(reason='MobileViTV2 does not support input and output embeddings')
   def     SCREAMING_SNAKE_CASE__					(self)			->      Tuple:
      """simple docstring"""






      pass







   @unittest.skip(reason='MobileViTV2 does not output attentions')
   def     SCREAMING_SNAKE_CASE__					(self)			->      List[Any]:
      """simple docstring"""






      pass







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






      pass







   def     SCREAMING_SNAKE_CASE__					(self)			->      int:
      """simple docstring"""






      __snake_case							,__snake_case				:      str     							= self.model_tester.prepare_config_and_inputs_for_common()

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

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







   def     SCREAMING_SNAKE_CASE__					(self)			->      Union[str, Any]:
      """simple docstring"""






      __snake_case				:      Optional[Any]     							= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*UpperCAmelCase_)







   def     SCREAMING_SNAKE_CASE__					(self)			->      Any:
      """simple docstring"""






      def check_hidden_states_output(__a					,   __a					,   __a):
         __snake_case				:      Union[str, Any]     							= model_class(UpperCAmelCase_)
         model.to(UpperCAmelCase_)
         model.eval()

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

         __snake_case				:      Tuple     							= outputs.hidden_states

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

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

      __snake_case							,__snake_case				:      List[str]     							= self.model_tester.prepare_config_and_inputs_for_common()

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

         # check that output_hidden_states also work using config
         del inputs_dict["output_hidden_states"]
         __snake_case				:      List[Any]     							= True

         check_hidden_states_output(UpperCAmelCase_					,   UpperCAmelCase_					,   UpperCAmelCase_)







   def     SCREAMING_SNAKE_CASE__					(self)			->      Dict:
      """simple docstring"""






      __snake_case				:      Optional[int]     							= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase_)







   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[int]:
      """simple docstring"""






      __snake_case				:      Optional[int]     							= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_semantic_segmentation(*UpperCAmelCase_)







   @slow
   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[int]:
      """simple docstring"""






      for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         __snake_case				:      str     							= MobileViTVaModel.from_pretrained(UpperCAmelCase_)
         self.assertIsNotNone(UpperCAmelCase_)




def        _SCREAMING_SNAKE_CASE  (				)    -> Tuple:







   """simple docstring"""
   __snake_case				:      Tuple     							= Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'  )
   return image





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







   @cached_property
   def     SCREAMING_SNAKE_CASE__					(self)			->      str:
      """simple docstring"""






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







   @slow
   def     SCREAMING_SNAKE_CASE__					(self)			->      Dict:
      """simple docstring"""






      __snake_case				:      int     							= MobileViTVaForImageClassification.from_pretrained('apple/mobilevitv2-1.0-imagenet1k-256').to(
          UpperCAmelCase_)

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

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

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

      __snake_case				:      Any     							= torch.tensor([-1.6_336E00, -7.3_204E-02, -5.1_883E-01]).to(UpperCAmelCase_)

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







   @slow
   def     SCREAMING_SNAKE_CASE__					(self)			->      Dict:
      """simple docstring"""






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

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

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

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

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

      __snake_case				:      Dict     							= 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=UpperCAmelCase_					,   )

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







   @slow
   def     SCREAMING_SNAKE_CASE__					(self)			->      Any:
      """simple docstring"""






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

      __snake_case				:      Union[str, Any]     							= MobileViTImageProcessor.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3')

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

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

      __snake_case				:      Dict     							= outputs.logits.detach().cpu()

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

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







'''simple docstring'''






import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()




class a_    (    UpperCamelCase_		,      UpperCamelCase_		,      unittest.TestCase							):
   _snake_case        =							VQModel
   _snake_case        =							"""sample"""







   @property
   def     SCREAMING_SNAKE_CASE__					(self					,   __a=(3_2, 3_2))			->      str:
      """simple docstring"""






      __snake_case				:      Dict     							= 4
      __snake_case				:      Optional[int]     							= 3

      __snake_case				:      str     							= floats_tensor((batch_size, num_channels) + sizes).to(__a)

      return {"sample": image}







   @property
   def     SCREAMING_SNAKE_CASE__					(self)			->      Union[str, Any]:
      """simple docstring"""






      return (3, 3_2, 3_2)







   @property
   def     SCREAMING_SNAKE_CASE__					(self)			->      List[Any]:
      """simple docstring"""






      return (3, 3_2, 3_2)







   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[int]:
      """simple docstring"""






      __snake_case				:      Optional[Any]     							= {
          'block_out_channels': [3_2, 6_4],
          'in_channels': 3,
          'out_channels': 3,
          'down_block_types': ['DownEncoderBlock2D', 'DownEncoderBlock2D'],
          'up_block_types': ['UpDecoderBlock2D', 'UpDecoderBlock2D'],
          'latent_channels': 3,
      }
      __snake_case				:      List[Any]     							= self.dummy_input
      return init_dict, inputs_dict







   def     SCREAMING_SNAKE_CASE__					(self)			->      Any:
      """simple docstring"""






      pass







   def     SCREAMING_SNAKE_CASE__					(self)			->      int:
      """simple docstring"""






      pass







   def     SCREAMING_SNAKE_CASE__					(self)			->      Union[str, Any]:
      """simple docstring"""






      __snake_case							,__snake_case				:      List[Any]     							= VQModel.from_pretrained('fusing/vqgan-dummy'					,   output_loading_info=__a)
      self.assertIsNotNone(__a)
      self.assertEqual(len(loading_info['missing_keys'])					,   0)

      model.to(__a)
      __snake_case				:      Any     							= model(**self.dummy_input)

      assert image is not None, "Make sure output is not None"







   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[int]:
      """simple docstring"""






      __snake_case				:      Union[str, Any]     							= VQModel.from_pretrained('fusing/vqgan-dummy')
      model.to(__a).eval()

      torch.manual_seed(0)
      if torch.cuda.is_available():
         torch.cuda.manual_seed_all(0)

      __snake_case				:      Tuple     							= torch.randn(1					,   model.config.in_channels					,   model.config.sample_size					,   model.config.sample_size)
      __snake_case				:      Optional[int]     							= image.to(__a)
      with torch.no_grad():
         __snake_case				:      List[Any]     							= model(__a).sample

      __snake_case				:      int     							= output[0, -1, -3:, -3:].flatten().cpu()
      # fmt: off
      __snake_case				:      int     							= torch.tensor([-0.0_153, -0.4_044, -0.1_880, -0.5_161, -0.2_418, -0.4_072, -0.1_612, -0.0_633, -0.0_143])
      # fmt: on
      self.assertTrue(torch.allclose(__a					,   __a					,   atol=1E-3))