Re-convert: store sam2_convs and boxes_pool_project in MLX (channels-last) layout

#2
by evo42 - opened

Re-converted from facebook/sam3 so the 12 conv weights that were published in PyTorch
(channels-first) layout are stored in MLX layout.

Why

sam3/convert.py in Deekshith-Dade/mlx_sam3
transposed the SAM-3 neck convs via a hard-coded key list. Anything outside that list shipped
channels-first: the 11 sam2_convs.* twins (the second neck, built only when
enable_inst_interactivity=True) and geometry_encoder.boxes_pool_project. The text-prompt path
never builds sam2_convs, which is why this only showed up on the box-prompt path, where
set_image crashed in mx.conv_transpose2d.

Fixed upstream in two steps:

  • PR #9 (merged) β€” a shape-guarded,
    idempotent _sanitize_conv_layout that repairs these weights at load time.
  • PR #10 β€” makes convert.py derive the
    permutations from the model's module types instead of a hard-coded list, so a re-convert produces
    correct weights. This file is the output of that converter.

Once this is merged, the load-time shim becomes a no-op (verified: it touches 0 keys on this file).

What changed

12 tensors, transposed only β€” nn.ConvTranspose2d via (1, 2, 3, 0), nn.Conv2d via
(0, 2, 3, 1):

tensor before after
backbone.vision_backbone.sam2_convs.0.conv_1x1.weight (256, 256, 1, 1) (256, 1, 1, 256)
backbone.vision_backbone.sam2_convs.0.conv_3x3.weight (256, 256, 3, 3) (256, 3, 3, 256)
backbone.vision_backbone.sam2_convs.0.dconv_2x2_0.weight (1024, 512, 2, 2) (512, 2, 2, 1024)
backbone.vision_backbone.sam2_convs.0.dconv_2x2_1.weight (512, 256, 2, 2) (256, 2, 2, 512)
backbone.vision_backbone.sam2_convs.1.conv_1x1.weight (256, 512, 1, 1) (256, 1, 1, 512)
backbone.vision_backbone.sam2_convs.1.conv_3x3.weight (256, 256, 3, 3) (256, 3, 3, 256)
backbone.vision_backbone.sam2_convs.1.dconv_2x2.weight (1024, 512, 2, 2) (512, 2, 2, 1024)
backbone.vision_backbone.sam2_convs.2.conv_1x1.weight (256, 1024, 1, 1) (256, 1, 1, 1024)
backbone.vision_backbone.sam2_convs.2.conv_3x3.weight (256, 256, 3, 3) (256, 3, 3, 256)
backbone.vision_backbone.sam2_convs.3.conv_1x1.weight (256, 1024, 1, 1) (256, 1, 1, 1024)
backbone.vision_backbone.sam2_convs.3.conv_3x3.weight (256, 256, 3, 3) (256, 3, 3, 256)
geometry_encoder.boxes_pool_project.weight (256, 256, 7, 7) (256, 7, 7, 256)

The other 1388 tensors are bit-identical to the current revision. Same 1400 keys, same dtypes
(1368x F32, 32x C64), same total file size (3 402 867 661 bytes), and
model.safetensors.index.json is unchanged, so it is not part of this PR.

Verification (Apple M5 Pro)

  • _sanitize_conv_layout touches 0 keys on this file, with enable_inst_interactivity both True
    and False
  • text prompt "person" -> 2 boxes @ 0.954 / 0.772
  • box prompt via Sam3Processor.add_geometric_prompt -> mask @ 0.9778

Both numbers are identical to what the current revision produces once the load-time shim has
repaired it, i.e. this changes the stored layout and nothing about the model's output.

Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment