
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__      :		Dict       							=  logging.get_logger(__name__)
def    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_=False						)     ->     Dict:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE      =				[]
   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      =				[(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_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_=False						)     ->     List[str]:




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




   """simple docstring"""
   _SCREAMING_SNAKE_CASE      =				["""head.weight""", """head.bias"""]
   for k in ignore_keys:
      state_dict.pop(SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)
def    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_						)     ->     Optional[Any]:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE      =				dct.pop(SCREAMING_SNAKE_CASE_						)
   _SCREAMING_SNAKE_CASE      =				val
def    lowerCAmelCase_	(				)     ->     Union[str, Any]:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE      =				"""http://images.cocodataset.org/val2017/000000039769.jpg"""
   _SCREAMING_SNAKE_CASE      =				Image.open(requests.get(SCREAMING_SNAKE_CASE_    ,      stream=SCREAMING_SNAKE_CASE_						).raw						)
   return im
@torch.no_grad()
def    lowerCAmelCase_	(				SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_=True						)     ->     List[Any]:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE      =				ViTConfig()
   # patch_size
   if model_name[-1] == "8":
      _SCREAMING_SNAKE_CASE      =				8
   # set labels if required
   if not base_model:
      _SCREAMING_SNAKE_CASE      =				10_00
      _SCREAMING_SNAKE_CASE      =				"""huggingface/label-files"""
      _SCREAMING_SNAKE_CASE      =				"""imagenet-1k-id2label.json"""
      _SCREAMING_SNAKE_CASE      =				json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE_    ,      SCREAMING_SNAKE_CASE_    ,      repo_type="""dataset"""						)    ,      """r"""						)						)
      _SCREAMING_SNAKE_CASE      =				{int(SCREAMING_SNAKE_CASE_						): v for k, v in idalabel.items()}
      _SCREAMING_SNAKE_CASE      =				idalabel
      _SCREAMING_SNAKE_CASE      =				{v: k for k, v in idalabel.items()}
   # size of the architecture
   if model_name in ["dino_vits8", "dino_vits16"]:
      _SCREAMING_SNAKE_CASE      =				3_84
      _SCREAMING_SNAKE_CASE      =				15_36
      _SCREAMING_SNAKE_CASE      =				12
      _SCREAMING_SNAKE_CASE      =				6

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

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

   # load HuggingFace model
   if base_model:
      _SCREAMING_SNAKE_CASE      =				ViTModel(SCREAMING_SNAKE_CASE_    ,      add_pooling_layer=SCREAMING_SNAKE_CASE_						).eval()
   else:
      _SCREAMING_SNAKE_CASE      =				ViTForImageClassification(SCREAMING_SNAKE_CASE_						).eval()
   model.load_state_dict(SCREAMING_SNAKE_CASE_						)

   # Check outputs on an image, prepared by ViTImageProcessor
   _SCREAMING_SNAKE_CASE      =				ViTImageProcessor()
   _SCREAMING_SNAKE_CASE      =				image_processor(images=prepare_img()    ,      return_tensors="""pt"""						)
   _SCREAMING_SNAKE_CASE      =				encoding["""pixel_values"""]
   _SCREAMING_SNAKE_CASE      =				model(SCREAMING_SNAKE_CASE_						)

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

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


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








'''simple docstring'''



import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding

class    _a				(_lowerCamelCase):





       """simple docstring"""




       SCREAMING_SNAKE_CASE      =  ['image_processor', 'tokenizer']
       SCREAMING_SNAKE_CASE      =  'ChineseCLIPImageProcessor'
       SCREAMING_SNAKE_CASE      =  ('BertTokenizer', 'BertTokenizerFast')






       def __init__(							self       ,							A__=None       ,							A__=None       ,							**A__      )    ->			int:
          _SCREAMING_SNAKE_CASE      =				None
          if "feature_extractor" in kwargs:
             warnings.warn(
                 """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
                 """ instead."""       ,							A__       ,							)
             _SCREAMING_SNAKE_CASE      =				kwargs.pop("""feature_extractor"""      )

          _SCREAMING_SNAKE_CASE      =				image_processor if image_processor is not None else feature_extractor
          if image_processor is None:
             raise ValueError("""You need to specify an `image_processor`."""      )
          if tokenizer is None:
             raise ValueError("""You need to specify a `tokenizer`."""      )

          super().__init__(A__       ,							A__      )
          _SCREAMING_SNAKE_CASE      =				self.image_processor






       def __call__(							self       ,							A__=None       ,							A__=None       ,							A__=None       ,							**A__      )    ->			Optional[int]:

          if text is None and images is None:
             raise ValueError("""You have to specify either text or images. Both cannot be none."""      )

          if text is not None:
             _SCREAMING_SNAKE_CASE      =				self.tokenizer(A__       ,							return_tensors=A__       ,							**A__      )

          if images is not None:
             _SCREAMING_SNAKE_CASE      =				self.image_processor(A__       ,							return_tensors=A__       ,							**A__      )

          if text is not None and images is not None:
             _SCREAMING_SNAKE_CASE      =				image_features.pixel_values
             return encoding
          elif text is not None:
             return encoding
          else:
             return BatchEncoding(data=dict(**A__      )       ,							tensor_type=A__      )






       def        UpperCamelCase				(							self       ,							*A__       ,							**A__      )    ->			Dict:
          return self.tokenizer.batch_decode(*A__       ,							**A__      )






       def        UpperCamelCase				(							self       ,							*A__       ,							**A__      )    ->			Optional[Any]:
          return self.tokenizer.decode(*A__       ,							**A__      )






       @property
       def        UpperCamelCase				(							self      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				self.tokenizer.model_input_names
          _SCREAMING_SNAKE_CASE      =				self.image_processor.model_input_names
          return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names      )      )






       @property
       def        UpperCamelCase				(							self      )    ->			Optional[int]:
          warnings.warn(
              """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""       ,							A__       ,							)
          return self.image_processor_class






