


"""simple docstring"""




import argparse
import collections
import json
from pathlib import Path

import requests
import torch
import yaml
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileViTImageProcessor,
    MobileViTVaConfig,
    MobileViTVaForImageClassification,
    MobileViTVaForSemanticSegmentation,
)
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCamelCase						=      logging.get_logger(__name__)


def     lowercase			(SCREAMING_SNAKE_CASE_		:				Dict  )   ->						Optional[Any]:
						print('Loading config file...'  )

						def flatten_yaml_as_dict(SCREAMING_SNAKE_CASE_		:				Union[str, Any] ,			SCREAMING_SNAKE_CASE_		:				Tuple="" ,			SCREAMING_SNAKE_CASE_		:				str="."  ):
												SCREAMING_SNAKE_CASE													=					[]
												for k, v in d.items():
																		SCREAMING_SNAKE_CASE													=					parent_key + sep + k if parent_key else k
																		if isinstance(SCREAMING_SNAKE_CASE_ ,			collections.abc.MutableMapping  ):
																								items.extend(flatten_yaml_as_dict(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			sep=SCREAMING_SNAKE_CASE_  ).items()  )
																		else:
																								items.append((new_key, v)  )
												return dict(SCREAMING_SNAKE_CASE_  )

						SCREAMING_SNAKE_CASE													=					argparse.Namespace()
						with open(SCREAMING_SNAKE_CASE_ ,			'r'  ) as yaml_file:
												try:
																		SCREAMING_SNAKE_CASE													=					yaml.load(SCREAMING_SNAKE_CASE_ ,			Loader=yaml.FullLoader  )

																		SCREAMING_SNAKE_CASE													=					flatten_yaml_as_dict(SCREAMING_SNAKE_CASE_  )
																		for k, v in flat_cfg.items():
																								setattr(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_  )
												except yaml.YAMLError as exc:
																		logger.error('Error while loading config file: {}. Error message: {}'.format(SCREAMING_SNAKE_CASE_ ,			str(SCREAMING_SNAKE_CASE_  )  )  )
						return config


def     lowercase			(SCREAMING_SNAKE_CASE_		:				int ,			SCREAMING_SNAKE_CASE_		:				Any  )   ->						List[Any]:
						SCREAMING_SNAKE_CASE													=					MobileViTVaConfig()

						SCREAMING_SNAKE_CASE													=					False

						# dataset
						if task_name.startswith('imagenet1k_'  ):
												SCREAMING_SNAKE_CASE													=					10_00
												if int(task_name.strip().split('_'  )[-1]  ) == 3_84:
																		SCREAMING_SNAKE_CASE													=					3_84
												else:
																		SCREAMING_SNAKE_CASE													=					2_56
												SCREAMING_SNAKE_CASE													=					'imagenet-1k-id2label.json'
						elif task_name.startswith('imagenet21k_to_1k_'  ):
												SCREAMING_SNAKE_CASE													=					2_10_00
												if int(task_name.strip().split('_'  )[-1]  ) == 3_84:
																		SCREAMING_SNAKE_CASE													=					3_84
												else:
																		SCREAMING_SNAKE_CASE													=					2_56
												SCREAMING_SNAKE_CASE													=					'imagenet-22k-id2label.json'
						elif task_name.startswith('ade20k_'  ):
												SCREAMING_SNAKE_CASE													=					1_51
												SCREAMING_SNAKE_CASE													=					5_12
												SCREAMING_SNAKE_CASE													=					'ade20k-id2label.json'
												SCREAMING_SNAKE_CASE													=					True
						elif task_name.startswith('voc_'  ):
												SCREAMING_SNAKE_CASE													=					21
												SCREAMING_SNAKE_CASE													=					5_12
												SCREAMING_SNAKE_CASE													=					'pascal-voc-id2label.json'
												SCREAMING_SNAKE_CASE													=					True

						# orig_config
						SCREAMING_SNAKE_CASE													=					load_orig_config_file(SCREAMING_SNAKE_CASE_  )
						assert getattr(SCREAMING_SNAKE_CASE_ ,			'model.classification.name' ,			-1  ) == "mobilevit_v2", "Invalid model"
						SCREAMING_SNAKE_CASE													=					getattr(SCREAMING_SNAKE_CASE_ ,			'model.classification.mitv2.width_multiplier' ,			1.0  )
						assert (
						    getattr(SCREAMING_SNAKE_CASE_ ,			'model.classification.mitv2.attn_norm_layer' ,			-1  ) == "layer_norm_2d"
						), "Norm layers other than layer_norm_2d is not supported"
						SCREAMING_SNAKE_CASE													=					getattr(SCREAMING_SNAKE_CASE_ ,			'model.classification.activation.name' ,			'swish'  )
						# config.image_size == getattr(orig_config,  'sampler.bs.crop_size_width', 256)

						if is_segmentation_model:
												SCREAMING_SNAKE_CASE													=					getattr(SCREAMING_SNAKE_CASE_ ,			'model.segmentation.output_stride' ,			16  )
												if "_deeplabv3" in task_name:
																		SCREAMING_SNAKE_CASE													=					getattr(SCREAMING_SNAKE_CASE_ ,			'model.segmentation.deeplabv3.aspp_rates' ,			[12, 24, 36]  )
																		SCREAMING_SNAKE_CASE													=					getattr(SCREAMING_SNAKE_CASE_ ,			'model.segmentation.deeplabv3.aspp_out_channels' ,			5_12  )
																		SCREAMING_SNAKE_CASE													=					getattr(SCREAMING_SNAKE_CASE_ ,			'model.segmentation.deeplabv3.aspp_dropout' ,			0.1  )

    # id2label
						SCREAMING_SNAKE_CASE													=					'huggingface/label-files'
						SCREAMING_SNAKE_CASE													=					json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			repo_type='dataset'  ) ,			'r'  )  )
						SCREAMING_SNAKE_CASE													=					{int(SCREAMING_SNAKE_CASE_  ): v for k, v in idalabel.items()}
						SCREAMING_SNAKE_CASE													=					idalabel
						SCREAMING_SNAKE_CASE													=					{v: k for k, v in idalabel.items()}

						return config


