





import collections
import inspect
import unittest

from transformers import FocalNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch
		from torch import nn

		from transformers import (
		    FocalNetBackbone,
		    FocalNetForImageClassification,
		    FocalNetForMaskedImageModeling,
		    FocalNetModel,
		)
		from transformers.models.focalnet.modeling_focalnet import FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor

class 		snake_case_   :







	'''simple docstring'''




	def __init__(						self							:     Optional[Any]					,      __lowerCamelCase							:     List[str]					,      __lowerCamelCase							:     List[str]=13					,      __lowerCamelCase							:     Optional[Any]=32					,      __lowerCamelCase							:     Any=2					,      __lowerCamelCase							:     List[Any]=3					,      __lowerCamelCase							:     Optional[Any]=16					,      __lowerCamelCase							:     Tuple=[32, 64, 128]					,      __lowerCamelCase							:     List[str]=[1, 2, 1]					,      __lowerCamelCase							:     str=[2, 2, 4]					,      __lowerCamelCase							:     str=2					,      __lowerCamelCase							:     str=2.0					,      __lowerCamelCase							:     str=True					,      __lowerCamelCase							:     Optional[int]=0.0					,      __lowerCamelCase							:     Any=0.0					,      __lowerCamelCase							:     int=0.1					,      __lowerCamelCase							:     int="gelu"					,      __lowerCamelCase							:     Union[str, Any]=False					,      __lowerCamelCase							:     Union[str, Any]=True					,      __lowerCamelCase							:     List[Any]=0.02					,      __lowerCamelCase							:     Tuple=1E-5					,      __lowerCamelCase							:     List[str]=True					,      __lowerCamelCase							:     Optional[Any]=None					,      __lowerCamelCase							:     Tuple=True					,      __lowerCamelCase							:     Any=10					,      __lowerCamelCase							:     Union[str, Any]=8					,      __lowerCamelCase							:     Tuple=["stage1", "stage2"]					,      __lowerCamelCase							:     str=[1, 2]					,      )       -> Optional[Any]:

			'''simple docstring'''




			__lowercase											=     parent
			__lowercase											=     batch_size
			__lowercase											=     image_size
			__lowercase											=     patch_size
			__lowercase											=     num_channels
			__lowercase											=     embed_dim
			__lowercase											=     hidden_sizes
			__lowercase											=     depths
			__lowercase											=     num_heads
			__lowercase											=     window_size
			__lowercase											=     mlp_ratio
			__lowercase											=     qkv_bias
			__lowercase											=     hidden_dropout_prob
			__lowercase											=     attention_probs_dropout_prob
			__lowercase											=     drop_path_rate
			__lowercase											=     hidden_act
			__lowercase											=     use_absolute_embeddings
			__lowercase											=     patch_norm
			__lowercase											=     layer_norm_eps
			__lowercase											=     initializer_range
			__lowercase											=     is_training
			__lowercase											=     scope
			__lowercase											=     use_labels
			__lowercase											=     type_sequence_label_size
			__lowercase											=     encoder_stride
			__lowercase											=     out_features
			__lowercase											=     out_indices




	def    UpperCAmelCase			(						self							:     int   )       -> int:

			'''simple docstring'''




			__lowercase											=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

			__lowercase											=     None
			if self.use_labels:
					__lowercase											=     ids_tensor([self.batch_size]					,      self.type_sequence_label_size   )

			__lowercase											=     self.get_config()

			return config, pixel_values, labels




	def    UpperCAmelCase			(						self							:     str   )       -> int:

			'''simple docstring'''




			return FocalNetConfig(
			    image_size=self.image_size					,      patch_size=self.patch_size					,      num_channels=self.num_channels					,      embed_dim=self.embed_dim					,      hidden_sizes=self.hidden_sizes					,      depths=self.depths					,      num_heads=self.num_heads					,      window_size=self.window_size					,      mlp_ratio=self.mlp_ratio					,      qkv_bias=self.qkv_bias					,      hidden_dropout_prob=self.hidden_dropout_prob					,      attention_probs_dropout_prob=self.attention_probs_dropout_prob					,      drop_path_rate=self.drop_path_rate					,      hidden_act=self.hidden_act					,      use_absolute_embeddings=self.use_absolute_embeddings					,      path_norm=self.patch_norm					,      layer_norm_eps=self.layer_norm_eps					,      initializer_range=self.initializer_range					,      encoder_stride=self.encoder_stride					,      out_features=self.out_features					,      out_indices=self.out_indices					,      )




	def    UpperCAmelCase			(						self							:     Any					,      __lowerCamelCase							:     Any					,      __lowerCamelCase							:     Dict					,      __lowerCamelCase							:     Tuple   )       -> Union[str, Any]:

			'''simple docstring'''




			__lowercase											=     FocalNetModel(config=__lowerCamelCase   )
			model.to(__lowerCamelCase   )
			model.eval()
			__lowercase											=     model(__lowerCamelCase   )

			__lowercase											=     ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths   ) - 1))
			__lowercase											=     int(config.embed_dim * 2 ** (len(config.depths   ) - 1)   )

			self.parent.assertEqual(result.last_hidden_state.shape					,      (self.batch_size, expected_seq_len, expected_dim)   )




	def    UpperCAmelCase			(						self							:     Optional[Any]					,      __lowerCamelCase							:     int					,      __lowerCamelCase							:     Optional[Any]					,      __lowerCamelCase							:     Optional[int]   )       -> int:

			'''simple docstring'''




			__lowercase											=     FocalNetBackbone(config=__lowerCamelCase   )
			model.to(__lowerCamelCase   )
			model.eval()
			__lowercase											=     model(__lowerCamelCase   )

			# verify feature maps
			self.parent.assertEqual(len(result.feature_maps   )					,      len(config.out_features   )   )
			self.parent.assertListEqual(list(result.feature_maps[0].shape   )					,      [self.batch_size, self.image_size, 8, 8]   )

			# verify channels
			self.parent.assertEqual(len(model.channels   )					,      len(config.out_features   )   )
			self.parent.assertListEqual(model.channels					,      config.hidden_sizes[:-1]   )

			# verify backbone works with out_features=None
			__lowercase											=     None
			__lowercase											=     FocalNetBackbone(config=__lowerCamelCase   )
			model.to(__lowerCamelCase   )
			model.eval()
			__lowercase											=     model(__lowerCamelCase   )

			# verify feature maps
			self.parent.assertEqual(len(result.feature_maps   )					,      1   )
			self.parent.assertListEqual(list(result.feature_maps[0].shape   )					,      [self.batch_size, self.image_size * 2, 4, 4]   )

			# verify channels
			self.parent.assertEqual(len(model.channels   )					,      1   )
			self.parent.assertListEqual(model.channels					,      [config.hidden_sizes[-1]]   )




	def    UpperCAmelCase			(						self							:     Tuple					,      __lowerCamelCase							:     List[Any]					,      __lowerCamelCase							:     Dict					,      __lowerCamelCase							:     Tuple   )       -> Optional[Any]:

			'''simple docstring'''




			__lowercase											=     FocalNetForMaskedImageModeling(config=__lowerCamelCase   )
			model.to(__lowerCamelCase   )
			model.eval()
			__lowercase											=     model(__lowerCamelCase   )
			self.parent.assertEqual(
			    result.reconstruction.shape					,      (self.batch_size, self.num_channels, self.image_size, self.image_size)   )

			# test greyscale images
			__lowercase											=     1
			__lowercase											=     FocalNetForMaskedImageModeling(__lowerCamelCase   )
			model.to(__lowerCamelCase   )
			model.eval()

			__lowercase											=     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
			__lowercase											=     model(__lowerCamelCase   )
			self.parent.assertEqual(result.reconstruction.shape					,      (self.batch_size, 1, self.image_size, self.image_size)   )




	def    UpperCAmelCase			(						self							:     List[Any]					,      __lowerCamelCase							:     str					,      __lowerCamelCase							:     List[str]					,      __lowerCamelCase							:     int   )       -> Optional[int]:

			'''simple docstring'''




			__lowercase											=     self.type_sequence_label_size
			__lowercase											=     FocalNetForImageClassification(__lowerCamelCase   )
			model.to(__lowerCamelCase   )
			model.eval()
			__lowercase											=     model(__lowerCamelCase					,      labels=__lowerCamelCase   )
			self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.type_sequence_label_size)   )

			# test greyscale images
			__lowercase											=     1
			__lowercase											=     FocalNetForImageClassification(__lowerCamelCase   )
			model.to(__lowerCamelCase   )
			model.eval()

			__lowercase											=     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
			__lowercase											=     model(__lowerCamelCase   )
			self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.type_sequence_label_size)   )






	def    UpperCAmelCase			(						self							:     str   )       -> Union[str, Any]:

			'''simple docstring'''




			__lowercase											=     self.prepare_config_and_inputs()

			__lowercase						,      __lowercase						,      __lowercase											=     config_and_inputs
			__lowercase											=     {'pixel_values': pixel_values}
			return config, inputs_dict





