

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


_lowerCAmelCase						=		{
    """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 		lowerCamelCase__		(					_lowerCamelCase   ):

    '''simple docstring'''



    _lowerCAmelCase       :   Optional[Any]  =    {}
    state_dict.pop('pixel_mean' ,   _lowerCamelCase   )
    state_dict.pop('pixel_std' ,   _lowerCamelCase   )

    _lowerCAmelCase       :   Optional[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       :   Optional[int]  =    key.replace(_lowerCamelCase ,   _lowerCamelCase   )

        if re.match(_lowerCamelCase ,   _lowerCamelCase   ):
            _lowerCAmelCase       :   Optional[int]  =    int(re.match(_lowerCamelCase ,   _lowerCamelCase   ).group(2   )   )
            if layer_nb == 0:
                _lowerCAmelCase       :   List[str]  =    key.replace('layers.0' ,   'proj_in'   )
            elif layer_nb == 1:
                _lowerCAmelCase       :   List[str]  =    key.replace('layers.1' ,   'layers.0'   )
            elif layer_nb == 2:
                _lowerCAmelCase       :   Dict  =    key.replace('layers.2' ,   'proj_out'   )

        _lowerCAmelCase       :   Tuple  =    value

    _lowerCAmelCase       :   int  =    model_state_dict[
        'prompt_encoder.shared_embedding.positional_embedding'
    ]

    return model_state_dict

def 		lowerCamelCase__		(					_lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase ,   _lowerCamelCase="ybelkada/segment-anything"   ):

    '''simple docstring'''



    _lowerCAmelCase       :   Optional[int]  =    hf_hub_download(_lowerCamelCase ,   f"""checkpoints/{model_name}.pth"""   )

    if "sam_vit_b" in model_name:
        _lowerCAmelCase       :   Optional[int]  =    SamConfig()
    elif "sam_vit_l" in model_name:
        _lowerCAmelCase       :   Any  =    SamVisionConfig(
            hidden_size=1024 ,   num_hidden_layers=24 ,   num_attention_heads=16 ,   global_attn_indexes=[5, 11, 17, 23] ,   )

        _lowerCAmelCase       :   Union[str, Any]  =    SamConfig(
            vision_config=_lowerCamelCase ,   )
    elif "sam_vit_h" in model_name:
        _lowerCAmelCase       :   Any  =    SamVisionConfig(
            hidden_size=1280 ,   num_hidden_layers=32 ,   num_attention_heads=16 ,   global_attn_indexes=[7, 15, 23, 31] ,   )

        _lowerCAmelCase       :   str  =    SamConfig(
            vision_config=_lowerCamelCase ,   )

    _lowerCAmelCase       :   Tuple  =    torch.load(_lowerCamelCase ,   map_location='cpu'   )
    _lowerCAmelCase       :   Optional[Any]  =    replace_keys(_lowerCamelCase   )

    _lowerCAmelCase       :   Dict  =    SamImageProcessor()

    _lowerCAmelCase       :   Optional[Any]  =    SamProcessor(image_processor=_lowerCamelCase   )
    _lowerCAmelCase       :   List[str]  =    SamModel(_lowerCamelCase   )

    hf_model.load_state_dict(_lowerCamelCase   )
    _lowerCAmelCase       :   Union[str, Any]  =    hf_model.to('cuda'   )

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

    _lowerCAmelCase       :   Optional[int]  =    [[[400, 650]]]
    _lowerCAmelCase       :   Union[str, Any]  =    [[1]]

    _lowerCAmelCase       :   Optional[int]  =    processor(images=np.array(_lowerCamelCase   ) ,   return_tensors='pt'   ).to('cuda'   )

    with torch.no_grad():
        _lowerCAmelCase       :   Union[str, Any]  =    hf_model(**_lowerCamelCase   )
    _lowerCAmelCase       :   Tuple  =    output.iou_scores.squeeze()

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

        _lowerCAmelCase       :   List[Any]  =    processor(
            images=np.array(_lowerCamelCase   ) ,   input_points=_lowerCamelCase ,   input_labels=_lowerCamelCase ,   return_tensors='pt'   ).to('cuda'   )

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

        assert scores[-1].item() == 0.9712603092193604

        _lowerCAmelCase       :   List[str]  =    ((75, 275, 1725, 850),)

        _lowerCAmelCase       :   List[Any]  =    processor(images=np.array(_lowerCamelCase   ) ,   input_boxes=_lowerCamelCase ,   return_tensors='pt'   ).to('cuda'   )

        with torch.no_grad():
            _lowerCAmelCase       :   Tuple  =    hf_model(**_lowerCamelCase   )
        _lowerCAmelCase       :   List[str]  =    output.iou_scores.squeeze()

        assert scores[-1].item() == 0.8686015605926514

        # Test with 2 points and 1 image.
        _lowerCAmelCase       :   Any  =    [[[400, 650], [800, 650]]]
        _lowerCAmelCase       :   Dict  =    [[1, 1]]

        _lowerCAmelCase       :   str  =    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.9936047792434692


if __name__ == "__main__":
 _lowerCAmelCase						=		argparse.ArgumentParser()
 _lowerCAmelCase						=		["""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""",
 )

 _lowerCAmelCase						=		parser.parse_args()

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








"""simple docstring"""






import json
import os
from typing import Dict, List, Optional, Tuple

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


_lowerCAmelCase						=		logging.get_logger(__name__)


_lowerCAmelCase						=		{
    """vocab_file""": """vocab.json""",
    """tokenizer_config_file""": """tokenizer_config.json""",
    """merges_file""": """merges.txt""",
}

_lowerCAmelCase						=		{
    """vocab_file""": {
        """facebook/s2t-wav2vec2-large-en-de""": (
            """https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/vocab.json"""
        ),
    },
    """tokenizer_config_file""": {
        """facebook/s2t-wav2vec2-large-en-de""": (
            """https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/tokenizer_config.json"""
        ),
    },
    """merges_file""": {
        """facebook/s2t-wav2vec2-large-en-de""": (
            """https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/merges.txt"""
        ),
    },
}

_lowerCAmelCase						=		"""</w>"""
_lowerCAmelCase						=		"""@@ """

def 		lowerCamelCase__		(					_lowerCamelCase   ):

    '''simple docstring'''



    _lowerCAmelCase       :   List[str]  =    set()
    _lowerCAmelCase       :   Dict  =    word[0]
    for char in word[1:]:
        pairs.add((prev_char, char)   )
        _lowerCAmelCase       :   Any  =    char
    return pairs


# Speech2Text2 has no max input length
_lowerCAmelCase						=		{"""facebook/s2t-wav2vec2-large-en-de""": 1_0_2_4}







class 	__UpperCamelCase    (    a__						):
  _UpperCAmelCase          =		VOCAB_FILES_NAMES
  _UpperCAmelCase          =		PRETRAINED_VOCAB_FILES_MAP
  _UpperCAmelCase          =		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  _UpperCAmelCase          =		["input_ids", "attention_mask"]

  def __init__(		self    ,_A    ,_A="<s>"    ,_A="<pad>"    ,_A="</s>"    ,_A="<unk>"    ,_A=False    ,_A=None    ,**_A    ,):
      '''simple docstring'''

      super().__init__(
          unk_token=_A    ,bos_token=_A    ,eos_token=_A    ,pad_token=_A    ,do_lower_case=_A    ,**_A    ,)

      _lowerCAmelCase       :   List[Any]  =    do_lower_case

      with open(_A    ,encoding='utf-8'      ) as vocab_handle:
          _lowerCAmelCase       :   Optional[int]  =    json.load(_A      )
      _lowerCAmelCase       :   Tuple  =    {v: k for k, v in self.encoder.items()}

      if merges_file is None:
          logger.info(F"""No merges files provided. {self.__class__.__name__} can only be used for decoding."""      )

          _lowerCAmelCase       :   Optional[Any]  =    None
          _lowerCAmelCase       :   Tuple  =    None
      else:
          with open(_A    ,encoding='utf-8'      ) as merges_handle:
              _lowerCAmelCase       :   Optional[Any]  =    merges_handle.read().split('\n'      )[:-1]

          _lowerCAmelCase       :   List[str]  =    [tuple(merge.split()[:2]      ) for merge in merges]
          _lowerCAmelCase       :   List[Any]  =    dict(zip(_A    ,range(len(_A      )      )      )      )
          _lowerCAmelCase       :   Union[str, Any]  =    {}

  @property
  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      return len(self.decoder      )

  def __lowerCamelCase		(		self      ):
      '''simple docstring'''

      return dict(self.encoder    ,**self.added_tokens_encoder      )

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   str  =    tuple(token[:-1]      ) + (token[-1] + BPE_TOKEN_MERGES,)
      if token in self.cache:
          return self.cache[token]
      _lowerCAmelCase       :   str  =    get_pairs(_A      )

      if not pairs:
          return token

      while True:
          _lowerCAmelCase       :   List[str]  =    min(_A    ,key=lambda _A      : self.bpe_ranks.get(_A    ,float('inf'      )      )      )
          if bigram not in self.bpe_ranks:
              break
          _lowerCAmelCase,					_lowerCAmelCase       :   Optional[int]  =    bigram
          _lowerCAmelCase       :   Union[str, Any]  =    []
          _lowerCAmelCase       :   Dict  =    0
          while i < len(_A      ):
              try:
                  _lowerCAmelCase       :   Dict  =    word.index(_A    ,_A      )
              except ValueError:
                  new_word.extend(word[i:]      )
                  break
              else:
                  new_word.extend(word[i:j]      )
                  _lowerCAmelCase       :   Optional[Any]  =    j

              if word[i] == first and i < len(_A      ) - 1 and word[i + 1] == second:
                  new_word.append(first + second      )
                  i += 2
              else:
                  new_word.append(word[i]      )
                  i += 1
          _lowerCAmelCase       :   Optional[Any]  =    tuple(_A      )
          _lowerCAmelCase       :   List[str]  =    new_word
          if len(_A      ) == 1:
              break
          else:
              _lowerCAmelCase       :   List[str]  =    get_pairs(_A      )
      _lowerCAmelCase       :   Any  =    ' '.join(_A      )
      if word == "\n  " + BPE_TOKEN_MERGES:
          _lowerCAmelCase       :   str  =    '\n' + BPE_TOKEN_MERGES

      if word.endswith(_A      ):
          _lowerCAmelCase       :   Dict  =    word.replace(_A    ,''      )

      _lowerCAmelCase       :   str  =    word.replace(' '    ,_A      )
      _lowerCAmelCase       :   str  =    word
      return word

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      if self.bpe_ranks is None:
          raise ValueError(
              'This tokenizer was instantiated without a `merges.txt` file, so'
              ' that it can only be used for decoding, not for encoding.'
              'Make sure to provide `merges.txt` file at instantiation to enable '
              'encoding.'      )

      if self.do_lower_case:
          _lowerCAmelCase       :   Optional[Any]  =    text.lower()

      _lowerCAmelCase       :   Tuple  =    text.split()

      _lowerCAmelCase       :   Union[str, Any]  =    []
      for token in text:
          if token:
              split_tokens.extend(list(self.bpe(_A      ).split(' '      )      )      )

      return split_tokens

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      return self.encoder.get(_A    ,self.encoder.get(self.unk_token      )      )

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   int  =    self.decoder.get(_A    ,self.unk_token      )
      return result

  def __lowerCamelCase		(		self    ,_A      ):
      '''simple docstring'''

      _lowerCAmelCase       :   Optional[Any]  =    ' '.join(_A      )

      # make sure @@ tokens are concatenated
      _lowerCAmelCase       :   int  =    ''.join(string.split(_A      )      )

      return string







  def __lowerCamelCase		(		self    ,_A    ,_A = None      ):
      '''simple docstring'''

      if not os.path.isdir(_A      ):
          logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""      )
          return
      _lowerCAmelCase       :   List[Any]  =    os.path.join(
          _A    ,(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']      )
      _lowerCAmelCase       :   str  =    os.path.join(
          _A    ,(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['merges_file']      )

      with open(_A    ,'w'    ,encoding='utf-8'      ) as f:
          f.write(json.dumps(self.encoder    ,indent=2    ,sort_keys=_A    ,ensure_ascii=_A      ) + '\n'      )

      _lowerCAmelCase       :   str  =    0
      if self.bpe_ranks is None:
          return (vocab_file,)

      with open(_A    ,'w'    ,encoding='utf-8'      ) as writer:
          for bpe_tokens, token_index in sorted(self.bpe_ranks.items()    ,key=lambda _A      : kv[1]      ):
              if index != token_index:
                  logger.warning(
                      F"""Saving vocabulary to {merges_file}: BPE merge indices are not consecutive."""
                      ' Please check that the tokenizer is not corrupted!'      )
                  _lowerCAmelCase       :   Dict  =    token_index
              writer.write(' '.join(_A      ) + '\n'      )
              index += 1

      return (vocab_file, merges_file)





