mirror of
https://github.com/storytold/effectcraft.git
synced 2026-10-31 21:48:20 -06:00
Kristian MyramandClaude Sonnet 5.5
b7c78065db
Sine position encodings depend only on the sizes: built once (150 ms of a frame). do_pool takes the maximum over grouped cells instead of permute + max_pool2d + permute (tested equal, odd sizes included). Encode 2.6 s -> 2.4 s per frame in release on 16 threads (3.4 s at the start of the speed work); reference tests unchanged. Chunked global attention was tried and did not help. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
62 lines
2.6 KiB
Rust
62 lines
2.6 KiB
Rust
//! Timing of the SAM 2 image path on this machine: load, encode one frame, decode a click.
|
|
//!
|
|
//! `cargo run -p effectcraft-sam2 --example sam2_bench -- [MODEL_DIR]` (default `sam2-tiny`).
|
|
//! Prints a message and exits when the weights are not there.
|
|
|
|
use std::time::Instant;
|
|
|
|
use candle_core::{DType, Device, Tensor};
|
|
use effectcraft_sam2::model::ImageModel;
|
|
use effectcraft_sam2::prompt::Prompt;
|
|
|
|
fn main() {
|
|
if let Err(e) = run() {
|
|
eprintln!("sam2_bench: {e}");
|
|
std::process::exit(1);
|
|
}
|
|
}
|
|
|
|
fn run() -> Result<(), Box<dyn std::error::Error>> {
|
|
let dir = std::env::args().nth(1).unwrap_or_else(|| "sam2-tiny".to_string());
|
|
let file = std::path::Path::new(&dir).join("model.safetensors");
|
|
if !file.exists() {
|
|
println!("no weights at {} (download facebook/sam2.1-hiera-tiny first)", file.display());
|
|
return Ok(());
|
|
}
|
|
let t = Instant::now();
|
|
let model = ImageModel::open(&file)?;
|
|
println!("load: {:.2} s", t.elapsed().as_secs_f64());
|
|
let pixels = Tensor::randn(0f32, 1.0, (1, 3, 1024, 1024), &Device::Cpu)?.to_dtype(DType::F32)?;
|
|
for run in 0..3 {
|
|
let t = Instant::now();
|
|
let enc = model.encode(&pixels)?;
|
|
let encode = t.elapsed().as_secs_f64();
|
|
let prompt = Prompt { points: vec![([512.0, 512.0], true)], bbox: None };
|
|
let t = Instant::now();
|
|
let out = model.segment(&enc, &prompt, true)?;
|
|
println!("run {run}: encode {encode:.2} s, click {:.0} ms, masks {:?}", t.elapsed().as_secs_f64() * 1e3, out.masks.dims());
|
|
}
|
|
// Where the encode time goes: patch embedding, then each Hiera block.
|
|
let hiera = &model.vision.hiera;
|
|
let t = Instant::now();
|
|
let mut x = hiera.patch_embed(&pixels)?;
|
|
println!("patch embed: {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
|
|
for (i, b) in hiera.blocks.iter().enumerate() {
|
|
let t = Instant::now();
|
|
x = b.forward(&x)?;
|
|
println!("block {i}: {:.0} ms, out {:?}", t.elapsed().as_secs_f64() * 1e3, x.dims());
|
|
}
|
|
// The rest of the encode: the neck (1x1 convolutions, top-down sums) and its sine position encodings.
|
|
let stages = hiera.forward(&pixels)?;
|
|
let t = Instant::now();
|
|
let _ = model.vision.neck.forward(&stages)?;
|
|
println!("neck (convolutions + positions): {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
|
|
let t = Instant::now();
|
|
for (h, w) in [(32, 32), (64, 64), (128, 128), (256, 256)] {
|
|
let _ = effectcraft_sam2::neck::sine_position(128, h, w)?;
|
|
}
|
|
println!("sine positions for the four levels: {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
|
|
println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get()));
|
|
Ok(())
|
|
}
|