




import argparse
import json
from pathlib import Path

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

from transformers import DeiTImageProcessor, ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


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







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_=False						)       ->	Optional[Any]:



       """simple docstring"""

       A__							=   []
       for i in range(config.num_hidden_layers						):
              # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
              rename_keys.append((f"""blocks.{i}.norm1.weight""", f"""vit.encoder.layer.{i}.layernorm_before.weight""")						)
              rename_keys.append((f"""blocks.{i}.norm1.bias""", f"""vit.encoder.layer.{i}.layernorm_before.bias""")						)
              rename_keys.append((f"""blocks.{i}.attn.proj.weight""", f"""vit.encoder.layer.{i}.attention.output.dense.weight""")						)
              rename_keys.append((f"""blocks.{i}.attn.proj.bias""", f"""vit.encoder.layer.{i}.attention.output.dense.bias""")						)
              rename_keys.append((f"""blocks.{i}.norm2.weight""", f"""vit.encoder.layer.{i}.layernorm_after.weight""")						)
              rename_keys.append((f"""blocks.{i}.norm2.bias""", f"""vit.encoder.layer.{i}.layernorm_after.bias""")						)
              rename_keys.append((f"""blocks.{i}.mlp.fc1.weight""", f"""vit.encoder.layer.{i}.intermediate.dense.weight""")						)
              rename_keys.append((f"""blocks.{i}.mlp.fc1.bias""", f"""vit.encoder.layer.{i}.intermediate.dense.bias""")						)
              rename_keys.append((f"""blocks.{i}.mlp.fc2.weight""", f"""vit.encoder.layer.{i}.output.dense.weight""")						)
              rename_keys.append((f"""blocks.{i}.mlp.fc2.bias""", f"""vit.encoder.layer.{i}.output.dense.bias""")						)

       # projection layer + position embeddings
       rename_keys.extend(
           [
               ('''cls_token''', '''vit.embeddings.cls_token'''),
               ('''patch_embed.proj.weight''', '''vit.embeddings.patch_embeddings.projection.weight'''),
               ('''patch_embed.proj.bias''', '''vit.embeddings.patch_embeddings.projection.bias'''),
               ('''pos_embed''', '''vit.embeddings.position_embeddings'''),
           ]						)

       if base_model:
              # layernorm + pooler
              rename_keys.extend(
                  [
                      ('''norm.weight''', '''layernorm.weight'''),
                      ('''norm.bias''', '''layernorm.bias'''),
                      ('''pre_logits.fc.weight''', '''pooler.dense.weight'''),
                      ('''pre_logits.fc.bias''', '''pooler.dense.bias'''),
                  ]						)

              # if just the base model, we should remove "vit" from all keys that start with "vit"
              A__							=   [(pair[0], pair[1][4:]) if pair[1].startswith('''vit'''						) else pair for pair in rename_keys]
       else:
              # layernorm + classification head
              rename_keys.extend(
                  [
                      ('''norm.weight''', '''vit.layernorm.weight'''),
                      ('''norm.bias''', '''vit.layernorm.bias'''),
                      ('''head.weight''', '''classifier.weight'''),
                      ('''head.bias''', '''classifier.bias'''),
                  ]						)

       return rename_keys







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_			,    lowercase_=False						)       ->	Optional[Any]:



       """simple docstring"""

       for i in range(config.num_hidden_layers						):
              if base_model:
                     A__							=   ''''''
              else:
                     A__							=   '''vit.'''
              # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
              A__							=   state_dict.pop(f"""blocks.{i}.attn.qkv.weight"""						)
              A__							=   state_dict.pop(f"""blocks.{i}.attn.qkv.bias"""						)
              # next, add query, keys and values (in that order) to the state dict
              A__							=   in_proj_weight[
                  : config.hidden_size, :
              ]
              A__							=   in_proj_bias[: config.hidden_size]
              A__							=   in_proj_weight[
                  config.hidden_size : config.hidden_size * 2, :
              ]
              A__							=   in_proj_bias[
                  config.hidden_size : config.hidden_size * 2
              ]
              A__							=   in_proj_weight[
                  -config.hidden_size :, :
              ]
              A__							=   in_proj_bias[-config.hidden_size :]