def     lowercase			(SCREAMING_SNAKE_CASE_		:				int ,			SCREAMING_SNAKE_CASE_		:				str ,			SCREAMING_SNAKE_CASE_		:				str  )   ->						List[str]:
						SCREAMING_SNAKE_CASE													=					dct.pop(SCREAMING_SNAKE_CASE_  )
						SCREAMING_SNAKE_CASE													=					val


def     lowercase			(SCREAMING_SNAKE_CASE_		:				Optional[int] ,			SCREAMING_SNAKE_CASE_		:				Dict=False  )   ->						List[Any]:
						if base_model:
												SCREAMING_SNAKE_CASE													=					''
						else:
												SCREAMING_SNAKE_CASE													=					'mobilevitv2.'

						SCREAMING_SNAKE_CASE													=					[]
						for k in state_dict.keys():
												if k[:8] == "encoder.":
																		SCREAMING_SNAKE_CASE													=					k[8:]
												else:
																		SCREAMING_SNAKE_CASE													=					k

												if ".block." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('.block.' ,			'.'  )
												if ".conv." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('.conv.' ,			'.convolution.'  )
												if ".norm." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('.norm.' ,			'.normalization.'  )

												if "conv_1." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('conv_1.' ,			F'{model_prefix}conv_stem.'  )
												for i in [1, 2]:
																		if F'layer_{i}.' in k:
																								SCREAMING_SNAKE_CASE													=					k_new.replace(F'layer_{i}.' ,			F'{model_prefix}encoder.layer.{i-1}.layer.'  )
												if ".exp_1x1." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('.exp_1x1.' ,			'.expand_1x1.'  )
												if ".red_1x1." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('.red_1x1.' ,			'.reduce_1x1.'  )

												for i in [3, 4, 5]:
																		if F'layer_{i}.0.' in k:
																								SCREAMING_SNAKE_CASE													=					k_new.replace(F'layer_{i}.0.' ,			F'{model_prefix}encoder.layer.{i-1}.downsampling_layer.'  )
																		if F'layer_{i}.1.local_rep.0.' in k:
																								SCREAMING_SNAKE_CASE													=					k_new.replace(F'layer_{i}.1.local_rep.0.' ,			F'{model_prefix}encoder.layer.{i-1}.conv_kxk.'  )
																		if F'layer_{i}.1.local_rep.1.' in k:
																								SCREAMING_SNAKE_CASE													=					k_new.replace(F'layer_{i}.1.local_rep.1.' ,			F'{model_prefix}encoder.layer.{i-1}.conv_1x1.'  )

												for i in [3, 4, 5]:
																		if i == 3:
																								SCREAMING_SNAKE_CASE													=					[0, 1]
																		elif i == 4:
																								SCREAMING_SNAKE_CASE													=					[0, 1, 2, 3]
																		elif i == 5:
																								SCREAMING_SNAKE_CASE													=					[0, 1, 2]

																		for j in j_in:
																								if F'layer_{i}.1.global_rep.{j}.' in k:
																														SCREAMING_SNAKE_CASE													=					k_new.replace(
																														    F'layer_{i}.1.global_rep.{j}.' ,			F'{model_prefix}encoder.layer.{i-1}.transformer.layer.{j}.'  )
																		if F'layer_{i}.1.global_rep.{j+1}.' in k:
																								SCREAMING_SNAKE_CASE													=					k_new.replace(
																								    F'layer_{i}.1.global_rep.{j+1}.' ,			F'{model_prefix}encoder.layer.{i-1}.layernorm.'  )

																		if F'layer_{i}.1.conv_proj.' in k:
																								SCREAMING_SNAKE_CASE													=					k_new.replace(F'layer_{i}.1.conv_proj.' ,			F'{model_prefix}encoder.layer.{i-1}.conv_projection.'  )

												if "pre_norm_attn.0." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('pre_norm_attn.0.' ,			'layernorm_before.'  )
												if "pre_norm_attn.1." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('pre_norm_attn.1.' ,			'attention.'  )
												if "pre_norm_ffn.0." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('pre_norm_ffn.0.' ,			'layernorm_after.'  )
												if "pre_norm_ffn.1." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('pre_norm_ffn.1.' ,			'ffn.conv1.'  )
												if "pre_norm_ffn.3." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('pre_norm_ffn.3.' ,			'ffn.conv2.'  )

												if "classifier.1." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('classifier.1.' ,			'classifier.'  )

												if "seg_head." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('seg_head.' ,			'segmentation_head.'  )
												if ".aspp_layer." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('.aspp_layer.' ,			'.'  )
												if ".aspp_pool." in k:
																		SCREAMING_SNAKE_CASE													=					k_new.replace('.aspp_pool.' ,			'.'  )

												rename_keys.append((k, k_new)  )
						return rename_keys


