



import collections
import os
import re
from pathlib import Path


SCREAMING_SNAKE_CASE__   			=   '''src/transformers'''


# Matches is_xxx_available()
SCREAMING_SNAKE_CASE__   			=   re.compile(r'''is\_([a-z_]*)_available()''')
# Catches a one-line _import_struct = {xxx}
SCREAMING_SNAKE_CASE__   			=   re.compile(r'''^_import_structure\s+=\s+\{([^\}]+)\}''')
# Catches a line with a key-values pattern: "bla": ["foo", "bar"]
SCREAMING_SNAKE_CASE__   			=   re.compile(r'''\s+"\S*":\s+\[([^\]]*)\]''')
# Catches a line if not is_foo_available
SCREAMING_SNAKE_CASE__   			=   re.compile(r'''^\s*if\s+not\s+is\_[a-z_]*\_available\(\)''')
# Catches a line _import_struct["bla"].append("foo")
SCREAMING_SNAKE_CASE__   			=   re.compile(r'''^\s*_import_structure\["\S*"\]\.append\("(\S*)"\)''')
# Catches a line _import_struct["bla"].extend(["foo", "bar"]) or _import_struct["bla"] = ["foo", "bar"]
SCREAMING_SNAKE_CASE__   			=   re.compile(r'''^\s*_import_structure\[\S*\](?:\.extend\(|\s*=\s+)\[([^\]]*)\]''')
# Catches a line with an object between quotes and a comma:     "MyModel",
SCREAMING_SNAKE_CASE__   			=   re.compile(r'''^\s+"([^"]+)",''')
# Catches a line with objects between brackets only:    ["foo", "bar"],
SCREAMING_SNAKE_CASE__   			=   re.compile(r'''^\s+\[([^\]]+)\]''')
# Catches a line with from foo import bar, bla, boo
SCREAMING_SNAKE_CASE__   			=   re.compile(r'''\s+from\s+\S*\s+import\s+([^\(\s].*)\n''')
# Catches a line with try:
SCREAMING_SNAKE_CASE__   			=   re.compile(r'''^\s*try:''')
# Catches a line with else:
SCREAMING_SNAKE_CASE__   			=   re.compile(r'''^\s*else:''')


def A    (							__UpperCamelCase       )    ->   Union[str, Any]:
			if _re_test_backend.search(__UpperCamelCase       ) is None:
						return None
			A__	  =     [b[0] for b in _re_backend.findall(__UpperCamelCase       )]
			backends.sort()
			return "_and_".join(__UpperCamelCase       )


