


import inspect
import unittest

from transformers import MobileViTConfig
from transformers.testing_utils import require_torch, 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 MobileViTForImageClassification, MobileViTForSemanticSegmentation, MobileViTModel
						from transformers.models.mobilevit.modeling_mobilevit import MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import MobileViTImageProcessor
class        _lowercase      (			lowerCamelCase_       ):






				'''simple docstring'''



				def 		__magic_name__(       self			:List[str]	)		->			Union[str, Any]:
										__SCREAMING_SNAKE_CASE						:     Any					=		self.config_class(**self.inputs_dict	)
										self.parent.assertTrue(hasattr(__a			,			'''hidden_sizes'''	)	)
										self.parent.assertTrue(hasattr(__a			,			'''neck_hidden_sizes'''	)	)
										self.parent.assertTrue(hasattr(__a			,			'''num_attention_heads'''	)	)




class        _lowercase      :






				'''simple docstring'''



				def __init__(       self			:str			,			lowerCAmelCase__			:Optional[Any]			,			lowerCAmelCase__			:Tuple=13			,			lowerCAmelCase__			:List[Any]=32			,			lowerCAmelCase__			:int=2			,			lowerCAmelCase__			:Dict=3			,			lowerCAmelCase__			:Optional[Any]=640			,			lowerCAmelCase__			:Optional[int]=4			,			lowerCAmelCase__			:List[str]="silu"			,			lowerCAmelCase__			:Tuple=3			,			lowerCAmelCase__			:List[Any]=32			,			lowerCAmelCase__			:int=0.1			,			lowerCAmelCase__			:Optional[int]=0.1			,			lowerCAmelCase__			:Tuple=0.1			,			lowerCAmelCase__			:Tuple=0.02			,			lowerCAmelCase__			:List[str]=True			,			lowerCAmelCase__			:Tuple=True			,			lowerCAmelCase__			:Any=10			,			lowerCAmelCase__			:str=None			,			)		->			List[Any]:
										__SCREAMING_SNAKE_CASE						:     int					=		parent
										__SCREAMING_SNAKE_CASE						:     int					=		batch_size
										__SCREAMING_SNAKE_CASE						:     str					=		image_size
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		patch_size
										__SCREAMING_SNAKE_CASE						:     Tuple					=		num_channels
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		last_hidden_size
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		num_attention_heads
										__SCREAMING_SNAKE_CASE						:     int					=		hidden_act
										__SCREAMING_SNAKE_CASE						:     int					=		conv_kernel_size
										__SCREAMING_SNAKE_CASE						:     List[str]					=		output_stride
										__SCREAMING_SNAKE_CASE						:     Tuple					=		hidden_dropout_prob
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		attention_probs_dropout_prob
										__SCREAMING_SNAKE_CASE						:     Any					=		classifier_dropout_prob
										__SCREAMING_SNAKE_CASE						:     Tuple					=		use_labels
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		is_training
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		num_labels
										__SCREAMING_SNAKE_CASE						:     Tuple					=		initializer_range
										__SCREAMING_SNAKE_CASE						:     Any					=		scope


				def 		__magic_name__(       self			:Optional[Any]	)		->			Union[str, Any]:
										__SCREAMING_SNAKE_CASE						:     List[str]					=		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

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

										__SCREAMING_SNAKE_CASE						:     Tuple					=		self.get_config()

										return config, pixel_values, labels, pixel_labels


				def 		__magic_name__(       self			:str	)		->			Dict:
										return MobileViTConfig(
										    image_size=self.image_size			,			patch_size=self.patch_size			,			num_channels=self.num_channels			,			num_attention_heads=self.num_attention_heads			,			hidden_act=self.hidden_act			,			conv_kernel_size=self.conv_kernel_size			,			output_stride=self.output_stride			,			hidden_dropout_prob=self.hidden_dropout_prob			,			attention_probs_dropout_prob=self.attention_probs_dropout_prob			,			classifier_dropout_prob=self.classifier_dropout_prob			,			initializer_range=self.initializer_range			,			)


				def 		__magic_name__(       self			:int			,			lowerCAmelCase__			:Dict			,			lowerCAmelCase__			:str			,			lowerCAmelCase__			:Any			,			lowerCAmelCase__			:Optional[Any]	)		->			Optional[int]:
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		MobileViTModel(config=__a	)
										model.to(__a	)
										model.eval()
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		model(__a	)
										self.parent.assertEqual(
										    result.last_hidden_state.shape			,			(
										        self.batch_size,
										        self.last_hidden_size,
										        self.image_size // self.output_stride,
										        self.image_size // self.output_stride,
										    )			,			)


				def 		__magic_name__(       self			:int			,			lowerCAmelCase__			:Tuple			,			lowerCAmelCase__			:Dict			,			lowerCAmelCase__			:List[Any]			,			lowerCAmelCase__			:Optional[int]	)		->			Optional[int]:
										__SCREAMING_SNAKE_CASE						:     int					=		self.num_labels
										__SCREAMING_SNAKE_CASE						:     int					=		MobileViTForImageClassification(__a	)
										model.to(__a	)
										model.eval()
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		model(__a			,			labels=__a	)
										self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.num_labels)	)


				def 		__magic_name__(       self			:Dict			,			lowerCAmelCase__			:Optional[Any]			,			lowerCAmelCase__			:Any			,			lowerCAmelCase__			:Dict			,			lowerCAmelCase__			:Optional[Any]	)		->			Union[str, Any]:
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		self.num_labels
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		MobileViTForSemanticSegmentation(__a	)
										model.to(__a	)
										model.eval()
										__SCREAMING_SNAKE_CASE						:     List[str]					=		model(__a	)
										self.parent.assertEqual(
										    result.logits.shape			,			(
										        self.batch_size,
										        self.num_labels,
										        self.image_size // self.output_stride,
										        self.image_size // self.output_stride,
										    )			,			)
										__SCREAMING_SNAKE_CASE						:     str					=		model(__a			,			labels=__a	)
										self.parent.assertEqual(
										    result.logits.shape			,			(
										        self.batch_size,
										        self.num_labels,
										        self.image_size // self.output_stride,
										        self.image_size // self.output_stride,
										    )			,			)



				def 		__magic_name__(       self			:Dict	)		->			Dict:
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		self.prepare_config_and_inputs()
										__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     Any					=		config_and_inputs
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		{'''pixel_values''': pixel_values}
										return config, inputs_dict




