

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      snake_case_				(__A					:						Optional[int]					)      ->     Optional[int]:
   __lowerCAmelCase  :      List[str]    =						{}
   state_dict.pop("""pixel_mean""" ,       __A					)
   state_dict.pop("""pixel_std""" ,       __A					)

   __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  :      Tuple    =						key.replace(__A ,       __A					)

      if re.match(__A ,       __A					):
         __lowerCAmelCase  :      List[Any]    =						int(re.match(__A ,       __A					).group(2					)					)
         if layer_nb == 0:
            __lowerCAmelCase  :      List[Any]    =						key.replace("""layers.0""" ,       """proj_in"""					)
         elif layer_nb == 1:
            __lowerCAmelCase  :      Tuple    =						key.replace("""layers.1""" ,       """layers.0"""					)
         elif layer_nb == 2:
            __lowerCAmelCase  :      Dict    =						key.replace("""layers.2""" ,       """proj_out"""					)

      __lowerCAmelCase  :      int    =						value

   __lowerCAmelCase  :      int    =						model_state_dict[
       """prompt_encoder.shared_embedding.positional_embedding"""
   ]

   return model_state_dict





def      snake_case_				(__A					:						Optional[Any] ,       __A					:						Tuple ,       __A					:						Any ,       __A					:						List[str]="ybelkada/segment-anything"					)      ->     Optional[int]:
   __lowerCAmelCase  :      Any    =						hf_hub_download(__A ,       f'''checkpoints/{model_name}.pth'''					)

   if "sam_vit_b" in model_name:
      __lowerCAmelCase  :      Tuple    =						SamConfig()
   elif "sam_vit_l" in model_name:
      __lowerCAmelCase  :      Any    =						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  :      str    =						SamConfig(
          vision_config=__A ,       )
   elif "sam_vit_h" in model_name:
      __lowerCAmelCase  :      Optional[Any]    =						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  :      Optional[int]    =						SamConfig(
          vision_config=__A ,       )

   __lowerCAmelCase  :      List[str]    =						torch.load(__A ,       map_location="""cpu"""					)
   __lowerCAmelCase  :      int    =						replace_keys(__A					)

   __lowerCAmelCase  :      Optional[int]    =						SamImageProcessor()

   __lowerCAmelCase  :      int    =						SamProcessor(image_processor=__A					)
   __lowerCAmelCase  :      Any    =						SamModel(__A					)

   hf_model.load_state_dict(__A					)
   __lowerCAmelCase  :      Union[str, Any]    =						hf_model.to("""cuda"""					)

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

   __lowerCAmelCase  :      Any    =						[[[4_0_0, 6_5_0]]]
   __lowerCAmelCase  :      Optional[int]    =						[[1]]

   __lowerCAmelCase  :      Optional[Any]    =						processor(images=np.array(__A					) ,       return_tensors="""pt"""					).to("""cuda"""					)

   with torch.no_grad():
      __lowerCAmelCase  :      Tuple    =						hf_model(**__A					)
   __lowerCAmelCase  :      str    =						output.iou_scores.squeeze()

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

      __lowerCAmelCase  :      Dict    =						processor(
          images=np.array(__A					) ,       input_points=__A ,       input_labels=__A ,       return_tensors="""pt"""					).to("""cuda"""					)

      with torch.no_grad():
         __lowerCAmelCase  :      List[Any]    =						hf_model(**__A					)
      __lowerCAmelCase  :      Dict    =						output.iou_scores.squeeze()

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

      __lowerCAmelCase  :      Union[str, Any]    =						((7_5, 2_7_5, 1_7_2_5, 8_5_0),)

      __lowerCAmelCase  :      List[str]    =						processor(images=np.array(__A					) ,       input_boxes=__A ,       return_tensors="""pt"""					).to("""cuda"""					)

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

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

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

      __lowerCAmelCase  :      Optional[Any]    =						processor(
          images=np.array(__A					) ,       input_points=__A ,       input_labels=__A ,       return_tensors="""pt"""					).to("""cuda"""					)

      with torch.no_grad():
         __lowerCAmelCase  :      Optional[int]    =						hf_model(**__A					)
      __lowerCAmelCase  :      List[Any]    =						output.iou_scores.squeeze()

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


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 unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
   import jax

   from transformers.models.big_bird.modeling_flax_big_bird import (
       FlaxBigBirdForCausalLM,
       FlaxBigBirdForMaskedLM,
       FlaxBigBirdForMultipleChoice,
       FlaxBigBirdForPreTraining,
       FlaxBigBirdForQuestionAnswering,
       FlaxBigBirdForSequenceClassification,
       FlaxBigBirdForTokenClassification,
       FlaxBigBirdModel,
   )