def A    (							__UpperCamelCase       )    ->   List[Any]:
			with open(__UpperCamelCase			,						'r'			,						encoding='utf-8'			,						newline='\n'       ) as f:
						A__	  =     f.readlines()

			A__	  =     0
			while line_index < len(__UpperCamelCase       ) and not lines[line_index].startswith('_import_structure = {'       ):
						line_index += 1

			# If this is a traditional init, just return.
			if line_index >= len(__UpperCamelCase       ):
						return None

			# First grab the objects without a specific backend in _import_structure
			A__	  =     []
			while not lines[line_index].startswith('if TYPE_CHECKING'       ) and find_backend(lines[line_index]       ) is None:
						A__	  =     lines[line_index]
						# If we have everything on a single line, let's deal with it.
						if _re_one_line_import_struct.search(__UpperCamelCase       ):
									A__	  =     _re_one_line_import_struct.search(__UpperCamelCase       ).groups()[0]
									A__	  =     re.findall(r'\[([^\]]+)\]'			,						__UpperCamelCase       )
									for imp in imports:
												objects.extend([obj[1:-1] for obj in imp.split(', '       )]       )
									line_index += 1
									continue
						A__	  =     _re_import_struct_key_value.search(__UpperCamelCase       )
						if single_line_import_search is not None:
									A__	  =     [obj[1:-1] for obj in single_line_import_search.groups()[0].split(', '       ) if len(__UpperCamelCase       ) > 0]
									objects.extend(__UpperCamelCase       )
						elif line.startswith(' ' * 8 + '"'       ):
									objects.append(line[9:-3]       )
						line_index += 1

			A__	  =     {'none': objects}
			# Let's continue with backend-specific objects in _import_structure
			while not lines[line_index].startswith('if TYPE_CHECKING'       ):
						# If the line is an if not is_backend_available, we grab all objects associated.
						A__	  =     find_backend(lines[line_index]       )
						# Check if the backend declaration is inside a try block:
						if _re_try.search(lines[line_index - 1]       ) is None:
									A__	  =     None

						if backend is not None:
									line_index += 1

									# Scroll until we hit the else block of try-except-else
									while _re_else.search(lines[line_index]       ) is None:
												line_index += 1

									line_index += 1

									A__	  =     []
									# Until we unindent, add backend objects to the list
									while len(lines[line_index]       ) <= 1 or lines[line_index].startswith(' ' * 4       ):
												A__	  =     lines[line_index]
												if _re_import_struct_add_one.search(__UpperCamelCase       ) is not None:
															objects.append(_re_import_struct_add_one.search(__UpperCamelCase       ).groups()[0]       )
												elif _re_import_struct_add_many.search(__UpperCamelCase       ) is not None:
															A__	  =     _re_import_struct_add_many.search(__UpperCamelCase       ).groups()[0].split(', '       )
															A__	  =     [obj[1:-1] for obj in imports if len(__UpperCamelCase       ) > 0]
															objects.extend(__UpperCamelCase       )
												elif _re_between_brackets.search(__UpperCamelCase       ) is not None:
															A__	  =     _re_between_brackets.search(__UpperCamelCase       ).groups()[0].split(', '       )
															A__	  =     [obj[1:-1] for obj in imports if len(__UpperCamelCase       ) > 0]
															objects.extend(__UpperCamelCase       )
												elif _re_quote_object.search(__UpperCamelCase       ) is not None:
															objects.append(_re_quote_object.search(__UpperCamelCase       ).groups()[0]       )
												elif line.startswith(' ' * 8 + '"'       ):
															objects.append(line[9:-3]       )
												elif line.startswith(' ' * 12 + '"'       ):
															objects.append(line[13:-3]       )
												line_index += 1

									A__	  =     objects
						else:
									line_index += 1

    # At this stage we are in the TYPE_CHECKING part, first grab the objects without a specific backend
			A__	  =     []
			while (
			    line_index < len(__UpperCamelCase       )
			    and find_backend(lines[line_index]       ) is None
			    and not lines[line_index].startswith('else'       )
			):
						A__	  =     lines[line_index]
						A__	  =     _re_import.search(__UpperCamelCase       )
						if single_line_import_search is not None:
									objects.extend(single_line_import_search.groups()[0].split(', '       )       )
						elif line.startswith(' ' * 8       ):
									objects.append(line[8:-2]       )
						line_index += 1

			A__	  =     {'none': objects}
			# Let's continue with backend-specific objects
			while line_index < len(__UpperCamelCase       ):
						# If the line is an if is_backend_available, we grab all objects associated.
						A__	  =     find_backend(lines[line_index]       )
						# Check if the backend declaration is inside a try block:
						if _re_try.search(lines[line_index - 1]       ) is None:
									A__	  =     None

						if backend is not None:
									line_index += 1

									# Scroll until we hit the else block of try-except-else
									while _re_else.search(lines[line_index]       ) is None:
												line_index += 1

									line_index += 1

									A__	  =     []
									# Until we unindent, add backend objects to the list
									while len(lines[line_index]       ) <= 1 or lines[line_index].startswith(' ' * 8       ):
												A__	  =     lines[line_index]
												A__	  =     _re_import.search(__UpperCamelCase       )
												if single_line_import_search is not None:
															objects.extend(single_line_import_search.groups()[0].split(', '       )       )
												elif line.startswith(' ' * 12       ):
															objects.append(line[12:-2]       )
												line_index += 1

									A__	  =     objects
						else:
									line_index += 1

			return import_dict_objects, type_hint_objects


def A    (							__UpperCamelCase			,						__UpperCamelCase       )    ->   Union[str, Any]:

			def find_duplicates(__UpperCamelCase       ):
						return [k for k, v in collections.Counter(__UpperCamelCase       ).items() if v > 1]

			if list(import_dict_objects.keys()       ) != list(type_hint_objects.keys()       ):
						return ["Both sides of the init do not have the same backends!"]

			A__	  =     []
			for key in import_dict_objects.keys():
						A__	  =     find_duplicates(import_dict_objects[key]       )
						if duplicate_imports:
									errors.append(f'''Duplicate _import_structure definitions for: {duplicate_imports}'''       )
						A__	  =     find_duplicates(type_hint_objects[key]       )
						if duplicate_type_hints:
									errors.append(f'''Duplicate TYPE_CHECKING objects for: {duplicate_type_hints}'''       )

						if sorted(set(import_dict_objects[key]       )       ) != sorted(set(type_hint_objects[key]       )       ):
									A__	  =     'base imports' if key == 'none' else f'''{key} backend'''
									errors.append(f'''Differences for {name}:'''       )
									for a in type_hint_objects[key]:
												if a not in import_dict_objects[key]:
															errors.append(f'''  {a} in TYPE_HINT but not in _import_structure.'''       )
									for a in import_dict_objects[key]:
												if a not in type_hint_objects[key]:
															errors.append(f'''  {a} in _import_structure but not in TYPE_HINT.'''       )
			return errors


