我正在将SAM模型和工具链移植到Apple Silicon上。在mlx上运行sam2.1-small模型时，推理速度已经提升了1.25倍。量化版本即将推出。仓库：https://github.com/avbiswas/sam2-mlx… 模型：https://huggingface.co/avbiswas/sam2.1-hiera-small-mlx-fp32…

avbiswas/sam2-mlx

来源：https://github.com/avbiswas/sam2-mlx

sam-mlx

Meta SAM 2.1 的 MLX 推理移植版本，当前目标为 facebook/sam2.1-hiera-small。

运行时包仅依赖 Python 3.14 + MLX，不安装 PyTorch。PyTorch 仅通过可选的 torch-parity 附加组件用于检查点转换和奇偶性测试用例。

当前检查点

预期的本地源检查点：

checkpoints/sam2.1_hiera_small.pt

转换后的 MLX 检查点：

checkpoints/sam2.1_hiera_small_image_segmenter.safetensors

该转换后的检查点包含：

• Hiera 图像编码器
• FPN 颈部模块
• 提示编码器
• 掩码解码器
• 对象指针投影
• 记忆编码器
• 记忆注意力模块

旧的仅图像编码器转换版本也可能存在于本地：

checkpoints/sam2.1_hiera_small_image_encoder.safetensors

生成的检查点会被 git 忽略。

环境设置

uv sync --python 3.14

如需 Torch 奇偶验证和转换脚本：

uv sync --python 3.14 --extra torch-parity

参考仓库预期存放于本地，但并非运行时依赖：

third_party/sam2
references/mlx-vlm

转换权重

uv run --extra torch-parity python scripts/convert_image_encoder_weights.py

该命令会写入：

checkpoints/sam2.1_hiera_small_image_segmenter.safetensors

奇偶验证测试用例

生成 Torch 图像嵌入验证用例：

uv run --extra torch-parity python scripts/export_torch_image_embeddings.py --frames 2
uv run python scripts/compare_image_embeddings.py

生成 Torch 提示掩码验证用例：

uv run --extra torch-parity python scripts/export_torch_prompt_mask.py
uv run python scripts/compare_prompt_mask.py

当前奇偶验证结果：

• 图像 vision_features 最大绝对误差：约 1.63e-05
• 提示低分辨率掩码最大绝对误差：约 4.67e-05
• 提示 IoU 最大绝对误差：约 4.77e-07

报告输出位置：

outputs/parity/

图像分割

运行一次提示帧并生成叠加图：

uv run python scripts/predict_image_mask.py \
  --point 500 610 \
  --output-video outputs/image_prompt_overlay.mp4 \
  --output-mask outputs/image_prompt_mask.npy

坐标基于缩放后的 1024x1024 SAM 输入空间。

视频追踪

掩码提示反馈基线：

uv run python scripts/propagate_video_masks.py --frames 30

SAM2 记忆追踪器：

uv run python scripts/track_video_memory.py --frames 289 \
  --point 500 610 \
  --output-video outputs/dog_memory_overlay_full_v3.mp4 \
  --output-mask outputs/dog_memory_masks_full_v3.npy \
  --report outputs/benchmarks/dog_memory_latency_full_v3.json

包 API，与官方 SAM2 视频预测器方法名一致：

import numpy as np

from sam_mlx import SAM2VideoPredictor

predictor = SAM2VideoPredictor(
    checkpoint="checkpoints/sam2.1_hiera_small_image_segmenter.safetensors"
)
state = predictor.init_state("third_party/sam2/demo/data/gallery/01_dog.mp4")

frame_idx, obj_ids, masks = predictor.add_new_points_or_box(
    state,
    frame_idx=0,
    obj_id=1,
    points=np.array([[625.0, 429.0]], dtype=np.float32),
    labels=np.array([1], dtype=np.int32),
)

for frame_idx, obj_ids, masks in predictor.propagate_in_video(state):
    # masks 是形状为 O,1,H,W 的 NumPy float32 数组，分辨率为原始视频分辨率。
    pass

已实现的 API 方法：

• SAM2VideoPredictor.from_pretrained(...)
• init_state(...)
• add_new_points_or_box(...)
• add_new_points(...)
• add_new_mask(...)
• propagate_in_video(...)
• clear_all_prompts_in_frame(...)
• reset_state(...)

当前记忆追踪器使用：

• 首帧点提示
• SAM2 记忆编码器
• SAM2 记忆注意力
• 对象指针
• 当单掩码追踪输出不稳定时的动态多掩码回退
• 条件帧记忆 + SAM2 风格的帧索引时序记忆选择
• 点击帧掩码在记忆编码前进行二值化，与官方后处理路径一致

狗画廊奇偶运行将 MLX 与官方 Torch SAM2VideoPredictor 在所有 289 帧上的输出进行比较：

• 所有帧的平均掩码 IoU：约 0.977
• 非空 Torch 帧的中位掩码 IoU：约 0.979
• 存在性匹配：289 / 289 帧
• 官方 Torch 叠加视频：outputs/torch_sam2_dog_overlay_full_twitter.mp4
• MLX 叠加视频：outputs/dog_memory_overlay_full_v3_twitter.mp4
• 比较报告：outputs/benchmarks/dog_memory_mlx_vs_torch_full_v3.json

重新生成官方 Torch 狗示例用例：

uv run --extra torch-parity python scripts/run_torch_sam2_dog_video.py

