



import math

import torch
from torch import nn
from torch.nn import init

from .binarizer import MagnitudeBinarizer, ThresholdBinarizer, TopKBinarizer





class 		UpperCAmelCase    (    nn.Linear					):





   def __init__(						self							,  _lowerCAmelCase							,  _lowerCAmelCase							,  _lowerCAmelCase = True							,  _lowerCAmelCase = "constant"							,  _lowerCAmelCase = 0.0							,  _lowerCAmelCase = "topK"							,  ):
         super(_lowerCAmelCase							,  self      ).__init__(in_features=_lowerCAmelCase							,  out_features=_lowerCAmelCase							,  bias=_lowerCAmelCase      )
         assert pruning_method in ["topK", "threshold", "sigmoied_threshold", "magnitude", "l0"]
         _lowerCAmelCase			      =	pruning_method

         if self.pruning_method in ["topK", "threshold", "sigmoied_threshold", "l0"]:
               _lowerCAmelCase			      =	mask_scale
               _lowerCAmelCase			      =	mask_init
               _lowerCAmelCase			      =	nn.Parameter(torch.empty(self.weight.size()      )      )
               self.init_mask()





   def 						__lowerCAmelCase				(						self      ):
         if self.mask_init == "constant":
               init.constant_(self.mask_scores							,  val=self.mask_scale      )
         elif self.mask_init == "uniform":
               init.uniform_(self.mask_scores							,  a=-self.mask_scale							,  b=self.mask_scale      )
         elif self.mask_init == "kaiming":
               init.kaiming_uniform_(self.mask_scores							,  a=math.sqrt(5      )      )





   def 						__lowerCAmelCase				(						self							,  _lowerCAmelCase							,  _lowerCAmelCase      ):
         # Get the mask
         if self.pruning_method == "topK":
               _lowerCAmelCase			      =	TopKBinarizer.apply(self.mask_scores							,  _lowerCAmelCase      )
         elif self.pruning_method in ["threshold", "sigmoied_threshold"]:
               _lowerCAmelCase			      =	'''sigmoied''' in self.pruning_method
               _lowerCAmelCase			      =	ThresholdBinarizer.apply(self.mask_scores							,  _lowerCAmelCase							,  _lowerCAmelCase      )
         elif self.pruning_method == "magnitude":
               _lowerCAmelCase			      =	MagnitudeBinarizer.apply(self.weight							,  _lowerCAmelCase      )
         elif self.pruning_method == "l0":
               _lowerCAmelCase      , _lowerCAmelCase      , _lowerCAmelCase			      =	-0.1, 1.1, 2 / 3
               if self.training:
                     _lowerCAmelCase			      =	torch.zeros_like(self.mask_scores      ).uniform_().clamp(0.0_001							,  0.9_999      )
                     _lowerCAmelCase			      =	torch.sigmoid((u.log() - (1 - u).log() + self.mask_scores) / b      )
               else:
                     _lowerCAmelCase			      =	torch.sigmoid(self.mask_scores      )
               _lowerCAmelCase			      =	s * (r - l) + l
               _lowerCAmelCase			      =	s_bar.clamp(min=0.0							,  max=1.0      )
         # Mask weights with computed mask
         _lowerCAmelCase			      =	mask * self.weight
         # Compute output (linear layer) with masked weights
         return nn.functional.linear(_lowerCAmelCase							,  _lowerCAmelCase							,  self.bias      )




import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase_          =				logging.get_logger(__name__)

UpperCAmelCase_          =				{
    "facebook/sam-vit-huge": "https://huggingface.co/facebook/sam-vit-huge/resolve/main/config.json",
    "facebook/sam-vit-large": "https://huggingface.co/facebook/sam-vit-large/resolve/main/config.json",
    "facebook/sam-vit-base": "https://huggingface.co/facebook/sam-vit-base/resolve/main/config.json",
}





class 		UpperCAmelCase    (    snake_case_					):





   def __init__(						self							,  _lowerCAmelCase=256							,  _lowerCAmelCase=1_024							,  _lowerCAmelCase=16							,  _lowerCAmelCase=16							,  _lowerCAmelCase=4							,  _lowerCAmelCase="gelu"							,  _lowerCAmelCase=1E-6							,  **_lowerCAmelCase							,  ):
         super().__init__(**_lowerCAmelCase      )
         _lowerCAmelCase			      =	hidden_size
         _lowerCAmelCase			      =	image_size
         _lowerCAmelCase			      =	patch_size
         _lowerCAmelCase			      =	image_size // patch_size
         _lowerCAmelCase			      =	mask_input_channels
         _lowerCAmelCase			      =	num_point_embeddings
         _lowerCAmelCase			      =	hidden_act
         _lowerCAmelCase			      =	layer_norm_eps