@require_torch
class        _lowercase      (			lowerCamelCase_						,			lowerCamelCase_						,			unittest.TestCase       ):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__				:			List[str]						=		(
				    (MobileViTModel, MobileViTForImageClassification, MobileViTForSemanticSegmentation)
				    if is_torch_available()
				    else ()
				)
				SCREAMING_SNAKE_CASE__				:			Tuple						=		(
				    {
				        '''feature-extraction''': MobileViTModel,
				        '''image-classification''': MobileViTForImageClassification,
				        '''image-segmentation''': MobileViTForSemanticSegmentation,
				    }
				    if is_torch_available()
				    else {}
				)

				SCREAMING_SNAKE_CASE__				:			Dict						=		False
				SCREAMING_SNAKE_CASE__				:			Union[str, Any]						=		False
				SCREAMING_SNAKE_CASE__				:			List[Any]						=		False
				SCREAMING_SNAKE_CASE__				:			Union[str, Any]						=		False


				def 		__magic_name__(       self			:List[str]	)		->			List[Any]:
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		MobileViTModelTester(self	)
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		MobileViTConfigTester(self			,			config_class=__a			,			has_text_modality=__a	)


				def 		__magic_name__(       self			:Optional[int]	)		->			List[Any]:
										self.config_tester.run_common_tests()


				@unittest.skip(reason='''MobileViT does not use inputs_embeds'''	)
				def 		__magic_name__(       self			:Any	)		->			Any:
										pass


				@unittest.skip(reason='''MobileViT does not support input and output embeddings'''	)
				def 		__magic_name__(       self			:Any	)		->			str:
										pass


				@unittest.skip(reason='''MobileViT does not output attentions'''	)
				def 		__magic_name__(       self			:List[Any]	)		->			Optional[int]:
										pass


				def 		__magic_name__(       self			:List[Any]	)		->			int:
										__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		self.model_tester.prepare_config_and_inputs_for_common()

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

																__SCREAMING_SNAKE_CASE						:     List[str]					=		['''pixel_values''']
																self.assertListEqual(arg_names[:1]			,			__a	)


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


				def 		__magic_name__(       self			:Dict	)		->			Any:
										__SCREAMING_SNAKE_CASE						:     Any					=		self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*__a	)


				def 		__magic_name__(       self			:Tuple	)		->			List[str]:
										def check_hidden_states_output(lowerCAmelCase__			:List[str]			,			lowerCAmelCase__			:Tuple			,			lowerCAmelCase__			:Optional[Any]	):
																__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		model_class(__a	)
																model.to(__a	)
																model.eval()

																with torch.no_grad():
																						__SCREAMING_SNAKE_CASE						:     Optional[int]					=		model(**self._prepare_for_class(__a			,			__a	)	)

																__SCREAMING_SNAKE_CASE						:     Dict					=		outputs.hidden_states

																__SCREAMING_SNAKE_CASE						:     Dict					=		5
																self.assertEqual(len(__a	)			,			__a	)

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

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

										__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		self.model_tester.prepare_config_and_inputs_for_common()

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

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

																check_hidden_states_output(__a			,			__a			,			__a	)


				def 		__magic_name__(       self			:int	)		->			List[str]:
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*__a	)


				def 		__magic_name__(       self			:Optional[int]	)		->			int:
										__SCREAMING_SNAKE_CASE						:     Dict					=		self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_semantic_segmentation(*__a	)



				@slow
				def 		__magic_name__(       self			:int	)		->			str:
										for model_name in MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__SCREAMING_SNAKE_CASE						:     str					=		MobileViTModel.from_pretrained(__a	)
																self.assertIsNotNone(__a	)




