mirror of
https://github.com/storytold/effectcraft.git
synced 2026-11-01 05:58:19 -06:00
Encode 5.0 s -> 3.4 s per 1024x1024 frame in release on 16 threads (fused softmax, queries scaled before the matmul, flat 2-D linear layers); a click on an encoded frame is about 120 ms. Reference tests unchanged. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
52 lines
2.1 KiB
Rust
52 lines
2.1 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());
|
|
}
|
|
println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get()));
|
|
Ok(())
|
|
}
|