


import argparse
import json
from pathlib import Path

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DeiTImageProcessor, ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase__							:						int							     =   logging.get_logger(__name__)







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase=False				)    ->				int:
  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'''),
            ('''pre_logits.fc.weight''', '''pooler.dense.weight'''),
            ('''pre_logits.fc.bias''', '''pooler.dense.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 SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase=False				)    ->				Dict:
  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 SCREAMING_SNAKE_CASE	(			__lowerCAmelCase				)    ->				Optional[Any]:
  snake_case__						= ['''head.weight''', '''head.bias''']
  for k in ignore_keys:
    state_dict.pop(__lowerCAmelCase							,   __lowerCAmelCase				)







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase				)    ->				Union[str, Any]:
  snake_case__						= dct.pop(__lowerCAmelCase				)
  snake_case__						= val







def SCREAMING_SNAKE_CASE	(			)    ->				str:
  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 SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase				)    ->				Dict:
  snake_case__						= ViTConfig()
  snake_case__						= False
  # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
  if vit_name[-5:] == "in21k":
    snake_case__						= True
    snake_case__						= int(vit_name[-12:-10]				)
    snake_case__						= int(vit_name[-9:-6]				)
  else:
    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()}
    snake_case__						= int(vit_name[-6:-4]				)
    snake_case__						= int(vit_name[-3:]				)
  # size of the architecture
  if "deit" in vit_name:
    if vit_name[9:].startswith('''tiny'''				):
      snake_case__						= 192
      snake_case__						= 768
      snake_case__						= 12
      snake_case__						= 3
    elif vit_name[9:].startswith('''small'''				):
      snake_case__						= 384
      snake_case__						= 1536
      snake_case__						= 12
      snake_case__						= 6
    else:
      pass
  else:
    if vit_name[4:].startswith('''small'''				):
      snake_case__						= 768
      snake_case__						= 2304
      snake_case__						= 8
      snake_case__						= 8
    elif vit_name[4:].startswith('''base'''				):
      pass
    elif vit_name[4:].startswith('''large'''				):
      snake_case__						= 1024
      snake_case__						= 4096
      snake_case__						= 24
      snake_case__						= 16
    elif vit_name[4:].startswith('''huge'''				):
      snake_case__						= 1280
      snake_case__						= 5120
      snake_case__						= 32
      snake_case__						= 16

    # load original model from timm
  snake_case__						= timm.create_model(__lowerCAmelCase							,   pretrained=__lowerCAmelCase				)
  timm_model.eval()

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

  # load HuggingFace model
  if vit_name[-5:] == "in21k":
    snake_case__						= ViTModel(__lowerCAmelCase				).eval()
  else:
    snake_case__						= ViTForImageClassification(__lowerCAmelCase				).eval()
  model.load_state_dict(__lowerCAmelCase				)

  # Check outputs on an image, prepared by ViTImageProcessor/DeiTImageProcessor
  if "deit" in vit_name:
    snake_case__						= DeiTImageProcessor(size=config.image_size				)
  else:
    snake_case__						= ViTImageProcessor(size=config.image_size				)
  snake_case__						= image_processor(images=prepare_img()							,   return_tensors='''pt'''				)
  snake_case__						= encoding['''pixel_values''']
  snake_case__						= model(__lowerCAmelCase				)

  if base_model:
    snake_case__						= timm_model.forward_features(__lowerCAmelCase				)
    assert timm_pooled_output.shape == outputs.pooler_output.shape
    assert torch.allclose(__lowerCAmelCase							,   outputs.pooler_output							,   atol=1e-3				)
  else:
    snake_case__						= timm_model(__lowerCAmelCase				)
    assert timm_logits.shape == outputs.logits.shape
    assert torch.allclose(__lowerCAmelCase							,   outputs.logits							,   atol=1e-3				)

  Path(__lowerCAmelCase				).mkdir(exist_ok=__lowerCAmelCase				)
  print(F"""Saving model {vit_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__							:						Optional[Any]							     =   argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       """--vit_name""",
       default="""vit_base_patch16_224""",
       type=str,
       help="""Name of the ViT timm model you'd like to convert.""",
   )
   parser.add_argument(
       """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
   )

   lowerCamelCase__							:						str							     =   parser.parse_args()
   convert_vit_checkpoint(args.vit_name, args.pytorch_dump_folder_path)








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,
)


SCREAMING_SNAKE_CASE_   :      str			=   {
    '''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  SCREAMING_SNAKE_CASE				(      snake_case							)		->    Tuple:
		__lowercase											=     {}
		state_dict.pop('pixel_mean'						,   snake_case							)
		state_dict.pop('pixel_std'						,   snake_case							)

		__lowercase											=     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:
								__lowercase											=     key.replace(snake_case						,   snake_case							)

				if re.match(snake_case						,   snake_case							):
						__lowercase											=     int(re.match(snake_case						,   snake_case							).group(2							)							)
						if layer_nb == 0:
								__lowercase											=     key.replace('layers.0'						,   'proj_in'							)
						elif layer_nb == 1:
								__lowercase											=     key.replace('layers.1'						,   'layers.0'							)
						elif layer_nb == 2:
								__lowercase											=     key.replace('layers.2'						,   'proj_out'							)

				__lowercase											=     value

		__lowercase											=     model_state_dict[
		    'prompt_encoder.shared_embedding.positional_embedding'
		]

		return model_state_dict



def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case						,   snake_case						,   snake_case="ybelkada/segment-anything"							)		->    int:
		__lowercase											=     hf_hub_download(snake_case						,   F"checkpoints/{model_name}.pth"							)

		if "sam_vit_b" in model_name:
				__lowercase											=     SamConfig()
		elif "sam_vit_l" in model_name:
				__lowercase											=     SamVisionConfig(
				    hidden_size=1_024						,   num_hidden_layers=24						,   num_attention_heads=16						,   global_attn_indexes=[5, 11, 17, 23]						,   )

				__lowercase											=     SamConfig(
				    vision_config=snake_case						,   )
		elif "sam_vit_h" in model_name:
				__lowercase											=     SamVisionConfig(
				    hidden_size=1_280						,   num_hidden_layers=32						,   num_attention_heads=16						,   global_attn_indexes=[7, 15, 23, 31]						,   )

				__lowercase											=     SamConfig(
				    vision_config=snake_case						,   )

		__lowercase											=     torch.load(snake_case						,   map_location='cpu'							)
		__lowercase											=     replace_keys(snake_case							)

		__lowercase											=     SamImageProcessor()

		__lowercase											=     SamProcessor(image_processor=snake_case							)
		__lowercase											=     SamModel(snake_case							)

		hf_model.load_state_dict(snake_case							)
		__lowercase											=     hf_model.to('cuda'							)

		__lowercase											=     'https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png'
		__lowercase											=     Image.open(requests.get(snake_case						,   stream=snake_case							).raw							).convert('RGB'							)

		__lowercase											=     [[[400, 650]]]
		__lowercase											=     [[1]]

		__lowercase											=     processor(images=np.array(snake_case							)						,   return_tensors='pt'							).to('cuda'							)

		with torch.no_grad():
				__lowercase											=     hf_model(**snake_case							)
		__lowercase											=     output.iou_scores.squeeze()

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

				__lowercase											=     processor(
				    images=np.array(snake_case							)						,   input_points=snake_case						,   input_labels=snake_case						,   return_tensors='pt'							).to('cuda'							)

				with torch.no_grad():
						__lowercase											=     hf_model(**snake_case							)
				__lowercase											=     output.iou_scores.squeeze()

				assert scores[-1].item() == 0.9712_6030_9219_3604

				__lowercase											=     ((75, 275, 1_725, 850),)

				__lowercase											=     processor(images=np.array(snake_case							)						,   input_boxes=snake_case						,   return_tensors='pt'							).to('cuda'							)

				with torch.no_grad():
						__lowercase											=     hf_model(**snake_case							)
				__lowercase											=     output.iou_scores.squeeze()

				assert scores[-1].item() == 0.8686_0156_0592_6514

				# Test with 2 points and 1 image.
				__lowercase											=     [[[400, 650], [800, 650]]]
				__lowercase											=     [[1, 1]]

				__lowercase											=     processor(
				    images=np.array(snake_case							)						,   input_points=snake_case						,   input_labels=snake_case						,   return_tensors='pt'							).to('cuda'							)

				with torch.no_grad():
						__lowercase											=     hf_model(**snake_case							)
				__lowercase											=     output.iou_scores.squeeze()

				assert scores[-1].item() == 0.9936_0477_9243_4692


if __name__ == "__main__":
		SCREAMING_SNAKE_CASE_   :      Optional[Any]			=   argparse.ArgumentParser()
		SCREAMING_SNAKE_CASE_   :      Union[str, Any]			=   ['''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''',
		)

		SCREAMING_SNAKE_CASE_   :      List[Any]			=   parser.parse_args()

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




