effectcraft/crates/sam2/tests/reference.rs

//! Stage-by-stage comparison with the PyTorch / transformers reference (`tools/sam2_reference.py`).
//!
//! Needs the `facebook/sam2.1-hiera-tiny` weights (default `<repo>/sam2-tiny`, or the folder in
//! `EFFECTCRAFT_SAM2_DIR`) and the fixture tensors (`<repo>/plan/sam2-ref/image.safetensors`, or
//! `EFFECTCRAFT_SAM2_REF`). Without them every test here skips, so CI stays green.
#![cfg(not(target_arch = "wasm32"))]

use std::collections::HashMap;
use std::path::PathBuf;

use candle_core::{Device, Tensor};
use effectcraft_sam2::hiera::{Hiera, HieraConfig};
use effectcraft_sam2::weights::Weights;

fn repo() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..")
}

fn weights_file() -> PathBuf {
    std::env::var_os("EFFECTCRAFT_SAM2_DIR").map_or_else(|| repo().join("sam2-tiny"), PathBuf::from).join("model.safetensors")
}

fn fixture_file() -> PathBuf {
    std::env::var_os("EFFECTCRAFT_SAM2_REF").map_or_else(|| repo().join("plan/sam2-ref/image.safetensors"), PathBuf::from)
}

/// The weights and the reference tensors, or `None` (the test then skips).
fn load() -> Option<(Weights, HashMap<String, Tensor>)> {
    let (wf, rf) = (weights_file(), fixture_file());
    if !wf.exists() || !rf.exists() {
        eprintln!("skipped: SAM 2 weights or reference tensors not found ({} / {})", wf.display(), rf.display());
        return None;
    }
    let w = Weights::open(&wf).unwrap();
    let r = candle_core::safetensors::load(&rf, &Device::Cpu).unwrap();
    Some((w, r))
}

fn max_abs(a: &Tensor, b: &Tensor) -> f32 {
    assert_eq!(a.dims(), b.dims(), "shape mismatch");
    (a - b).unwrap().abs().unwrap().flatten_all().unwrap().max(0).unwrap().to_scalar::<f32>().unwrap()
}

/// f32 activations differ from PyTorch by accumulation order only; the Hiera blocks stay well
/// inside this bound (measured max 2.6e-5, 2026-10-15).
const HIERA_TOL: f32 = 1e-4;

#[test]
fn hiera_matches_pytorch_block_by_block() {
    let Some((w, r)) = load() else { return };
    let hiera = Hiera::load(&w, &HieraConfig::tiny()).unwrap();
    let pixels = r.get("input.pixel_values").unwrap();
    let (emb, blocks) = hiera.forward_all(pixels).unwrap();
    let e = max_abs(&emb, r.get("vision_encoder.backbone.patch_embed.0").unwrap());
    eprintln!("patch_embed max abs error {e:e}");
    assert!(e < 1e-4, "patch embed error {e}");
    assert_eq!(blocks.len(), 12);
    for (i, b) in blocks.iter().enumerate() {
        let want = r.get(&format!("vision_encoder.backbone.blocks.{i}.0")).unwrap();
        let e = max_abs(b, want);
        eprintln!("block {i}: max abs error {e:e}");
        assert!(e < HIERA_TOL, "block {i} error {e} (tolerance {HIERA_TOL})");
    }
}

#[test]
fn hiera_is_deterministic_and_rejects_bad_sizes() {
    let Some((w, r)) = load() else { return };
    let hiera = Hiera::load(&w, &HieraConfig::tiny()).unwrap();
    let pixels = r.get("input.pixel_values").unwrap();
    let a = hiera.forward(pixels).unwrap();
    let b = hiera.forward(pixels).unwrap();
    assert_eq!(a.len(), 4);
    for (x, y) in a.iter().zip(&b) {
        assert_eq!(max_abs(x, y), 0.0);
    }
    let odd = pixels.narrow(2, 0, 1022).unwrap();
    assert!(hiera.forward(&odd).is_err());
}

/// The neck adds 1x1 convolutions and nearest-neighbour sums to the Hiera output; the sine
/// position encodings are exact functions of the grid size.
const NECK_TOL: f32 = 1e-4;

#[test]
fn neck_and_vision_encoder_match_pytorch() {
    let Some((w, r)) = load() else { return };
    let enc = effectcraft_sam2::neck::VisionEncoder::load_tiny(&w).unwrap();
    let pixels = r.get("input.pixel_values").unwrap();
    let stages = enc.hiera.forward(pixels).unwrap();
    let out = enc.neck.forward(&stages).unwrap();
    assert_eq!((out.features.len(), out.positions.len()), (4, 4));
    for (i, f) in out.features.iter().enumerate() {
        let e = max_abs(f, r.get(&format!("vision_encoder.neck.{i}")).unwrap());
        eprintln!("neck feature {i}: max abs error {e:e}");
        assert!(e < NECK_TOL, "neck feature {i} error {e}");
    }
    for (i, p) in out.positions.iter().enumerate() {
        let e = max_abs(p, r.get(&format!("vision_encoder.neck.{}", i + 4)).unwrap());
        eprintln!("neck position {i}: max abs error {e:e}");
        assert!(e < NECK_TOL, "neck position {i} error {e}");
    }
    // The vision encoder keeps the three highest-resolution levels, highest first.
    let feats = enc.forward(pixels).unwrap();
    let dims: Vec<usize> = feats.features.iter().map(|f| f.dim(2).unwrap()).collect();
    assert_eq!(dims, vec![256, 128, 64]);
    let e = max_abs(feats.features.first().unwrap(), r.get("vision_encoder.neck.3").unwrap());
    assert!(e < NECK_TOL);
}