def SCREAMING_SNAKE_CASE    (	lowercase_						)       ->	List[str]:



       """simple docstring"""

       A__							=   ['''head.weight''', '''head.bias''']
       for k in ignore_keys:
              state_dict.pop(lowercase_			,    lowercase_						)







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_			,    lowercase_						)       ->	Union[str, Any]:



       """simple docstring"""

       A__							=   dct.pop(lowercase_						)
       A__							=   val







def SCREAMING_SNAKE_CASE    (	)       ->	Optional[Any]:



       """simple docstring"""

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







@torch.no_grad()
def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_						)       ->	Optional[Any]:



       """simple docstring"""

       A__							=   ViTConfig()
       A__							=   False
       # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
       if vit_name[-5:] == "in21k":
              A__							=   True
              A__							=   int(vit_name[-12:-10]						)
              A__							=   int(vit_name[-9:-6]						)
       else:
              A__							=   1_000
              A__							=   '''huggingface/label-files'''
              A__							=   '''imagenet-1k-id2label.json'''
              A__							=   json.load(open(hf_hub_download(lowercase_			,    lowercase_			,    repo_type='''dataset'''						)			,    '''r'''						)						)
              A__							=   {int(lowercase_						): v for k, v in idalabel.items()}
              A__							=   idalabel
              A__							=   {v: k for k, v in idalabel.items()}
              A__							=   int(vit_name[-6:-4]						)
              A__							=   int(vit_name[-3:]						)
       # size of the architecture
       if "deit" in vit_name:
              if vit_name[9:].startswith('''tiny'''						):
                     A__							=   192
                     A__							=   768
                     A__							=   12
                     A__							=   3
              elif vit_name[9:].startswith('''small'''						):
                     A__							=   384
                     A__							=   1_536
                     A__							=   12
                     A__							=   6
              else:
                     pass
       else:
              if vit_name[4:].startswith('''small'''						):
                     A__							=   768
                     A__							=   2_304
                     A__							=   8
                     A__							=   8
              elif vit_name[4:].startswith('''base'''						):
                     pass
              elif vit_name[4:].startswith('''large'''						):
                     A__							=   1_024
                     A__							=   4_096
                     A__							=   24
                     A__							=   16
              elif vit_name[4:].startswith('''huge'''						):
                     A__							=   1_280
                     A__							=   5_120
                     A__							=   32
                     A__							=   16

    # load original model from timm
       A__							=   timm.create_model(lowercase_			,    pretrained=lowercase_						)
       timm_model.eval()

       # load state_dict of original model, remove and rename some keys
       A__							=   timm_model.state_dict()
       if base_model:
              remove_classification_head_(lowercase_						)
       A__							=   create_rename_keys(lowercase_			,    lowercase_						)
       for src, dest in rename_keys:
              rename_key(lowercase_			,    lowercase_			,    lowercase_						)
       read_in_q_k_v(lowercase_			,    lowercase_			,    lowercase_						)

       # load HuggingFace model
       if vit_name[-5:] == "in21k":
              A__							=   ViTModel(lowercase_						).eval()
       else:
              A__							=   ViTForImageClassification(lowercase_						).eval()
       model.load_state_dict(lowercase_						)

       # Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
       if "deit" in vit_name:
              A__							=   DeiTImageProcessor(size=config.image_size						)
       else:
              A__							=   ViTImageProcessor(size=config.image_size						)
       A__							=   image_processor(images=prepare_img()			,    return_tensors='''pt'''						)
       A__							=   encoding['''pixel_values''']
       A__							=   model(lowercase_						)

       if base_model:
              A__							=   timm_model.forward_features(lowercase_						)
              assert timm_pooled_output.shape == outputs.pooler_output.shape
              assert torch.allclose(lowercase_			,    outputs.pooler_output			,    atol=1E-3						)
       else:
              A__							=   timm_model(lowercase_						)
              assert timm_logits.shape == outputs.logits.shape
              assert torch.allclose(lowercase_			,    outputs.logits			,    atol=1E-3						)

       Path(lowercase_						).mkdir(exist_ok=lowercase_						)
       print(f"""Saving model {vit_name} to {pytorch_dump_folder_path}"""						)
       model.save_pretrained(lowercase_						)
       print(f"""Saving image processor to {pytorch_dump_folder_path}"""						)
       image_processor.save_pretrained(lowercase_						)


