






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


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


def 	SCREAMING_SNAKE_CASE__       (  snake_case       : List[str]			,					snake_case       : Optional[Any]=False				)->		Optional[Any]:






     '''simple docstring'''



     UpperCAmelCase__						:    Optional[Any]							=							[]
     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"
          UpperCAmelCase__						:    Optional[Any]							=							[(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__       (  snake_case       : int			,					snake_case       : int			,					snake_case       : Union[str, Any]=False				)->		str:






     '''simple docstring'''



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


def 	SCREAMING_SNAKE_CASE__       (  snake_case       : Optional[Any]				)->		Optional[Any]:






     '''simple docstring'''



     UpperCAmelCase__						:    int							=							["head.weight", "head.bias"]
     for k in ignore_keys:
          state_dict.pop(snake_case			,					snake_case				)


def 	SCREAMING_SNAKE_CASE__       (  snake_case       : Optional[int]			,					snake_case       : List[Any]			,					snake_case       : List[str]				)->		int:






     '''simple docstring'''



     UpperCAmelCase__						:    List[str]							=							dct.pop(snake_case				)
     UpperCAmelCase__						:    Union[str, Any]							=							val


def 	SCREAMING_SNAKE_CASE__       (  )->		int:






     '''simple docstring'''



     UpperCAmelCase__						:    str							=							"http://images.cocodataset.org/val2017/000000039769.jpg"
     UpperCAmelCase__						:    Tuple							=							Image.open(requests.get(snake_case			,					stream=snake_case				).raw				)
     return im


@torch.no_grad()
def 	SCREAMING_SNAKE_CASE__       (  snake_case       : int			,					snake_case       : Optional[Any]			,					snake_case       : Any=True				)->		Union[str, Any]:






     '''simple docstring'''



     UpperCAmelCase__						:    List[str]							=							ViTConfig()
     # patch_size
     if model_name[-1] == "8":
          UpperCAmelCase__						:    Tuple							=							8
     # set labels if required
     if not base_model:
          UpperCAmelCase__						:    Optional[Any]							=							1000
          UpperCAmelCase__						:    Union[str, Any]							=							"huggingface/label-files"
          UpperCAmelCase__						:    List[Any]							=							"imagenet-1k-id2label.json"
          UpperCAmelCase__						:    List[Any]							=							json.load(open(hf_hub_download(snake_case			,					snake_case			,					repo_type="dataset"				)			,					"r"				)				)
          UpperCAmelCase__						:    Dict							=							{int(snake_case				): v for k, v in idalabel.items()}
          UpperCAmelCase__						:    Union[str, Any]							=							idalabel
          UpperCAmelCase__						:    Any							=							{v: k for k, v in idalabel.items()}
     # size of the architecture
     if model_name in ["dino_vits8", "dino_vits16"]:
          UpperCAmelCase__						:    Union[str, Any]							=							384
          UpperCAmelCase__						:    List[Any]							=							1536
          UpperCAmelCase__						:    Union[str, Any]							=							12
          UpperCAmelCase__						:    Optional[int]							=							6

     # load original model from torch hub
     UpperCAmelCase__						:    Union[str, Any]							=							torch.hub.load("facebookresearch/dino:main"			,					snake_case				)
     original_model.eval()

     # load state_dict of original model, remove and rename some keys
     UpperCAmelCase__						:    Tuple							=							original_model.state_dict()
     if base_model:
          remove_classification_head_(snake_case				)
     UpperCAmelCase__						:    Dict							=							create_rename_keys(snake_case			,					base_model=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			,					snake_case				)

     # load HuggingFace model
     if base_model:
          UpperCAmelCase__						:    List[str]							=							ViTModel(snake_case			,					add_pooling_layer=snake_case				).eval()
     else:
          UpperCAmelCase__						:    List[str]							=							ViTForImageClassification(snake_case				).eval()
     model.load_state_dict(snake_case				)

     # Check outputs on an image, prepared by ViTImageProcessor
     UpperCAmelCase__						:    Tuple							=							ViTImageProcessor()
     UpperCAmelCase__						:    Dict							=							image_processor(images=prepare_img()			,					return_tensors="pt"				)
     UpperCAmelCase__						:    List[str]							=							encoding["pixel_values"]
     UpperCAmelCase__						:    int							=							model(snake_case				)

     if base_model:
          UpperCAmelCase__						:    Dict							=							original_model(snake_case				)
          assert torch.allclose(snake_case			,					outputs.last_hidden_state[:, 0, :]			,					atol=1E-1				)
     else:
          UpperCAmelCase__						:    int							=							original_model(snake_case				)
          assert logits.shape == outputs.logits.shape
          assert torch.allclose(snake_case			,					outputs.logits			,					atol=1E-3				)

     Path(snake_case				).mkdir(exist_ok=snake_case				)
     print(f'Saving model {model_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 __name__ == "__main__":
 _lowerCAmelCase	:	Any				=      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)










"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_lowerCAmelCase	:	Tuple				=      {
    """configuration_lxmert""": ["""LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """LxmertConfig"""],
    """tokenization_lxmert""": ["""LxmertTokenizer"""],
}

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowerCAmelCase	:	int				=      ["""LxmertTokenizerFast"""]

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowerCAmelCase	:	Tuple				=      [
     """LxmertEncoder""",
     """LxmertForPreTraining""",
     """LxmertForQuestionAnswering""",
     """LxmertModel""",
     """LxmertPreTrainedModel""",
     """LxmertVisualFeatureEncoder""",
     """LxmertXLayer""",
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowerCAmelCase	:	Optional[int]				=      [
     """TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """TFLxmertForPreTraining""",
     """TFLxmertMainLayer""",
     """TFLxmertModel""",
     """TFLxmertPreTrainedModel""",
     """TFLxmertVisualFeatureEncoder""",
 ]


if TYPE_CHECKING:
 from .configuration_lxmert import LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP, LxmertConfig
 from .tokenization_lxmert import LxmertTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_lxmert_fast import LxmertTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_lxmert import (
      LxmertEncoder,
      LxmertForPreTraining,
      LxmertForQuestionAnswering,
      LxmertModel,
      LxmertPreTrainedModel,
      LxmertVisualFeatureEncoder,
      LxmertXLayer,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_lxmert import (
      TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFLxmertForPreTraining,
      TFLxmertMainLayer,
      TFLxmertModel,
      TFLxmertPreTrainedModel,
      TFLxmertVisualFeatureEncoder,
  )

else:
 import sys

 _lowerCAmelCase	:	Any				=      _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