@require_torch
class 		snake_case_   (							UpperCAmelCase_     ,     UpperCAmelCase_     ,     unittest.TestCase						):







	'''simple docstring'''
	__UpperCamelCase   =		(
	    (
	        FocalNetModel,
	        FocalNetForImageClassification,
	        FocalNetForMaskedImageModeling,
	        FocalNetBackbone,
	    )
	    if is_torch_available()
	    else ()
	)
	__UpperCamelCase   =		(
	    {'''feature-extraction''': FocalNetModel, '''image-classification''': FocalNetForImageClassification}
	    if is_torch_available()
	    else {}
	)
	__UpperCamelCase   =		False

	__UpperCamelCase   =		False
	__UpperCamelCase   =		False
	__UpperCamelCase   =		False
	__UpperCamelCase   =		False




	def    UpperCAmelCase			(						self							:     Optional[Any]   )       -> Tuple:

			'''simple docstring'''




			__lowercase											=     FocalNetModelTester(self   )
			__lowercase											=     ConfigTester(self					,      config_class=__lowerCamelCase					,      embed_dim=37					,      has_text_modality=__lowerCamelCase   )




	def    UpperCAmelCase			(						self							:     Tuple   )       -> Any:

			'''simple docstring'''




			self.create_and_test_config_common_properties()
			self.config_tester.create_and_test_config_to_json_string()
			self.config_tester.create_and_test_config_to_json_file()
			self.config_tester.create_and_test_config_from_and_save_pretrained()
			self.config_tester.create_and_test_config_with_num_labels()
			self.config_tester.check_config_can_be_init_without_params()
			self.config_tester.check_config_arguments_init()




	def    UpperCAmelCase			(						self							:     Optional[int]   )       -> List[Any]:

			'''simple docstring'''




			return




	def    UpperCAmelCase			(						self							:     Any   )       -> Tuple:

			'''simple docstring'''




			__lowercase											=     self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*__lowerCamelCase   )




	def    UpperCAmelCase			(						self							:     int   )       -> int:

			'''simple docstring'''




			__lowercase											=     self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_backbone(*__lowerCamelCase   )




	def    UpperCAmelCase			(						self							:     Tuple   )       -> int:

			'''simple docstring'''




			__lowercase											=     self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_masked_image_modeling(*__lowerCamelCase   )




	def    UpperCAmelCase			(						self							:     int   )       -> Union[str, Any]:

			'''simple docstring'''




			__lowercase											=     self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_image_classification(*__lowerCamelCase   )




	@unittest.skip(reason='FocalNet does not use inputs_embeds'   )
	def    UpperCAmelCase			(						self							:     Optional[Any]   )       -> Optional[int]:

			'''simple docstring'''




			pass




	@unittest.skip(reason='FocalNet does not use feedforward chunking'   )
	def    UpperCAmelCase			(						self							:     str   )       -> Tuple:

			'''simple docstring'''




			pass




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

			'''simple docstring'''




			__lowercase						,      __lowercase											=     self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes[:-1]:
					__lowercase											=     model_class(__lowerCamelCase   )
					self.assertIsInstance(model.get_input_embeddings()					,      (nn.Module)   )
					__lowercase											=     model.get_output_embeddings()
					self.assertTrue(x is None or isinstance(__lowerCamelCase					,      nn.Linear   )   )




	def    UpperCAmelCase			(						self							:     str   )       -> int:

			'''simple docstring'''




			__lowercase						,      __lowercase											=     self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes[:-1]:
					__lowercase											=     model_class(__lowerCamelCase   )
					__lowercase											=     inspect.signature(model.forward   )
					# signature.parameters is an OrderedDict => so arg_names order is deterministic
					__lowercase											=     [*signature.parameters.keys()]

					__lowercase											=     ['pixel_values']
					self.assertListEqual(arg_names[:1]					,      __lowerCamelCase   )




	def    UpperCAmelCase			(						self							:     Union[str, Any]					,      __lowerCamelCase							:     List[Any]					,      __lowerCamelCase							:     Tuple					,      __lowerCamelCase							:     Optional[int]					,      __lowerCamelCase							:     Dict   )       -> Union[str, Any]:

			'''simple docstring'''




			__lowercase											=     model_class(__lowerCamelCase   )
			model.to(__lowerCamelCase   )
			model.eval()

			with torch.no_grad():
					__lowercase											=     model(**self._prepare_for_class(__lowerCamelCase					,      __lowerCamelCase   )   )

			__lowercase											=     outputs.hidden_states

			__lowercase											=     getattr(
			    self.model_tester					,      'expected_num_hidden_layers'					,      len(self.model_tester.depths   ) + 1   )
			self.assertEqual(len(__lowerCamelCase   )					,      __lowerCamelCase   )

			# FocalNet has a different seq_length
			__lowercase											=     (
			    config.patch_size
			    if isinstance(config.patch_size					,      collections.abc.Iterable   )
			    else (config.patch_size, config.patch_size)
			)

			__lowercase											=     (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

			self.assertListEqual(
			    list(hidden_states[0].shape[-2:]   )					,      [num_patches, self.model_tester.embed_dim]					,      )

			__lowercase											=     outputs.reshaped_hidden_states
			self.assertEqual(len(__lowerCamelCase   )					,      __lowerCamelCase   )

			__lowercase						,      __lowercase						,      __lowercase						,      __lowercase											=     reshaped_hidden_states[0].shape
			__lowercase											=     (
			    reshaped_hidden_states[0].view(__lowerCamelCase					,      __lowerCamelCase					,      height * width   ).permute(0					,      2					,      1   )
			)
			self.assertListEqual(
			    list(reshaped_hidden_states.shape[-2:]   )					,      [num_patches, self.model_tester.embed_dim]					,      )




	def    UpperCAmelCase			(						self							:     Dict   )       -> List[str]:

			'''simple docstring'''




			__lowercase						,      __lowercase											=     self.model_tester.prepare_config_and_inputs_for_common()

			__lowercase											=     (
			    self.model_tester.image_size
			    if isinstance(self.model_tester.image_size					,      collections.abc.Iterable   )
			    else (self.model_tester.image_size, self.model_tester.image_size)
			)

			for model_class in self.all_model_classes[:-1]:
					__lowercase											=     True
					self.check_hidden_states_output(__lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase   )

					# check that output_hidden_states also work using config
					del inputs_dict["output_hidden_states"]
					__lowercase											=     True

					self.check_hidden_states_output(__lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase   )




	def    UpperCAmelCase			(						self							:     Optional[Any]   )       -> str:

			'''simple docstring'''




			__lowercase						,      __lowercase											=     self.model_tester.prepare_config_and_inputs_for_common()
			__lowercase											=     3

			__lowercase											=     (
			    self.model_tester.image_size
			    if isinstance(self.model_tester.image_size					,      collections.abc.Iterable   )
			    else (self.model_tester.image_size, self.model_tester.image_size)
			)
			__lowercase											=     (
			    config.patch_size
			    if isinstance(config.patch_size					,      collections.abc.Iterable   )
			    else (config.patch_size, config.patch_size)
			)

			__lowercase											=     image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
			__lowercase											=     image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

			for model_class in self.all_model_classes[:-1]:
					__lowercase											=     True
					self.check_hidden_states_output(__lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      (padded_height, padded_width)   )

					# check that output_hidden_states also work using config
					del inputs_dict["output_hidden_states"]
					__lowercase											=     True
					self.check_hidden_states_output(__lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      (padded_height, padded_width)   )




	@slow
	def    UpperCAmelCase			(						self							:     Any   )       -> Optional[Any]:

			'''simple docstring'''




			for model_name in FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					__lowercase											=     FocalNetModel.from_pretrained(__lowerCamelCase   )
					self.assertIsNotNone(__lowerCamelCase   )






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

			'''simple docstring'''




			__lowercase						,      __lowercase											=     self.model_tester.prepare_config_and_inputs_for_common()

			__lowercase											=     _config_zero_init(__lowerCamelCase   )
			for model_class in self.all_model_classes:
					__lowercase											=     model_class(config=__lowerCamelCase   )
					for name, param in model.named_parameters():
							if "embeddings" not in name and param.requires_grad:
									self.assertIn(
									    ((param.data.mean() * 1E9).round() / 1E9).item()					,      [0.0, 1.0]					,      msg=F"Parameter {name} of model {model_class} seems not properly initialized"					,      )





@require_vision
@require_torch
class 		snake_case_   (							unittest.TestCase						):







	'''simple docstring'''




	@cached_property
	def    UpperCAmelCase			(						self							:     Optional[Any]   )       -> Optional[int]:

			'''simple docstring'''




			return AutoImageProcessor.from_pretrained('microsoft/focalnet-tiny'   ) if is_vision_available() else None






	@slow
	def    UpperCAmelCase			(						self							:     str   )       -> List[Any]:

			'''simple docstring'''




			__lowercase											=     FocalNetForImageClassification.from_pretrained('microsoft/focalnet-tiny'   ).to(__lowerCamelCase   )
			__lowercase											=     self.default_image_processor

			__lowercase											=     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
			__lowercase											=     image_processor(images=__lowerCamelCase					,      return_tensors='pt'   ).to(__lowerCamelCase   )

			# forward pass
			with torch.no_grad():
					__lowercase											=     model(**__lowerCamelCase   )

			# verify the logits
			__lowercase											=     torch.Size((1, 1_000)   )
			self.assertEqual(outputs.logits.shape					,      __lowerCamelCase   )
			__lowercase											=     torch.tensor([0.2166, -0.4368, 0.2191]   ).to(__lowerCamelCase   )
			self.assertTrue(torch.allclose(outputs.logits[0, :3]					,      __lowerCamelCase					,      atol=1E-4   )   )
			self.assertTrue(outputs.logits.argmax(dim=-1   ).item()					,      281   )





@require_torch
class 		snake_case_   (							UpperCAmelCase_     ,     unittest.TestCase						):







	'''simple docstring'''
	__UpperCamelCase   =		(FocalNetBackbone,) if is_torch_available() else ()
	__UpperCamelCase   =		FocalNetConfig

	__UpperCamelCase   =		False




	def    UpperCAmelCase			(						self							:     Union[str, Any]   )       -> Tuple:

			'''simple docstring'''




			__lowercase											=     FocalNetModelTester(self   )











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)




