

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 ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase     :   List[str]						     =  logging.get_logger(__name__)

def    A_  (							a					,   a=False  ):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_				:  int           =     []
      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'),
                ]  )

            # if just the base model, we should remove "vit" from all keys that start with "vit"
            SCREAMING_SNAKE_CASE_				:  int           =     [(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    A_  (							a					,   a					,   a=False  ):

      """simple docstring"""



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

def    A_  (							a  ):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_				:  List[str]           =     ['head.weight', 'head.bias']
      for k in ignore_keys:
            state_dict.pop(a					,   a  )

def    A_  (							a					,   a					,   a  ):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_				:  int           =     dct.pop(a  )
      SCREAMING_SNAKE_CASE_				:  str           =     val

def    A_  (							):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_				:  str           =     'http://images.cocodataset.org/val2017/000000039769.jpg'
      SCREAMING_SNAKE_CASE_				:  List[str]           =     Image.open(requests.get(a					,   stream=a  ).raw  )
      return im

@torch.no_grad()
def    A_  (							a					,   a					,   a=True  ):

      """simple docstring"""



      SCREAMING_SNAKE_CASE_				:  List[Any]           =     ViTConfig()
      # patch_size
      if model_name[-1] == "8":
            SCREAMING_SNAKE_CASE_				:  List[Any]           =     8
      # set labels if required
      if not base_model:
            SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     1_0_0_0
            SCREAMING_SNAKE_CASE_				:  int           =     'huggingface/label-files'
            SCREAMING_SNAKE_CASE_				:  List[Any]           =     'imagenet-1k-id2label.json'
            SCREAMING_SNAKE_CASE_				:  str           =     json.load(open(hf_hub_download(a					,   a					,   repo_type='dataset'  )					,   'r'  )  )
            SCREAMING_SNAKE_CASE_				:  List[Any]           =     {int(a  ): v for k, v in idalabel.items()}
            SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     idalabel
            SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     {v: k for k, v in idalabel.items()}
      # size of the architecture
      if model_name in ["dino_vits8", "dino_vits16"]:
            SCREAMING_SNAKE_CASE_				:  str           =     3_8_4
            SCREAMING_SNAKE_CASE_				:  List[str]           =     1_5_3_6
            SCREAMING_SNAKE_CASE_				:  Dict           =     1_2
            SCREAMING_SNAKE_CASE_				:  Dict           =     6

      # load original model from torch hub
      SCREAMING_SNAKE_CASE_				:  Optional[int]           =     torch.hub.load('facebookresearch/dino:main'					,   a  )
      original_model.eval()

      # load state_dict of original model, remove and rename some keys
      SCREAMING_SNAKE_CASE_				:  Any           =     original_model.state_dict()
      if base_model:
            remove_classification_head_(a  )
      SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     create_rename_keys(a					,   base_model=a  )
      for src, dest in rename_keys:
            rename_key(a					,   a					,   a  )
      read_in_q_k_v(a					,   a					,   a  )

      # load HuggingFace model
      if base_model:
            SCREAMING_SNAKE_CASE_				:  List[Any]           =     ViTModel(a					,   add_pooling_layer=a  ).eval()
      else:
            SCREAMING_SNAKE_CASE_				:  Dict           =     ViTForImageClassification(a  ).eval()
      model.load_state_dict(a  )

      # Check outputs on an image, prepared by ViTImageProcessor
      SCREAMING_SNAKE_CASE_				:  List[str]           =     ViTImageProcessor()
      SCREAMING_SNAKE_CASE_				:  Tuple           =     image_processor(images=prepare_img()					,   return_tensors='pt'  )
      SCREAMING_SNAKE_CASE_				:  List[Any]           =     encoding['pixel_values']
      SCREAMING_SNAKE_CASE_				:  Optional[int]           =     model(a  )

      if base_model:
            SCREAMING_SNAKE_CASE_				:  Tuple           =     original_model(a  )
            assert torch.allclose(a					,   outputs.last_hidden_state[:, 0, :]					,   atol=1e-1  )
      else:
            SCREAMING_SNAKE_CASE_				:  List[str]           =     original_model(a  )
            assert logits.shape == outputs.logits.shape
            assert torch.allclose(a					,   outputs.logits					,   atol=1e-3  )

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


if __name__ == "__main__":
    lowerCAmelCase     :   Optional[int]						     =  argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '--model_name',
        default='dino_vitb16',
        type=str,
        help='Name of the model trained with DINO you\'d like to convert.',
    )
    parser.add_argument(
        '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
    )
    parser.add_argument(
        '--base_model',
        action='store_true',
        help='Whether to only convert the base model (no projection head weights).',
    )

    parser.set_defaults(base_model=True)
    lowerCAmelCase     :   int						     =  parser.parse_args()
    convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)









import unittest

from transformers import is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device


if is_torch_available():
    import torch

    from transformers import AutoModelForImageClassification

if is_vision_available():
    from transformers import AutoImageProcessor
@require_torch
@require_vision
class  _A		(  unittest.TestCase):

  @slow
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  List[str]           =     AutoImageProcessor.from_pretrained('microsoft/dit-base-finetuned-rvlcdip'					)
        SCREAMING_SNAKE_CASE_				:  Any           =     AutoModelForImageClassification.from_pretrained('microsoft/dit-base-finetuned-rvlcdip'					)
        model.to(_SCREAMING_SNAKE_CASE					)

        from datasets import load_dataset

        SCREAMING_SNAKE_CASE_				:  str           =     load_dataset('nielsr/rvlcdip-demo'					)

        SCREAMING_SNAKE_CASE_				:  List[Any]           =     dataset['train'][0]['image'].convert('RGB'					)

        SCREAMING_SNAKE_CASE_				:  List[Any]           =     image_processor(_SCREAMING_SNAKE_CASE ,					return_tensors='pt'					).to(_SCREAMING_SNAKE_CASE					)

        # forward pass
        with torch.no_grad():
              SCREAMING_SNAKE_CASE_				:  int           =     model(**_SCREAMING_SNAKE_CASE					)
              SCREAMING_SNAKE_CASE_				:  Any           =     outputs.logits

        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     torch.Size((1, 16)					)
        self.assertEqual(logits.shape ,					_SCREAMING_SNAKE_CASE					)

        SCREAMING_SNAKE_CASE_				:  List[str]           =     torch.tensor(
            [-0.4158, -0.4092, -0.4347] ,					device=_SCREAMING_SNAKE_CASE ,					dtype=torch.float ,					)
        self.assertTrue(torch.allclose(logits[0, :3] ,					_SCREAMING_SNAKE_CASE ,					atol=1e-4					)					)