/// Click-prompted single image: prompt embeddings and decoder outputs against `Sam2Model`.
const DECODER_TOL: f32 = 1e-3;

#[test]
fn prompt_encoder_and_decoder_match_pytorch() {
    use effectcraft_sam2::model::ImageModel;
    use effectcraft_sam2::prompt::Prompt;
    let Some((w, r)) = load() else { return };
    let model = ImageModel::load(&w).unwrap();
    let pixels = r.get("input.pixel_values").unwrap();
    let prompt = Prompt { points: vec![([460.0, 512.0], true)], bbox: None };

    let (sparse, dense) = model.embed_prompt(&prompt).unwrap();
    let want_sparse = r.get("prompt_encoder.0").unwrap().squeeze(0).unwrap();
    let e = max_abs(&sparse, &want_sparse);
    eprintln!("sparse prompt: max abs error {e:e}");
    assert!(e < 1e-5, "sparse prompt error {e}");
    let e = max_abs(&dense, r.get("prompt_encoder.1").unwrap());
    assert!(e < 1e-6, "dense prompt error {e}");

    let enc = model.encode(pixels).unwrap();
    let out = model.segment(&enc, &prompt, true).unwrap();
    let want_masks = r.get("mask_decoder.0").unwrap().squeeze(0).unwrap();
    let e = max_abs(&out.masks, &want_masks);
    eprintln!("masks: max abs error {e:e} (logit range {:?})", want_masks.flatten_all().unwrap().max(0).unwrap().to_scalar::<f32>());
    assert!(e < DECODER_TOL, "mask logits error {e}");
    let e = max_abs(&out.iou, &r.get("mask_decoder.1").unwrap().squeeze(0).unwrap());
    eprintln!("iou: max abs error {e:e}");
    assert!(e < 1e-4, "iou error {e}");
    let e = max_abs(&out.object_score, &r.get("mask_decoder.3").unwrap().squeeze(0).unwrap());
    eprintln!("object score: max abs error {e:e}");
    assert!(e < 1e-4, "object score error {e}");
    let e = max_abs(&out.sam_tokens, &r.get("mask_decoder.2").unwrap().squeeze(0).unwrap());
    assert!(e < 1e-4, "sam token error {e}");
}

fn video_fixture_file() -> PathBuf {
    std::env::var_os("EFFECTCRAFT_SAM2_VIDEO_REF").map_or_else(|| repo().join("plan/sam2-ref/video.safetensors"), PathBuf::from)
}

/// The weights and the video reference tensors (`tools/sam2_video_reference.py`), or `None`.
fn load_video() -> Option<(Weights, HashMap<String, Tensor>)> {
    let (wf, rf) = (weights_file(), video_fixture_file());
    if !wf.exists() || !rf.exists() {
        eprintln!("skipped: SAM 2 weights or video reference tensors not found ({} / {})", wf.display(), rf.display());
        return None;
    }
    Some((Weights::open(&wf).unwrap(), candle_core::safetensors::load(&rf, &Device::Cpu).unwrap()))
}

/// Measured max 3.5e-6 (memory features), 2.6e-5 (memory attention), 2026-10-15.
const MEMORY_TOL: f32 = 1e-4;

#[test]
fn memory_encoder_matches_pytorch() {
    let Some((w, r)) = load_video() else { return };
    let enc = effectcraft_sam2::memory::MemoryEncoder::load(&w).unwrap();
    let (feat, pos) = enc.encode(r.get("mem.enc.vision_features").unwrap(), r.get("mem.enc.masks").unwrap()).unwrap();
    let e = max_abs(&feat, r.get("mem.enc.out_features").unwrap());
    eprintln!("memory features: max abs error {e:e}");
    assert!(e < MEMORY_TOL, "memory feature error {e}");
    let e = max_abs(&pos, r.get("mem.enc.out_pos").unwrap());
    eprintln!("memory position encoding: max abs error {e:e}");
    assert!(e < 1e-5, "memory position error {e}");
}

#[test]
fn memory_attention_matches_pytorch() {
    let Some((w, r)) = load_video() else { return };
    let att = effectcraft_sam2::memattn::MemoryAttention::load(&w).unwrap();
    let get = |k: &str| r.get(k).unwrap().squeeze(1).unwrap();
    let pointers = r.get("mem.attn.num_pointers").unwrap().to_vec1::<f32>().unwrap()[0] as usize;
    let out = att
        .forward(
            &get("mem.attn.current_vision_features"),
            &get("mem.attn.current_vision_position_embeddings"),
            &get("mem.attn.memory"),
            &get("mem.attn.memory_posision_embeddings"),
            pointers,
        )
        .unwrap();
    let want = r.get("mem.attn.out").unwrap().squeeze(0).unwrap().squeeze(0).unwrap();
    let e = max_abs(&out, &want);
    eprintln!("memory attention: max abs error {e:e} (pointers {pointers})");
    assert!(e < MEMORY_TOL, "memory attention error {e}");
}