


'''simple docstring'''





import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD






def 						UpperCAmelCase__(				):






			"""simple docstring"""




			__A=	'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
			__A=	Image.open(requests.get(_SCREAMING_SNAKE_CASE,stream=_SCREAMING_SNAKE_CASE ).raw ).convert('RGB' )

			return image






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Union[str, Any] ):






			"""simple docstring"""




			__A=	[]
			# fmt: off

			# vision encoder
			rename_keys.append(('visual_encoder.cls_token', 'vision_model.embeddings.class_embedding') )
			rename_keys.append(('visual_encoder.pos_embed', 'vision_model.embeddings.position_embedding') )
			rename_keys.append(('visual_encoder.patch_embed.proj.weight', 'vision_model.embeddings.patch_embedding.weight') )
			rename_keys.append(('visual_encoder.patch_embed.proj.bias', 'vision_model.embeddings.patch_embedding.bias') )
			rename_keys.append(('ln_vision.weight', 'vision_model.post_layernorm.weight') )
			rename_keys.append(('ln_vision.bias', 'vision_model.post_layernorm.bias') )

			for i in range(config.vision_config.num_hidden_layers ):
						rename_keys.append((f"""visual_encoder.blocks.{i}.norm1.weight""", f"""vision_model.encoder.layers.{i}.layer_norm1.weight""") )
						rename_keys.append((f"""visual_encoder.blocks.{i}.norm1.bias""", f"""vision_model.encoder.layers.{i}.layer_norm1.bias""") )
						rename_keys.append((f"""visual_encoder.blocks.{i}.norm2.weight""", f"""vision_model.encoder.layers.{i}.layer_norm2.weight""") )
						rename_keys.append((f"""visual_encoder.blocks.{i}.norm2.bias""", f"""vision_model.encoder.layers.{i}.layer_norm2.bias""") )
						rename_keys.append((f"""visual_encoder.blocks.{i}.attn.qkv.weight""", f"""vision_model.encoder.layers.{i}.self_attn.qkv.weight""") )
						rename_keys.append((f"""visual_encoder.blocks.{i}.attn.proj.weight""", f"""vision_model.encoder.layers.{i}.self_attn.projection.weight""",) )
						rename_keys.append((f"""visual_encoder.blocks.{i}.attn.proj.bias""", f"""vision_model.encoder.layers.{i}.self_attn.projection.bias""") )
						rename_keys.append((f"""visual_encoder.blocks.{i}.mlp.fc1.weight""", f"""vision_model.encoder.layers.{i}.mlp.fc1.weight""") )
						rename_keys.append((f"""visual_encoder.blocks.{i}.mlp.fc1.bias""", f"""vision_model.encoder.layers.{i}.mlp.fc1.bias""") )
						rename_keys.append((f"""visual_encoder.blocks.{i}.mlp.fc2.weight""", f"""vision_model.encoder.layers.{i}.mlp.fc2.weight""") )
						rename_keys.append((f"""visual_encoder.blocks.{i}.mlp.fc2.bias""", f"""vision_model.encoder.layers.{i}.mlp.fc2.bias""") )

			# QFormer
			rename_keys.append(('Qformer.bert.embeddings.LayerNorm.weight', 'qformer.layernorm.weight') )
			rename_keys.append(('Qformer.bert.embeddings.LayerNorm.bias', 'qformer.layernorm.bias') )

			# fmt: on
			return rename_keys






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Optional[Any],_SCREAMING_SNAKE_CASE					: Dict,_SCREAMING_SNAKE_CASE					: Tuple ):






			"""simple docstring"""




			__A=	dct.pop(_SCREAMING_SNAKE_CASE )
			__A=	val






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Dict,_SCREAMING_SNAKE_CASE					: List[str] ):






			"""simple docstring"""




			for i in range(config.vision_config.num_hidden_layers ):
						# read in original q and v biases
						__A=	state_dict.pop(f"""visual_encoder.blocks.{i}.attn.q_bias""" )
						__A=	state_dict.pop(f"""visual_encoder.blocks.{i}.attn.v_bias""" )

						# next, set bias in the state dict
						__A=	torch.cat((q_bias, torch.zeros_like(_SCREAMING_SNAKE_CASE,requires_grad=_SCREAMING_SNAKE_CASE ), v_bias) )
						__A=	qkv_bias






