



"""simple docstring"""





import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, PNDMScheduler, StableDiffusionInpaintPipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class    lowercase__     (					SCREAMING_SNAKE_CASE       ,      SCREAMING_SNAKE_CASE       ,      SCREAMING_SNAKE_CASE       ,      unittest.TestCase   ):







       '''simple docstring'''


       UpperCamelCase		=  StableDiffusionInpaintPipeline
       UpperCamelCase		=  TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
       UpperCamelCase		=  TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
       UpperCamelCase		=  frozenset(
           []   )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
       UpperCamelCase		=  frozenset([]   )



       def   lowercase__       (							self    :						Optional[Any]					)				->		Any:


             '''simple docstring'''




             torch.manual_seed(0					)
             UpperCAmelCase_    				=							UNetaDConditionModel(
                 block_out_channels=(32, 64)	,       layers_per_block=2	,       sample_size=32	,       in_channels=9	,       out_channels=4	,       down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")	,       up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")	,       cross_attention_dim=32	,       attention_head_dim=(2, 4)	,       use_linear_projection=_UpperCAmelCase	,       )
             UpperCAmelCase_    				=							PNDMScheduler(skip_prk_steps=_UpperCAmelCase					)
             torch.manual_seed(0					)
             UpperCAmelCase_    				=							AutoencoderKL(
                 block_out_channels=[32, 64]	,       in_channels=3	,       out_channels=3	,       down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]	,       up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]	,       latent_channels=4	,       sample_size=128	,       )
             torch.manual_seed(0					)
             UpperCAmelCase_    				=							CLIPTextConfig(
                 bos_token_id=0	,       eos_token_id=2	,       hidden_size=32	,       intermediate_size=37	,       layer_norm_eps=1e-05	,       num_attention_heads=4	,       num_hidden_layers=5	,       pad_token_id=1	,       vocab_size=1000	,       hidden_act="gelu"	,       projection_dim=512	,       )
             UpperCAmelCase_    				=							CLIPTextModel(_UpperCAmelCase					)
             UpperCAmelCase_    				=							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"					)

             UpperCAmelCase_    				=							{
                 "unet": unet,
                 "scheduler": scheduler,
                 "vae": vae,
                 "text_encoder": text_encoder,
                 "tokenizer": tokenizer,
                 "safety_checker": None,
                 "feature_extractor": None,
             }
             return components



       def   lowercase__       (							self    :						Dict	,       _UpperCAmelCase    :						Dict	,       _UpperCAmelCase    :						Union[str, Any]=0					)				->		Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_    				=							floats_tensor((1, 3, 32, 32)	,       rng=random.Random(_UpperCAmelCase					)					).to(_UpperCAmelCase					)
             UpperCAmelCase_    				=							image.cpu().permute(0	,       2	,       3	,       1					)[0]
             UpperCAmelCase_    				=							Image.fromarray(np.uinta(_UpperCAmelCase					)					).convert("RGB"					).resize((64, 64)					)
             UpperCAmelCase_    				=							Image.fromarray(np.uinta(image + 4					)					).convert("RGB"					).resize((64, 64)					)
             if str(_UpperCAmelCase					).startswith("mps"					):
                   UpperCAmelCase_    				=							torch.manual_seed(_UpperCAmelCase					)
             else:
                   UpperCAmelCase_    				=							torch.Generator(device=_UpperCAmelCase					).manual_seed(_UpperCAmelCase					)
             UpperCAmelCase_    				=							{
                 "prompt": "A painting of a squirrel eating a burger",
                 "image": init_image,
                 "mask_image": mask_image,
                 "generator": generator,
                 "num_inference_steps": 2,
                 "guidance_scale": 6.0,
                 "output_type": "numpy",
             }
             return inputs



       def   lowercase__       (							self    :						Any					)				->		int:


             '''simple docstring'''




             UpperCAmelCase_    				=							"cpu"  # ensure determinism for the device-dependent torch.Generator
             UpperCAmelCase_    				=							self.get_dummy_components()
             UpperCAmelCase_    				=							StableDiffusionInpaintPipeline(**_UpperCAmelCase					)
             UpperCAmelCase_    				=							sd_pipe.to(_UpperCAmelCase					)
             sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase					)

             UpperCAmelCase_    				=							self.get_dummy_inputs(_UpperCAmelCase					)
             UpperCAmelCase_    				=							sd_pipe(**_UpperCAmelCase					).images
             UpperCAmelCase_    				=							image[0, -3:, -3:, -1]

             assert image.shape == (1, 64, 64, 3)
             UpperCAmelCase_    				=							np.array([0.4727, 0.5735, 0.3941, 0.5446, 0.5926, 0.4394, 0.5062, 0.4654, 0.4476]					)

             assert np.abs(image_slice.flatten() - expected_slice					).max() < 1e-2







       def   lowercase__       (							self    :						Optional[Any]					)				->		str:


             '''simple docstring'''




             super().test_inference_batch_single_identical(expected_max_diff=3e-3					)