def        _UpperCamelCase							(      ):
						__SCREAMING_SNAKE_CASE						:     str					=		Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''      )
						return image
@require_torch
@require_vision
class        _lowercase      (			unittest.TestCase       ):






				'''simple docstring'''



				@cached_property
				def 		__magic_name__(       self			:Union[str, Any]	)		->			List[str]:
										return MobileViTImageProcessor.from_pretrained('''apple/mobilevit-xx-small'''	) if is_vision_available() else None


				@slow
				def 		__magic_name__(       self			:Tuple	)		->			int:
										__SCREAMING_SNAKE_CASE						:     int					=		MobileViTForImageClassification.from_pretrained('''apple/mobilevit-xx-small'''	).to(__a	)

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

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

										# verify the logits
										__SCREAMING_SNAKE_CASE						:     Any					=		torch.Size((1, 1_000)	)
										self.assertEqual(outputs.logits.shape			,			__a	)

										__SCREAMING_SNAKE_CASE						:     Any					=		torch.tensor([-1.9364, -1.2327, -0.4653]	).to(__a	)

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


				@slow
				def 		__magic_name__(       self			:Dict	)		->			Tuple:
										__SCREAMING_SNAKE_CASE						:     str					=		MobileViTForSemanticSegmentation.from_pretrained('''apple/deeplabv3-mobilevit-xx-small'''	)
										__SCREAMING_SNAKE_CASE						:     List[str]					=		model.to(__a	)

										__SCREAMING_SNAKE_CASE						:     Tuple					=		MobileViTImageProcessor.from_pretrained('''apple/deeplabv3-mobilevit-xx-small'''	)

										__SCREAMING_SNAKE_CASE						:     str					=		prepare_img()
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		image_processor(images=__a			,			return_tensors='''pt'''	).to(__a	)

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

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

										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		torch.tensor(
										    [
										        [[6.9713, 6.9786, 7.2422], [7.2893, 7.2825, 7.4446], [7.6580, 7.8797, 7.9420]],
										        [[-10.6869, -10.3250, -10.3471], [-10.4228, -9.9868, -9.7132], [-11.0405, -11.0221, -10.7318]],
										        [[-3.3089, -2.8539, -2.6740], [-3.2706, -2.5621, -2.5108], [-3.2534, -2.6615, -2.6651]],
										    ]			,			device=__a			,			)

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



				@slow
				def 		__magic_name__(       self			:List[str]	)		->			int:
										__SCREAMING_SNAKE_CASE						:     int					=		MobileViTForSemanticSegmentation.from_pretrained('''apple/deeplabv3-mobilevit-xx-small'''	)
										__SCREAMING_SNAKE_CASE						:     int					=		model.to(__a	)

										__SCREAMING_SNAKE_CASE						:     List[str]					=		MobileViTImageProcessor.from_pretrained('''apple/deeplabv3-mobilevit-xx-small'''	)

										__SCREAMING_SNAKE_CASE						:     Tuple					=		prepare_img()
										__SCREAMING_SNAKE_CASE						:     List[str]					=		image_processor(images=__a			,			return_tensors='''pt'''	).to(__a	)

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

										__SCREAMING_SNAKE_CASE						:     List[Any]					=		outputs.logits.detach().cpu()

										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		image_processor.post_process_semantic_segmentation(outputs=__a			,			target_sizes=[(50, 60)]	)
										__SCREAMING_SNAKE_CASE						:     str					=		torch.Size((50, 60)	)
										self.assertEqual(segmentation[0].shape			,			__a	)

										__SCREAMING_SNAKE_CASE						:     Any					=		image_processor.post_process_semantic_segmentation(outputs=__a	)
										__SCREAMING_SNAKE_CASE						:     int					=		torch.Size((32, 32)	)
										self.assertEqual(segmentation[0].shape			,			__a	)