def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Any,_SCREAMING_SNAKE_CASE					: str ):






			"""simple docstring"""




			__A=	364 if 'coco' in model_name else 224
			__A=	BlipaVisionConfig(image_size=_SCREAMING_SNAKE_CASE ).to_dict()

			# make sure the models have proper bos_token_id and eos_token_id set (important for generation)
			# seems like flan-T5 models don't have bos_token_id properly set?
			if "opt-2.7b" in model_name:
						__A=	OPTConfig.from_pretrained('facebook/opt-2.7b',eos_token_id=_SCREAMING_SNAKE_CASE ).to_dict()
			elif "opt-6.7b" in model_name:
						__A=	OPTConfig.from_pretrained('facebook/opt-6.7b',eos_token_id=_SCREAMING_SNAKE_CASE ).to_dict()
			elif "t5-xl" in model_name:
						__A=	TaConfig.from_pretrained('google/flan-t5-xl',dense_act_fn='gelu',bos_token_id=1 ).to_dict()
			elif "t5-xxl" in model_name:
						__A=	TaConfig.from_pretrained('google/flan-t5-xxl',dense_act_fn='gelu',bos_token_id=1 ).to_dict()

			__A=	BlipaConfig(vision_config=_SCREAMING_SNAKE_CASE,text_config=_SCREAMING_SNAKE_CASE )

			return config, image_size






