





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	:	Optional[Any]		     =       logging.get_logger(__name__)
def      _A       (		SCREAMING_SNAKE_CASE     :   Tuple       ,       SCREAMING_SNAKE_CASE     :   Any=False						):





  """simple docstring"""



  a__	:					Optional[int]										=[]
  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__	:					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      _A       (		SCREAMING_SNAKE_CASE     :   Optional[Any]       ,       SCREAMING_SNAKE_CASE     :   List[str]       ,       SCREAMING_SNAKE_CASE     :   List[Any]=False						):





  """simple docstring"""



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





  """simple docstring"""



  a__	:					Optional[int]										=["head.weight", "head.bias"]
  for k in ignore_keys:
    state_dict.pop(SCREAMING_SNAKE_CASE       ,       SCREAMING_SNAKE_CASE						)
def      _A       (		SCREAMING_SNAKE_CASE     :   Any       ,       SCREAMING_SNAKE_CASE     :   Union[str, Any]       ,       SCREAMING_SNAKE_CASE     :   Dict						):





  """simple docstring"""



  a__	:					int										=dct.pop(SCREAMING_SNAKE_CASE						)
  a__	:					Optional[int]										=val
def      _A       (		):





  """simple docstring"""



  a__	:					List[Any]										="http://images.cocodataset.org/val2017/000000039769.jpg"
  a__	:					Union[str, Any]										=Image.open(requests.get(SCREAMING_SNAKE_CASE       ,       stream=SCREAMING_SNAKE_CASE						).raw						)
  return im
@torch.no_grad()
def      _A       (		SCREAMING_SNAKE_CASE     :   str       ,       SCREAMING_SNAKE_CASE     :   Any       ,       SCREAMING_SNAKE_CASE     :   Tuple=True						):





  """simple docstring"""



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

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

  # load state_dict of original model, remove and rename some keys
  a__	:					List[Any]										=original_model.state_dict()
  if base_model:
    remove_classification_head_(SCREAMING_SNAKE_CASE						)
  a__	:					Optional[Any]										=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:
    a__	:					Optional[int]										=ViTModel(SCREAMING_SNAKE_CASE       ,       add_pooling_layer=SCREAMING_SNAKE_CASE						).eval()
  else:
    a__	:					str										=ViTForImageClassification(SCREAMING_SNAKE_CASE						).eval()
  model.load_state_dict(SCREAMING_SNAKE_CASE						)

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

  if base_model:
    a__	:					List[str]										=original_model(SCREAMING_SNAKE_CASE						)
    assert torch.allclose(SCREAMING_SNAKE_CASE       ,       outputs.last_hidden_state[:, 0, :]       ,       atol=1e-1						)
  else:
    a__	:					Optional[int]										=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	:	List[str]		     =       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"""





from math import isqrt, loga



def 						_lowerCAmelCase				(				UpperCamelCase_				):
		__SCREAMING_SNAKE_CASE            = [True] * max_number
		for i in range(2  ,				isqrt(max_number - 1				) + 1				):
				if is_prime[i]:
						for j in range(i**2  ,				UpperCamelCase_  ,				UpperCamelCase_				):
								__SCREAMING_SNAKE_CASE            = False

		return [i for i in range(2  ,				UpperCamelCase_				) if is_prime[i]]



def 						_lowerCAmelCase				(				UpperCamelCase_ = 80_0800  ,				UpperCamelCase_ = 80_0800				):
		__SCREAMING_SNAKE_CASE            = degree * loga(UpperCamelCase_				)
		__SCREAMING_SNAKE_CASE            = int(UpperCamelCase_				)
		__SCREAMING_SNAKE_CASE            = calculate_prime_numbers(UpperCamelCase_				)

		__SCREAMING_SNAKE_CASE            = 0
		__SCREAMING_SNAKE_CASE            = 0
		__SCREAMING_SNAKE_CASE            = len(UpperCamelCase_				) - 1
		while left < right:
				while (
				    prime_numbers[right] * loga(prime_numbers[left]				)
				    + prime_numbers[left] * loga(prime_numbers[right]				)
				    > upper_bound
				):
						right -= 1
				hybrid_integers_count += right - left
				left += 1

		return hybrid_integers_count


if __name__ == "__main__":
							print(F"""{solution() = }""")