if __name__ == "__main__":
  _lowerCamelCase     :		Union[str, Any]								=       argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--vit_name""",
      default="""vit_base_patch16_224""",
      type=str,
      help="""Name of the ViT timm model you'd like to convert.""",
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
  )

  _lowerCamelCase     :		Optional[Any]								=       parser.parse_args()
  convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)








import argparse

import requests
import torch
from PIL import Image

from transformers import ViTMAEConfig, ViTMAEForPreTraining, ViTMAEImageProcessor







def SCREAMING_SNAKE_CASE    (	lowercase_						)       ->	str:



       """simple docstring"""

       if "cls_token" in name:
              A__							=   name.replace('''cls_token'''			,    '''vit.embeddings.cls_token'''						)
       if "mask_token" in name:
              A__							=   name.replace('''mask_token'''			,    '''decoder.mask_token'''						)
       if "decoder_pos_embed" in name:
              A__							=   name.replace('''decoder_pos_embed'''			,    '''decoder.decoder_pos_embed'''						)
       if "pos_embed" in name and "decoder" not in name:
              A__							=   name.replace('''pos_embed'''			,    '''vit.embeddings.position_embeddings'''						)
       if "patch_embed.proj" in name:
              A__							=   name.replace('''patch_embed.proj'''			,    '''vit.embeddings.patch_embeddings.projection'''						)
       if "patch_embed.norm" in name:
              A__							=   name.replace('''patch_embed.norm'''			,    '''vit.embeddings.norm'''						)
       if "decoder_blocks" in name:
              A__							=   name.replace('''decoder_blocks'''			,    '''decoder.decoder_layers'''						)
       if "blocks" in name:
              A__							=   name.replace('''blocks'''			,    '''vit.encoder.layer'''						)
       if "attn.proj" in name:
              A__							=   name.replace('''attn.proj'''			,    '''attention.output.dense'''						)
       if "attn" in name:
              A__							=   name.replace('''attn'''			,    '''attention.self'''						)
       if "norm1" in name:
              A__							=   name.replace('''norm1'''			,    '''layernorm_before'''						)
       if "norm2" in name:
              A__							=   name.replace('''norm2'''			,    '''layernorm_after'''						)
       if "mlp.fc1" in name:
              A__							=   name.replace('''mlp.fc1'''			,    '''intermediate.dense'''						)
       if "mlp.fc2" in name:
              A__							=   name.replace('''mlp.fc2'''			,    '''output.dense'''						)
       if "decoder_embed" in name:
              A__							=   name.replace('''decoder_embed'''			,    '''decoder.decoder_embed'''						)
       if "decoder_norm" in name:
              A__							=   name.replace('''decoder_norm'''			,    '''decoder.decoder_norm'''						)
       if "decoder_pred" in name:
              A__							=   name.replace('''decoder_pred'''			,    '''decoder.decoder_pred'''						)
       if "norm.weight" in name and "decoder" not in name:
              A__							=   name.replace('''norm.weight'''			,    '''vit.layernorm.weight'''						)
       if "norm.bias" in name and "decoder" not in name:
              A__							=   name.replace('''norm.bias'''			,    '''vit.layernorm.bias'''						)

       return name







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_						)       ->	str:



       """simple docstring"""

       for key in orig_state_dict.copy().keys():
              A__							=   orig_state_dict.pop(lowercase_						)

              if "qkv" in key:
                     A__							=   key.split('''.'''						)
                     A__							=   int(key_split[1]						)
                     if "decoder_blocks" in key:
                            A__							=   config.decoder_hidden_size
                            A__							=   '''decoder.decoder_layers.'''
                            if "weight" in key:
                                   A__							=   val[:dim, :]
                                   A__							=   val[dim : dim * 2, :]
                                   A__							=   val[-dim:, :]
                            elif "bias" in key:
                                   A__							=   val[:dim]
                                   A__							=   val[dim : dim * 2]
                                   A__							=   val[-dim:]
                     else:
                            A__							=   config.hidden_size
                            A__							=   '''vit.encoder.layer.'''
                            if "weight" in key:
                                   A__							=   val[:dim, :]
                                   A__							=   val[dim : dim * 2, :]
                                   A__							=   val[-dim:, :]
                            elif "bias" in key:
                                   A__							=   val[:dim]
                                   A__							=   val[dim : dim * 2]
                                   A__							=   val[-dim:]

              else:
                     A__							=   val

       return orig_state_dict







