configuration_sam2.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453
  1. # coding=utf-8
  2. # Copyright 2025 The HuggingFace Inc. team. All rights reserved.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. """SAM2 model configuration"""
  16. from ...configuration_utils import PretrainedConfig
  17. from ...utils import logging
  18. from ..auto import CONFIG_MAPPING, AutoConfig
  19. logger = logging.get_logger(__name__)
  20. class Sam2HieraDetConfig(PretrainedConfig):
  21. r"""
  22. This is the configuration class to store the configuration of a [`Sam2HieraDetModel`]. It is used to instantiate
  23. a HieraDet model as defined in the original sam2 repo according to the specified arguments, defining the model architecture.
  24. Instantiating a configuration defaults will yield a similar configuration to that of SAM 2.1 Hiera-tiny
  25. [facebook/sam2.1-hiera-tiny](https://huggingface.co/facebook/sam2.1-hiera-tiny) architecture.
  26. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
  27. documentation from [`PretrainedConfig`] for more information.
  28. Args:
  29. hidden_size (`int`, *optional*, defaults to 96):
  30. The hidden dimension of the image encoder.
  31. num_attention_heads (`int`, *optional*, defaults to 1):
  32. Number of attention heads for each attention layer in the Transformer encoder.
  33. num_channels (`int`, *optional*, defaults to 3):
  34. The number of channels in the image.
  35. image_size (`list[int]`, *optional*, defaults to `[1024, 1024]`):
  36. The size of the image.
  37. patch_kernel_size (`list[int]`, *optional*, defaults to `[7, 7]`):
  38. The kernel size of the patch.
  39. patch_stride (`list[int]`, *optional*, defaults to `[4, 4]`):
  40. The stride of the patch.
  41. patch_padding (`list[int]`, *optional*, defaults to `[3, 3]`):
  42. The padding of the patch.
  43. query_stride (`list[int]`, *optional*, defaults to `[2, 2]`):
  44. The downsample stride between stages.
  45. window_positional_embedding_background_size (`list[int]`, *optional*, defaults to `[7, 7]`):
  46. The window size per stage when not using global attention.
  47. num_query_pool_stages (`int`, *optional*, defaults to 3):
  48. The number of query pool stages.
  49. blocks_per_stage (`list[int]`, *optional*, defaults to `[1, 2, 7, 2]`):
  50. The number of blocks per stage.
  51. embed_dim_per_stage (`list[int]`, *optional*, defaults to `[96, 192, 384, 768]`):
  52. The embedding dimension per stage.
  53. num_attention_heads_per_stage (`list[int]`, *optional*, defaults to `[1, 2, 4, 8]`):
  54. The number of attention heads per stage.
  55. window_size_per_stage (`list[int]`, *optional*, defaults to `[8, 4, 14, 7]`):
  56. The window size per stage.
  57. global_attention_blocks (`list[int]`, *optional*, defaults to `[5, 7, 9]`):
  58. The blocks where global attention is used.
  59. mlp_ratio (`float`, *optional*, defaults to 4.0):
  60. The ratio of the MLP hidden dimension to the embedding dimension.
  61. hidden_act (`str`, *optional*, defaults to `"gelu"`):
  62. The non-linear activation function in the neck.
  63. layer_norm_eps (`float`, *optional*, defaults to 1e-06):
  64. The epsilon for the layer normalization.
  65. initializer_range (`float`, *optional*, defaults to 0.02):
  66. The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
  67. """
  68. base_config_key = "backbone_config"
  69. model_type = "sam2_hiera_det_model"
  70. def __init__(
  71. self,
  72. hidden_size=96,
  73. num_attention_heads=1,
  74. num_channels=3,
  75. image_size=None,
  76. patch_kernel_size=None,
  77. patch_stride=None,
  78. patch_padding=None,
  79. query_stride=None,
  80. window_positional_embedding_background_size=None,
  81. num_query_pool_stages=3,
  82. blocks_per_stage=None,
  83. embed_dim_per_stage=None,
  84. num_attention_heads_per_stage=None,
  85. window_size_per_stage=None,
  86. global_attention_blocks=None,
  87. mlp_ratio=4.0,
  88. hidden_act="gelu",
  89. layer_norm_eps=1e-6,
  90. initializer_range=0.02,
  91. **kwargs,
  92. ):
  93. super().__init__(**kwargs)
  94. image_size = image_size if image_size is not None else [1024, 1024]
  95. patch_kernel_size = patch_kernel_size if patch_kernel_size is not None else [7, 7]
  96. patch_stride = patch_stride if patch_stride is not None else [4, 4]
  97. patch_padding = patch_padding if patch_padding is not None else [3, 3]
  98. query_stride = query_stride if query_stride is not None else [2, 2]
  99. window_positional_embedding_background_size = (
  100. window_positional_embedding_background_size
  101. if window_positional_embedding_background_size is not None
  102. else [7, 7]
  103. )
  104. blocks_per_stage = blocks_per_stage if blocks_per_stage is not None else [1, 2, 7, 2]
  105. embed_dim_per_stage = embed_dim_per_stage if embed_dim_per_stage is not None else [96, 192, 384, 768]
  106. num_attention_heads_per_stage = (
  107. num_attention_heads_per_stage if num_attention_heads_per_stage is not None else [1, 2, 4, 8]
  108. )
  109. window_size_per_stage = window_size_per_stage if window_size_per_stage is not None else [8, 4, 14, 7]
  110. global_attention_blocks = global_attention_blocks if global_attention_blocks is not None else [5, 7, 9]
  111. self.hidden_size = hidden_size
  112. self.num_attention_heads = num_attention_heads
  113. self.num_channels = num_channels
  114. self.image_size = image_size
  115. self.patch_kernel_size = patch_kernel_size
  116. self.patch_stride = patch_stride
  117. self.patch_padding = patch_padding
  118. self.query_stride = query_stride
  119. self.window_positional_embedding_background_size = window_positional_embedding_background_size
  120. self.num_query_pool_stages = num_query_pool_stages
  121. self.blocks_per_stage = blocks_per_stage
  122. self.embed_dim_per_stage = embed_dim_per_stage
  123. self.num_attention_heads_per_stage = num_attention_heads_per_stage
  124. self.window_size_per_stage = window_size_per_stage
  125. self.global_attention_blocks = global_attention_blocks
  126. self.mlp_ratio = mlp_ratio
  127. self.hidden_act = hidden_act
  128. self.layer_norm_eps = layer_norm_eps
  129. self.initializer_range = initializer_range
  130. class Sam2VisionConfig(PretrainedConfig):
  131. r"""
  132. This is the configuration class to store the configuration of a [`Sam2VisionModel`]. It is used to instantiate a SAM
  133. vision encoder according to the specified arguments, defining the model architecture. Instantiating a configuration
  134. defaults will yield a similar configuration to that of SAM 2.1 Hiera-tiny
  135. [facebook/sam2.1-hiera-tiny](https://huggingface.co/facebook/sam2.1-hiera-tiny) architecture.
  136. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
  137. documentation from [`PretrainedConfig`] for more information.
  138. Args:
  139. backbone_config (`Union[dict, "PretrainedConfig"]`, *optional*):
  140. Configuration for the vision backbone. This is used to instantiate the backbone using
  141. `AutoModel.from_config`.
  142. backbone_channel_list (`List[int]`, *optional*, defaults to `[768, 384, 192, 96]`):
  143. The list of channel dimensions for the backbone.
  144. backbone_feature_sizes (`List[List[int]]`, *optional*, defaults to `[[256, 256], [128, 128], [64, 64]]`):
  145. The spatial sizes of the feature maps from the backbone.
  146. fpn_hidden_size (`int`, *optional*, defaults to 256):
  147. The hidden dimension of the FPN.
  148. fpn_kernel_size (`int`, *optional*, defaults to 1):
  149. The kernel size for the convolutions in the neck.
  150. fpn_stride (`int`, *optional*, defaults to 1):
  151. The stride for the convolutions in the neck.
  152. fpn_padding (`int`, *optional*, defaults to 0):
  153. The padding for the convolutions in the neck.
  154. fpn_top_down_levels (`List[int]`, *optional*, defaults to `[2, 3]`):
  155. The levels for the top-down FPN connections.
  156. num_feature_levels (`int`, *optional*, defaults to 3):
  157. The number of feature levels from the FPN to use.
  158. hidden_act (`str`, *optional*, defaults to `"gelu"`):
  159. The non-linear activation function in the neck.
  160. layer_norm_eps (`float`, *optional*, defaults to 1e-06):
  161. The epsilon for the layer normalization.
  162. initializer_range (`float`, *optional*, defaults to 0.02):
  163. The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
  164. """
  165. base_config_key = "vision_config"
  166. model_type = "sam2_vision_model"
  167. sub_configs = {
  168. "backbone_config": AutoConfig,
  169. }
  170. def __init__(
  171. self,
  172. backbone_config=None,
  173. backbone_channel_list=None,
  174. backbone_feature_sizes=None,
  175. fpn_hidden_size=256,
  176. fpn_kernel_size=1,
  177. fpn_stride=1,
  178. fpn_padding=0,
  179. fpn_top_down_levels=None,
  180. num_feature_levels=3,
  181. hidden_act="gelu",
  182. layer_norm_eps=1e-6,
  183. initializer_range=0.02,
  184. **kwargs,
  185. ):
  186. super().__init__(**kwargs)
  187. backbone_channel_list = [768, 384, 192, 96] if backbone_channel_list is None else backbone_channel_list
  188. backbone_feature_sizes = (
  189. [[256, 256], [128, 128], [64, 64]] if backbone_feature_sizes is None else backbone_feature_sizes
  190. )
  191. fpn_top_down_levels = [2, 3] if fpn_top_down_levels is None else fpn_top_down_levels
  192. if isinstance(backbone_config, dict):
  193. backbone_config["model_type"] = backbone_config.get("model_type", "sam2_hiera_det_model")
  194. backbone_config = CONFIG_MAPPING[backbone_config["model_type"]](**backbone_config)
  195. elif isinstance(backbone_config, Sam2HieraDetConfig):
  196. pass
  197. elif backbone_config is None:
  198. backbone_config = Sam2HieraDetConfig()
  199. self.backbone_config = backbone_config
  200. # Neck
  201. self.backbone_channel_list = backbone_channel_list
  202. self.backbone_feature_sizes = backbone_feature_sizes
  203. self.fpn_hidden_size = fpn_hidden_size
  204. self.fpn_kernel_size = fpn_kernel_size
  205. self.fpn_stride = fpn_stride
  206. self.fpn_padding = fpn_padding
  207. self.fpn_top_down_levels = fpn_top_down_levels
  208. self.num_feature_levels = num_feature_levels
  209. self.hidden_act = hidden_act
  210. self.layer_norm_eps = layer_norm_eps
  211. self.initializer_range = initializer_range
  212. class Sam2PromptEncoderConfig(PretrainedConfig):
  213. r"""
  214. This is the configuration class to store the configuration of a [`Sam2PromptEncoder`]. The [`Sam2PromptEncoder`]
  215. module is used to encode the input 2D points and bounding boxes.
  216. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
  217. documentation from [`PretrainedConfig`] for more information.
  218. Args:
  219. hidden_size (`int`, *optional*, defaults to 256):
  220. Dimensionality of the hidden states.
  221. image_size (`int`, *optional*, defaults to 1024):
  222. The expected output resolution of the image.
  223. patch_size (`int`, *optional*, defaults to 16):
  224. The size (resolution) of each patch.
  225. mask_input_channels (`int`, *optional*, defaults to 16):
  226. The number of channels to be fed to the `MaskDecoder` module.
  227. num_point_embeddings (`int`, *optional*, defaults to 4):
  228. The number of point embeddings to be used.
  229. hidden_act (`str`, *optional*, defaults to `"gelu"`):
  230. The non-linear activation function in the encoder and pooler.
  231. layer_norm_eps (`float`, *optional*, defaults to 1e-06):
  232. The epsilon used by the layer normalization layers.
  233. scale (`float`, *optional*, defaults to 1):
  234. The scale factor for the prompt encoder.
  235. """
  236. base_config_key = "prompt_encoder_config"
  237. def __init__(
  238. self,
  239. hidden_size=256,
  240. image_size=1024,
  241. patch_size=16,
  242. mask_input_channels=16,
  243. num_point_embeddings=4,
  244. hidden_act="gelu",
  245. layer_norm_eps=1e-6,
  246. scale=1,
  247. **kwargs,
  248. ):
  249. super().__init__(**kwargs)
  250. self.hidden_size = hidden_size
  251. self.image_size = image_size
  252. self.patch_size = patch_size
  253. self.mask_input_channels = mask_input_channels
  254. self.num_point_embeddings = num_point_embeddings
  255. self.hidden_act = hidden_act
  256. self.layer_norm_eps = layer_norm_eps
  257. self.scale = scale
  258. class Sam2MaskDecoderConfig(PretrainedConfig):
  259. r"""
  260. This is the configuration class to store the configuration of a [`Sam2MaskDecoder`]. It is used to instantiate a SAM2
  261. memory encoder according to the specified arguments, defining the model architecture.
  262. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
  263. documentation from [`PretrainedConfig`] for more information.
  264. Args:
  265. hidden_size (`int`, *optional*, defaults to 256):
  266. Dimensionality of the hidden states.
  267. hidden_act (`str`, *optional*, defaults to `"gelu"`):
  268. The non-linear activation function in the SAM2 mask decoder.
  269. mlp_dim (`int`, *optional*, defaults to 2048):
  270. The dimension of the MLP in the two-way transformer.
  271. num_hidden_layers (`int`, *optional*, defaults to 2):
  272. The number of hidden layers in the two-way transformer.
  273. num_attention_heads (`int`, *optional*, defaults to 8):
  274. The number of attention heads in the two-way transformer.
  275. attention_downsample_rate (`int`, *optional*, defaults to 2):
  276. The downsample rate for the attention layers.
  277. num_multimask_outputs (`int`, *optional*, defaults to 3):
  278. The number of multimask outputs.
  279. iou_head_depth (`int`, *optional*, defaults to 3):
  280. The depth of the IoU head.
  281. iou_head_hidden_dim (`int`, *optional*, defaults to 256):
  282. The hidden dimension of the IoU head.
  283. dynamic_multimask_via_stability (`bool`, *optional*, defaults to `True`):
  284. Whether to use dynamic multimask via stability.
  285. dynamic_multimask_stability_delta (`float`, *optional*, defaults to 0.05):
  286. The stability delta for the dynamic multimask.
  287. dynamic_multimask_stability_thresh (`float`, *optional*, defaults to 0.98):
  288. The stability threshold for the dynamic multimask.
  289. """
  290. base_config_key = "mask_decoder_config"
  291. def __init__(
  292. self,
  293. hidden_size=256,
  294. hidden_act="gelu",
  295. mlp_dim=2048,
  296. num_hidden_layers=2,
  297. num_attention_heads=8,
  298. attention_downsample_rate=2,
  299. num_multimask_outputs=3,
  300. iou_head_depth=3,
  301. iou_head_hidden_dim=256,
  302. dynamic_multimask_via_stability=True,
  303. dynamic_multimask_stability_delta=0.05,
  304. dynamic_multimask_stability_thresh=0.98,
  305. **kwargs,
  306. ):
  307. super().__init__(**kwargs)
  308. self.hidden_size = hidden_size
  309. self.num_multimask_outputs = num_multimask_outputs
  310. self.hidden_act = hidden_act
  311. self.iou_head_depth = iou_head_depth
  312. self.iou_head_hidden_dim = iou_head_hidden_dim
  313. self.dynamic_multimask_via_stability = dynamic_multimask_via_stability
  314. self.dynamic_multimask_stability_delta = dynamic_multimask_stability_delta
  315. self.dynamic_multimask_stability_thresh = dynamic_multimask_stability_thresh
  316. # TwoWayTransformer configuration
  317. self.num_hidden_layers = num_hidden_layers
  318. self.hidden_size = hidden_size
  319. self.num_attention_heads = num_attention_heads
  320. self.mlp_dim = mlp_dim
  321. self.attention_downsample_rate = attention_downsample_rate
  322. class Sam2Config(PretrainedConfig):
  323. r"""
  324. [`Sam2Config`] is the configuration class to store the configuration of a [`Sam2Model`]. It is used to instantiate a
  325. SAM2 model according to the specified arguments, defining the memory attention, memory encoder, and image encoder
  326. configs. Instantiating a configuration defaults will yield a similar configuration to that of the SAM 2.1 Hiera-tiny
  327. [facebook/sam2.1-hiera-tiny](https://huggingface.co/facebook/sam2.1-hiera-tiny) architecture.
  328. Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
  329. documentation from [`PretrainedConfig`] for more information.
  330. Args:
  331. vision_config (Union[`dict`, `Sam2VisionConfig`], *optional*):
  332. Dictionary of configuration options used to initialize [`Sam2VisionConfig`].
  333. prompt_encoder_config (Union[`dict`, `Sam2PromptEncoderConfig`], *optional*):
  334. Dictionary of configuration options used to initialize [`Sam2PromptEncoderConfig`].
  335. mask_decoder_config (Union[`dict`, `Sam2MaskDecoderConfig`], *optional*):
  336. Dictionary of configuration options used to initialize [`Sam2MaskDecoderConfig`].
  337. initializer_range (`float`, *optional*, defaults to 0.02):
  338. Standard deviation for parameter initialization.
  339. Example:
  340. ```python
  341. >>> from transformers import (
  342. ... Sam2VisionConfig,
  343. ... Sam2PromptEncoderConfig,
  344. ... Sam2MaskDecoderConfig,
  345. ... Sam2Model,
  346. ... )
  347. >>> # Initializing a Sam2Config with `"facebook/sam2.1_hiera_tiny"` style configuration
  348. >>> configuration = Sam2config()
  349. >>> # Initializing a Sam2Model (with random weights) from the `"facebook/sam2.1_hiera_tiny"` style configuration
  350. >>> model = Sam2Model(configuration)
  351. >>> # Accessing the model configuration
  352. >>> configuration = model.config
  353. >>> # We can also initialize a Sam2Config from a Sam2VisionConfig, Sam2PromptEncoderConfig, and Sam2MaskDecoderConfig
  354. >>> # Initializing SAM2 vision encoder, memory attention, and memory encoder configurations
  355. >>> vision_config = Sam2VisionConfig()
  356. >>> prompt_encoder_config = Sam2PromptEncoderConfig()
  357. >>> mask_decoder_config = Sam2MaskDecoderConfig()
  358. >>> config = Sam2Config(vision_config, prompt_encoder_config, mask_decoder_config)
  359. ```"""
  360. model_type = "sam2"
  361. sub_configs = {
  362. "vision_config": AutoConfig,
  363. "prompt_encoder_config": Sam2PromptEncoderConfig,
  364. "mask_decoder_config": Sam2MaskDecoderConfig,
  365. }
  366. def __init__(
  367. self,
  368. vision_config=None,
  369. prompt_encoder_config=None,
  370. mask_decoder_config=None,
  371. initializer_range=0.02,
  372. **kwargs,
  373. ):
  374. super().__init__(**kwargs)
  375. vision_config = vision_config if vision_config is not None else {}
  376. prompt_encoder_config = prompt_encoder_config if prompt_encoder_config is not None else {}
  377. mask_decoder_config = mask_decoder_config if mask_decoder_config is not None else {}
  378. if isinstance(vision_config, dict):
  379. vision_config["model_type"] = vision_config.get("model_type", "sam2_vision_model")
  380. vision_config = CONFIG_MAPPING[vision_config["model_type"]](**vision_config)
  381. if isinstance(prompt_encoder_config, Sam2PromptEncoderConfig):
  382. prompt_encoder_config = prompt_encoder_config.to_dict()
  383. if isinstance(mask_decoder_config, Sam2MaskDecoderConfig):
  384. mask_decoder_config = mask_decoder_config.to_dict()
  385. self.vision_config = vision_config
  386. self.prompt_encoder_config = Sam2PromptEncoderConfig(**prompt_encoder_config)
  387. self.mask_decoder_config = Sam2MaskDecoderConfig(**mask_decoder_config)
  388. self.initializer_range = initializer_range
  389. __all__ = [
  390. "Sam2Config",
  391. "Sam2HieraDetConfig",
  392. "Sam2VisionConfig",
  393. "Sam2PromptEncoderConfig",
  394. "Sam2MaskDecoderConfig",
  395. ]