



import argparse
import json
from pathlib import Path

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

from transformers import (
    BeitConfig,
    BeitForImageClassification,
    BeitForMaskedImageModeling,
    BeitForSemanticSegmentation,
    BeitImageProcessor,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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




def       UpperCAmelCase__	(							__snake_case  ,			__snake_case=False  ,			__snake_case=False							)				->	List[str]:
  _A          =      '''backbone.''' if is_semantic else ''''''

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

  # projection layer + position embeddings
  rename_keys.extend(
      [
          (F'''{prefix}cls_token''', '''beit.embeddings.cls_token'''),
          (F'''{prefix}patch_embed.proj.weight''', '''beit.embeddings.patch_embeddings.projection.weight'''),
          (F'''{prefix}patch_embed.proj.bias''', '''beit.embeddings.patch_embeddings.projection.bias'''),
      ]							)

  if has_lm_head:
    # mask token + shared relative position bias + layernorm
    rename_keys.extend(
        [
            ('''mask_token''', '''beit.embeddings.mask_token'''),
            (
                '''rel_pos_bias.relative_position_bias_table''',
                '''beit.encoder.relative_position_bias.relative_position_bias_table''',
            ),
            (
                '''rel_pos_bias.relative_position_index''',
                '''beit.encoder.relative_position_bias.relative_position_index''',
            ),
            ('''norm.weight''', '''layernorm.weight'''),
            ('''norm.bias''', '''layernorm.bias'''),
        ]							)
  elif is_semantic:
    # semantic segmentation classification heads
    rename_keys.extend(
        [
            ('''decode_head.conv_seg.weight''', '''decode_head.classifier.weight'''),
            ('''decode_head.conv_seg.bias''', '''decode_head.classifier.bias'''),
            ('''auxiliary_head.conv_seg.weight''', '''auxiliary_head.classifier.weight'''),
            ('''auxiliary_head.conv_seg.bias''', '''auxiliary_head.classifier.bias'''),
        ]							)
  else:
    # layernorm + classification head
    rename_keys.extend(
        [
            ('''fc_norm.weight''', '''beit.pooler.layernorm.weight'''),
            ('''fc_norm.bias''', '''beit.pooler.layernorm.bias'''),
            ('''head.weight''', '''classifier.weight'''),
            ('''head.bias''', '''classifier.bias'''),
        ]							)

  return rename_keys




def       UpperCAmelCase__	(							__snake_case  ,			__snake_case  ,			__snake_case=False  ,			__snake_case=False							)				->	List[Any]:
  for i in range(config.num_hidden_layers							):
    _A          =      '''backbone.''' if is_semantic else ''''''
    # queries, keys and values
    _A          =      state_dict.pop(F'''{prefix}blocks.{i}.attn.qkv.weight'''							)
    _A          =      state_dict.pop(F'''{prefix}blocks.{i}.attn.q_bias'''							)
    _A          =      state_dict.pop(F'''{prefix}blocks.{i}.attn.v_bias'''							)

    _A          =      in_proj_weight[
        : config.hidden_size, :
    ]
    _A          =      q_bias
    _A          =      in_proj_weight[
        config.hidden_size : config.hidden_size * 2, :
    ]
    _A          =      in_proj_weight[
        -config.hidden_size :, :
    ]
    _A          =      v_bias

    # gamma_1 and gamma_2
    # we call them lambda because otherwise they are renamed when using .from_pretrained
    _A          =      state_dict.pop(F'''{prefix}blocks.{i}.gamma_1'''							)
    _A          =      state_dict.pop(F'''{prefix}blocks.{i}.gamma_2'''							)

    _A          =      gamma_a
    _A          =      gamma_a

    # relative_position bias table + index
    if not has_lm_head:
      # each layer has its own relative position bias
      _A          =      state_dict.pop(F'''{prefix}blocks.{i}.attn.relative_position_bias_table'''							)
      _A          =      state_dict.pop(F'''{prefix}blocks.{i}.attn.relative_position_index'''							)

      _A          =      table
      _A          =      index




def       UpperCAmelCase__	(							__snake_case  ,			__snake_case  ,			__snake_case							)				->	Union[str, Any]:
  _A          =      dct.pop(__snake_case							)
  _A          =      val




def       UpperCAmelCase__	(							)				->	Tuple:
  _A          =      '''http://images.cocodataset.org/val2017/000000039769.jpg'''
  _A          =      Image.open(requests.get(__snake_case  ,			stream=__snake_case							).raw							)
  return im




@torch.no_grad()
def       UpperCAmelCase__	(							__snake_case  ,			__snake_case							)				->	List[str]:
  _A          =      BeitConfig()
  _A          =      False
  _A          =      False
  _A          =      '''huggingface/label-files'''
  # set config parameters based on URL
  if checkpoint_url[-9:-4] == "pt22k":
    # masked image modeling
    _A          =      True
    _A          =      True
    _A          =      True
  elif checkpoint_url[-9:-4] == "ft22k":
    # intermediate fine-tuning on ImageNet-22k
    _A          =      True
    _A          =      21_841
    _A          =      '''imagenet-22k-id2label.json'''
    _A          =      json.load(open(hf_hub_download(__snake_case  ,			__snake_case  ,			repo_type='''dataset'''							)  ,			'''r'''							)							)
    _A          =      {int(__snake_case							): v for k, v in idalabel.items()}
    # this dataset contains 21843 labels but the model only has 21841
    # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
    del idalabel[9_205]
    del idalabel[15_027]
    _A          =      idalabel
    _A          =      {v: k for k, v in idalabel.items()}
  elif checkpoint_url[-8:-4] == "to1k":
    # fine-tuning on ImageNet-1k
    _A          =      True
    _A          =      1_000
    _A          =      '''imagenet-1k-id2label.json'''
    _A          =      json.load(open(hf_hub_download(__snake_case  ,			__snake_case  ,			repo_type='''dataset'''							)  ,			'''r'''							)							)
    _A          =      {int(__snake_case							): v for k, v in idalabel.items()}
    _A          =      idalabel
    _A          =      {v: k for k, v in idalabel.items()}
    if "384" in checkpoint_url:
      _A          =      384
    if "512" in checkpoint_url:
      _A          =      512
  elif "ade20k" in checkpoint_url:
    # fine-tuning
    _A          =      True
    _A          =      150
    _A          =      '''ade20k-id2label.json'''
    _A          =      json.load(open(hf_hub_download(__snake_case  ,			__snake_case  ,			repo_type='''dataset'''							)  ,			'''r'''							)							)
    _A          =      {int(__snake_case							): v for k, v in idalabel.items()}
    _A          =      idalabel
    _A          =      {v: k for k, v in idalabel.items()}
    _A          =      640
    _A          =      True
  else:
    raise ValueError('''Checkpoint not supported, URL should either end with \'pt22k\', \'ft22k\', \'to1k\' or \'ade20k\''''							)

  # size of the architecture
  if "base" in checkpoint_url:
    pass
  elif "large" in checkpoint_url:
    _A          =      1_024
    _A          =      4_096
    _A          =      24
    _A          =      16
    if "ade20k" in checkpoint_url:
      _A          =      640
      _A          =      [7, 11, 15, 23]
  else:
    raise ValueError('''Should either find \'base\' or \'large\' in checkpoint URL'''							)

  # load state_dict of original model, remove and rename some keys
  _A          =      torch.hub.load_state_dict_from_url(__snake_case  ,			map_location='''cpu'''  ,			check_hash=__snake_case							)
  _A          =      state_dict['''model'''] if '''ade20k''' not in checkpoint_url else state_dict['''state_dict''']

  _A          =      create_rename_keys(__snake_case  ,			has_lm_head=__snake_case  ,			is_semantic=__snake_case							)
  for src, dest in rename_keys:
    rename_key(__snake_case  ,			__snake_case  ,			__snake_case							)
  read_in_q_k_v(__snake_case  ,			__snake_case  ,			has_lm_head=__snake_case  ,			is_semantic=__snake_case							)
  if is_semantic:
    # add prefix to decoder keys
    for key, val in state_dict.copy().items():
      _A          =      state_dict.pop(__snake_case							)
      if key.startswith('''backbone.fpn'''							):
        _A          =      key.replace('''backbone.fpn'''  ,			'''fpn'''							)
      _A          =      val

    # load HuggingFace model
  if checkpoint_url[-9:-4] == "pt22k":
    _A          =      BeitForMaskedImageModeling(__snake_case							)
  elif "ade20k" in checkpoint_url:
    _A          =      BeitForSemanticSegmentation(__snake_case							)
  else:
    _A          =      BeitForImageClassification(__snake_case							)
  model.eval()
  model.load_state_dict(__snake_case							)

  # Check outputs on an image
  if is_semantic:
    _A          =      BeitImageProcessor(size=config.image_size  ,			do_center_crop=__snake_case							)
    _A          =      load_dataset('''hf-internal-testing/fixtures_ade20k'''  ,			split='''test'''							)
    _A          =      Image.open(ds[0]['''file''']							)
  else:
    _A          =      BeitImageProcessor(
        size=config.image_size  ,			resample=PILImageResampling.BILINEAR  ,			do_center_crop=__snake_case							)
    _A          =      prepare_img()

  _A          =      image_processor(images=__snake_case  ,			return_tensors='''pt'''							)
  _A          =      encoding['''pixel_values''']

  _A          =      model(__snake_case							)
  _A          =      outputs.logits

  # verify logits
  _A          =      torch.Size([1, 1_000]							)
  if checkpoint_url[:-4].endswith('''beit_base_patch16_224_pt22k'''							):
    _A          =      torch.Size([1, 196, 8_192]							)
  elif checkpoint_url[:-4].endswith('''beit_large_patch16_224_pt22k'''							):
    _A          =      torch.Size([1, 196, 8_192]							)
  elif checkpoint_url[:-4].endswith('''beit_base_patch16_224_pt22k_ft22k'''							):
    _A          =      torch.Size([1, 21_841]							)
    _A          =      torch.tensor([2.22_88, 2.46_71, 0.73_95]							)
    _A          =      2_397
  elif checkpoint_url[:-4].endswith('''beit_large_patch16_224_pt22k_ft22k'''							):
    _A          =      torch.Size([1, 21_841]							)
    _A          =      torch.tensor([1.68_81, -0.27_87, 0.59_01]							)
    _A          =      2_396
  elif checkpoint_url[:-4].endswith('''beit_base_patch16_224_pt22k_ft1k'''							):
    _A          =      torch.tensor([0.12_41, 0.07_98, -0.65_69]							)
    _A          =      285
  elif checkpoint_url[:-4].endswith('''beit_base_patch16_224_pt22k_ft22kto1k'''							):
    _A          =      torch.tensor([-1.23_85, -1.09_87, -1.01_08]							)
    _A          =      281
  elif checkpoint_url[:-4].endswith('''beit_base_patch16_384_pt22k_ft22kto1k'''							):
    _A          =      torch.tensor([-1.53_03, -0.94_84, -0.31_47]							)
    _A          =      761
  elif checkpoint_url[:-4].endswith('''beit_large_patch16_224_pt22k_ft1k'''							):
    _A          =      torch.tensor([0.46_10, -0.09_28, 0.20_86]							)
    _A          =      761
  elif checkpoint_url[:-4].endswith('''beit_large_patch16_224_pt22k_ft22kto1k'''							):
    _A          =      torch.tensor([-0.48_04, 0.62_57, -0.18_37]							)
    _A          =      761
  elif checkpoint_url[:-4].endswith('''beit_large_patch16_384_pt22k_ft22kto1k'''							):
    _A          =      torch.tensor([[-0.51_22, 0.51_17, -0.21_13]]							)
    _A          =      761
  elif checkpoint_url[:-4].endswith('''beit_large_patch16_512_pt22k_ft22kto1k'''							):
    _A          =      torch.tensor([-0.30_62, 0.72_61, 0.48_52]							)
    _A          =      761
  elif checkpoint_url[:-4].endswith('''beit_base_patch16_640_pt22k_ft22ktoade20k'''							):
    _A          =      (1, 150, 160, 160)
    _A          =      torch.tensor(
        [
            [[-4.92_25, -2.39_54, -3.05_22], [-2.88_22, -1.00_46, -1.75_61], [-2.95_49, -1.32_28, -2.13_47]],
            [[-5.81_68, -3.41_29, -4.07_78], [-3.86_51, -2.22_14, -3.02_77], [-3.83_56, -2.46_43, -3.35_35]],
            [[-0.00_78, 3.99_52, 4.07_54], [2.98_56, 4.69_44, 5.00_35], [3.24_13, 4.78_13, 4.99_69]],
        ]							)
  elif checkpoint_url[:-4].endswith('''beit_large_patch16_640_pt22k_ft22ktoade20k'''							):
    _A          =      (1, 150, 160, 160)
    _A          =      torch.tensor(
        [
            [[-4.33_05, -2.30_49, -3.01_61], [-2.95_91, -1.53_05, -2.22_51], [-3.41_98, -1.80_04, -2.90_62]],
            [[-5.89_22, -3.74_35, -4.39_78], [-4.20_63, -2.78_72, -3.47_55], [-4.27_91, -3.18_74, -4.16_81]],
            [[0.98_95, 4.34_67, 4.76_63], [4.24_76, 5.68_30, 6.15_18], [4.55_50, 6.24_95, 6.51_54]],
        ]							)
  else:
    raise ValueError('''Can\'t verify logits as model is not supported'''							)

  assert logits.shape == expected_shape, "Shape of logits not as expected"
  if not has_lm_head:
    if is_semantic:
      assert torch.allclose(
          logits[0, :3, :3, :3]  ,			__snake_case  ,			atol=1e-3							), "First elements of logits not as expected"
    else:
      print('''Predicted class idx:'''  ,			logits.argmax(-1							).item()							)
      assert torch.allclose(
          logits[0, :3]  ,			__snake_case  ,			atol=1e-3							), "First elements of logits not as expected"
      assert logits.argmax(-1							).item() == expected_class_idx, "Predicted class index not as expected"

  Path(__snake_case							).mkdir(exist_ok=__snake_case							)
  print(F'''Saving model 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 __name__ == "__main__":
    __lowerCamelCase								=   argparse.ArgumentParser()

    parser.add_argument(
        """--checkpoint_url""",
        default="""https://conversationhub.blob.core.windows.net/beit-share-public/beit/beit_base_patch16_224_pt22k_ft22kto1k.pth""",
        type=str,
        help="""URL to the original PyTorch checkpoint (.pth file).""",
    )
    parser.add_argument(
        """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the folder to output PyTorch model."""
    )
    __lowerCamelCase								=   parser.parse_args()
    convert_beit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)




import importlib.metadata
import subprocess
import sys




def       UpperCAmelCase__	(							__snake_case = False							)				->	List[Any]:
  _A          =      False
  if "IPython" in sys.modules:
    _A          =      '''google.colab''' in str(sys.modules['''IPython'''].get_ipython()							)

  if in_colab:
    if upgrade:
      _A          =      ['''pip''', '''install''', '''-U''', '''torch''']
      subprocess.run(__snake_case  ,			check=__snake_case							)
    # get the current version of torch
    _A          =      importlib.metadata.version('''torch'''							)
    _A          =      torch_version[: torch_version.rindex('''.'''							)]
    _A          =      F'''https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-{torch_version_trunc}-cp37-cp37m-linux_x86_64.whl'''
    _A          =      ['''pip''', '''install''', xla_wheel]
    subprocess.run(__snake_case  ,			check=__snake_case							)
  else:
    raise RuntimeError('''`install_xla` utility works only on google colab.'''							)