class SCREAMING_SNAKE_CASE     (  unittest.TestCase			):
  """simple docstring"""






  def __init__(      self			:   Dict	,							lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   str=2	,							lowerCAmelCase			:   Optional[Any]=56	,							lowerCAmelCase			:   Any=True	,							lowerCAmelCase			:   Optional[int]=True	,							lowerCAmelCase			:   Any=True	,							lowerCAmelCase			:   Union[str, Any]=True	,							lowerCAmelCase			:   Tuple=99	,							lowerCAmelCase			:   Optional[Any]=32	,							lowerCAmelCase			:   List[str]=2	,							lowerCAmelCase			:   int=2	,							lowerCAmelCase			:   str=7	,							lowerCAmelCase			:   List[Any]="gelu_new"	,							lowerCAmelCase			:   Optional[int]=0.1	,							lowerCAmelCase			:   Optional[Any]=0.1	,							lowerCAmelCase			:   Optional[Any]=5_12	,							lowerCAmelCase			:   Dict=16	,							lowerCAmelCase			:   int=2	,							lowerCAmelCase			:   Optional[Any]=0.02	,							lowerCAmelCase			:   Tuple=4	,							lowerCAmelCase			:   Union[str, Any]="block_sparse"	,							lowerCAmelCase			:   List[str]=True	,							lowerCAmelCase			:   Optional[int]=False	,							lowerCAmelCase			:   Optional[Any]=2	,							lowerCAmelCase			:   List[Any]=3	,							)							->     Tuple:






     """simple docstring"""







     __lowerCAmelCase  :      Tuple    =						parent
     __lowerCAmelCase  :      Union[str, Any]    =						batch_size
     __lowerCAmelCase  :      List[Any]    =						seq_length
     __lowerCAmelCase  :      int    =						is_training
     __lowerCAmelCase  :      Union[str, Any]    =						use_attention_mask
     __lowerCAmelCase  :      Tuple    =						use_token_type_ids
     __lowerCAmelCase  :      Union[str, Any]    =						use_labels
     __lowerCAmelCase  :      List[str]    =						vocab_size
     __lowerCAmelCase  :      int    =						hidden_size
     __lowerCAmelCase  :      Tuple    =						num_hidden_layers
     __lowerCAmelCase  :      List[str]    =						num_attention_heads
     __lowerCAmelCase  :      Optional[Any]    =						intermediate_size
     __lowerCAmelCase  :      List[Any]    =						hidden_act
     __lowerCAmelCase  :      Optional[Any]    =						hidden_dropout_prob
     __lowerCAmelCase  :      List[str]    =						attention_probs_dropout_prob
     __lowerCAmelCase  :      Optional[int]    =						max_position_embeddings
     __lowerCAmelCase  :      str    =						type_vocab_size
     __lowerCAmelCase  :      List[Any]    =						type_sequence_label_size
     __lowerCAmelCase  :      Tuple    =						initializer_range
     __lowerCAmelCase  :      str    =						num_choices

     __lowerCAmelCase  :      Any    =						rescale_embeddings
     __lowerCAmelCase  :      str    =						attention_type
     __lowerCAmelCase  :      List[Any]    =						use_bias
     __lowerCAmelCase  :      List[str]    =						block_size
     __lowerCAmelCase  :      Union[str, Any]    =						num_random_blocks







  def 	SCREAMING_SNAKE_CASE   (      self			:   List[Any]   )							->     str:






     """simple docstring"""







     __lowerCAmelCase  :      str    =						ids_tensor([self.batch_size, self.seq_length]	,							self.vocab_size   )

     __lowerCAmelCase  :      Optional[Any]    =						None
     if self.use_attention_mask:
        __lowerCAmelCase  :      Optional[Any]    =						random_attention_mask([self.batch_size, self.seq_length]   )

     __lowerCAmelCase  :      Any    =						None
     if self.use_token_type_ids:
        __lowerCAmelCase  :      Optional[int]    =						ids_tensor([self.batch_size, self.seq_length]	,							self.type_vocab_size   )

     __lowerCAmelCase  :      List[Any]    =						BigBirdConfig(
         vocab_size=self.vocab_size	,							hidden_size=self.hidden_size	,							num_hidden_layers=self.num_hidden_layers	,							num_attention_heads=self.num_attention_heads	,							intermediate_size=self.intermediate_size	,							hidden_act=self.hidden_act	,							hidden_dropout_prob=self.hidden_dropout_prob	,							attention_probs_dropout_prob=self.attention_probs_dropout_prob	,							max_position_embeddings=self.max_position_embeddings	,							type_vocab_size=self.type_vocab_size	,							is_decoder=lowerCAmelCase	,							initializer_range=self.initializer_range	,							attention_type=self.attention_type	,							block_size=self.block_size	,							num_random_blocks=self.num_random_blocks	,							use_bias=self.use_bias	,							rescale_embeddings=self.rescale_embeddings	,							)

     return config, input_ids, token_type_ids, attention_mask







  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple   )							->     int:






     """simple docstring"""







     __lowerCAmelCase  :      List[str]    =						self.prepare_config_and_inputs()
     __lowerCAmelCase   ,__lowerCAmelCase   ,__lowerCAmelCase   ,__lowerCAmelCase  :      Any    =						config_and_inputs
     __lowerCAmelCase  :      Union[str, Any]    =						{
         """input_ids""": input_ids,
         """token_type_ids""": token_type_ids,
         """attention_mask""": attention_mask,
     }
     return config, inputs_dict