@torch.no_grad()
def 						UpperCAmelCase__(				_SCREAMING_SNAKE_CASE					: Optional[int],_SCREAMING_SNAKE_CASE					: List[Any]=None,_SCREAMING_SNAKE_CASE					: Tuple=False ):






			"""simple docstring"""




			__A=	(
			    AutoTokenizer.from_pretrained('facebook/opt-2.7b' )
			    if 'opt' in model_name
			    else AutoTokenizer.from_pretrained('google/flan-t5-xl' )
			)
			__A=	tokenizer('\n',add_special_tokens=_SCREAMING_SNAKE_CASE ).input_ids[0]
			__A,	__A=	get_blipa_config(_SCREAMING_SNAKE_CASE,eos_token_id=_SCREAMING_SNAKE_CASE )

			__A=	BlipaForConditionalGeneration(_SCREAMING_SNAKE_CASE ).eval()

			__A=	{
			    'blip2-opt-2.7b': ('blip2_opt', 'pretrain_opt2.7b'),
			    'blip2-opt-6.7b': ('blip2_opt', 'pretrain_opt6.7b'),
			    'blip2-opt-2.7b-coco': ('blip2_opt', 'caption_coco_opt2.7b'),
			    'blip2-opt-6.7b-coco': ('blip2_opt', 'caption_coco_opt6.7b'),
			    'blip2-flan-t5-xl': ('blip2_t5', 'pretrain_flant5xl'),
			    'blip2-flan-t5-xl-coco': ('blip2_t5', 'caption_coco_flant5xl'),
			    'blip2-flan-t5-xxl': ('blip2_t5', 'pretrain_flant5xxl'),
			}

			__A,	__A=	model_name_to_original[model_name]

			# load original model
			print('Loading original model...' )
			__A=	'cuda' if torch.cuda.is_available() else 'cpu'
			__A,	__A,	__A=	load_model_and_preprocess(
			    name=_SCREAMING_SNAKE_CASE,model_type=_SCREAMING_SNAKE_CASE,is_eval=_SCREAMING_SNAKE_CASE,device=_SCREAMING_SNAKE_CASE )
			original_model.eval()
			print('Done!' )

			# update state dict keys
			__A=	original_model.state_dict()
			__A=	create_rename_keys(_SCREAMING_SNAKE_CASE )
			for src, dest in rename_keys:
						rename_key(_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE )

			# some keys can be renamed efficiently
			for key, val in state_dict.copy().items():
						__A=	state_dict.pop(_SCREAMING_SNAKE_CASE )
						if key.startswith('Qformer.bert' ):
									__A=	key.replace('Qformer.bert','qformer' )
						if "attention.self" in key:
									__A=	key.replace('self','attention' )
						if "opt_proj" in key:
									__A=	key.replace('opt_proj','language_projection' )
						if "t5_proj" in key:
									__A=	key.replace('t5_proj','language_projection' )
						if key.startswith('opt' ):
									__A=	key.replace('opt','language' )
						if key.startswith('t5' ):
									__A=	key.replace('t5','language' )
						__A=	val

			# read in qv biases
			read_in_q_v_bias(_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE )

			__A,	__A=	hf_model.load_state_dict(_SCREAMING_SNAKE_CASE,strict=_SCREAMING_SNAKE_CASE )
			assert len(_SCREAMING_SNAKE_CASE ) == 0
			assert unexpected_keys == ["qformer.embeddings.position_ids"]

			__A=	load_demo_image()
			__A=	vis_processors['eval'](_SCREAMING_SNAKE_CASE ).unsqueeze(0 ).to(_SCREAMING_SNAKE_CASE )
			__A=	tokenizer(['\n'],return_tensors='pt' ).input_ids.to(_SCREAMING_SNAKE_CASE )

			# create processor
			__A=	BlipImageProcessor(
			    size={'height': image_size, 'width': image_size},image_mean=_SCREAMING_SNAKE_CASE,image_std=_SCREAMING_SNAKE_CASE )
			__A=	BlipaProcessor(image_processor=_SCREAMING_SNAKE_CASE,tokenizer=_SCREAMING_SNAKE_CASE )
			__A=	processor(images=_SCREAMING_SNAKE_CASE,return_tensors='pt' ).pixel_values.to(_SCREAMING_SNAKE_CASE )

			# make sure processor creates exact same pixel values
			assert torch.allclose(_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE )

			original_model.to(_SCREAMING_SNAKE_CASE )
			hf_model.to(_SCREAMING_SNAKE_CASE )
			with torch.no_grad():
						if "opt" in model_name:
									__A=	original_model({'image': original_pixel_values, 'text_input': ['']} ).logits
									__A=	hf_model(_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE ).logits
						else:
									__A=	original_model(
									    {'image': original_pixel_values, 'text_input': ['\n'], 'text_output': ['\n']} ).logits
									__A=	input_ids.masked_fill(input_ids == tokenizer.pad_token_id,-100 )
									__A=	hf_model(_SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE,labels=_SCREAMING_SNAKE_CASE ).logits

			assert original_logits.shape == logits.shape
			print('First values of original logits:',original_logits[0, :3, :3] )
			print('First values of HF logits:',logits[0, :3, :3] )

			# assert values
			if model_name == "blip2-flan-t5-xl":
						__A=	torch.tensor(
						    [[-41.58_50, -4.44_40, -8.99_22], [-47.43_22, -5.91_43, -1.73_40]],device=_SCREAMING_SNAKE_CASE )
						assert torch.allclose(logits[0, :3, :3],_SCREAMING_SNAKE_CASE,atol=1e-4 )
			elif model_name == "blip2-flan-t5-xl-coco":
						__A=	torch.tensor(
						    [[-57.01_09, -9.89_67, -12.62_80], [-68.65_78, -12.71_91, -10.50_65]],device=_SCREAMING_SNAKE_CASE )
			else:
						# cast to same type
						__A=	logits.dtype
						assert torch.allclose(original_logits.to(_SCREAMING_SNAKE_CASE ),_SCREAMING_SNAKE_CASE,atol=1e-2 )
			print('Looks ok!' )

			print('Generating a caption...' )
			__A=	''
			__A=	tokenizer(_SCREAMING_SNAKE_CASE,return_tensors='pt' ).input_ids.to(_SCREAMING_SNAKE_CASE )

			__A=	original_model.generate({'image': original_pixel_values} )
			__A=	hf_model.generate(
			    _SCREAMING_SNAKE_CASE,_SCREAMING_SNAKE_CASE,do_sample=_SCREAMING_SNAKE_CASE,num_beams=5,max_length=30,min_length=1,top_p=0.9,repetition_penalty=1.0,length_penalty=1.0,temperature=1,)
			print('Original generation:',_SCREAMING_SNAKE_CASE )
			__A=	input_ids.shape[1]
			__A=	processor.batch_decode(outputs[:, prompt_length:],skip_special_tokens=_SCREAMING_SNAKE_CASE )
			__A=	[text.strip() for text in output_text]
			print('HF generation:',_SCREAMING_SNAKE_CASE )

			if pytorch_dump_folder_path is not None:
						processor.save_pretrained(_SCREAMING_SNAKE_CASE )
						hf_model.save_pretrained(_SCREAMING_SNAKE_CASE )

			if push_to_hub:
						processor.push_to_hub(f"""nielsr/{model_name}""" )
						hf_model.push_to_hub(f"""nielsr/{model_name}""" )