重新生成 MLX 与 Torch 狗比较报告：

uv run python scripts/compare_video_masks.py \
  --reference outputs/torch_sam2_dog_masks_full.npy \
  --candidate outputs/dog_memory_masks_full_v3.npy \
  --output outputs/benchmarks/dog_memory_mlx_vs_torch_full_v3.json

功能奇偶基准测试

生成官方 Torch 测试用例，用于我们仍需复现的视频 UX/状态功能：

uv run --extra torch-parity python scripts/run_torch_video_feature_benchmarks.py \
  --scenario all \
  --frames 130 \
  --frames-dir outputs/torch_feature_benchmark_frames_130f \
  --output-dir outputs/feature_benchmarks_130f

该命令会写入形状为 T,O,H,W 的掩码用例及各场景报告，涵盖：

• multi_object
• box_prompt
• negative_clicks
• cross_frame_corrections
• bidirectional_middle

当前生成的用例汇总：

outputs/feature_benchmarks_130f/torch_feature_benchmarks_summary.json

运行 MLX 并与 Torch 用例比较所有追踪功能：

uv run python scripts/run_feature_regression.py --frames 130

如需先重新生成 Torch 用例，使用：

uv run python scripts/run_feature_regression.py --refresh-torch --frames 130

在内循环开发期间，可仅比较已有输出而不重新运行 MLX：

uv run python scripts/run_feature_regression.py --skip-mlx --frames 130

直接使用相同比较器处理单个场景：

uv run python scripts/track_video_features_mlx.py \
  --scenario multi_object \
  --frames 130 \
  --output-dir outputs/feature_benchmarks_130f

uv run python scripts/compare_video_masks.py \
  --reference outputs/feature_benchmarks_130f/multi_object_torch_masks.npy \
  --candidate outputs/feature_benchmarks_130f/multi_object_mlx_masks.npy \
  --output outputs/feature_benchmarks_130f/multi_object_mlx_vs_torch.json

当前 MLX 与 Torch 功能基准测试结果（基于 130 帧狗画廊用例）：

• multi_object：平均 IoU 0.973，存在性 260 / 260
• box_prompt：平均 IoU 0.953，存在性 129 / 130
• negative_clicks：平均 IoU 0.972，存在性 130 / 130
• cross_frame_corrections：平均 IoU 0.974，存在性 130 / 130
• bidirectional_middle：平均 IoU 0.924，存在性 128 / 130

双向场景目前功能正确，但指标仍低于其他场景。剩余差异在于靠近遮挡处的两帧对象存在性边界：Torch 在第 76-77 帧保留了极小掩码，而 MLX 将这些帧判定为无对象。

在另一个画廊视频 02_cups.mp4 上，相同的双向基准测试（使用第 120 帧的中心杯子提示）结果大幅改善：

• 杯子上的 bidirectional_middle：平均 IoU 0.979，存在性 130 / 130
• 报告：outputs/cups_feature_benchmarks_130f/bidirectional_middle_mlx_vs_torch.json
• 叠加视频：outputs/cups_feature_benchmarks_130f/bidirectional_middle_mlx_overlay.mp4

MLX 功能运行器现在使用共享视频状态处理多个对象。每帧仅编码一次，为每个对象维护条件/非条件记忆库，并支持 NLE 风格的修正重放：

• 从修正帧开始正向重放，用于正/负点击编辑
• 从中间帧修正开始双向重放
• 编辑重放范围外的掩码保持不变

提供了一个仅 MLX 的检查场景，用于中间帧编辑器修正：

uv run python scripts/track_video_features_mlx.py \
  --scenario nle_bidirectional_correction \
  --frames 90 \
  --output-dir outputs/feature_benchmarks_130f

功能报告和检查叠加视频写入以下目录：

outputs/feature_benchmarks_130f/

叠加工具

将掩码渲染到视频上：

uv run python scripts/overlay_masks.py \
  --masks outputs/dog_memory_masks_full_v3.npy \
  --output outputs/dog_memory_overlay_from_masks.mp4

叠加脚本接受 .npy 或 .npz 格式的掩码，形状为 T,H,W 或 T,1,H,W。合成叠加仅用于写入冒烟测试，需要：

uv run python scripts/overlay_masks.py --synthetic-smoke-test

基准测试

图像编码器：

uv run --extra torch-parity python scripts/benchmark_image_encoder.py --warmup 3 --runs 10

提示分割：

uv run python scripts/benchmark_prompt_segmenter.py --warmup 3 --runs 20

视频记忆追踪：

uv run python scripts/track_video_memory.py --frames 150 \
  --report outputs/benchmarks/video_memory_latency_150f.json

本机当前参考数据：

• 图像编码器 MLX：约 81 ms/帧
• 图像编码器 Torch/MPS：约 104 ms/帧
• MLX 图像编码器加速比：约 1.28x
• 缓存提示解码：约 4 ms
• 完整图像 + 提示：约 85 ms
• 狗完整视频记忆追踪器：约 269 ms/帧（289 帧运行）

基准测试报告写入以下目录：

outputs/benchmarks/

运行时依赖边界

默认运行时不应包含 Torch：

uv sync --python 3.14
uv run python - <<'PY'
import importlib.util as u
print({m: bool(u.find_spec(m)) for m in ["torch", "torchvision", "hydra", "iopath", "mlx", "cv2"]})
PY

预期输出：

torch=False, torchvision=False, hydra=False, iopath=False, mlx=True, cv2=True