def A    (							)    ->   Optional[Any]:
			A__	  =     []
			for root, _, files in os.walk(__UpperCamelCase       ):
						if "__init__.py" in files:
									A__	  =     os.path.join(__UpperCamelCase			,						'__init__.py'       )
									A__	  =     parse_init(__UpperCamelCase       )
									if objects is not None:
												A__	  =     analyze_results(*__UpperCamelCase       )
												if len(__UpperCamelCase       ) > 0:
															A__	  =     f'''Problem in {fname}, both halves do not define the same objects.\n{errors[0]}'''
															failures.append('\n'.join(__UpperCamelCase       )       )
			if len(__UpperCamelCase       ) > 0:
						raise ValueError('\n\n'.join(__UpperCamelCase       )       )


def A    (							)    ->   str:
			A__	  =     []
			for path, directories, files in os.walk(__UpperCamelCase       ):
						for folder in directories:
									# Ignore private modules
									if folder.startswith('_'       ):
												directories.remove(__UpperCamelCase       )
												continue
									# Ignore leftovers from branches (empty folders apart from pycache)
									if len(list((Path(__UpperCamelCase       ) / folder).glob('*.py'       )       )       ) == 0:
												continue
									A__	  =     str((Path(__UpperCamelCase       ) / folder).relative_to(__UpperCamelCase       )       )
									A__	  =     short_path.replace(os.path.sep			,						'.'       )
									submodules.append(__UpperCamelCase       )
						for fname in files:
									if fname == "__init__.py":
												continue
									A__	  =     str((Path(__UpperCamelCase       ) / fname).relative_to(__UpperCamelCase       )       )
									A__	  =     short_path.replace('.py'			,						''       ).replace(os.path.sep			,						'.'       )
									if len(submodule.split('.'       )       ) == 1:
												submodules.append(__UpperCamelCase       )
			return submodules


SCREAMING_SNAKE_CASE__   			=   [
    '''convert_pytorch_checkpoint_to_tf2''',
    '''modeling_flax_pytorch_utils''',
    '''models.esm.openfold_utils''',
]


def A    (							)    ->   Any:
			# This is to make sure the transformers module imported is the one in the repo.
			from transformers.utils import direct_transformers_import

			A__	  =     direct_transformers_import(__UpperCamelCase       )

			A__	  =     set(transformers._import_structure.keys()       )
			# This contains all the base keys of the _import_structure object defined in the init, but if the user is missing
			# some optional dependencies, they may not have all of them. Thus we read the init to read all additions and
			# (potentiall re-) add them.
			with open(os.path.join(__UpperCamelCase			,						'__init__.py'       )			,						'r'       ) as f:
						A__	  =     f.read()
			import_structure_keys.update(set(re.findall(r'import_structure\[\"([^\"]*)\"\]'			,						__UpperCamelCase       )       )       )

			A__	  =     [
			    module
			    for module in get_transformers_submodules()
			    if module not in IGNORE_SUBMODULES and module not in import_structure_keys
			]

			if len(__UpperCamelCase       ) > 0:
						A__	  =     '\n'.join(f'''- {module}''' for module in module_not_registered       )
						raise ValueError(
						    'The following submodules are not properly registed in the main init of Transformers:\n'
						    f'''{list_of_modules}\n'''
						    'Make sure they appear somewhere in the keys of `_import_structure` with an empty list as value.'       )


if __name__ == "__main__":
							check_all_inits()
							check_submodules()









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()
SCREAMING_SNAKE_CASE__   			=   logging.get_logger(__name__)


