





'''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 		UpperCamelCase							(				_lowerCamelCase   :   List[Any]		,			_lowerCamelCase   :   Any=False			):
    A__							=     []
    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"
        A__							=     [(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 		UpperCamelCase							(				_lowerCamelCase   :   str		,			_lowerCamelCase   :   int		,			_lowerCamelCase   :   Optional[int]=False			):
    for i in range(config.num_hidden_layers			):
        if base_model:
            A__							=     ""
        else:
            A__							=     "vit."
        # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
        A__							=     state_dict.pop(F"blocks.{i}.attn.qkv.weight"			)
        A__							=     state_dict.pop(F"blocks.{i}.attn.qkv.bias"			)
        # next, add query, keys and values (in that order) to the state dict
        A__							=     in_proj_weight[
            : config.hidden_size, :
        ]
        A__							=     in_proj_bias[: config.hidden_size]
        A__							=     in_proj_weight[
            config.hidden_size : config.hidden_size * 2, :
        ]
        A__							=     in_proj_bias[
            config.hidden_size : config.hidden_size * 2
        ]
        A__							=     in_proj_weight[
            -config.hidden_size :, :
        ]
        A__							=     in_proj_bias[-config.hidden_size :]





def 		UpperCamelCase							(				_lowerCamelCase   :   List[str]			):
    A__							=     ["head.weight", "head.bias"]
    for k in ignore_keys:
        state_dict.pop(_lowerCamelCase		,			_lowerCamelCase			)





def 		UpperCamelCase							(				_lowerCamelCase   :   str		,			_lowerCamelCase   :   Any		,			_lowerCamelCase   :   Optional[int]			):
    A__							=     dct.pop(_lowerCamelCase			)
    A__							=     val





def 		UpperCamelCase							(				):
    A__							=     "http://images.cocodataset.org/val2017/000000039769.jpg"
    A__							=     Image.open(requests.get(_lowerCamelCase		,			stream=_lowerCamelCase			).raw			)
    return im





@torch.no_grad()
def 		UpperCamelCase							(				_lowerCamelCase   :   str		,			_lowerCamelCase   :   Optional[int]		,			_lowerCamelCase   :   Dict=True			):
    A__							=     ViTConfig()
    # patch_size
    if model_name[-1] == "8":
        A__							=     8
    # set labels if required
    if not base_model:
        A__							=     10_00
        A__							=     "huggingface/label-files"
        A__							=     "imagenet-1k-id2label.json"
        A__							=     json.load(open(hf_hub_download(_lowerCamelCase		,			_lowerCamelCase		,			repo_type="dataset"			)		,			"r"			)			)
        A__							=     {int(_lowerCamelCase			): v for k, v in idalabel.items()}
        A__							=     idalabel
        A__							=     {v: k for k, v in idalabel.items()}
    # size of the architecture
    if model_name in ["dino_vits8", "dino_vits16"]:
        A__							=     3_84
        A__							=     15_36
        A__							=     12
        A__							=     6

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

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

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

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

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

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


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










'''simple docstring'''

import os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
  import torch
  from torch import nn

  from transformers import (
      Adafactor,
      AdamW,
      get_constant_schedule,
      get_constant_schedule_with_warmup,
      get_cosine_schedule_with_warmup,
      get_cosine_with_hard_restarts_schedule_with_warmup,
      get_inverse_sqrt_schedule,
      get_linear_schedule_with_warmup,
      get_polynomial_decay_schedule_with_warmup,
  )





def 		UpperCamelCase							(				_lowerCamelCase   :   Any		,			_lowerCamelCase   :   Optional[Any]=10			):
    A__							=     []
    for _ in range(_lowerCamelCase			):
        lrs.append(scheduler.get_lr()[0]			)
        scheduler.step()
    return lrs





def 		UpperCamelCase							(				_lowerCamelCase   :   Tuple		,			_lowerCamelCase   :   List[str]=10			):
    A__							=     []
    for step in range(_lowerCamelCase			):
        lrs.append(scheduler.get_lr()[0]			)
        scheduler.step()
        if step == num_steps // 2:
            with tempfile.TemporaryDirectory() as tmpdirname:
                A__							=     os.path.join(_lowerCamelCase		,			"schedule.bin"			)
                torch.save(scheduler.state_dict()		,			_lowerCamelCase			)

                A__							=     torch.load(_lowerCamelCase			)
                scheduler.load_state_dict(_lowerCamelCase			)
    return lrs






@require_torch
class   UpperCAmelCase      (			unittest.TestCase							):
 def 	UpperCAmelCase_				(							self :int		,			lowercase_ :Any		,			lowercase_ :List[Any]		,			lowercase_ :Optional[Any]				)->				int:
     self.assertEqual(len(lowercase_				)		,			len(lowercase_				)				)
     for a, b in zip(lowercase_		,			lowercase_				):
         self.assertAlmostEqual(lowercase_		,			lowercase_		,			delta=lowercase_				)
 def 	UpperCAmelCase_				(							self :Union[str, Any]				)->				Dict:
     A__							=     torch.tensor([0.1, -0.2, -0.1]		,			requires_grad=lowercase_				)
     A__							=     torch.tensor([0.4, 0.2, -0.5]				)
     A__							=     nn.MSELoss()
     # No warmup, constant schedule, no gradient clipping
     A__							=     AdamW(params=[w]		,			lr=2E-1		,			weight_decay=0.0				)
     for _ in range(1_00				):
         A__							=     criterion(lowercase_		,			lowercase_				)
         loss.backward()
         optimizer.step()
         w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
         w.grad.zero_()
     self.assertListAlmostEqual(w.tolist()		,			[0.4, 0.2, -0.5]		,			tol=1E-2				)



 def 	UpperCAmelCase_				(							self :Tuple				)->				List[str]:
     A__							=     torch.tensor([0.1, -0.2, -0.1]		,			requires_grad=lowercase_				)
     A__							=     torch.tensor([0.4, 0.2, -0.5]				)
     A__							=     nn.MSELoss()
     # No warmup, constant schedule, no gradient clipping
     A__							=     Adafactor(
         params=[w]		,			lr=1E-2		,			eps=(1E-30, 1E-3)		,			clip_threshold=1.0		,			decay_rate=-0.8		,			betaa=lowercase_		,			weight_decay=0.0		,			relative_step=lowercase_		,			scale_parameter=lowercase_		,			warmup_init=lowercase_		,			)
     for _ in range(10_00				):
         A__							=     criterion(lowercase_		,			lowercase_				)
         loss.backward()
         optimizer.step()
         w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
         w.grad.zero_()
     self.assertListAlmostEqual(w.tolist()		,			[0.4, 0.2, -0.5]		,			tol=1E-2				)






@require_torch
class   UpperCAmelCase      (			unittest.TestCase							):
 __lowercase  						=							nn.Linear(50     ,			50							) if is_torch_available() else None
 __lowercase  						=							AdamW(m.parameters()     ,			lr=1_0.0							) if is_torch_available() else None
 __lowercase  						=							10
 def 	UpperCAmelCase_				(							self :Tuple		,			lowercase_ :Any		,			lowercase_ :List[Any]		,			lowercase_ :List[Any]		,			lowercase_ :str=None				)->				Optional[int]:
     self.assertEqual(len(lowercase_				)		,			len(lowercase_				)				)
     for a, b in zip(lowercase_		,			lowercase_				):
         self.assertAlmostEqual(lowercase_		,			lowercase_		,			delta=lowercase_		,			msg=lowercase_				)



 def 	UpperCAmelCase_				(							self :Optional[Any]				)->				Any:
     A__							=     {"num_warmup_steps": 2, "num_training_steps": 10}
     # schedulers doct format
     # function: (sched_args_dict, expected_learning_rates)
     A__							=     {
         get_constant_schedule: ({}, [1_0.0] * self.num_steps),
         get_constant_schedule_with_warmup: (
             {"num_warmup_steps": 4},
             [0.0, 2.5, 5.0, 7.5, 1_0.0, 1_0.0, 1_0.0, 1_0.0, 1_0.0, 1_0.0],
         ),
         get_linear_schedule_with_warmup: (
             {**common_kwargs},
             [0.0, 5.0, 1_0.0, 8.7_5, 7.5, 6.2_5, 5.0, 3.7_5, 2.5, 1.2_5],
         ),
         get_cosine_schedule_with_warmup: (
             {**common_kwargs},
             [0.0, 5.0, 1_0.0, 9.6_1, 8.5_3, 6.9_1, 5.0, 3.0_8, 1.4_6, 0.3_8],
         ),
         get_cosine_with_hard_restarts_schedule_with_warmup: (
             {**common_kwargs, "num_cycles": 2},
             [0.0, 5.0, 1_0.0, 8.5_3, 5.0, 1.4_6, 1_0.0, 8.5_3, 5.0, 1.4_6],
         ),
         get_polynomial_decay_schedule_with_warmup: (
             {**common_kwargs, "power": 2.0, "lr_end": 1E-7},
             [0.0, 5.0, 1_0.0, 7.6_5_6, 5.6_2_5, 3.9_0_6, 2.5, 1.4_0_6, 0.6_2_5, 0.1_5_6],
         ),
         get_inverse_sqrt_schedule: (
             {"num_warmup_steps": 2},
             [0.0, 5.0, 1_0.0, 8.1_6_5, 7.0_7_1, 6.3_2_5, 5.7_7_4, 5.3_4_5, 5.0, 4.7_1_4],
         ),
     }

     for scheduler_func, data in scheds.items():
         A__,      A__							=     data

         A__							=     scheduler_func(self.optimizer		,			**lowercase_				)
         self.assertEqual(len([scheduler.get_lr()[0]]				)		,			1				)
         A__							=     unwrap_schedule(lowercase_		,			self.num_steps				)
         self.assertListAlmostEqual(
             lowercase_		,			lowercase_		,			tol=1E-2		,			msg=F"failed for {scheduler_func} in normal scheduler"		,			)

         A__							=     scheduler_func(self.optimizer		,			**lowercase_				)
         if scheduler_func.__name__ != "get_constant_schedule":
             LambdaScheduleWrapper.wrap_scheduler(lowercase_				)  # wrap to test picklability of the schedule
         A__							=     unwrap_and_save_reload_schedule(lowercase_		,			self.num_steps				)
         self.assertListEqual(lowercase_		,			lowercase_		,			msg=F"failed for {scheduler_func} in save and reload"				)






class   UpperCAmelCase      :
 def __init__(							self :str		,			lowercase_ :List[str]				)->				Tuple:
     A__							=     fn
 def __call__(							self :List[Any]		,			*lowercase_ :Dict		,			**lowercase_ :Dict				)->				Tuple:
     return self.fn(*lowercase_		,			**lowercase_				)



 @classmethod
 def 	UpperCAmelCase_				(							self :Any		,			lowercase_ :Tuple				)->				List[Any]:
     A__							=     list(map(self		,			scheduler.lr_lambdas				)				)