def     lowercase			(SCREAMING_SNAKE_CASE_		:				List[str]  )   ->						Tuple:
						SCREAMING_SNAKE_CASE													=					[]
						for k in state_dict.keys():
												if k.startswith('seg_head.aux_head.'  ):
																		keys_to_ignore.append(SCREAMING_SNAKE_CASE_  )
						for k in keys_to_ignore:
												state_dict.pop(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_  )


def     lowercase			()   ->						Any:
						SCREAMING_SNAKE_CASE													=					'http://images.cocodataset.org/val2017/000000039769.jpg'
						# url = "https://cdn.britannica.com/86/141086-050-9D7C75EE/Gulfstream-G450-business-jet-passengers.jpg"
						SCREAMING_SNAKE_CASE													=					Image.open(requests.get(SCREAMING_SNAKE_CASE_ ,			stream=SCREAMING_SNAKE_CASE_  ).raw  )
						return im


@torch.no_grad()
def     lowercase			(SCREAMING_SNAKE_CASE_		:				str ,			SCREAMING_SNAKE_CASE_		:				Union[str, Any] ,			SCREAMING_SNAKE_CASE_		:				Tuple ,			SCREAMING_SNAKE_CASE_		:				str  )   ->						List[Any]:
						SCREAMING_SNAKE_CASE													=					get_mobilevitva_config(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_  )

						# load original state_dict
						SCREAMING_SNAKE_CASE													=					torch.load(SCREAMING_SNAKE_CASE_ ,			map_location='cpu'  )

						# load huggingface model
						if task_name.startswith('ade20k_'  ) or task_name.startswith('voc_'  ):
												SCREAMING_SNAKE_CASE													=					MobileViTVaForSemanticSegmentation(SCREAMING_SNAKE_CASE_  ).eval()
												SCREAMING_SNAKE_CASE													=					False
						else:
												SCREAMING_SNAKE_CASE													=					MobileViTVaForImageClassification(SCREAMING_SNAKE_CASE_  ).eval()
												SCREAMING_SNAKE_CASE													=					False

						# remove and rename some keys of load the original model
						SCREAMING_SNAKE_CASE													=					checkpoint
						remove_unused_keys(SCREAMING_SNAKE_CASE_  )
						SCREAMING_SNAKE_CASE													=					create_rename_keys(SCREAMING_SNAKE_CASE_ ,			base_model=SCREAMING_SNAKE_CASE_  )
						for rename_key_src, rename_key_dest in rename_keys:
												rename_key(SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_ ,			SCREAMING_SNAKE_CASE_  )

						# load modified state_dict
						model.load_state_dict(SCREAMING_SNAKE_CASE_  )

						# Check outputs on an image, prepared by MobileViTImageProcessor
						SCREAMING_SNAKE_CASE													=					MobileViTImageProcessor(crop_size=config.image_size ,			size=config.image_size + 32  )
						SCREAMING_SNAKE_CASE													=					image_processor(images=prepare_img() ,			return_tensors='pt'  )
						SCREAMING_SNAKE_CASE													=					model(**SCREAMING_SNAKE_CASE_  )

						# verify classification model
						if task_name.startswith('imagenet'  ):
												SCREAMING_SNAKE_CASE													=					outputs.logits
												SCREAMING_SNAKE_CASE													=					logits.argmax(-1  ).item()
												print('Predicted class:' ,			model.config.idalabel[predicted_class_idx]  )
												if task_name.startswith('imagenet1k_256'  ) and config.width_multiplier == 1.0:
																		# expected_logits for base variant
																		SCREAMING_SNAKE_CASE													=					torch.tensor([-1.6336E00, -7.3204E-02, -5.1883E-01]  )
																		assert torch.allclose(logits[0, :3] ,			SCREAMING_SNAKE_CASE_ ,			atol=1E-4  )

						Path(SCREAMING_SNAKE_CASE_  ).mkdir(exist_ok=SCREAMING_SNAKE_CASE_  )
						print(F'Saving model {task_name} to {pytorch_dump_folder_path}'  )
						model.save_pretrained(SCREAMING_SNAKE_CASE_  )
						print(F'Saving image processor to {pytorch_dump_folder_path}'  )
						image_processor.save_pretrained(SCREAMING_SNAKE_CASE_  )