def SCREAMING_SNAKE_CASE    (	lowercase_			,    lowercase_						)       ->	Tuple:



       """simple docstring"""

       A__							=   ViTMAEConfig()
       if "large" in checkpoint_url:
              A__							=   1_024
              A__							=   4_096
              A__							=   24
              A__							=   16
       elif "huge" in checkpoint_url:
              A__							=   14
              A__							=   1_280
              A__							=   5_120
              A__							=   32
              A__							=   16

       A__							=   ViTMAEForPreTraining(lowercase_						)

       A__							=   torch.hub.load_state_dict_from_url(lowercase_			,    map_location='''cpu'''						)['''model''']

       A__							=   ViTMAEImageProcessor(size=config.image_size						)

       A__							=   convert_state_dict(lowercase_			,    lowercase_						)

       model.load_state_dict(lowercase_						)
       model.eval()

       A__							=   '''https://user-images.githubusercontent.com/11435359/147738734-196fd92f-9260-48d5-ba7e-bf103d29364d.jpg'''

       A__							=   Image.open(requests.get(lowercase_			,    stream=lowercase_						).raw						)
       A__							=   ViTMAEImageProcessor(size=config.image_size						)
       A__							=   image_processor(images=lowercase_			,    return_tensors='''pt'''						)

       # forward pass
       torch.manual_seed(2						)
       A__							=   model(**lowercase_						)
       A__							=   outputs.logits

       if "large" in checkpoint_url:
              A__							=   torch.tensor(
                  [[-0.73_09, -0.71_28, -1.01_69], [-1.01_61, -0.90_58, -1.18_78], [-1.04_78, -0.94_11, -1.19_11]]						)
       elif "huge" in checkpoint_url:
              A__							=   torch.tensor(
                  [[-1.15_99, -0.91_99, -1.22_21], [-1.19_52, -0.92_69, -1.23_07], [-1.21_43, -0.93_37, -1.22_62]]						)
       else:
              A__							=   torch.tensor(
                  [[-0.91_92, -0.84_81, -1.12_59], [-1.13_49, -1.00_34, -1.25_99], [-1.17_57, -1.04_29, -1.27_26]]						)

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

       print(f"""Saving model to {pytorch_dump_folder_path}"""						)
       model.save_pretrained(lowercase_						)

       print(f"""Saving image processor to {pytorch_dump_folder_path}"""						)
       image_processor.save_pretrained(lowercase_						)


if __name__ == "__main__":
  _lowerCamelCase     :		Any								=       argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--checkpoint_url""",
      default="""https://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_base.pth""",
      type=str,
      help="""URL of the checkpoint you'd like to convert.""",
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
  )

  _lowerCamelCase     :		Optional[int]								=       parser.parse_args()
  convert_vit_mae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)