class 		UpperCAmelCase    (    snake_case_					):





   def __init__(						self							,  _lowerCAmelCase=256							,  _lowerCAmelCase="relu"							,  _lowerCAmelCase=2_048							,  _lowerCAmelCase=2							,  _lowerCAmelCase=8							,  _lowerCAmelCase=2							,  _lowerCAmelCase=3							,  _lowerCAmelCase=3							,  _lowerCAmelCase=256							,  _lowerCAmelCase=1E-6							,  **_lowerCAmelCase							,  ):
         super().__init__(**_lowerCAmelCase      )
         _lowerCAmelCase			      =	hidden_size
         _lowerCAmelCase			      =	hidden_act
         _lowerCAmelCase			      =	mlp_dim
         _lowerCAmelCase			      =	num_hidden_layers
         _lowerCAmelCase			      =	num_attention_heads
         _lowerCAmelCase			      =	attention_downsample_rate
         _lowerCAmelCase			      =	num_multimask_outputs
         _lowerCAmelCase			      =	iou_head_depth
         _lowerCAmelCase			      =	iou_head_hidden_dim
         _lowerCAmelCase			      =	layer_norm_eps






class 		UpperCAmelCase    (    snake_case_					):





   def __init__(						self							,  _lowerCAmelCase=768							,  _lowerCAmelCase=256							,  _lowerCAmelCase=12							,  _lowerCAmelCase=12							,  _lowerCAmelCase=3							,  _lowerCAmelCase=1_024							,  _lowerCAmelCase=16							,  _lowerCAmelCase="gelu"							,  _lowerCAmelCase=1E-06							,  _lowerCAmelCase=0.0							,  _lowerCAmelCase=1E-10							,  _lowerCAmelCase=True							,  _lowerCAmelCase=4.0							,  _lowerCAmelCase=True							,  _lowerCAmelCase=True							,  _lowerCAmelCase=14							,  _lowerCAmelCase=[2, 5, 8, 11]							,  _lowerCAmelCase=128							,  _lowerCAmelCase=None							,  **_lowerCAmelCase							,  ):
         super().__init__(**_lowerCAmelCase      )

         _lowerCAmelCase			      =	hidden_size
         _lowerCAmelCase			      =	output_channels
         _lowerCAmelCase			      =	num_hidden_layers
         _lowerCAmelCase			      =	num_attention_heads
         _lowerCAmelCase			      =	num_channels
         _lowerCAmelCase			      =	image_size
         _lowerCAmelCase			      =	patch_size
         _lowerCAmelCase			      =	hidden_act
         _lowerCAmelCase			      =	layer_norm_eps
         _lowerCAmelCase			      =	attention_dropout
         _lowerCAmelCase			      =	initializer_range
         _lowerCAmelCase			      =	qkv_bias
         _lowerCAmelCase			      =	mlp_ratio
         _lowerCAmelCase			      =	use_abs_pos
         _lowerCAmelCase			      =	use_rel_pos
         _lowerCAmelCase			      =	window_size
         _lowerCAmelCase			      =	global_attn_indexes
         _lowerCAmelCase			      =	num_pos_feats
         _lowerCAmelCase			      =	int(hidden_size * mlp_ratio      ) if mlp_dim is None else mlp_dim






class 		UpperCAmelCase    (    snake_case_					):
   SCREAMING_SNAKE_CASE__					      =  '''sam'''
   SCREAMING_SNAKE_CASE__					      =  True





   def __init__(						self							,  _lowerCAmelCase=None							,  _lowerCAmelCase=None							,  _lowerCAmelCase=None							,  _lowerCAmelCase=0.02							,  **_lowerCAmelCase							,  ):
         super().__init__(**_lowerCAmelCase      )
         _lowerCAmelCase			      =	vision_config if vision_config is not None else {}
         _lowerCAmelCase			      =	prompt_encoder_config if prompt_encoder_config is not None else {}
         _lowerCAmelCase			      =	mask_decoder_config if mask_decoder_config is not None else {}

         if isinstance(_lowerCAmelCase							,  _lowerCAmelCase      ):
               _lowerCAmelCase			      =	vision_config.to_dict()
         if isinstance(_lowerCAmelCase							,  _lowerCAmelCase      ):
               _lowerCAmelCase			      =	prompt_encoder_config.to_dict()
         if isinstance(_lowerCAmelCase							,  _lowerCAmelCase      ):
               _lowerCAmelCase			      =	mask_decoder_config.to_dict()

         _lowerCAmelCase			      =	SamVisionConfig(**_lowerCAmelCase      )
         _lowerCAmelCase			      =	SamPromptEncoderConfig(**_lowerCAmelCase      )
         _lowerCAmelCase			      =	SamMaskDecoderConfig(**_lowerCAmelCase      )
         _lowerCAmelCase			      =	initializer_range





   def 						__lowerCAmelCase				(						self      ):
         _lowerCAmelCase			      =	copy.deepcopy(self.__dict__      )
         _lowerCAmelCase			      =	self.vision_config.to_dict()
         _lowerCAmelCase			      =	self.prompt_encoder_config.to_dict()
         _lowerCAmelCase			      =	self.mask_decoder_config.to_dict()
         _lowerCAmelCase			      =	self.__class__.model_type
         return output