


"""simple docstring"""





import argparse
import re

import numpy as np
import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SamConfig,
    SamImageProcessor,
    SamModel,
    SamProcessor,
    SamVisionConfig,
)


__UpperCAmelCase     					=   {
    'iou_prediction_head.layers.0': 'iou_prediction_head.proj_in',
    'iou_prediction_head.layers.1': 'iou_prediction_head.layers.0',
    'iou_prediction_head.layers.2': 'iou_prediction_head.proj_out',
    'mask_decoder.output_upscaling.0': 'mask_decoder.upscale_conv1',
    'mask_decoder.output_upscaling.1': 'mask_decoder.upscale_layer_norm',
    'mask_decoder.output_upscaling.3': 'mask_decoder.upscale_conv2',
    'mask_downscaling.0': 'mask_embed.conv1',
    'mask_downscaling.1': 'mask_embed.layer_norm1',
    'mask_downscaling.3': 'mask_embed.conv2',
    'mask_downscaling.4': 'mask_embed.layer_norm2',
    'mask_downscaling.6': 'mask_embed.conv3',
    'point_embeddings': 'point_embed',
    'pe_layer.positional_encoding_gaussian_matrix': 'shared_embedding.positional_embedding',
    'image_encoder': 'vision_encoder',
    'neck.0': 'neck.conv1',
    'neck.1': 'neck.layer_norm1',
    'neck.2': 'neck.conv2',
    'neck.3': 'neck.layer_norm2',
    'patch_embed.proj': 'patch_embed.projection',
    '.norm': '.layer_norm',
    'blocks': 'layers',
}





def      lowercase__						(	lowerCamelCase  :	str					)       ->  Any:
  lowerCAmelCase__     :     Dict										=       {}
  state_dict.pop("pixel_mean"					,						_lowerCAmelCase					)
  state_dict.pop("pixel_std"					,						_lowerCAmelCase					)

  lowerCAmelCase__     :     Any										=       R".*.output_hypernetworks_mlps.(\d+).layers.(\d+).*"

  for key, value in state_dict.items():
    for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
      if key_to_modify in key:
        lowerCAmelCase__     :     List[str]										=       key.replace(_lowerCAmelCase					,						_lowerCAmelCase					)

    if re.match(_lowerCAmelCase					,						_lowerCAmelCase					):
      lowerCAmelCase__     :     Union[str, Any]										=       int(re.match(_lowerCAmelCase					,						_lowerCAmelCase					).group(2					)					)
      if layer_nb == 0:
        lowerCAmelCase__     :     Optional[Any]										=       key.replace("layers.0"					,						"proj_in"					)
      elif layer_nb == 1:
        lowerCAmelCase__     :     List[str]										=       key.replace("layers.1"					,						"layers.0"					)
      elif layer_nb == 2:
        lowerCAmelCase__     :     Any										=       key.replace("layers.2"					,						"proj_out"					)

    lowerCAmelCase__     :     str										=       value

  lowerCAmelCase__     :     Optional[int]										=       model_state_dict[
      "prompt_encoder.shared_embedding.positional_embedding"
  ]

  return model_state_dict