@slow
@require_torch_gpu
class    lowercase__     (					unittest.TestCase   ):







       '''simple docstring'''


       def   lowercase__       (							self    :						Tuple					)				->		Optional[int]:


             '''simple docstring'''




             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()



       def   lowercase__       (							self    :						Tuple					)				->		str:


             '''simple docstring'''




             UpperCAmelCase_    				=							load_image(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                 "/sd2-inpaint/init_image.png"					)
             UpperCAmelCase_    				=							load_image(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"					)
             UpperCAmelCase_    				=							load_numpy(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"
                 "/yellow_cat_sitting_on_a_park_bench.npy"					)

             UpperCAmelCase_    				=							"stabilityai/stable-diffusion-2-inpainting"
             UpperCAmelCase_    				=							StableDiffusionInpaintPipeline.from_pretrained(_UpperCAmelCase	,       safety_checker=_UpperCAmelCase					)
             pipe.to(_UpperCAmelCase					)
             pipe.set_progress_bar_config(disable=_UpperCAmelCase					)
             pipe.enable_attention_slicing()

             UpperCAmelCase_    				=							"Face of a yellow cat, high resolution, sitting on a park bench"

             UpperCAmelCase_    				=							torch.manual_seed(0					)
             UpperCAmelCase_    				=							pipe(
                 prompt=_UpperCAmelCase	,       image=_UpperCAmelCase	,       mask_image=_UpperCAmelCase	,       generator=_UpperCAmelCase	,       output_type="np"	,       )
             UpperCAmelCase_    				=							output.images[0]

             assert image.shape == (512, 512, 3)
             assert np.abs(expected_image - image					).max() < 9e-3



       def   lowercase__       (							self    :						str					)				->		List[Any]:


             '''simple docstring'''




             UpperCAmelCase_    				=							load_image(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                 "/sd2-inpaint/init_image.png"					)
             UpperCAmelCase_    				=							load_image(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"					)
             UpperCAmelCase_    				=							load_numpy(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint"
                 "/yellow_cat_sitting_on_a_park_bench_fp16.npy"					)

             UpperCAmelCase_    				=							"stabilityai/stable-diffusion-2-inpainting"
             UpperCAmelCase_    				=							StableDiffusionInpaintPipeline.from_pretrained(
                 _UpperCAmelCase	,       torch_dtype=torch.floataa	,       safety_checker=_UpperCAmelCase	,       )
             pipe.to(_UpperCAmelCase					)
             pipe.set_progress_bar_config(disable=_UpperCAmelCase					)
             pipe.enable_attention_slicing()

             UpperCAmelCase_    				=							"Face of a yellow cat, high resolution, sitting on a park bench"

             UpperCAmelCase_    				=							torch.manual_seed(0					)
             UpperCAmelCase_    				=							pipe(
                 prompt=_UpperCAmelCase	,       image=_UpperCAmelCase	,       mask_image=_UpperCAmelCase	,       generator=_UpperCAmelCase	,       output_type="np"	,       )
             UpperCAmelCase_    				=							output.images[0]

             assert image.shape == (512, 512, 3)
             assert np.abs(expected_image - image					).max() < 5e-1







       def   lowercase__       (							self    :						Dict					)				->		Dict:


             '''simple docstring'''




             torch.cuda.empty_cache()
             torch.cuda.reset_max_memory_allocated()
             torch.cuda.reset_peak_memory_stats()

             UpperCAmelCase_    				=							load_image(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                 "/sd2-inpaint/init_image.png"					)
             UpperCAmelCase_    				=							load_image(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"					)

             UpperCAmelCase_    				=							"stabilityai/stable-diffusion-2-inpainting"
             UpperCAmelCase_    				=							PNDMScheduler.from_pretrained(_UpperCAmelCase	,       subfolder="scheduler"					)
             UpperCAmelCase_    				=							StableDiffusionInpaintPipeline.from_pretrained(
                 _UpperCAmelCase	,       safety_checker=_UpperCAmelCase	,       scheduler=_UpperCAmelCase	,       torch_dtype=torch.floataa	,       )
             pipe.to(_UpperCAmelCase					)
             pipe.set_progress_bar_config(disable=_UpperCAmelCase					)
             pipe.enable_attention_slicing(1					)
             pipe.enable_sequential_cpu_offload()

             UpperCAmelCase_    				=							"Face of a yellow cat, high resolution, sitting on a park bench"

             UpperCAmelCase_    				=							torch.manual_seed(0					)
             UpperCAmelCase_    				=							pipe(
                 prompt=_UpperCAmelCase	,       image=_UpperCAmelCase	,       mask_image=_UpperCAmelCase	,       generator=_UpperCAmelCase	,       num_inference_steps=2	,       output_type="np"	,       )

             UpperCAmelCase_    				=							torch.cuda.max_memory_allocated()
             # make sure that less than 2.65 GB is allocated
             assert mem_bytes < 2.65 * 10**9









"""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    lowercase__     (					SCREAMING_SNAKE_CASE   ):







       '''simple docstring'''


       def   lowercase__       (							self    :						Union[str, Any]					)				->		Dict:


             '''simple docstring'''




             UpperCAmelCase_    				=							self.config_class(**self.inputs_dict					)
             self.parent.assertTrue(hasattr(_UpperCAmelCase	,       "width_multiplier"					)					)



class    lowercase__     :







       '''simple docstring'''


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


             '''simple docstring'''




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



       def   lowercase__       (							self    :						Optional[int]					)				->		List[str]:


             '''simple docstring'''




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

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

             UpperCAmelCase_    				=							self.get_config()

             return config, pixel_values, labels, pixel_labels



       def   lowercase__       (							self    :						Tuple					)				->		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   lowercase__       (							self    :						List[str]	,       _UpperCAmelCase    :						str	,       _UpperCAmelCase    :						Optional[Any]	,       _UpperCAmelCase    :						Optional[Any]	,       _UpperCAmelCase    :						int					)				->		int:


             '''simple docstring'''




             UpperCAmelCase_    				=							MobileViTVaModel(config=_UpperCAmelCase					)
             model.to(_UpperCAmelCase					)
             model.eval()
             UpperCAmelCase_    				=							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   lowercase__       (							self    :						Optional[int]	,       _UpperCAmelCase    :						int	,       _UpperCAmelCase    :						Any	,       _UpperCAmelCase    :						str	,       _UpperCAmelCase    :						Tuple					)				->		Tuple:


             '''simple docstring'''




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



       def   lowercase__       (							self    :						Any	,       _UpperCAmelCase    :						int	,       _UpperCAmelCase    :						Optional[int]	,       _UpperCAmelCase    :						List[Any]	,       _UpperCAmelCase    :						Tuple					)				->		Dict:


             '''simple docstring'''




             UpperCAmelCase_    				=							self.num_labels
             UpperCAmelCase_    				=							MobileViTVaForSemanticSegmentation(_UpperCAmelCase					)
             model.to(_UpperCAmelCase					)
             model.eval()
             UpperCAmelCase_    				=							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,
                 )	,       )
             UpperCAmelCase_    				=							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   lowercase__       (							self    :						List[str]					)				->		Tuple:


             '''simple docstring'''




             UpperCAmelCase_    				=							self.prepare_config_and_inputs()
             UpperCAmelCase_				,					UpperCAmelCase_				,					UpperCAmelCase_				,					UpperCAmelCase_    				=							config_and_inputs
             UpperCAmelCase_    				=							{"pixel_values": pixel_values}
             return config, inputs_dict



@require_torch
class    lowercase__     (					SCREAMING_SNAKE_CASE       ,      SCREAMING_SNAKE_CASE       ,      unittest.TestCase   ):







       '''simple docstring'''


       UpperCamelCase		=  (
           (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
           if is_torch_available()
           else ()
       )

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

       UpperCamelCase		=  False
       UpperCamelCase		=  False
       UpperCamelCase		=  False
       UpperCamelCase		=  False



       def   lowercase__       (							self    :						str					)				->		Dict:


             '''simple docstring'''




             UpperCAmelCase_    				=							MobileViTVaModelTester(self					)
             UpperCAmelCase_    				=							MobileViTVaConfigTester(self	,       config_class=_UpperCAmelCase	,       has_text_modality=_UpperCAmelCase					)



       def   lowercase__       (							self    :						int					)				->		List[str]:


             '''simple docstring'''




             self.config_tester.run_common_tests()



       @unittest.skip(reason="MobileViTV2 does not use inputs_embeds"					)
       def   lowercase__       (							self    :						List[Any]					)				->		List[str]:


             '''simple docstring'''




             pass



       @unittest.skip(reason="MobileViTV2 does not support input and output embeddings"					)
       def   lowercase__       (							self    :						Tuple					)				->		Dict:


             '''simple docstring'''




             pass



       @unittest.skip(reason="MobileViTV2 does not output attentions"					)
       def   lowercase__       (							self    :						Optional[int]					)				->		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   lowercase__       (							self    :						int					)				->		int:


             '''simple docstring'''




             pass



       @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."					)
       def   lowercase__       (							self    :						Tuple					)				->		Optional[int]:


             '''simple docstring'''




             pass



       def   lowercase__       (							self    :						int					)				->		int:


             '''simple docstring'''




             UpperCAmelCase_				,					UpperCAmelCase_    				=							self.model_tester.prepare_config_and_inputs_for_common()

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

                   UpperCAmelCase_    				=							["pixel_values"]
                   self.assertListEqual(arg_names[:1]	,       _UpperCAmelCase					)



       def   lowercase__       (							self    :						Union[str, Any]					)				->		int:


             '''simple docstring'''




             UpperCAmelCase_    				=							self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*_UpperCAmelCase					)



       def   lowercase__       (							self    :						List[Any]					)				->		Optional[int]:


             '''simple docstring'''




             def check_hidden_states_output(_UpperCAmelCase    :						List[Any]	,       _UpperCAmelCase    :						Union[str, Any]	,       _UpperCAmelCase    :						int					):
                   UpperCAmelCase_    				=							model_class(_UpperCAmelCase					)
                   model.to(_UpperCAmelCase					)
                   model.eval()

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

                   UpperCAmelCase_    				=							outputs.hidden_states

                   UpperCAmelCase_    				=							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.
                   UpperCAmelCase_    				=							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					)

             UpperCAmelCase_				,					UpperCAmelCase_    				=							self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   UpperCAmelCase_    				=							True
                   check_hidden_states_output(_UpperCAmelCase	,       _UpperCAmelCase	,       _UpperCAmelCase					)

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

                   check_hidden_states_output(_UpperCAmelCase	,       _UpperCAmelCase	,       _UpperCAmelCase					)



       def   lowercase__       (							self    :						int					)				->		int:


             '''simple docstring'''




             UpperCAmelCase_    				=							self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*_UpperCAmelCase					)



       def   lowercase__       (							self    :						int					)				->		Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_    				=							self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_semantic_segmentation(*_UpperCAmelCase					)







       @slow
       def   lowercase__       (							self    :						Optional[int]					)				->		Union[str, Any]:


             '''simple docstring'''




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



def      a__			(   ):
      UpperCAmelCase_    				=							Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"   )
      return image


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







       '''simple docstring'''


       @cached_property
       def   lowercase__       (							self    :						int					)				->		List[Any]:


             '''simple docstring'''




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



       @slow
       def   lowercase__       (							self    :						Union[str, Any]					)				->		Tuple:


             '''simple docstring'''




             UpperCAmelCase_    				=							MobileViTVaForImageClassification.from_pretrained("apple/mobilevitv2-1.0-imagenet1k-256"					).to(
                 _UpperCAmelCase					)

             UpperCAmelCase_    				=							self.default_image_processor
             UpperCAmelCase_    				=							prepare_img()
             UpperCAmelCase_    				=							image_processor(images=_UpperCAmelCase	,       return_tensors="pt"					).to(_UpperCAmelCase					)

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

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

             UpperCAmelCase_    				=							torch.tensor([-1.6336e00, -7.3204e-02, -5.1883e-01]					).to(_UpperCAmelCase					)

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



       @slow
       def   lowercase__       (							self    :						List[str]					)				->		Dict:


             '''simple docstring'''




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

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

             UpperCAmelCase_    				=							prepare_img()
             UpperCAmelCase_    				=							image_processor(images=_UpperCAmelCase	,       return_tensors="pt"					).to(_UpperCAmelCase					)

             # forward pass
             with torch.no_grad():
                   UpperCAmelCase_    				=							model(**_UpperCAmelCase					)
             UpperCAmelCase_    				=							outputs.logits

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

             UpperCAmelCase_    				=							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=_UpperCAmelCase	,       )

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







       @slow
       def   lowercase__       (							self    :						Optional[int]					)				->		Tuple:


             '''simple docstring'''




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

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

             UpperCAmelCase_    				=							prepare_img()
             UpperCAmelCase_    				=							image_processor(images=_UpperCAmelCase	,       return_tensors="pt"					).to(_UpperCAmelCase					)

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

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

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

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




