
"""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()
lowercase_         =       logging.get_logger(__name__)




def       A_      (       lowercase    ,			lowercase=False   )  ->				Optional[int]:


     """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_    :     Dict            =	[(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_      (       lowercase    ,			lowercase    ,			lowercase=False   )  ->				str:


     """simple docstring"""







     for i in range(config.num_hidden_layers   ):
          if base_model:
               UpperCAmelCase_    :     Dict            =	""""""
          else:
               UpperCAmelCase_    :     Dict            =	"""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_    :     Any            =	state_dict.pop(f'''blocks.{i}.attn.qkv.bias'''   )
          # next, add query, keys and values (in that order) to the state dict
          UpperCAmelCase_    :     int            =	in_proj_weight[
              : config.hidden_size, :
          ]
          UpperCAmelCase_    :     Union[str, Any]            =	in_proj_bias[: config.hidden_size]
          UpperCAmelCase_    :     str            =	in_proj_weight[
              config.hidden_size : config.hidden_size * 2, :
          ]
          UpperCAmelCase_    :     Union[str, Any]            =	in_proj_bias[
              config.hidden_size : config.hidden_size * 2
          ]
          UpperCAmelCase_    :     Any            =	in_proj_weight[
              -config.hidden_size :, :
          ]
          UpperCAmelCase_    :     Union[str, Any]            =	in_proj_bias[-config.hidden_size :]




def       A_      (       lowercase   )  ->				Optional[int]:


     """simple docstring"""







     UpperCAmelCase_    :     List[str]            =	["""head.weight""", """head.bias"""]
     for k in ignore_keys:
          state_dict.pop(lowercase    ,			lowercase   )




def       A_      (       lowercase    ,			lowercase    ,			lowercase   )  ->				Optional[Any]:


     """simple docstring"""







     UpperCAmelCase_    :     List[Any]            =	dct.pop(lowercase   )
     UpperCAmelCase_    :     Dict            =	val




def       A_      (       )  ->				Any:


     """simple docstring"""







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




@torch.no_grad()
def       A_      (       lowercase    ,			lowercase    ,			lowercase=True   )  ->				List[str]:


     """simple docstring"""







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

     # load original model from torch hub
     UpperCAmelCase_    :     List[Any]            =	torch.hub.load("""facebookresearch/dino:main"""    ,			lowercase   )
     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_(lowercase   )
     UpperCAmelCase_    :     Dict            =	create_rename_keys(lowercase    ,			base_model=lowercase   )
     for src, dest in rename_keys:
          rename_key(lowercase    ,			lowercase    ,			lowercase   )
     read_in_q_k_v(lowercase    ,			lowercase    ,			lowercase   )

     # load HuggingFace model
     if base_model:
          UpperCAmelCase_    :     Dict            =	ViTModel(lowercase    ,			add_pooling_layer=lowercase   ).eval()
     else:
          UpperCAmelCase_    :     Optional[int]            =	ViTForImageClassification(lowercase   ).eval()
     model.load_state_dict(lowercase   )

     # Check outputs on an image, prepared by ViTImageProcessor
     UpperCAmelCase_    :     Optional[Any]            =	ViTImageProcessor()
     UpperCAmelCase_    :     List[Any]            =	image_processor(images=prepare_img()    ,			return_tensors="""pt"""   )
     UpperCAmelCase_    :     List[Any]            =	encoding["""pixel_values"""]
     UpperCAmelCase_    :     Tuple            =	model(lowercase   )

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

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


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



"""simple docstring"""







from . import __version__

# Backward compatibility imports, to make sure all those objects can be found in file_utils
from .utils import (
    CLOUDFRONT_DISTRIB_PREFIX,
    CONFIG_NAME,
    DISABLE_TELEMETRY,
    DUMMY_INPUTS,
    DUMMY_MASK,
    ENV_VARS_TRUE_AND_AUTO_VALUES,
    ENV_VARS_TRUE_VALUES,
    FEATURE_EXTRACTOR_NAME,
    FLAX_WEIGHTS_NAME,
    HF_MODULES_CACHE,
    HUGGINGFACE_CO_PREFIX,
    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
    MODEL_CARD_NAME,
    MULTIPLE_CHOICE_DUMMY_INPUTS,
    PYTORCH_PRETRAINED_BERT_CACHE,
    PYTORCH_TRANSFORMERS_CACHE,
    S3_BUCKET_PREFIX,
    SENTENCEPIECE_UNDERLINE,
    SPIECE_UNDERLINE,
    TF2_WEIGHTS_NAME,
    TF_WEIGHTS_NAME,
    TORCH_FX_REQUIRED_VERSION,
    TRANSFORMERS_CACHE,
    TRANSFORMERS_DYNAMIC_MODULE_NAME,
    USE_JAX,
    USE_TF,
    USE_TORCH,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
    ContextManagers,
    DummyObject,
    EntryNotFoundError,
    ExplicitEnum,
    ModelOutput,
    PaddingStrategy,
    PushToHubMixin,
    RepositoryNotFoundError,
    RevisionNotFoundError,
    TensorType,
    _LazyModule,
    add_code_sample_docstrings,
    add_end_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    cached_property,
    copy_func,
    default_cache_path,
    define_sagemaker_information,
    get_cached_models,
    get_file_from_repo,
    get_full_repo_name,
    get_torch_version,
    has_file,
    http_user_agent,
    is_apex_available,
    is_bsa_available,
    is_coloredlogs_available,
    is_datasets_available,
    is_detectrona_available,
    is_faiss_available,
    is_flax_available,
    is_ftfy_available,
    is_in_notebook,
    is_ipex_available,
    is_librosa_available,
    is_offline_mode,
    is_onnx_available,
    is_pandas_available,
    is_phonemizer_available,
    is_protobuf_available,
    is_psutil_available,
    is_pyanvml_available,
    is_pyctcdecode_available,
    is_pytesseract_available,
    is_pytorch_quantization_available,
    is_rjieba_available,
    is_sagemaker_dp_enabled,
    is_sagemaker_mp_enabled,
    is_scipy_available,
    is_sentencepiece_available,
    is_seqio_available,
    is_sklearn_available,
    is_soundfile_availble,
    is_spacy_available,
    is_speech_available,
    is_tensor,
    is_tensorflow_probability_available,
    is_tfaonnx_available,
    is_tf_available,
    is_timm_available,
    is_tokenizers_available,
    is_torch_available,
    is_torch_bfaa_available,
    is_torch_cuda_available,
    is_torch_fx_available,
    is_torch_fx_proxy,
    is_torch_mps_available,
    is_torch_tfaa_available,
    is_torch_tpu_available,
    is_torchaudio_available,
    is_training_run_on_sagemaker,
    is_vision_available,
    replace_return_docstrings,
    requires_backends,
    to_numpy,
    to_py_obj,
    torch_only_method,
)

