mirror of
https://github.com/storytold/effectcraft.git
synced 2026-10-31 21:48:20 -06:00
S2.13: prompt encoder and two-way mask decoder match PyTorch
A click on the synthetic image: mask logits within 1.1e-4 (range 11.6), IoU head 1e-6, object score 2e-5, prompt embeddings 1e-7. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
Kristian Myram
co-authored by
Claude Sonnet 5.5
Claude Sonnet 5.5
parent
d2291926f0
commit
8a6ad6f5e9
@@ -0,0 +1,258 @@
|
||||
// Modified work: ported from Hugging Face `transformers` `models/sam2/modeling_sam2.py`
|
||||
// (Sam2Attention, Sam2TwoWayAttentionBlock, Sam2TwoWayTransformer, Sam2MaskDecoder; Apache-2.0,
|
||||
// Copyright The HuggingFace Inc. team) to Rust on candle. See NOTICE.
|
||||
//! The two-way mask decoder: output tokens and prompt tokens attend to the image and back, then
|
||||
//! hyper-networks turn the mask tokens into masks over the up-scaled image embedding.
|
||||
//! One object per call (batch 1, one point batch).
|
||||
|
||||
use candle_core::{D, Tensor};
|
||||
use candle_nn::{LayerNorm, Linear, Module, ops::sigmoid, ops::softmax};
|
||||
|
||||
use crate::nn::{layer_norm, linear};
|
||||
use crate::weights::Weights;
|
||||
use crate::{Error, Result};
|
||||
|
||||
const DIM: usize = 256;
|
||||
const MASK_TOKENS: usize = 4;
|
||||
|
||||
struct Attention {
|
||||
q: Linear,
|
||||
k: Linear,
|
||||
v: Linear,
|
||||
o: Linear,
|
||||
heads: usize,
|
||||
scale: f64,
|
||||
}
|
||||
|
||||
impl Attention {
|
||||
fn load(w: &Weights, prefix: &str, downsample: usize, heads: usize) -> Result<Attention> {
|
||||
let inner = DIM / downsample;
|
||||
Ok(Attention {
|
||||
q: linear(w, &format!("{prefix}.q_proj"), DIM, inner)?,
|
||||
k: linear(w, &format!("{prefix}.k_proj"), DIM, inner)?,
|
||||
v: linear(w, &format!("{prefix}.v_proj"), DIM, inner)?,
|
||||
o: linear(w, &format!("{prefix}.o_proj"), inner, DIM)?,
|
||||
heads,
|
||||
scale: ((inner / heads) as f64).powf(-0.5),
|
||||
})
|
||||
}
|
||||
|
||||
/// `q`: `[1, nq, 256]`, `k` and `v`: `[1, nk, 256]`.
|
||||
fn forward(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
|
||||
let split = |t: Tensor| -> Result<Tensor> { Ok(t.reshape((1, (), self.heads, t.dim(2)? / self.heads))?.transpose(1, 2)?.contiguous()?) };
|
||||
let (q, k, v) = (split(self.q.forward(q)?)?, split(self.k.forward(k)?)?, split(self.v.forward(v)?)?);
|
||||
let att = softmax(&(q.matmul(&k.transpose(2, 3)?.contiguous()?)? * self.scale)?, D::Minus1)?;
|
||||
let out = att.matmul(&v)?.transpose(1, 2)?.contiguous()?;
|
||||
let (n, h, d) = (out.dim(1)?, out.dim(2)?, out.dim(3)?);
|
||||
Ok(self.o.forward(&out.reshape((1, n, h * d))?)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// Feed-forward with ReLU, `Sam2FeedForward`: `proj_in`, `layers`, `proj_out`.
|
||||
pub struct FeedForward {
|
||||
proj_in: Linear,
|
||||
layers: Vec<Linear>,
|
||||
proj_out: Linear,
|
||||
sigmoid_output: bool,
|
||||
}
|
||||
|
||||
impl FeedForward {
|
||||
pub fn load(w: &Weights, prefix: &str, input: usize, hidden: usize, output: usize, num_layers: usize, sigmoid_output: bool) -> Result<FeedForward> {
|
||||
Ok(FeedForward {
|
||||
proj_in: linear(w, &format!("{prefix}.proj_in"), input, hidden)?,
|
||||
layers: (0..num_layers.saturating_sub(2)).map(|i| linear(w, &format!("{prefix}.layers.{i}"), hidden, hidden)).collect::<Result<_>>()?,
|
||||
proj_out: linear(w, &format!("{prefix}.proj_out"), hidden, output)?,
|
||||
sigmoid_output,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let mut x = self.proj_in.forward(x)?.relu()?;
|
||||
for l in &self.layers {
|
||||
x = l.forward(&x)?.relu()?;
|
||||
}
|
||||
let x = self.proj_out.forward(&x)?;
|
||||
Ok(if self.sigmoid_output { sigmoid(&x)? } else { x })
|
||||
}
|
||||
}
|
||||
|
||||
struct Block {
|
||||
self_attn: Attention,
|
||||
ln1: LayerNorm,
|
||||
token_to_image: Attention,
|
||||
ln2: LayerNorm,
|
||||
mlp: FeedForward,
|
||||
ln3: LayerNorm,
|
||||
ln4: LayerNorm,
|
||||
image_to_token: Attention,
|
||||
skip_first_pe: bool,
|
||||
}
|
||||
|
||||
impl Block {
|
||||
fn forward(&self, queries: &Tensor, keys: &Tensor, query_pe: &Tensor, key_pe: &Tensor) -> Result<(Tensor, Tensor)> {
|
||||
let mut queries = if self.skip_first_pe {
|
||||
self.self_attn.forward(queries, queries, queries)?
|
||||
} else {
|
||||
let q = (queries + query_pe)?;
|
||||
(queries + self.self_attn.forward(&q, &q, queries)?)?
|
||||
};
|
||||
queries = self.ln1.forward(&queries)?;
|
||||
let (q, k) = ((&queries + query_pe)?, (keys + key_pe)?);
|
||||
queries = self.ln2.forward(&(&queries + self.token_to_image.forward(&q, &k, keys)?)?)?;
|
||||
queries = self.ln3.forward(&(&queries + self.mlp.forward(&queries)?)?)?;
|
||||
let (q, k) = ((&queries + query_pe)?, (keys + key_pe)?);
|
||||
let keys = self.ln4.forward(&(keys + self.image_to_token.forward(&k, &q, &queries)?)?)?;
|
||||
Ok((queries, keys))
|
||||
}
|
||||
}
|
||||
|
||||
/// What the decoder returns for one object.
|
||||
pub struct DecoderOutput {
|
||||
/// Mask logits `[1, k, 4 * GRID, 4 * GRID]` (`k` is 3 for multimask output, else 1).
|
||||
pub masks: Tensor,
|
||||
/// Predicted IoU per mask `[1, k]`.
|
||||
pub iou: Tensor,
|
||||
/// The output tokens of the chosen masks `[1, k, 256]` (the object pointer is made from these).
|
||||
pub sam_tokens: Tensor,
|
||||
/// Object-presence logit `[1, 1]` (negative: the object is not in the frame).
|
||||
pub object_score: Tensor,
|
||||
}
|
||||
|
||||
pub struct MaskDecoder {
|
||||
obj_token: Tensor,
|
||||
iou_token: Tensor,
|
||||
mask_tokens: Tensor,
|
||||
layers: Vec<Block>,
|
||||
final_attn: Attention,
|
||||
final_ln: LayerNorm,
|
||||
up1: (Tensor, Tensor),
|
||||
up2: (Tensor, Tensor),
|
||||
up_ln: LayerNorm,
|
||||
conv_s0: (Tensor, Tensor),
|
||||
conv_s1: (Tensor, Tensor),
|
||||
hyper: Vec<FeedForward>,
|
||||
iou_head: FeedForward,
|
||||
obj_head: FeedForward,
|
||||
}
|
||||
|
||||
fn conv_w(w: &Weights, name: &str, i: usize, o: usize, k: usize, transposed: bool) -> Result<(Tensor, Tensor)> {
|
||||
let shape = if transposed { [i, o, k, k] } else { [o, i, k, k] };
|
||||
Ok((w.get(&format!("{name}.weight"), &shape)?, w.get(&format!("{name}.bias"), &[o])?))
|
||||
}
|
||||
|
||||
fn add_bias(x: Tensor, b: &Tensor) -> Result<Tensor> {
|
||||
Ok(x.broadcast_add(&b.reshape((1, b.dim(0)?, 1, 1))?)?)
|
||||
}
|
||||
|
||||
impl MaskDecoder {
|
||||
pub fn load(w: &Weights) -> Result<MaskDecoder> {
|
||||
let layers = (0..2)
|
||||
.map(|i| {
|
||||
let p = format!("mask_decoder.transformer.layers.{i}");
|
||||
Ok(Block {
|
||||
self_attn: Attention::load(w, &format!("{p}.self_attn"), 1, 8)?,
|
||||
ln1: layer_norm(w, &format!("{p}.layer_norm1"), DIM, 1e-5)?,
|
||||
token_to_image: Attention::load(w, &format!("{p}.cross_attn_token_to_image"), 2, 8)?,
|
||||
ln2: layer_norm(w, &format!("{p}.layer_norm2"), DIM, 1e-5)?,
|
||||
mlp: FeedForward::load(w, &format!("{p}.mlp"), DIM, 2048, DIM, 2, false)?,
|
||||
ln3: layer_norm(w, &format!("{p}.layer_norm3"), DIM, 1e-5)?,
|
||||
ln4: layer_norm(w, &format!("{p}.layer_norm4"), DIM, 1e-5)?,
|
||||
image_to_token: Attention::load(w, &format!("{p}.cross_attn_image_to_token"), 2, 8)?,
|
||||
skip_first_pe: i == 0,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
Ok(MaskDecoder {
|
||||
obj_token: w.get("mask_decoder.obj_score_token.weight", &[1, DIM])?,
|
||||
iou_token: w.get("mask_decoder.iou_token.weight", &[1, DIM])?,
|
||||
mask_tokens: w.get("mask_decoder.mask_tokens.weight", &[MASK_TOKENS, DIM])?,
|
||||
layers,
|
||||
final_attn: Attention::load(w, "mask_decoder.transformer.final_attn_token_to_image", 2, 8)?,
|
||||
final_ln: layer_norm(w, "mask_decoder.transformer.layer_norm_final_attn", DIM, 1e-5)?,
|
||||
up1: conv_w(w, "mask_decoder.upscale_conv1", DIM, DIM / 4, 2, true)?,
|
||||
up2: conv_w(w, "mask_decoder.upscale_conv2", DIM / 4, DIM / 8, 2, true)?,
|
||||
up_ln: layer_norm(w, "mask_decoder.upscale_layer_norm", DIM / 4, 1e-6)?,
|
||||
conv_s0: conv_w(w, "mask_decoder.conv_s0", DIM, DIM / 8, 1, false)?,
|
||||
conv_s1: conv_w(w, "mask_decoder.conv_s1", DIM, DIM / 4, 1, false)?,
|
||||
hyper: (0..MASK_TOKENS)
|
||||
.map(|i| FeedForward::load(w, &format!("mask_decoder.output_hypernetworks_mlps.{i}"), DIM, DIM, DIM / 8, 3, false))
|
||||
.collect::<Result<_>>()?,
|
||||
iou_head: FeedForward::load(w, "mask_decoder.iou_prediction_head", DIM, DIM, MASK_TOKENS, 3, true)?,
|
||||
obj_head: FeedForward::load(w, "mask_decoder.pred_obj_score_head", DIM, DIM, 1, 3, false)?,
|
||||
})
|
||||
}
|
||||
|
||||
/// Project the two highest-resolution neck levels to the decoder's skip features
|
||||
/// (`conv_s0` on the 4x level, `conv_s1` on the 2x level), `[1, 32, 4g, 4g]` and `[1, 64, 2g, 2g]`.
|
||||
pub fn high_res(&self, level0: &Tensor, level1: &Tensor) -> Result<(Tensor, Tensor)> {
|
||||
let s0 = add_bias(level0.conv2d(&self.conv_s0.0, 0, 1, 1, 1)?, &self.conv_s0.1)?;
|
||||
let s1 = add_bias(level1.conv2d(&self.conv_s1.0, 0, 1, 1, 1)?, &self.conv_s1.1)?;
|
||||
Ok((s0, s1))
|
||||
}
|
||||
|
||||
/// Predict masks. `image` and `pos`: `[1, 256, g, g]` (the lowest-resolution level with the
|
||||
/// no-memory embedding already added when there is no memory), `sparse`: `[1, n, 256]`,
|
||||
/// `dense`: `[1, 256, g, g]`, `high`: the skip features from [`MaskDecoder::high_res`].
|
||||
pub fn forward(&self, image: &Tensor, pos: &Tensor, sparse: &Tensor, dense: &Tensor, high: &(Tensor, Tensor), multimask: bool) -> Result<DecoderOutput> {
|
||||
let (_, c, g, gw) = image.dims4()?;
|
||||
if c != DIM || g != gw {
|
||||
return Err(Error::Input(format!("image embedding must be 256 x g x g, got {:?}", image.dims())));
|
||||
}
|
||||
let output_tokens = Tensor::cat(&[&self.obj_token, &self.iou_token, &self.mask_tokens], 0)?.unsqueeze(0)?;
|
||||
let tokens = Tensor::cat(&[&output_tokens, sparse], 1)?;
|
||||
let image = (image + dense)?;
|
||||
let flat = |t: &Tensor| -> Result<Tensor> { Ok(t.flatten_from(2)?.transpose(1, 2)?.contiguous()?) };
|
||||
let (mut keys, key_pe) = (flat(&image)?, flat(pos)?);
|
||||
let mut queries = tokens.clone();
|
||||
for l in &self.layers {
|
||||
(queries, keys) = l.forward(&queries, &keys, &tokens, &key_pe)?;
|
||||
}
|
||||
let (q, k) = ((&queries + &tokens)?, (&keys + &key_pe)?);
|
||||
let queries = self.final_ln.forward(&(&queries + self.final_attn.forward(&q, &k, &keys)?)?)?;
|
||||
|
||||
let iou_out = queries.narrow(1, 1, 1)?.squeeze(1)?;
|
||||
let mask_tokens_out = queries.narrow(1, 2, MASK_TOKENS)?;
|
||||
let obj_out = queries.narrow(1, 0, 1)?.squeeze(1)?;
|
||||
|
||||
let emb = keys.transpose(1, 2)?.contiguous()?.reshape((1, DIM, g, g))?;
|
||||
let (s0, s1) = high;
|
||||
let up = add_bias(emb.conv_transpose2d(&self.up1.0, 0, 0, 2, 1)?, &self.up1.1)?;
|
||||
let up = (up + s1)?;
|
||||
let up = self.up_ln.forward(&up.permute((0, 2, 3, 1))?)?.permute((0, 3, 1, 2))?.contiguous()?.gelu_erf()?;
|
||||
let up = (add_bias(up.conv_transpose2d(&self.up2.0, 0, 0, 2, 1)?, &self.up2.1)? + s0)?.gelu_erf()?;
|
||||
|
||||
let hyper: Vec<Tensor> = self.hyper.iter().enumerate().map(|(i, m)| m.forward(&mask_tokens_out.narrow(1, i, 1)?.squeeze(1)?)).collect::<Result<_>>()?;
|
||||
let hyper = Tensor::stack(&hyper, 1)?;
|
||||
let (_, uc, uh, uw) = up.dims4()?;
|
||||
let masks = hyper.matmul(&up.reshape((1, uc, uh * uw))?)?.reshape((1, MASK_TOKENS, uh, uw))?;
|
||||
let iou = self.iou_head.forward(&iou_out)?;
|
||||
let object_score = self.obj_head.forward(&obj_out)?;
|
||||
|
||||
let (masks, iou, sam_tokens) = if multimask {
|
||||
(masks.narrow(1, 1, 3)?, iou.narrow(1, 1, 3)?, mask_tokens_out.narrow(1, 1, 3)?)
|
||||
} else {
|
||||
let (m, i) = dynamic_single(&masks, &iou)?;
|
||||
(m, i, mask_tokens_out.narrow(1, 0, 1)?)
|
||||
};
|
||||
Ok(DecoderOutput { masks, iou, sam_tokens, object_score })
|
||||
}
|
||||
}
|
||||
|
||||
/// Single-mask output with the stability fallback (`_dynamic_multimask_via_stability`): if the
|
||||
/// first mask token's mask is unstable, use the best of the other three.
|
||||
fn dynamic_single(masks: &Tensor, iou: &Tensor) -> Result<(Tensor, Tensor)> {
|
||||
const DELTA: f32 = 0.05;
|
||||
const THRESH: f32 = 0.98;
|
||||
let single = masks.narrow(1, 0, 1)?;
|
||||
let single_iou = iou.narrow(1, 0, 1)?;
|
||||
let logits = single.flatten_all()?.to_vec1::<f32>()?;
|
||||
let area_i = logits.iter().filter(|v| **v > DELTA).count() as f32;
|
||||
let area_u = logits.iter().filter(|v| **v > -DELTA).count() as f32;
|
||||
let stability = if area_u > 0.0 { area_i / area_u } else { 1.0 };
|
||||
if stability >= THRESH {
|
||||
return Ok((single, single_iou));
|
||||
}
|
||||
let multi_iou = iou.narrow(1, 1, 3)?.flatten_all()?.to_vec1::<f32>()?;
|
||||
let best = multi_iou.iter().enumerate().fold((0usize, f32::NEG_INFINITY), |b, (i, v)| if *v > b.1 { (i, *v) } else { b }).0;
|
||||
Ok((masks.narrow(1, 1 + best, 1)?, iou.narrow(1, 1 + best, 1)?))
|
||||
}
|
||||
@@ -11,9 +11,12 @@
|
||||
|
||||
use std::fmt;
|
||||
|
||||
pub mod decoder;
|
||||
pub mod hiera;
|
||||
pub mod model;
|
||||
pub mod neck;
|
||||
pub mod nn;
|
||||
pub mod prompt;
|
||||
pub mod weights;
|
||||
|
||||
/// Everything that can go wrong in this crate. Never a panic.
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
// Modified work: ported from Hugging Face `transformers` `Sam2Model` (Apache-2.0, Copyright The
|
||||
// HuggingFace Inc. team) to Rust on candle. See NOTICE.
|
||||
//! The single-image model: encode a frame once, then decode any number of prompts against it.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use candle_core::{Device, Tensor};
|
||||
|
||||
use crate::decoder::{DecoderOutput, MaskDecoder};
|
||||
use crate::neck::{ImageFeatures, VisionEncoder};
|
||||
use crate::prompt::{INPUT_SIZE, Prompt, PromptEncoder};
|
||||
use crate::weights::Weights;
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// A frame after the vision encoder: reusable for every prompt on that frame.
|
||||
pub struct Encoded {
|
||||
/// The three neck levels, highest resolution first (`[1, 256, 256, 256]`, `128`, `64`).
|
||||
pub features: ImageFeatures,
|
||||
/// The decoder's skip features from the two highest-resolution levels.
|
||||
pub high: (Tensor, Tensor),
|
||||
}
|
||||
|
||||
pub struct ImageModel {
|
||||
pub vision: VisionEncoder,
|
||||
pub prompt: PromptEncoder,
|
||||
pub decoder: MaskDecoder,
|
||||
no_memory: Tensor,
|
||||
image_pos: Tensor,
|
||||
}
|
||||
|
||||
impl ImageModel {
|
||||
pub fn load(w: &Weights) -> Result<ImageModel> {
|
||||
let prompt = PromptEncoder::load(w)?;
|
||||
Ok(ImageModel {
|
||||
vision: VisionEncoder::load_tiny(w)?,
|
||||
decoder: MaskDecoder::load(w)?,
|
||||
no_memory: w.get("no_memory_embedding", &[1, 1, 256])?.reshape((1, 256, 1, 1))?,
|
||||
image_pos: prompt.image_positions()?,
|
||||
prompt,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn open(path: &Path) -> Result<ImageModel> {
|
||||
ImageModel::load(&Weights::open(path)?)
|
||||
}
|
||||
|
||||
/// Encode a normalised `[1, 3, 1024, 1024]` frame.
|
||||
pub fn encode(&self, pixels: &Tensor) -> Result<Encoded> {
|
||||
if pixels.dims() != [1, 3, INPUT_SIZE, INPUT_SIZE] {
|
||||
return Err(Error::Input(format!("pixels must be [1, 3, {INPUT_SIZE}, {INPUT_SIZE}], got {:?}", pixels.dims())));
|
||||
}
|
||||
let features = self.vision.forward(pixels)?;
|
||||
let (l0, l1) = match (features.features.first(), features.features.get(1)) {
|
||||
(Some(a), Some(b)) => (a, b),
|
||||
_ => return Err(Error::Compute("vision encoder returned too few levels".into())),
|
||||
};
|
||||
let high = self.decoder.high_res(l0, l1)?;
|
||||
Ok(Encoded { features, high })
|
||||
}
|
||||
|
||||
/// Prompt embeddings (sparse tokens, empty dense map) for a click/box prompt.
|
||||
pub fn embed_prompt(&self, p: &Prompt) -> Result<(Tensor, Tensor)> {
|
||||
Ok((self.prompt.sparse(p)?, self.prompt.dense_empty()?))
|
||||
}
|
||||
|
||||
/// Decode masks for a prompt on an encoded frame.
|
||||
pub fn segment(&self, enc: &Encoded, p: &Prompt, multimask: bool) -> Result<DecoderOutput> {
|
||||
let low = enc.features.features.get(2).ok_or_else(|| Error::Compute("vision encoder returned too few levels".into()))?;
|
||||
let image = low.broadcast_add(&self.no_memory)?;
|
||||
let (sparse, dense) = self.embed_prompt(p)?;
|
||||
self.decoder.forward(&image, &self.image_pos, &sparse, &dense, &enc.high, multimask)
|
||||
}
|
||||
|
||||
pub fn device(&self) -> Device {
|
||||
Device::Cpu
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
// Modified work: ported from Hugging Face `transformers` `models/sam2/modeling_sam2.py`
|
||||
// (Sam2PositionalEmbedding, Sam2MaskEmbedding, Sam2PromptEncoder; Apache-2.0, Copyright The
|
||||
// HuggingFace Inc. team) to Rust on candle. See NOTICE.
|
||||
//! The prompt encoder: clicks and boxes become sparse tokens, a mask or "no mask" becomes a dense map.
|
||||
|
||||
use candle_core::{Device, Tensor};
|
||||
use candle_nn::{LayerNorm, Module};
|
||||
|
||||
use crate::nn::layer_norm;
|
||||
use crate::weights::Weights;
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// Side of the (square) model input in pixels; prompts are given in these coordinates.
|
||||
pub const INPUT_SIZE: usize = 1024;
|
||||
/// Side of the image-embedding grid (input / 16).
|
||||
pub const GRID: usize = INPUT_SIZE / 16;
|
||||
const DIM: usize = 256;
|
||||
|
||||
/// What the user pointed at, in model-input pixels (0..1024).
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct Prompt {
|
||||
/// Clicks and whether each is foreground.
|
||||
pub points: Vec<([f32; 2], bool)>,
|
||||
/// A box `[x0, y0, x1, y1]`.
|
||||
pub bbox: Option<[f32; 4]>,
|
||||
}
|
||||
|
||||
impl Prompt {
|
||||
/// Reject non-finite or absurd coordinates (they come from the user or an agent).
|
||||
pub fn validate(&self) -> Result<()> {
|
||||
let bad = |v: f32| !v.is_finite() || v.abs() > 4.0 * INPUT_SIZE as f32;
|
||||
if self.points.iter().any(|(p, _)| p.iter().any(|v| bad(*v))) || self.bbox.is_some_and(|b| b.iter().any(|v| bad(*v))) {
|
||||
return Err(Error::Input("prompt coordinates must be finite and near the frame".into()));
|
||||
}
|
||||
if self.points.len() > 64 {
|
||||
return Err(Error::Input("at most 64 points per prompt".into()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
struct MaskEmbed {
|
||||
conv: [(Tensor, Tensor, usize, usize); 3],
|
||||
ln1: LayerNorm,
|
||||
ln2: LayerNorm,
|
||||
}
|
||||
|
||||
/// Sparse and dense prompt embeddings.
|
||||
pub struct PromptEncoder {
|
||||
gauss: Vec<f32>,
|
||||
point_embed: Vec<f32>,
|
||||
not_a_point: Vec<f32>,
|
||||
no_mask: Tensor,
|
||||
mask: MaskEmbed,
|
||||
}
|
||||
|
||||
fn rows(t: Tensor) -> Result<Vec<f32>> {
|
||||
Ok(t.flatten_all()?.to_vec1::<f32>()?)
|
||||
}
|
||||
|
||||
impl PromptEncoder {
|
||||
pub fn load(w: &Weights) -> Result<PromptEncoder> {
|
||||
let conv = |name: &str, i: usize, o: usize, k: usize, stride: usize| -> Result<(Tensor, Tensor, usize, usize)> {
|
||||
Ok((
|
||||
w.get(&format!("prompt_encoder.mask_embed.{name}.weight"), &[o, i, k, k])?,
|
||||
w.get(&format!("prompt_encoder.mask_embed.{name}.bias"), &[o])?,
|
||||
k,
|
||||
stride,
|
||||
))
|
||||
};
|
||||
Ok(PromptEncoder {
|
||||
gauss: rows(w.get("prompt_encoder.shared_embedding.positional_embedding", &[2, DIM / 2])?)?,
|
||||
point_embed: rows(w.get("prompt_encoder.point_embed.weight", &[4, DIM])?)?,
|
||||
not_a_point: rows(w.get("prompt_encoder.not_a_point_embed.weight", &[1, DIM])?)?,
|
||||
no_mask: w.get("prompt_encoder.no_mask_embed.weight", &[1, DIM])?,
|
||||
mask: MaskEmbed {
|
||||
conv: [conv("conv1", 1, 4, 2, 2)?, conv("conv2", 4, 16, 2, 2)?, conv("conv3", 16, DIM, 1, 1)?],
|
||||
ln1: layer_norm(w, "prompt_encoder.mask_embed.layer_norm1", 4, 1e-6)?,
|
||||
ln2: layer_norm(w, "prompt_encoder.mask_embed.layer_norm2", 16, 1e-6)?,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
/// Random-Fourier position encoding of coordinates already normalised to 0..1: `[n * 256]`.
|
||||
fn fourier(&self, coords01: &[[f32; 2]]) -> Vec<f32> {
|
||||
let half = DIM / 2;
|
||||
let mut out = vec![0f32; coords01.len() * DIM];
|
||||
for (n, c) in coords01.iter().enumerate() {
|
||||
let (x, y) = (2.0 * c[0] - 1.0, 2.0 * c[1] - 1.0);
|
||||
for j in 0..half {
|
||||
let g0 = self.gauss.get(j).copied().unwrap_or(0.0);
|
||||
let g1 = self.gauss.get(half + j).copied().unwrap_or(0.0);
|
||||
let v = 2.0 * std::f32::consts::PI * (x * g0 + y * g1);
|
||||
if let Some(o) = out.get_mut(n * DIM + j) {
|
||||
*o = v.sin();
|
||||
}
|
||||
if let Some(o) = out.get_mut(n * DIM + half + j) {
|
||||
*o = v.cos();
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// The position encoding of the whole `GRID x GRID` embedding, `[1, 256, GRID, GRID]`.
|
||||
pub fn image_positions(&self) -> Result<Tensor> {
|
||||
let coords: Vec<[f32; 2]> = (0..GRID * GRID).map(|i| [((i % GRID) as f32 + 0.5) / GRID as f32, ((i / GRID) as f32 + 0.5) / GRID as f32]).collect();
|
||||
let flat = self.fourier(&coords);
|
||||
Ok(Tensor::from_vec(flat, (GRID * GRID, DIM), &Device::Cpu)?.t()?.reshape((1, DIM, GRID, GRID))?.contiguous()?)
|
||||
}
|
||||
|
||||
fn point_row(&self, i: usize) -> &[f32] {
|
||||
self.point_embed.get(i * DIM..(i + 1) * DIM).unwrap_or(&[])
|
||||
}
|
||||
|
||||
/// Sparse tokens `[1, n, 256]` for a prompt (a click-only prompt gets one padding token).
|
||||
pub fn sparse(&self, p: &Prompt) -> Result<Tensor> {
|
||||
p.validate()?;
|
||||
let norm = |c: [f32; 2]| [(c[0] + 0.5) / INPUT_SIZE as f32, (c[1] + 0.5) / INPUT_SIZE as f32];
|
||||
let mut tokens: Vec<f32> = vec![];
|
||||
let mut n = 0;
|
||||
let mut push = |v: &[f32]| {
|
||||
tokens.extend_from_slice(v);
|
||||
n += 1;
|
||||
};
|
||||
let pts: Vec<[f32; 2]> = p.points.iter().map(|(c, _)| norm(*c)).collect();
|
||||
let enc = self.fourier(&pts);
|
||||
for (k, (_, fg)) in p.points.iter().enumerate() {
|
||||
let e = enc.get(k * DIM..(k + 1) * DIM).unwrap_or(&[]);
|
||||
let label_row = self.point_row(usize::from(*fg));
|
||||
let v: Vec<f32> = e.iter().zip(label_row).map(|(a, b)| a + b).collect();
|
||||
push(&v);
|
||||
}
|
||||
match p.bbox {
|
||||
None => push(&self.not_a_point.clone()),
|
||||
Some(b) => {
|
||||
let corners = [norm([b[0], b[1]]), norm([b[2], b[3]])];
|
||||
let enc = self.fourier(&corners);
|
||||
for (k, row) in [2usize, 3].into_iter().enumerate() {
|
||||
let e = enc.get(k * DIM..(k + 1) * DIM).unwrap_or(&[]);
|
||||
let v: Vec<f32> = e.iter().zip(self.point_row(row)).map(|(a, b)| a + b).collect();
|
||||
push(&v);
|
||||
}
|
||||
push(&self.not_a_point.clone());
|
||||
}
|
||||
}
|
||||
Ok(Tensor::from_vec(tokens, (1, n, DIM), &Device::Cpu)?)
|
||||
}
|
||||
|
||||
/// The dense embedding for "no mask prompt", `[1, 256, GRID, GRID]`.
|
||||
pub fn dense_empty(&self) -> Result<Tensor> {
|
||||
Ok(self.no_mask.reshape((1, DIM, 1, 1))?.broadcast_as((1, DIM, GRID, GRID))?.contiguous()?)
|
||||
}
|
||||
|
||||
/// The dense embedding of a mask prompt `[1, 1, 4 * GRID, 4 * GRID]` (low-resolution logits).
|
||||
pub fn dense_mask(&self, mask: &Tensor) -> Result<Tensor> {
|
||||
let (c1, c2, c3) = (&self.mask.conv[0], &self.mask.conv[1], &self.mask.conv[2]);
|
||||
let step = |x: &Tensor, c: &(Tensor, Tensor, usize, usize)| -> Result<Tensor> {
|
||||
Ok(x.conv2d(&c.0, 0, c.3, 1, 1)?.broadcast_add(&c.1.reshape((1, c.1.dim(0)?, 1, 1))?)?)
|
||||
};
|
||||
let cl = |x: &Tensor, ln: &LayerNorm| -> Result<Tensor> { Ok(ln.forward(&x.permute((0, 2, 3, 1))?)?.permute((0, 3, 1, 2))?.contiguous()?) };
|
||||
let x = cl(&step(mask, c1)?, &self.mask.ln1)?.gelu_erf()?;
|
||||
let x = cl(&step(&x, c2)?, &self.mask.ln2)?.gelu_erf()?;
|
||||
step(&x, c3)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn hostile_prompts_are_rejected() {
|
||||
let ok = Prompt { points: vec![([10.0, 20.0], true)], bbox: None };
|
||||
assert!(ok.validate().is_ok());
|
||||
assert!(Prompt { points: vec![([f32::NAN, 0.0], true)], bbox: None }.validate().is_err());
|
||||
assert!(Prompt { points: vec![], bbox: Some([0.0, 0.0, f32::INFINITY, 5.0]) }.validate().is_err());
|
||||
assert!(Prompt { points: vec![([1e9, 0.0], false)], bbox: None }.validate().is_err());
|
||||
assert!(Prompt { points: vec![([1.0, 1.0], true); 65], bbox: None }.validate().is_err());
|
||||
}
|
||||
}
|
||||
@@ -107,3 +107,39 @@ fn neck_and_vision_encoder_match_pytorch() {
|
||||
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}");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user