def      lowercase__						(	lowerCamelCase  :	Tuple					,						lowerCamelCase  :	Optional[Any]					,						lowerCamelCase  :	Tuple					,						lowerCamelCase  :	List[Any]="ybelkada/segment-anything"					)       ->  Any:
  lowerCAmelCase__     :     Union[str, Any]										=       hf_hub_download(_lowerCAmelCase					,						F"checkpoints/{model_name}.pth"					)

  if "sam_vit_b" in model_name:
    lowerCAmelCase__     :     Dict										=       SamConfig()
  elif "sam_vit_l" in model_name:
    lowerCAmelCase__     :     Tuple										=       SamVisionConfig(
        hidden_size=1_0_2_4					,						num_hidden_layers=2_4					,						num_attention_heads=1_6					,						global_attn_indexes=[5, 1_1, 1_7, 2_3]					,						)

    lowerCAmelCase__     :     Union[str, Any]										=       SamConfig(
        vision_config=_lowerCAmelCase					,						)
  elif "sam_vit_h" in model_name:
    lowerCAmelCase__     :     Dict										=       SamVisionConfig(
        hidden_size=1_2_8_0					,						num_hidden_layers=3_2					,						num_attention_heads=1_6					,						global_attn_indexes=[7, 1_5, 2_3, 3_1]					,						)

    lowerCAmelCase__     :     int										=       SamConfig(
        vision_config=_lowerCAmelCase					,						)

  lowerCAmelCase__     :     Optional[Any]										=       torch.load(_lowerCAmelCase					,						map_location="cpu"					)
  lowerCAmelCase__     :     str										=       replace_keys(_lowerCAmelCase					)

  lowerCAmelCase__     :     Union[str, Any]										=       SamImageProcessor()

  lowerCAmelCase__     :     Dict										=       SamProcessor(image_processor=_lowerCAmelCase					)
  lowerCAmelCase__     :     Tuple										=       SamModel(_lowerCAmelCase					)

  hf_model.load_state_dict(_lowerCAmelCase					)
  lowerCAmelCase__     :     str										=       hf_model.to("cuda"					)

  lowerCAmelCase__     :     Optional[Any]										=       "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
  lowerCAmelCase__     :     str										=       Image.open(requests.get(_lowerCAmelCase					,						stream=_lowerCAmelCase					).raw					).convert("RGB"					)

  lowerCAmelCase__     :     Any										=       [[[4_0_0, 6_5_0]]]
  lowerCAmelCase__     :     Any										=       [[1]]

  lowerCAmelCase__     :     Dict										=       processor(images=np.array(_lowerCAmelCase					)					,						return_tensors="pt"					).to("cuda"					)

  with torch.no_grad():
    lowerCAmelCase__     :     Optional[Any]										=       hf_model(**_lowerCAmelCase					)
  lowerCAmelCase__     :     Dict										=       output.iou_scores.squeeze()

  if model_name == "sam_vit_h_4b8939":
    assert scores[-1].item() == 0.5_79_89_02_51_15_96_68

    lowerCAmelCase__     :     Tuple										=       processor(
        images=np.array(_lowerCAmelCase					)					,						input_points=_lowerCAmelCase					,						input_labels=_lowerCAmelCase					,						return_tensors="pt"					).to("cuda"					)

    with torch.no_grad():
      lowerCAmelCase__     :     Any										=       hf_model(**_lowerCAmelCase					)
    lowerCAmelCase__     :     Tuple										=       output.iou_scores.squeeze()

    assert scores[-1].item() == 0.97_12_60_30_92_19_36_04

    lowerCAmelCase__     :     Dict										=       ((7_5, 2_7_5, 1_7_2_5, 8_5_0),)

    lowerCAmelCase__     :     Any										=       processor(images=np.array(_lowerCAmelCase					)					,						input_boxes=_lowerCAmelCase					,						return_tensors="pt"					).to("cuda"					)

    with torch.no_grad():
      lowerCAmelCase__     :     List[Any]										=       hf_model(**_lowerCAmelCase					)
    lowerCAmelCase__     :     int										=       output.iou_scores.squeeze()

    assert scores[-1].item() == 0.86_86_01_56_05_92_65_14

    # Test with 2 points and 1 image.
    lowerCAmelCase__     :     Union[str, Any]										=       [[[4_0_0, 6_5_0], [8_0_0, 6_5_0]]]
    lowerCAmelCase__     :     int										=       [[1, 1]]

    lowerCAmelCase__     :     Tuple										=       processor(
        images=np.array(_lowerCAmelCase					)					,						input_points=_lowerCAmelCase					,						input_labels=_lowerCAmelCase					,						return_tensors="pt"					).to("cuda"					)

    with torch.no_grad():
      lowerCAmelCase__     :     str										=       hf_model(**_lowerCAmelCase					)
    lowerCAmelCase__     :     Dict										=       output.iou_scores.squeeze()

    assert scores[-1].item() == 0.99_36_04_77_92_43_46_92


if __name__ == "__main__":
   __UpperCAmelCase     					=   argparse.ArgumentParser()
   __UpperCAmelCase     					=   ['sam_vit_b_01ec64', 'sam_vit_h_4b8939', 'sam_vit_l_0b3195']
   parser.add_argument(
       "--model_name",
       default="sam_vit_h_4b8939",
       choices=choices,
       type=str,
       help="Path to hf config.json of model to convert",
   )
   parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
   parser.add_argument(
       "--push_to_hub",
       action="store_true",
       help="Whether to push the model and processor to the hub after converting",
   )
   parser.add_argument(
       "--model_hub_id",
       default="ybelkada/segment-anything",
       choices=choices,
       type=str,
       help="Path to hf config.json of model to convert",
   )

   __UpperCAmelCase     					=   parser.parse_args()

   convert_sam_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub, args.model_hub_id)




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()
A				:       int      =  logging.get_logger(__name__)





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


						'''simple docstring'''






						__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"
												__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							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase=False			)    ->						int:


						'''simple docstring'''






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





def 						_lowerCAmelCase							(     _lowerCAmelCase			)    ->						List[str]:


						'''simple docstring'''






						__snake_case				=		["head.weight", "head.bias"]
						for k in ignore_keys:
												state_dict.pop(_lowerCAmelCase   ,						_lowerCAmelCase			)





def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase			)    ->						Dict:


						'''simple docstring'''






						__snake_case				=		dct.pop(_lowerCAmelCase			)
						__snake_case				=		val





def 						_lowerCAmelCase							(     )    ->						Any:


						'''simple docstring'''






						__snake_case				=		"http://images.cocodataset.org/val2017/000000039769.jpg"
						__snake_case				=		Image.open(requests.get(_lowerCAmelCase   ,						stream=_lowerCAmelCase			).raw			)
						return im





@torch.no_grad()
def 						_lowerCAmelCase							(     _lowerCAmelCase   ,						_lowerCAmelCase   ,						_lowerCAmelCase=True			)    ->						int:


						'''simple docstring'''






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

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

						# load state_dict of original model, remove and rename some keys
						__snake_case				=		original_model.state_dict()
						if base_model:
												remove_classification_head_(_lowerCAmelCase			)
						__snake_case				=		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:
												__snake_case				=		ViTModel(_lowerCAmelCase   ,						add_pooling_layer=_lowerCAmelCase			).eval()
						else:
												__snake_case				=		ViTForImageClassification(_lowerCAmelCase			).eval()
						model.load_state_dict(_lowerCAmelCase			)

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

						if base_model:
												__snake_case				=		original_model(_lowerCAmelCase			)
												assert torch.allclose(_lowerCAmelCase   ,						outputs.last_hidden_state[:, 0, :]   ,						atol=1E-1			)
						else:
												__snake_case				=		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__":
	A				:       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)
	A				:       Any      =  parser.parse_args()
	convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)