@require_flax
class SCREAMING_SNAKE_CASE     (  a_       ,							unittest.TestCase			):
  """simple docstring"""






  lowerCamelCase   :      List[Any]										=(
      (
          FlaxBigBirdForCausalLM,
          FlaxBigBirdModel,
          FlaxBigBirdForPreTraining,
          FlaxBigBirdForMaskedLM,
          FlaxBigBirdForMultipleChoice,
          FlaxBigBirdForQuestionAnswering,
          FlaxBigBirdForSequenceClassification,
          FlaxBigBirdForTokenClassification,
      )
      if is_flax_available()
      else ()
  )

  lowerCamelCase   :      List[str]										=False
  lowerCamelCase   :      Union[str, Any]										=False







  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple   )							->     Any:






     """simple docstring"""







     __lowerCAmelCase  :      str    =						FlaxBigBirdModelTester(self   )







  @slow
  # copied from `test_modeling_flax_common` because it takes much longer than other models
  def 	SCREAMING_SNAKE_CASE   (      self			:   Dict   )							->     Optional[Any]:






     """simple docstring"""







     super().test_from_pretrained_save_pretrained()







  @slow
  # copied from `test_modeling_flax_common` because it takes much longer than other models
  def 	SCREAMING_SNAKE_CASE   (      self			:   Dict   )							->     List[str]:






     """simple docstring"""







     super().test_from_pretrained_with_no_automatic_init()







  @slow
  # copied from `test_modeling_flax_common` because it takes much longer than other models
  def 	SCREAMING_SNAKE_CASE   (      self			:   List[str]   )							->     Optional[int]:






     """simple docstring"""







     super().test_no_automatic_init()







  @slow
  # copied from `test_modeling_flax_common` because it takes much longer than other models
  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple   )							->     Optional[int]:






     """simple docstring"""







     super().test_hidden_states_output()







  @slow
  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple   )							->     Optional[int]:






     """simple docstring"""







     for model_class_name in self.all_model_classes:
        __lowerCAmelCase  :      List[str]    =						model_class_name.from_pretrained("""google/bigbird-roberta-base"""   )
        self.assertIsNotNone(lowerCAmelCase   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   List[Any]   )							->     List[str]:






     """simple docstring"""







     if self.test_attn_probs:
        super().test_attention_outputs()







  @slow
  # copied from `test_modeling_flax_common` because it takes much longer than other models
  def 	SCREAMING_SNAKE_CASE   (      self			:   List[Any]   )							->     List[Any]:






     """simple docstring"""







     __lowerCAmelCase   ,__lowerCAmelCase  :      Tuple    =						self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        with self.subTest(model_class.__name__   ):
           __lowerCAmelCase  :      List[str]    =						self._prepare_for_class(lowerCAmelCase	,							lowerCAmelCase   )
           __lowerCAmelCase  :      Dict    =						model_class(lowerCAmelCase   )

           @jax.jit
           def model_jitted(lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   Optional[Any]=None	,							**lowerCAmelCase			:   Union[str, Any]   ):
              return model(input_ids=lowerCAmelCase	,							attention_mask=lowerCAmelCase	,							**lowerCAmelCase   )

           with self.subTest("""JIT Enabled"""   ):
              __lowerCAmelCase  :      str    =						model_jitted(**lowerCAmelCase   ).to_tuple()

           with self.subTest("""JIT Disabled"""   ):
              with jax.disable_jit():
                 __lowerCAmelCase  :      List[Any]    =						model_jitted(**lowerCAmelCase   ).to_tuple()

           self.assertEqual(len(lowerCAmelCase   )	,							len(lowerCAmelCase   )   )
           for jitted_output, output in zip(lowerCAmelCase	,							lowerCAmelCase   ):
              self.assertEqual(jitted_output.shape	,							output.shape   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[Any]	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   Optional[int]	,							lowerCAmelCase			:   Union[str, Any]=1e-5	,							lowerCAmelCase			:   Union[str, Any]="outputs"	,							lowerCAmelCase			:   Union[str, Any]=None   )							->     Optional[int]:






     """simple docstring"""







     if name.startswith("""outputs.attentions"""   ):
        return
     else:
        super().check_pt_flax_outputs(lowerCAmelCase	,							lowerCAmelCase	,							lowerCAmelCase	,							lowerCAmelCase	,							lowerCAmelCase	,							lowerCAmelCase   )