def A    (							__UpperCamelCase       )    ->   List[Any]:
			print('Loading config file...'       )

			def flatten_yaml_as_dict(__UpperCamelCase			,						__UpperCamelCase=""			,						__UpperCamelCase="."       ):
						A__	  =     []
						for k, v in d.items():
									A__	  =     parent_key + sep + k if parent_key else k
									if isinstance(__UpperCamelCase			,						collections.abc.MutableMapping       ):
												items.extend(flatten_yaml_as_dict(__UpperCamelCase			,						__UpperCamelCase			,						sep=__UpperCamelCase       ).items()       )
									else:
												items.append((new_key, v)       )
						return dict(__UpperCamelCase       )

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

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


def A    (							__UpperCamelCase			,						__UpperCamelCase       )    ->   Optional[Any]:
			A__	  =     MobileViTVaConfig()

			A__	  =     False

			# dataset
			if task_name.startswith('imagenet1k_'       ):
						A__	  =     1_000
						if int(task_name.strip().split('_'       )[-1]       ) == 384:
									A__	  =     384
						else:
									A__	  =     256
						A__	  =     'imagenet-1k-id2label.json'
			elif task_name.startswith('imagenet21k_to_1k_'       ):
						A__	  =     21_000
						if int(task_name.strip().split('_'       )[-1]       ) == 384:
									A__	  =     384
						else:
									A__	  =     256
						A__	  =     'imagenet-22k-id2label.json'
			elif task_name.startswith('ade20k_'       ):
						A__	  =     151
						A__	  =     512
						A__	  =     'ade20k-id2label.json'
						A__	  =     True
			elif task_name.startswith('voc_'       ):
						A__	  =     21
						A__	  =     512
						A__	  =     'pascal-voc-id2label.json'
						A__	  =     True

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

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

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

			return config


def A    (							__UpperCamelCase			,						__UpperCamelCase			,						__UpperCamelCase       )    ->   List[str]:
			A__	  =     dct.pop(__UpperCamelCase       )
			A__	  =     val


def A    (							__UpperCamelCase			,						__UpperCamelCase=False       )    ->   Dict:
			if base_model:
						A__	  =     ''
			else:
						A__	  =     'mobilevitv2.'

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

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

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

						for i in [3, 4, 5]:
									if f'''layer_{i}.0.''' in k:
												A__	  =     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:
												A__	  =     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:
												A__	  =     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:
												A__	  =     [0, 1]
									elif i == 4:
												A__	  =     [0, 1, 2, 3]
									elif i == 5:
												A__	  =     [0, 1, 2]

									for j in j_in:
												if f'''layer_{i}.1.global_rep.{j}.''' in k:
															A__	  =     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:
												A__	  =     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:
												A__	  =     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:
									A__	  =     k_new.replace('pre_norm_attn.0.'			,						'layernorm_before.'       )
						if "pre_norm_attn.1." in k:
									A__	  =     k_new.replace('pre_norm_attn.1.'			,						'attention.'       )
						if "pre_norm_ffn.0." in k:
									A__	  =     k_new.replace('pre_norm_ffn.0.'			,						'layernorm_after.'       )
						if "pre_norm_ffn.1." in k:
									A__	  =     k_new.replace('pre_norm_ffn.1.'			,						'ffn.conv1.'       )
						if "pre_norm_ffn.3." in k:
									A__	  =     k_new.replace('pre_norm_ffn.3.'			,						'ffn.conv2.'       )

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

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

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


def A    (							__UpperCamelCase       )    ->   Tuple:
			A__	  =     []
			for k in state_dict.keys():
						if k.startswith('seg_head.aux_head.'       ):
									keys_to_ignore.append(__UpperCamelCase       )
			for k in keys_to_ignore:
						state_dict.pop(__UpperCamelCase			,						__UpperCamelCase       )


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


@torch.no_grad()
def A    (							__UpperCamelCase			,						__UpperCamelCase			,						__UpperCamelCase			,						__UpperCamelCase       )    ->   Optional[Any]:
			A__	  =     get_mobilevitva_config(__UpperCamelCase			,						__UpperCamelCase       )

			# load original state_dict
			A__	  =     torch.load(__UpperCamelCase			,						map_location='cpu'       )

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

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

			# load modified state_dict
			model.load_state_dict(__UpperCamelCase       )

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

			# verify classification model
			if task_name.startswith('imagenet'       ):
						A__	  =     outputs.logits
						A__	  =     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
									A__	  =     torch.tensor([-1.6336E00, -7.3204E-02, -5.1883E-01]       )
									assert torch.allclose(logits[0, :3]			,						__UpperCamelCase			,						atol=1E-4       )

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


if __name__ == "__main__":
							SCREAMING_SNAKE_CASE__   			=   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.'''
							)

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




