Load the SAM 3.1 multiplex checkpoint once, without spurious key warnings (fixes #506)

sam3/model/vitdet.py

@@ -551,8 +551,9 @@ def _setup_rope_freqs(self) -> None:
 
         self.register_buffer("freqs_cis", freqs_cis)
         if self.use_rope_real:
-            self.register_buffer("freqs_cis_real", freqs_cis.real)
-            self.register_buffer("freqs_cis_imag", freqs_cis.imag)
+            # Derived from freqs_cis, so not saved in (or expected from) checkpoints
+            self.register_buffer("freqs_cis_real", freqs_cis.real, persistent=False)
+            self.register_buffer("freqs_cis_imag", freqs_cis.imag, persistent=False)
 
     def _apply_rope(self, q, k) -> Tuple[Tensor, Tensor]:
         if not self.use_rope:

sam3/model_builder.py

@@ -1117,9 +1117,10 @@ def build_sam3_multiplex_video_predictor(
     )
     from sam3.model.sam3_multiplex_video_predictor import Sam3MultiplexVideoPredictor
 
-    # Build tracker
+    # Build tracker. Its weights come from the merged checkpoint, which is loaded
+    # into the full demo model below (keys are prefixed with "tracker.model.").
     tracker_model = build_sam3_multiplex_video_model(
-        checkpoint_path=checkpoint_path,
+        checkpoint_path=None,
         load_from_HF=False,
         multiplex_count=multiplex_count,
         use_fa3=use_fa3,

test/test_multiplex_checkpoint_loading.py

@@ -0,0 +1,66 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
+
+"""Tests for SAM 3.1 multiplex checkpoint loading (#506)."""
+
+import unittest
+from unittest.mock import patch
+
+from sam3 import model_builder
+from sam3.model.vitdet import ViT
+
+
+class _StopBuild(Exception):
+    pass
+
+
+class TestMultiplexCheckpointLoading(unittest.TestCase):
+    def test_tracker_does_not_load_merged_checkpoint(self) -> None:
+        # The merged checkpoint is loaded once into the full demo model; loading it
+        # into the bare tracker first matches no keys and only prints warnings.
+        with patch.object(
+            model_builder,
+            "build_sam3_multiplex_video_model",
+            side_effect=_StopBuild,
+        ) as build_tracker:
+            with self.assertRaises(_StopBuild):
+                model_builder.build_sam3_multiplex_video_predictor(
+                    checkpoint_path="sam3.1_multiplex.pt"
+                )
+        self.assertIsNone(build_tracker.call_args.kwargs["checkpoint_path"])
+        self.assertFalse(build_tracker.call_args.kwargs["load_from_HF"])
+
+    def test_real_rope_buffers_are_not_persistent(self) -> None:
+        def make_vit(use_rope_real):
+            return ViT(
+                img_size=56,
+                pretrain_img_size=56,
+                patch_size=14,
+                embed_dim=32,
+                depth=1,
+                num_heads=2,
+                use_rope=True,
+                use_abs_pos=False,
+                tile_abs_pos=False,
+                window_size=0,
+                global_att_blocks=(0,),
+                rel_pos_blocks=(),
+                retain_cls_token=False,
+                use_act_checkpoint=False,
+                use_rope_real=use_rope_real,
+            )
+
+        vit = make_vit(use_rope_real=True)
+        attn = vit.blocks[0].attn
+        self.assertIsNotNone(attn.freqs_cis_real)
+        self.assertIsNotNone(attn.freqs_cis_imag)
+
+        keys = vit.state_dict().keys()
+        self.assertFalse(
+            any(k.endswith(("freqs_cis_real", "freqs_cis_imag")) for k in keys)
+        )
+        # A checkpoint saved without real-valued RoPE loads strictly
+        vit.load_state_dict(make_vit(use_rope_real=False).state_dict())
+
+
+if __name__ == "__main__":
+    unittest.main()