if __name__ == "__main__":
	__UpperCamelCase						=      argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '''--task''',
	    default='''imagenet1k_256''',
	    type=str,
	    help=(
	        '''Name of the task for which the MobileViTV2 model you\'d like to convert is trained on . '''
	        '''
                Classification (ImageNet-1k)
                    - MobileViTV2 (256x256) : imagenet1k_256
                    - MobileViTV2 (Trained on 256x256 and Finetuned on 384x384) : imagenet1k_384
                    - MobileViTV2 (Trained on ImageNet-21k and Finetuned on ImageNet-1k 256x256) :
                      imagenet21k_to_1k_256
                    - MobileViTV2 (Trained on ImageNet-21k, Finetuned on ImageNet-1k 256x256, and Finetuned on
                      ImageNet-1k 384x384) : imagenet21k_to_1k_384
                Segmentation
                    - ADE20K Dataset : ade20k_deeplabv3
                    - Pascal VOC 2012 Dataset: voc_deeplabv3
            '''
	    ),
	    choices=[
	        '''imagenet1k_256''',
	        '''imagenet1k_384''',
	        '''imagenet21k_to_1k_256''',
	        '''imagenet21k_to_1k_384''',
	        '''ade20k_deeplabv3''',
	        '''voc_deeplabv3''',
	    ],
	)

	parser.add_argument(
	    '''--orig_checkpoint_path''', required=True, type=str, help='''Path to the original state dict (.pt file).'''
	)
	parser.add_argument('''--orig_config_path''', required=True, type=str, help='''Path to the original config file.''')
	parser.add_argument(
	    '''--pytorch_dump_folder_path''', required=True, type=str, help='''Path to the output PyTorch model directory.'''
	)

	__UpperCamelCase						=      parser.parse_args()
	convert_mobilevitva_checkpoint(
	    args.task, args.orig_checkpoint_path, args.orig_config_path, args.pytorch_dump_folder_path
	)






"""simple docstring"""




def     lowercase			(SCREAMING_SNAKE_CASE_		:				int = 10_00  )   ->						int:
						SCREAMING_SNAKE_CASE													=					2**power
						SCREAMING_SNAKE_CASE													=					0
						while n:
												SCREAMING_SNAKE_CASE							,				SCREAMING_SNAKE_CASE													=					r + n % 10, n // 10
						return r


if __name__ == "__main__":
	print(solution(int(str(input()).strip())))


