






"""simple docstring"""
import argparse
import json
from pathlib import Path

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

from transformers import (
    MobileViTConfig,
    MobileViTForImageClassification,
    MobileViTForSemanticSegmentation,
    MobileViTImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
snake_case : str     		=		logging.get_logger(__name__)


def 							A						(				__snake_case: Optional[int]  )    ->		Any:

       """simple docstring"""





       __magic_name__      				=    MobileViTConfig()

       # size of the architecture
       if "mobilevit_s" in mobilevit_name:
              __magic_name__      				=    [1_4_4, 1_9_2, 2_4_0]
              __magic_name__      				=    [1_6, 3_2, 6_4, 9_6, 1_2_8, 1_6_0, 6_4_0]
       elif "mobilevit_xs" in mobilevit_name:
              __magic_name__      				=    [9_6, 1_2_0, 1_4_4]
              __magic_name__      				=    [1_6, 3_2, 4_8, 6_4, 8_0, 9_6, 3_8_4]
       elif "mobilevit_xxs" in mobilevit_name:
              __magic_name__      				=    [6_4, 8_0, 9_6]
              __magic_name__      				=    [1_6, 1_6, 2_4, 4_8, 6_4, 8_0, 3_2_0]
              __magic_name__      				=    0.05
              __magic_name__      				=    2.0

       if mobilevit_name.startswith('deeplabv3_'  ):
              __magic_name__      				=    5_1_2
              __magic_name__      				=    1_6
              __magic_name__      				=    2_1
              __magic_name__      				=    'pascal-voc-id2label.json'
       else:
              __magic_name__      				=    1_0_0_0
              __magic_name__      				=    'imagenet-1k-id2label.json'

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

       return config


def 							A						(				__snake_case: Optional[int] ,     __snake_case: Optional[int]=False  )    ->		str:

       """simple docstring"""





       for i in range(1 ,     6  ):
              if F"""layer_{i}.""" in name:
                     __magic_name__      				=    name.replace(F"""layer_{i}.""" ,     F"""encoder.layer.{i - 1}."""  )

       if "conv_1." in name:
              __magic_name__      				=    name.replace('conv_1.' ,     'conv_stem.'  )
       if ".block." in name:
              __magic_name__      				=    name.replace('.block.' ,     '.'  )

       if "exp_1x1" in name:
              __magic_name__      				=    name.replace('exp_1x1' ,     'expand_1x1'  )
       if "red_1x1" in name:
              __magic_name__      				=    name.replace('red_1x1' ,     'reduce_1x1'  )
       if ".local_rep.conv_3x3." in name:
              __magic_name__      				=    name.replace('.local_rep.conv_3x3.' ,     '.conv_kxk.'  )
       if ".local_rep.conv_1x1." in name:
              __magic_name__      				=    name.replace('.local_rep.conv_1x1.' ,     '.conv_1x1.'  )
       if ".norm." in name:
              __magic_name__      				=    name.replace('.norm.' ,     '.normalization.'  )
       if ".conv." in name:
              __magic_name__      				=    name.replace('.conv.' ,     '.convolution.'  )
       if ".conv_proj." in name:
              __magic_name__      				=    name.replace('.conv_proj.' ,     '.conv_projection.'  )

       for i in range(0 ,     2  ):
              for j in range(0 ,     4  ):
                     if F""".{i}.{j}.""" in name:
                            __magic_name__      				=    name.replace(F""".{i}.{j}.""" ,     F""".{i}.layer.{j}."""  )

       for i in range(2 ,     6  ):
              for j in range(0 ,     4  ):
                     if F""".{i}.{j}.""" in name:
                            __magic_name__      				=    name.replace(F""".{i}.{j}.""" ,     F""".{i}."""  )
                            if "expand_1x1" in name:
                                   __magic_name__      				=    name.replace('expand_1x1' ,     'downsampling_layer.expand_1x1'  )
                            if "conv_3x3" in name:
                                   __magic_name__      				=    name.replace('conv_3x3' ,     'downsampling_layer.conv_3x3'  )
                            if "reduce_1x1" in name:
                                   __magic_name__      				=    name.replace('reduce_1x1' ,     'downsampling_layer.reduce_1x1'  )

       for i in range(2 ,     5  ):
              if F""".global_rep.{i}.weight""" in name:
                     __magic_name__      				=    name.replace(F""".global_rep.{i}.weight""" ,     '.layernorm.weight'  )
              if F""".global_rep.{i}.bias""" in name:
                     __magic_name__      				=    name.replace(F""".global_rep.{i}.bias""" ,     '.layernorm.bias'  )

       if ".global_rep." in name:
              __magic_name__      				=    name.replace('.global_rep.' ,     '.transformer.'  )
       if ".pre_norm_mha.0." in name:
              __magic_name__      				=    name.replace('.pre_norm_mha.0.' ,     '.layernorm_before.'  )
       if ".pre_norm_mha.1.out_proj." in name:
              __magic_name__      				=    name.replace('.pre_norm_mha.1.out_proj.' ,     '.attention.output.dense.'  )
       if ".pre_norm_ffn.0." in name:
              __magic_name__      				=    name.replace('.pre_norm_ffn.0.' ,     '.layernorm_after.'  )
       if ".pre_norm_ffn.1." in name:
              __magic_name__      				=    name.replace('.pre_norm_ffn.1.' ,     '.intermediate.dense.'  )
       if ".pre_norm_ffn.4." in name:
              __magic_name__      				=    name.replace('.pre_norm_ffn.4.' ,     '.output.dense.'  )
       if ".transformer." in name:
              __magic_name__      				=    name.replace('.transformer.' ,     '.transformer.layer.'  )

       if ".aspp_layer." in name:
              __magic_name__      				=    name.replace('.aspp_layer.' ,     '.'  )
       if ".aspp_pool." in name:
              __magic_name__      				=    name.replace('.aspp_pool.' ,     '.'  )
       if "seg_head." in name:
              __magic_name__      				=    name.replace('seg_head.' ,     'segmentation_head.'  )
       if "segmentation_head.classifier.classifier." in name:
              __magic_name__      				=    name.replace('segmentation_head.classifier.classifier.' ,     'segmentation_head.classifier.'  )

       if "classifier.fc." in name:
              __magic_name__      				=    name.replace('classifier.fc.' ,     'classifier.'  )
       elif (not base_model) and ("segmentation_head." not in name):
              __magic_name__      				=    'mobilevit.' + name

       return name


def 							A						(				__snake_case: List[Any] ,     __snake_case: int ,     __snake_case: str=False  )    ->		Optional[int]:

       """simple docstring"""





       if base_model:
              __magic_name__      				=    ''
       else:
              __magic_name__      				=    'mobilevit.'

       for key in orig_state_dict.copy().keys():
              __magic_name__      				=    orig_state_dict.pop(__snake_case  )

              if key[:8] == "encoder.":
                     __magic_name__      				=    key[8:]

              if "qkv" in key:
                     __magic_name__      				=    key.split('.'  )
                     __magic_name__      				=    int(key_split[0][6:]  ) - 1
                     __magic_name__      				=    int(key_split[3]  )
                     __magic_name__      				=    model.get_submodule(F"""{model_prefix}encoder.layer.{layer_num}"""  )
                     __magic_name__      				=    layer.transformer.layer[transformer_num].attention.attention.all_head_size
                     __magic_name__      				=    (
                         F"""{model_prefix}encoder.layer.{layer_num}.transformer.layer.{transformer_num}.attention.attention."""
                     )
                     if "weight" in key:
                            __magic_name__      				=    val[:dim, :]
                            __magic_name__      				=    val[dim : dim * 2, :]
                            __magic_name__      				=    val[-dim:, :]
                     else:
                            __magic_name__      				=    val[:dim]
                            __magic_name__      				=    val[dim : dim * 2]
                            __magic_name__      				=    val[-dim:]
              else:
                     __magic_name__      				=    val

       return orig_state_dict


def 							A						(				)    ->		Optional[Any]:

       """simple docstring"""





       __magic_name__      				=    'http://images.cocodataset.org/val2017/000000039769.jpg'
       __magic_name__      				=    Image.open(requests.get(__snake_case ,     stream=__snake_case  ).raw  )
       return im


@torch.no_grad()
def 							A						(				__snake_case: Union[str, Any] ,     __snake_case: Any ,     __snake_case: Optional[Any] ,     __snake_case: Tuple=False  )    ->		Dict:

       """simple docstring"""





       __magic_name__      				=    get_mobilevit_config(__snake_case  )

       # load original state_dict
       __magic_name__      				=    torch.load(__snake_case ,     map_location='cpu'  )

       # load 🤗 model
       if mobilevit_name.startswith('deeplabv3_'  ):
              __magic_name__      				=    MobileViTForSemanticSegmentation(__snake_case  ).eval()
       else:
              __magic_name__      				=    MobileViTForImageClassification(__snake_case  ).eval()

       __magic_name__      				=    convert_state_dict(__snake_case ,     __snake_case  )
       model.load_state_dict(__snake_case  )

       # Check outputs on an image, prepared by MobileViTImageProcessor
       __magic_name__      				=    MobileViTImageProcessor(crop_size=config.image_size ,     size=config.image_size + 3_2  )
       __magic_name__      				=    image_processor(images=prepare_img() ,     return_tensors='pt'  )
       __magic_name__      				=    model(**__snake_case  )
       __magic_name__      				=    outputs.logits

       if mobilevit_name.startswith('deeplabv3_'  ):
              assert logits.shape == (1, 2_1, 3_2, 3_2)

              if mobilevit_name == "deeplabv3_mobilevit_s":
                     __magic_name__      				=    torch.tensor(
                         [
                             [[6.2065, 6.1292, 6.2070], [6.1079, 6.1254, 6.1747], [6.0042, 6.1071, 6.1034]],
                             [[-6.9253, -6.8653, -7.0398], [-7.3218, -7.3983, -7.3670], [-7.1961, -7.2482, -7.1569]],
                             [[-4.4723, -4.4348, -4.3769], [-5.3629, -5.4632, -5.4598], [-5.1587, -5.3402, -5.5059]],
                         ]  )
              elif mobilevit_name == "deeplabv3_mobilevit_xs":
                     __magic_name__      				=    torch.tensor(
                         [
                             [[5.4449, 5.5733, 5.6314], [5.1815, 5.3930, 5.5963], [5.1656, 5.4333, 5.4853]],
                             [[-9.4423, -9.7766, -9.6714], [-9.1581, -9.5720, -9.5519], [-9.1006, -9.6458, -9.5703]],
                             [[-7.7721, -7.3716, -7.1583], [-8.4599, -8.0624, -7.7944], [-8.4172, -7.8366, -7.5025]],
                         ]  )
              elif mobilevit_name == "deeplabv3_mobilevit_xxs":
                     __magic_name__      				=    torch.tensor(
                         [
                             [[6.9811, 6.9743, 7.3123], [7.1777, 7.1931, 7.3938], [7.5633, 7.8050, 7.8901]],
                             [[-10.5536, -10.2332, -10.2924], [-10.2336, -9.8624, -9.5964], [-10.8840, -10.8158, -10.6659]],
                             [[-3.4938, -3.0631, -2.8620], [-3.4205, -2.8135, -2.6875], [-3.4179, -2.7945, -2.8750]],
                         ]  )
              else:
                     raise ValueError(F"""Unknown mobilevit_name: {mobilevit_name}"""  )

              assert torch.allclose(logits[0, :3, :3, :3] ,     __snake_case ,     atol=1E-4  )
       else:
              assert logits.shape == (1, 1_0_0_0)

              if mobilevit_name == "mobilevit_s":
                     __magic_name__      				=    torch.tensor([-0.9866, 0.2392, -1.1241]  )
              elif mobilevit_name == "mobilevit_xs":
                     __magic_name__      				=    torch.tensor([-2.4761, -0.9399, -1.9587]  )
              elif mobilevit_name == "mobilevit_xxs":
                     __magic_name__      				=    torch.tensor([-1.9364, -1.2327, -0.4653]  )
              else:
                     raise ValueError(F"""Unknown mobilevit_name: {mobilevit_name}"""  )

              assert torch.allclose(logits[0, :3] ,     __snake_case ,     atol=1E-4  )

       Path(__snake_case  ).mkdir(exist_ok=__snake_case  )
       print(F"""Saving model {mobilevit_name} to {pytorch_dump_folder_path}"""  )
       model.save_pretrained(__snake_case  )
       print(F"""Saving image processor to {pytorch_dump_folder_path}"""  )
       image_processor.save_pretrained(__snake_case  )

       if push_to_hub:
              __magic_name__      				=    {
                  'mobilevit_s': 'mobilevit-small',
                  'mobilevit_xs': 'mobilevit-x-small',
                  'mobilevit_xxs': 'mobilevit-xx-small',
                  'deeplabv3_mobilevit_s': 'deeplabv3-mobilevit-small',
                  'deeplabv3_mobilevit_xs': 'deeplabv3-mobilevit-x-small',
                  'deeplabv3_mobilevit_xxs': 'deeplabv3-mobilevit-xx-small',
              }

              print('Pushing to the hub...'  )
              __magic_name__      				=    model_mapping[mobilevit_name]
              image_processor.push_to_hub(__snake_case ,     organization='apple'  )
              model.push_to_hub(__snake_case ,     organization='apple'  )


if __name__ == "__main__":
 snake_case : Optional[int]     		=		argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     """--mobilevit_name""",
     default="""mobilevit_s""",
     type=str,
     help=(
         """Name of the MobileViT model you'd like to convert. Should be one of 'mobilevit_s', 'mobilevit_xs',"""
         """ 'mobilevit_xxs', 'deeplabv3_mobilevit_s', 'deeplabv3_mobilevit_xs', 'deeplabv3_mobilevit_xxs'."""
     ),
 )
 parser.add_argument(
     """--checkpoint_path""", required=True, type=str, help="""Path to the original state dict (.pt file)."""
 )
 parser.add_argument(
     """--pytorch_dump_folder_path""", required=True, type=str, help="""Path to the output PyTorch model directory."""
 )
 parser.add_argument(
     """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
 )

 snake_case : int     		=		parser.parse_args()
 convert_movilevit_checkpoint(
     args.mobilevit_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
 )







"""simple docstring"""
def 							A						(				__snake_case: int  )    ->		int:

       """simple docstring"""





       if divisor % 5 == 0 or divisor % 2 == 0:
              return 0
       __magic_name__      				=    1
       __magic_name__      				=    1
       while repunit:
              __magic_name__      				=    (1_0 * repunit + 1) % divisor
              repunit_index += 1
       return repunit_index


def 							A						(				__snake_case: int = 1_0_0_0_0_0_0  )    ->		int:

       """simple docstring"""





       __magic_name__      				=    limit - 1
       if divisor % 2 == 0:
              divisor += 1
       while least_divisible_repunit(__snake_case  ) <= limit:
              divisor += 2
       return divisor


if __name__ == "__main__":
 print(f"""{solution() = }""")