'''simple docstring'''







import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin
class 				__SCREAMING_SNAKE_CASE  (lowerCamelCase_       ,							lowerCamelCase_      ):






				"""simple docstring"""







				@register_to_config
				def __init__(  self     :					Dict					,				*,
				    __a     :					int = 4					,				__a     :					int = 7_68					,				__a     :					int					,				__a     :					int					,				):
								super().__init__()

								_a								= nn.Parameter(torch.zeros(__a					)					)

								# parameters for additional clip time embeddings
								_a								= nn.Linear(__a					,				__a					)
								_a								= nn.Linear(__a					,				__a					)

								# parameters for encoder hidden states
								_a								= clip_extra_context_tokens
								_a								= nn.Linear(
								    __a					,				self.clip_extra_context_tokens * cross_attention_dim					)
								_a								= nn.Linear(__a					,				__a					)
								_a								= nn.LayerNorm(__a					)




				def   UpperCamelCase__   (  self     :					Optional[Any]					,				*, __a     :					Tuple					,				__a     :					Union[str, Any]					,				__a     :					Any					,				__a     :					List[Any]					):
								if do_classifier_free_guidance:
												# Add the classifier free guidance embeddings to the image embeddings
												_a								= image_embeddings.shape[0]
												_a								= self.learned_classifier_free_guidance_embeddings.unsqueeze(0					)
												_a								= classifier_free_guidance_embeddings.expand(
												    __a					,				-1					)
												_a								= torch.cat([classifier_free_guidance_embeddings, image_embeddings]					,				dim=0					)

								# The image embeddings batch size and the text embeddings batch size are equal
								assert image_embeddings.shape[0] == prompt_embeds.shape[0]

								_a								= prompt_embeds.shape[0]

								# "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
								# adding CLIP embeddings to the existing timestep embedding, ...
								_a								= self.embedding_proj(__a					)
								_a								= self.clip_image_embeddings_project_to_time_embeddings(__a					)
								_a								= time_projected_image_embeddings + time_projected_prompt_embeds

								# ... and by projecting CLIP embeddings into four
								# extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
								_a								= self.clip_extra_context_tokens_proj(__a					)
								_a								= clip_extra_context_tokens.reshape(__a					,				-1					,				self.clip_extra_context_tokens					)
								_a								= clip_extra_context_tokens.permute(0					,				2					,				1					)

								_a								= self.encoder_hidden_states_proj(__a					)
								_a								= self.text_encoder_hidden_states_norm(__a					)
								_a								= torch.cat([clip_extra_context_tokens, text_encoder_hidden_states]					,				dim=1					)

								return text_encoder_hidden_states, additive_clip_time_embeddings