if __name__ == "__main__":
					UpperCAmelCase__  			=				argparse.ArgumentParser()
					UpperCAmelCase__  			=				[
					    '''blip2-opt-2.7b''',
					    '''blip2-opt-6.7b''',
					    '''blip2-opt-2.7b-coco''',
					    '''blip2-opt-6.7b-coco''',
					    '''blip2-flan-t5-xl''',
					    '''blip2-flan-t5-xl-coco''',
					    '''blip2-flan-t5-xxl''',
					]
					parser.add_argument(
					    '''--model_name''',
					    default='''blip2-opt-2.7b''',
					    choices=choices,
					    type=str,
					    help='''Path to hf config.json of model to convert''',
					)
					parser.add_argument('''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
					parser.add_argument(
					    '''--push_to_hub''',
					    action='''store_true''',
					    help='''Whether to push the model and processor to the hub after converting''',
					)

					UpperCAmelCase__  			=				parser.parse_args()

					convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)







'''simple docstring'''





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





			'''simple docstring'''

			def 					lowerCAmelCase     (  self							:			Tuple	)      ->		Optional[int]:
						__A=	self.config_class(**self.inputs_dict	)
						self.parent.assertTrue(hasattr(lowerCAmelCase_			,    'hidden_sizes'	)	)
						self.parent.assertTrue(hasattr(lowerCAmelCase_			,    'neck_hidden_sizes'	)	)
						self.parent.assertTrue(hasattr(lowerCAmelCase_			,    'num_attention_heads'	)	)


class 						a__					:





			'''simple docstring'''

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

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

						__A=	self.get_config()

						return config, pixel_values, labels, pixel_labels
			def 					lowerCAmelCase     (  self							:			List[Any]	)      ->		List[Any]:
						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 					lowerCAmelCase     (  self							:			Any			,    lowerCAmelCase_							:			List[Any]			,    lowerCAmelCase_							:			Any			,    lowerCAmelCase_							:			Tuple			,    lowerCAmelCase_							:			Optional[Any]	)      ->		Dict:
						__A=	MobileViTModel(config=lowerCAmelCase_	)
						model.to(lowerCAmelCase_	)
						model.eval()
						__A=	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 					lowerCAmelCase     (  self							:			List[str]			,    lowerCAmelCase_							:			List[str]			,    lowerCAmelCase_							:			Optional[int]			,    lowerCAmelCase_							:			List[Any]			,    lowerCAmelCase_							:			Any	)      ->		Optional[Any]:
						__A=	self.num_labels
						__A=	MobileViTForImageClassification(lowerCAmelCase_	)
						model.to(lowerCAmelCase_	)
						model.eval()
						__A=	model(lowerCAmelCase_			,    labels=lowerCAmelCase_	)
						self.parent.assertEqual(result.logits.shape			,    (self.batch_size, self.num_labels)	)
			def 					lowerCAmelCase     (  self							:			Optional[int]			,    lowerCAmelCase_							:			Optional[Any]			,    lowerCAmelCase_							:			Optional[Any]			,    lowerCAmelCase_							:			Optional[int]			,    lowerCAmelCase_							:			int	)      ->		int:
						__A=	self.num_labels
						__A=	MobileViTForSemanticSegmentation(lowerCAmelCase_	)
						model.to(lowerCAmelCase_	)
						model.eval()
						__A=	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,
						    )			,    )
						__A=	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 					lowerCAmelCase     (  self							:			Optional[int]	)      ->		Optional[int]:
						__A=	self.prepare_config_and_inputs()
						__A,	__A,	__A,	__A=	config_and_inputs
						__A=	{'pixel_values': pixel_values}
						return config, inputs_dict


@require_torch
class 						a__					(					a_      ,  a_      ,  unittest.TestCase   ):





			'''simple docstring'''

			A						:				Union[str, Any]					      =			(
			    (MobileViTModel, MobileViTForImageClassification, MobileViTForSemanticSegmentation)
			    if is_torch_available()
			    else ()
			)
			A						:				Any					      =			(
			    {
			        '''feature-extraction''': MobileViTModel,
			        '''image-classification''': MobileViTForImageClassification,
			        '''image-segmentation''': MobileViTForSemanticSegmentation,
			    }
			    if is_torch_available()
			    else {}
			)

			A						:				int					      =			False
			A						:				Dict					      =			False
			A						:				Dict					      =			False
			A						:				str					      =			False
			def 					lowerCAmelCase     (  self							:			Union[str, Any]	)      ->		Union[str, Any]:
						__A=	MobileViTModelTester(self	)
						__A=	MobileViTConfigTester(self			,    config_class=lowerCAmelCase_			,    has_text_modality=lowerCAmelCase_	)
			def 					lowerCAmelCase     (  self							:			Any	)      ->		str:
						self.config_tester.run_common_tests()
			@unittest.skip(reason='MobileViT does not use inputs_embeds'	)
			def 					lowerCAmelCase     (  self							:			int	)      ->		Any:
						pass
			@unittest.skip(reason='MobileViT does not support input and output embeddings'	)
			def 					lowerCAmelCase     (  self							:			List[str]	)      ->		List[Any]:
						pass
			@unittest.skip(reason='MobileViT does not output attentions'	)
			def 					lowerCAmelCase     (  self							:			List[str]	)      ->		Optional[int]:
						pass
			def 					lowerCAmelCase     (  self							:			Optional[int]	)      ->		Optional[Any]:
						__A,	__A=	self.model_tester.prepare_config_and_inputs_for_common()

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

									__A=	['pixel_values']
									self.assertListEqual(arg_names[:1]			,    lowerCAmelCase_	)
			@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'	)
			def 					lowerCAmelCase     (  self							:			str	)      ->		int:
						pass
			def 					lowerCAmelCase     (  self							:			str	)      ->		Any:
						__A=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*lowerCAmelCase_	)
			def 					lowerCAmelCase     (  self							:			Optional[Any]	)      ->		str:
						def check_hidden_states_output(lowerCAmelCase_							:			Optional[int]			,    lowerCAmelCase_							:			int			,    lowerCAmelCase_							:			List[str]	):
									__A=	model_class(lowerCAmelCase_	)
									model.to(lowerCAmelCase_	)
									model.eval()

									with torch.no_grad():
												__A=	model(**self._prepare_for_class(lowerCAmelCase_			,    lowerCAmelCase_	)	)

									__A=	outputs.hidden_states

									__A=	5
									self.assertEqual(len(lowerCAmelCase_	)			,    lowerCAmelCase_	)

									# MobileViT's feature maps are of shape (batch_size, num_channels, height, width)
									# with the width and height being successively divided by 2.
									__A=	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	)

						__A,	__A=	self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									__A=	True
									check_hidden_states_output(lowerCAmelCase_			,    lowerCAmelCase_			,    lowerCAmelCase_	)

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

									check_hidden_states_output(lowerCAmelCase_			,    lowerCAmelCase_			,    lowerCAmelCase_	)
			def 					lowerCAmelCase     (  self							:			Tuple	)      ->		Optional[Any]:
						__A=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*lowerCAmelCase_	)
			def 					lowerCAmelCase     (  self							:			Union[str, Any]	)      ->		Union[str, Any]:
						__A=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_semantic_segmentation(*lowerCAmelCase_	)

			@slow
			def 					lowerCAmelCase     (  self							:			Union[str, Any]	)      ->		Optional[Any]:
						for model_name in MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									__A=	MobileViTModel.from_pretrained(lowerCAmelCase_	)
									self.assertIsNotNone(lowerCAmelCase_	)






def 						UpperCAmelCase__(				):






			"""simple docstring"""




			__A=	Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png' )
			return image






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





			'''simple docstring'''

			@cached_property
			def 					lowerCAmelCase     (  self							:			Optional[Any]	)      ->		Tuple:
						return MobileViTImageProcessor.from_pretrained('apple/mobilevit-xx-small'	) if is_vision_available() else None
			@slow
			def 					lowerCAmelCase     (  self							:			int	)      ->		str:
						__A=	MobileViTForImageClassification.from_pretrained('apple/mobilevit-xx-small'	).to(lowerCAmelCase_	)

						__A=	self.default_image_processor
						__A=	prepare_img()
						__A=	image_processor(images=lowerCAmelCase_			,    return_tensors='pt'	).to(lowerCAmelCase_	)

						# forward pass
						with torch.no_grad():
									__A=	model(**lowerCAmelCase_	)

						# verify the logits
						__A=	torch.Size((1, 1_000)	)
						self.assertEqual(outputs.logits.shape			,    lowerCAmelCase_	)

						__A=	torch.tensor([-1.93_64, -1.23_27, -0.46_53]	).to(lowerCAmelCase_	)

						self.assertTrue(torch.allclose(outputs.logits[0, :3]			,    lowerCAmelCase_			,    atol=1E-4	)	)
			@slow
			def 					lowerCAmelCase     (  self							:			Union[str, Any]	)      ->		List[str]:
						__A=	MobileViTForSemanticSegmentation.from_pretrained('apple/deeplabv3-mobilevit-xx-small'	)
						__A=	model.to(lowerCAmelCase_	)

						__A=	MobileViTImageProcessor.from_pretrained('apple/deeplabv3-mobilevit-xx-small'	)

						__A=	prepare_img()
						__A=	image_processor(images=lowerCAmelCase_			,    return_tensors='pt'	).to(lowerCAmelCase_	)

						# forward pass
						with torch.no_grad():
									__A=	model(**lowerCAmelCase_	)
						__A=	outputs.logits

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

						__A=	torch.tensor(
						    [
						        [[6.97_13, 6.97_86, 7.24_22], [7.28_93, 7.28_25, 7.44_46], [7.65_80, 7.87_97, 7.94_20]],
						        [[-10.68_69, -10.32_50, -10.34_71], [-10.42_28, -9.98_68, -9.71_32], [-11.04_05, -11.02_21, -10.73_18]],
						        [[-3.30_89, -2.85_39, -2.67_40], [-3.27_06, -2.56_21, -2.51_08], [-3.25_34, -2.66_15, -2.66_51]],
						    ]			,    device=lowerCAmelCase_			,    )

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

			@slow
			def 					lowerCAmelCase     (  self							:			Tuple	)      ->		str:
						__A=	MobileViTForSemanticSegmentation.from_pretrained('apple/deeplabv3-mobilevit-xx-small'	)
						__A=	model.to(lowerCAmelCase_	)

						__A=	MobileViTImageProcessor.from_pretrained('apple/deeplabv3-mobilevit-xx-small'	)

						__A=	prepare_img()
						__A=	image_processor(images=lowerCAmelCase_			,    return_tensors='pt'	).to(lowerCAmelCase_	)

						# forward pass
						with torch.no_grad():
									__A=	model(**lowerCAmelCase_	)

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

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

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



