

"""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()
__UpperCAmelCase =				logging.get_logger(__name__)
def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase=False      ):
 '''simple docstring'''





 UpperCAmelCase__					:  Union[str, 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[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   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase							,    __UpperCamelCase=False      ):
 '''simple docstring'''





 for i in range(config.num_hidden_layers      ):
  if base_model:
   UpperCAmelCase__					:  Tuple					   =					""""""
  else:
   UpperCAmelCase__					:  Any					   =					"""vit."""
  # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
  UpperCAmelCase__					:  Optional[Any]					   =					state_dict.pop(F"blocks.{i}.attn.qkv.weight"      )
  UpperCAmelCase__					:  Dict					   =					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__					:  str					   =					in_proj_bias[: config.hidden_size]
  UpperCAmelCase__					:  Optional[Any]					   =					in_proj_weight[
      config.hidden_size : config.hidden_size * 2, :
  ]
  UpperCAmelCase__					:  Optional[Any]					   =					in_proj_bias[
      config.hidden_size : config.hidden_size * 2
  ]
  UpperCAmelCase__					:  Union[str, Any]					   =					in_proj_weight[
      -config.hidden_size :, :
  ]
  UpperCAmelCase__					:  str					   =					in_proj_bias[-config.hidden_size :]
def   lowerCAmelCase						(					__UpperCamelCase      ):
 '''simple docstring'''





 UpperCAmelCase__					:  int					   =					["""head.weight""", """head.bias"""]
 for k in ignore_keys:
  state_dict.pop(__UpperCamelCase							,    __UpperCamelCase      )
def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase							,    __UpperCamelCase      ):
 '''simple docstring'''





 UpperCAmelCase__					:  Any					   =					dct.pop(__UpperCamelCase      )
 UpperCAmelCase__					:  List[str]					   =					val
def   lowerCAmelCase						(					):
 '''simple docstring'''





 UpperCAmelCase__					:  int					   =					"""http://images.cocodataset.org/val2017/000000039769.jpg"""
 UpperCAmelCase__					:  int					   =					Image.open(requests.get(__UpperCamelCase							,    stream=__UpperCamelCase      ).raw      )
 return im
@torch.no_grad()
def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase							,    __UpperCamelCase=True      ):
 '''simple docstring'''





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

 # load original model from torch hub
 UpperCAmelCase__					:  Any					   =					torch.hub.load("""facebookresearch/dino:main"""							,    __UpperCamelCase      )
 original_model.eval()

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

 # load HuggingFace model
 if base_model:
  UpperCAmelCase__					:  Dict					   =					ViTModel(__UpperCamelCase							,    add_pooling_layer=__UpperCamelCase      ).eval()
 else:
  UpperCAmelCase__					:  Tuple					   =					ViTForImageClassification(__UpperCamelCase      ).eval()
 model.load_state_dict(__UpperCamelCase      )

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

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

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


if __name__ == "__main__":
      __UpperCAmelCase =				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)
      __UpperCAmelCase =				parser.parse_args()
      convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)








"""simple docstring"""




from typing import Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
      from PIL import Image

      from ..image_utils import load_image

if is_torch_available():
      from ..models.auto.modeling_auto import MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING

__UpperCAmelCase =				logging.get_logger(__name__)







@add_end_docstrings(__lowerCamelCase					)
class      __lowercase			(    __lowerCamelCase					):

  def __init__(  self		:			List[str]					,*A		:			Optional[int]					,**A		:			List[str]    ):






   '''simple docstring'''
   super().__init__(*A					,**A    )
   self.check_model_type(A    )

  def    __lowercase       (  self		:			Optional[int]					,A		:			Dict=None					,A		:			Optional[int]=None					,A		:			Tuple=None					,**A		:			str    ):






   '''simple docstring'''
   UpperCAmelCase__   ,      UpperCAmelCase__					:  Any					   =					{}, {}
   if padding is not None:
    UpperCAmelCase__					:  Union[str, Any]					   =					padding
   if truncation is not None:
    UpperCAmelCase__					:  List[str]					   =					truncation
   if top_k is not None:
    UpperCAmelCase__					:  List[Any]					   =					top_k
   return preprocess_params, {}, postprocess_params

  def __call__(  self		:			List[Any]					,A		:			Union["Image.Image", str]					,A		:			str = None					,**A		:			Union[str, Any]    ):






   '''simple docstring'''
   if isinstance(A					,(Image.Image, str)    ) and isinstance(A					,A    ):
    UpperCAmelCase__					:  Any					   =					{"""image""": image, """question""": question}
   else:
    UpperCAmelCase__					:  Dict					   =					image
   UpperCAmelCase__					:  Dict					   =					super().__call__(A					,**A    )
   return results

  def    __lowercase       (  self		:			Optional[int]					,A		:			Optional[Any]					,A		:			Optional[Any]=False					,A		:			List[Any]=False    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					load_image(inputs["""image"""]    )
   UpperCAmelCase__					:  Optional[Any]					   =					self.tokenizer(
       inputs["""question"""]					,return_tensors=self.framework					,padding=A					,truncation=A    )
   UpperCAmelCase__					:  int					   =					self.image_processor(images=A					,return_tensors=self.framework    )
   model_inputs.update(A    )
   return model_inputs

  def    __lowercase       (  self		:			Dict					,A		:			Tuple    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					self.model(**A    )
   return model_outputs




  def    __lowercase       (  self		:			int					,A		:			List[Any]					,A		:			List[str]=5    ):






   '''simple docstring'''
   if top_k > self.model.config.num_labels:
    UpperCAmelCase__					:  Any					   =					self.model.config.num_labels

   if self.framework == "pt":
    UpperCAmelCase__					:  Optional[int]					   =					model_outputs.logits.sigmoid()[0]
    UpperCAmelCase__   ,      UpperCAmelCase__					:  Optional[int]					   =					probs.topk(A    )
   else:
    raise ValueError(f"Unsupported framework: {self.framework}"    )

   UpperCAmelCase__					:  List[str]					   =					scores.tolist()
   UpperCAmelCase__					:  Optional[int]					   =					ids.tolist()
   return [{"score": score, "answer": self.model.config.idalabel[_id]} for score, _id in zip(A					,A    )]





