309 lines
15 KiB
Rust
309 lines
15 KiB
Rust
//! Text prompts (SAM 3's concept detector): the prompt's CLIP features fused into the image
|
||
//! features by a DETR encoder, 200 object queries refined by a DETR decoder (with box
|
||
//! relative-position bias and a presence token), per-query masks from a pixel decoder.
|
||
//!
|
||
//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from
|
||
//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace
|
||
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
|
||
|
||
use candle_core::{D, Device, IndexOp, Module, Tensor};
|
||
use candle_nn::{Conv2d, Conv2dConfig, GroupNorm, LayerNorm, Linear, VarBuilder};
|
||
|
||
use crate::clip::TextEncoder;
|
||
use crate::neck::{FPN, Neck};
|
||
use crate::nn::{Attention, DecoderMlp, Mlp, TORCH_LN_EPS, cxcywh_to_xyxy, inverse_sigmoid, linear, ln, padding_bias, sine_grid, to_vec};
|
||
use crate::{Error, Prediction, Result};
|
||
|
||
const HIDDEN: usize = FPN;
|
||
const HEADS: usize = 8;
|
||
const QUERIES: usize = 200;
|
||
const LAYERS: usize = 6;
|
||
|
||
/// Detector image features, computed once per image.
|
||
pub struct Features {
|
||
/// Levels 0..=2 (288, 144, 72 per side).
|
||
levels: Vec<Tensor>,
|
||
/// Level 2 flattened `[1, 5184, 256]` and its sine position encoding.
|
||
flat: Tensor,
|
||
pos: Tensor,
|
||
side: usize,
|
||
}
|
||
|
||
/// Everything the detector predicts for one prompt.
|
||
pub struct Output {
|
||
/// Per-query mask logits `[200, 288, 288]` flattened.
|
||
pub masks: Tensor,
|
||
/// Per-query logits (before the presence score).
|
||
pub logits: Vec<f32>,
|
||
/// Whether the concept is in the image at all (logit).
|
||
pub presence: f32,
|
||
/// Boxes (x0, y0, x1, y1), normalized.
|
||
pub boxes: Vec<[f32; 4]>,
|
||
/// The text features `[1, 32, 256]` (tests).
|
||
pub text: Tensor,
|
||
}
|
||
|
||
struct EncoderLayer {
|
||
ln1: LayerNorm,
|
||
self_attn: Attention,
|
||
ln2: LayerNorm,
|
||
cross_attn: Attention,
|
||
ln3: LayerNorm,
|
||
mlp: Mlp,
|
||
}
|
||
|
||
struct DecoderLayer {
|
||
self_attn: Attention,
|
||
self_ln: LayerNorm,
|
||
text_attn: Attention,
|
||
text_ln: LayerNorm,
|
||
vision_attn: Attention,
|
||
vision_ln: LayerNorm,
|
||
mlp: Mlp,
|
||
mlp_ln: LayerNorm,
|
||
}
|
||
|
||
pub struct Detector {
|
||
device: Device,
|
||
neck: Neck,
|
||
text: TextEncoder,
|
||
text_proj: Linear,
|
||
encoder: Vec<EncoderLayer>,
|
||
decoder: Vec<DecoderLayer>,
|
||
out_ln: LayerNorm,
|
||
box_head: DecoderMlp,
|
||
query_embed: Tensor,
|
||
reference_points: Tensor,
|
||
presence_token: Tensor,
|
||
presence_head: DecoderMlp,
|
||
presence_ln: LayerNorm,
|
||
ref_point_head: DecoderMlp,
|
||
rpb_x: DecoderMlp,
|
||
rpb_y: DecoderMlp,
|
||
// dot-product scoring
|
||
score_text_mlp: DecoderMlp,
|
||
score_text_ln: LayerNorm,
|
||
score_text_proj: Linear,
|
||
score_query_proj: Linear,
|
||
// mask decoder
|
||
prompt_attn: Attention,
|
||
prompt_ln: LayerNorm,
|
||
pixel_convs: Vec<Conv2d>,
|
||
pixel_norms: Vec<GroupNorm>,
|
||
instance_proj: Conv2d,
|
||
mask_embedder: Vec<Linear>,
|
||
}
|
||
|
||
impl Detector {
|
||
/// `vb` at `detector_model`.
|
||
pub fn new(vb: VarBuilder) -> Result<Self> {
|
||
let e = vb.pp("detr_encoder");
|
||
let mut encoder = Vec::with_capacity(LAYERS);
|
||
for i in 0..LAYERS {
|
||
let l = e.pp(format!("layers.{i}"));
|
||
encoder.push(EncoderLayer {
|
||
ln1: ln(HIDDEN, TORCH_LN_EPS, l.pp("layer_norm1"))?,
|
||
self_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("self_attn"))?,
|
||
ln2: ln(HIDDEN, TORCH_LN_EPS, l.pp("layer_norm2"))?,
|
||
cross_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("cross_attn"))?,
|
||
ln3: ln(HIDDEN, TORCH_LN_EPS, l.pp("layer_norm3"))?,
|
||
mlp: Mlp::new(HIDDEN, 2048, false, l.pp("mlp"))?,
|
||
});
|
||
}
|
||
let d = vb.pp("detr_decoder");
|
||
let mut decoder = Vec::with_capacity(LAYERS);
|
||
for i in 0..LAYERS {
|
||
let l = d.pp(format!("layers.{i}"));
|
||
decoder.push(DecoderLayer {
|
||
self_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("self_attn"))?,
|
||
self_ln: ln(HIDDEN, TORCH_LN_EPS, l.pp("self_attn_layer_norm"))?,
|
||
text_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("text_cross_attn"))?,
|
||
text_ln: ln(HIDDEN, TORCH_LN_EPS, l.pp("text_cross_attn_layer_norm"))?,
|
||
vision_attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("vision_cross_attn"))?,
|
||
vision_ln: ln(HIDDEN, TORCH_LN_EPS, l.pp("vision_cross_attn_layer_norm"))?,
|
||
mlp: Mlp::new(HIDDEN, 2048, false, l.pp("mlp"))?,
|
||
mlp_ln: ln(HIDDEN, TORCH_LN_EPS, l.pp("mlp_layer_norm"))?,
|
||
});
|
||
}
|
||
let s = vb.pp("dot_product_scoring");
|
||
let m = vb.pp("mask_decoder");
|
||
let p = m.pp("pixel_decoder");
|
||
let conv3 = Conv2dConfig { padding: 1, ..Default::default() };
|
||
Ok(Self {
|
||
device: vb.device().clone(),
|
||
neck: Neck::new(vb.pp("vision_encoder.neck"))?,
|
||
text: TextEncoder::new(vb.pp("text_encoder.text_model"))?,
|
||
text_proj: linear(1024, HIDDEN, vb.pp("text_projection"))?,
|
||
encoder,
|
||
decoder,
|
||
out_ln: ln(HIDDEN, TORCH_LN_EPS, d.pp("output_layer_norm"))?,
|
||
box_head: DecoderMlp::new(HIDDEN, HIDDEN, 4, 3, d.pp("box_head"))?,
|
||
query_embed: d.get((QUERIES, HIDDEN), "query_embed.weight")?.unsqueeze(0)?,
|
||
reference_points: d.get((QUERIES, 4), "reference_points.weight")?.unsqueeze(0)?,
|
||
presence_token: d.get((1, HIDDEN), "presence_token.weight")?.unsqueeze(0)?,
|
||
presence_head: DecoderMlp::new(HIDDEN, HIDDEN, 1, 3, d.pp("presence_head"))?,
|
||
presence_ln: ln(HIDDEN, TORCH_LN_EPS, d.pp("presence_layer_norm"))?,
|
||
ref_point_head: DecoderMlp::new(2 * HIDDEN, HIDDEN, HIDDEN, 2, d.pp("ref_point_head"))?,
|
||
rpb_x: DecoderMlp::new(2, HIDDEN, HEADS, 2, d.pp("box_rpb_embed_x"))?,
|
||
rpb_y: DecoderMlp::new(2, HIDDEN, HEADS, 2, d.pp("box_rpb_embed_y"))?,
|
||
score_text_mlp: DecoderMlp::new(HIDDEN, 2048, HIDDEN, 2, s.pp("text_mlp"))?,
|
||
score_text_ln: ln(HIDDEN, TORCH_LN_EPS, s.pp("text_mlp_out_norm"))?,
|
||
score_text_proj: linear(HIDDEN, HIDDEN, s.pp("text_proj"))?,
|
||
score_query_proj: linear(HIDDEN, HIDDEN, s.pp("query_proj"))?,
|
||
prompt_attn: Attention::new(HIDDEN, HIDDEN, HEADS, m.pp("prompt_cross_attn"))?,
|
||
prompt_ln: ln(HIDDEN, TORCH_LN_EPS, m.pp("prompt_cross_attn_norm"))?,
|
||
pixel_convs: (0..3).map(|i| Ok(candle_nn::conv2d(HIDDEN, HIDDEN, 3, conv3, p.pp(format!("conv_layers.{i}")))?)).collect::<Result<_>>()?,
|
||
pixel_norms: (0..3).map(|i| Ok(candle_nn::group_norm(8, HIDDEN, 1e-5, p.pp(format!("norms.{i}")))?)).collect::<Result<_>>()?,
|
||
instance_proj: candle_nn::conv2d(HIDDEN, HIDDEN, 1, Conv2dConfig::default(), m.pp("instance_projection"))?,
|
||
mask_embedder: (0..3).map(|i| linear(HIDDEN, HIDDEN, m.pp(format!("mask_embedder.layers.{i}")))).collect::<Result<_>>()?,
|
||
})
|
||
}
|
||
|
||
/// The detector's feature pyramid for the backbone output `[1, 1024, 72, 72]`.
|
||
pub fn features(&self, backbone: &Tensor) -> Result<Features> {
|
||
let levels = self.neck.forward(backbone, 3)?;
|
||
let last = levels.last().ok_or_else(|| Error::Model("detector neck".into()))?;
|
||
let (_, _, side, _) = last.dims4()?;
|
||
let flat = last.flatten_from(2)?.transpose(1, 2)?.contiguous()?;
|
||
let pos = sine_grid(side, side, HIDDEN / 2, &self.device)?.flatten_from(2)?.transpose(1, 2)?.contiguous()?;
|
||
Ok(Features { levels, flat, pos, side })
|
||
}
|
||
|
||
/// Box relative-position bias `[1, heads, 1 + 200, side²]` for reference boxes (cx, cy, w, h).
|
||
fn rpb(&self, boxes: &Tensor, side: usize) -> Result<Tensor> {
|
||
let xyxy = cxcywh_to_xyxy(boxes)?.squeeze(0)?; // [200, 4]
|
||
let coords = Tensor::arange(0u32, side as u32, &self.device)?.to_dtype(candle_core::DType::F32)?.affine(1.0 / side as f64, 0.0)?;
|
||
let coords = coords.reshape((1, side, 1))?;
|
||
let log_scale = |d: Tensor| -> Result<Tensor> {
|
||
let d = (d * 8.0)?;
|
||
Ok(d.sign()?.mul(&((d.abs()? + 1.0)?.log()? / (8f64.ln()))?)?)
|
||
};
|
||
let pick = |a: usize, b: usize| -> Result<Tensor> { Ok(Tensor::cat(&[xyxy.narrow(1, a, 1)?, xyxy.narrow(1, b, 1)?], 1)?.unsqueeze(1)?) };
|
||
let dy = log_scale(coords.broadcast_sub(&pick(1, 3)?)?)?; // [200, side, 2]
|
||
let dx = log_scale(coords.broadcast_sub(&pick(0, 2)?)?)?;
|
||
let ey = self.rpb_y.forward(&dy)?; // [200, side, heads]
|
||
let ex = self.rpb_x.forward(&dx)?;
|
||
let m = ey.unsqueeze(2)?.broadcast_add(&ex.unsqueeze(1)?)?; // [200, side, side, heads]
|
||
let m = m.reshape((QUERIES, side * side, HEADS))?.permute((2, 0, 1))?; // [heads, 200, side²]
|
||
let m = m.pad_with_zeros(1, 1, 0)?; // the presence token attends everywhere
|
||
Ok(m.unsqueeze(0)?.contiguous()?)
|
||
}
|
||
|
||
/// Sine embedding of boxes (cx, cy, w, h) `[1, 200, 512]` (y, x, w, h features).
|
||
fn box_sine(&self, boxes: &Tensor) -> Result<Tensor> {
|
||
let feats = HIDDEN / 2;
|
||
let dim_t: Vec<f32> = (0..feats).map(|i| 10000f64.powf(2.0 * (i / 2) as f64 / feats as f64) as f32).collect();
|
||
let dim_t = Tensor::from_vec(dim_t, (1, 1, feats), &self.device)?;
|
||
let even: Vec<f32> = (0..feats).map(|i| if i % 2 == 0 { 1.0 } else { 0.0 }).collect();
|
||
let even = Tensor::from_vec(even, (1, 1, feats), &self.device)?;
|
||
let one = |i: usize| -> Result<Tensor> {
|
||
let p = (boxes.narrow(D::Minus1, i, 1)? * (2.0 * std::f64::consts::PI))?.broadcast_div(&dim_t)?;
|
||
Ok((p.sin()?.broadcast_mul(&even)? + p.cos()?.broadcast_mul(&(1.0 - &even)?)?)?)
|
||
};
|
||
Ok(Tensor::cat(&[one(1)?, one(0)?, one(2)?, one(3)?], D::Minus1)?)
|
||
}
|
||
|
||
/// Run the detector for one tokenized prompt.
|
||
pub fn predict(&self, f: &Features, ids: &[u32], valid: &[bool]) -> Result<Output> {
|
||
let text = self.text_proj.forward(&self.text.forward(ids, valid)?)?; // [1, 32, 256]
|
||
let text_bias = padding_bias(valid, &self.device)?;
|
||
|
||
// DETR encoder: image tokens attend to themselves and to the prompt
|
||
let mut x = f.flat.clone();
|
||
for l in &self.encoder {
|
||
let h = l.ln1.forward(&x)?;
|
||
let hp = (&h + &f.pos)?;
|
||
x = (&x + l.self_attn.forward(&hp, &hp, &h, None)?)?;
|
||
let h = l.ln2.forward(&x)?;
|
||
x = (&x + l.cross_attn.forward(&h, &text, &text, Some(&text_bias))?)?;
|
||
let h = l.ln3.forward(&x)?;
|
||
x = (&x + l.mlp.forward(&h)?)?;
|
||
}
|
||
let memory = x;
|
||
|
||
// DETR decoder: presence token + 200 queries with iterative box refinement
|
||
let mut boxes = candle_nn::ops::sigmoid(&self.reference_points)?;
|
||
let mut h = Tensor::cat(&[self.presence_token.clone(), self.query_embed.clone()], 1)?;
|
||
let key_pos = (&memory + &f.pos)?;
|
||
let mut queries = None;
|
||
let mut presence = None;
|
||
for l in &self.decoder {
|
||
let qpos = self.ref_point_head.forward(&self.box_sine(&boxes)?)?.pad_with_zeros(1, 1, 0)?;
|
||
let bias = self.rpb(&boxes, f.side)?;
|
||
let q = (&h + &qpos)?;
|
||
h = l.self_ln.forward(&(&h + l.self_attn.forward(&q, &q, &h, None)?)?)?;
|
||
let q = (&h + &qpos)?;
|
||
h = l.text_ln.forward(&(&h + l.text_attn.forward(&q, &text, &text, Some(&text_bias))?)?)?;
|
||
let q = (&h + &qpos)?;
|
||
h = l.vision_ln.forward(&(&h + l.vision_attn.forward(&q, &key_pos, &memory, Some(&bias))?)?)?;
|
||
h = l.mlp_ln.forward(&(&h + l.mlp.forward(&h)?)?)?;
|
||
let qh = self.out_ln.forward(&h.narrow(1, 1, QUERIES)?)?;
|
||
boxes = candle_nn::ops::sigmoid(&(self.box_head.forward(&qh)? + inverse_sigmoid(&boxes)?)?)?;
|
||
let p = self.presence_head.forward(&self.presence_ln.forward(&h.narrow(1, 0, 1)?)?)?;
|
||
presence = Some(p.clamp(-10f32, 10f32)?);
|
||
queries = Some(qh);
|
||
}
|
||
let (Some(queries), Some(presence)) = (queries, presence) else { return Err(Error::Model("decoder has no layers".into())) };
|
||
|
||
// dot-product scoring against the mean of the (refined) prompt tokens
|
||
let t = self.score_text_ln.forward(&(self.score_text_mlp.forward(&text)? + &text)?)?;
|
||
let w: Vec<f32> = valid.iter().map(|v| if *v { 1.0 } else { 0.0 }).collect();
|
||
let count = w.iter().sum::<f32>().max(1.0);
|
||
let w = Tensor::from_vec(w, (1, valid.len(), 1), &self.device)?;
|
||
let pooled = (t.broadcast_mul(&w)?.sum(1)? / f64::from(count))?; // [1, 256]
|
||
let pt = self.score_text_proj.forward(&pooled)?; // [1, 256]
|
||
let pq = self.score_query_proj.forward(&queries)?; // [1, 200, 256]
|
||
let logits = (pq.matmul(&pt.unsqueeze(2)?)? / (HIDDEN as f64).sqrt())?.clamp(-12f32, 12f32)?;
|
||
|
||
// masks: the encoder memory attends to the prompt, then a pixel decoder up to 288²
|
||
let mem = (&memory + self.prompt_attn.forward(&self.prompt_ln.forward(&memory)?, &text, &text, Some(&text_bias))?)?;
|
||
let [l0, l1, _] = &f.levels[..] else { return Err(Error::Model("detector levels".into())) };
|
||
let mut prev = mem.transpose(1, 2)?.reshape((1, HIDDEN, f.side, f.side))?;
|
||
for (i, skip) in [l1, l0].into_iter().enumerate() {
|
||
let (_, _, sh, sw) = skip.dims4()?;
|
||
prev = (prev.upsample_nearest2d(sh, sw)? + skip)?;
|
||
let (Some(conv), Some(norm)) = (self.pixel_convs.get(i), self.pixel_norms.get(i)) else {
|
||
return Err(Error::Model("pixel decoder".into()));
|
||
};
|
||
prev = norm.forward(&conv.forward(&prev)?)?.relu()?;
|
||
}
|
||
let inst = self.instance_proj.forward(&prev)?; // [1, 256, 288, 288]
|
||
let (_, c, mh, mw) = inst.dims4()?;
|
||
let mut e = queries.clone();
|
||
for (i, l) in self.mask_embedder.iter().enumerate() {
|
||
e = l.forward(&e)?;
|
||
if i + 1 < self.mask_embedder.len() {
|
||
e = e.relu()?;
|
||
}
|
||
}
|
||
let masks = e.matmul(&inst.reshape((1, c, mh * mw))?)?.reshape((QUERIES, mh, mw))?;
|
||
let boxes = cxcywh_to_xyxy(&boxes)?;
|
||
let bv = to_vec(&boxes)?;
|
||
Ok(Output {
|
||
masks,
|
||
logits: to_vec(&logits)?,
|
||
presence: to_vec(&presence)?.first().copied().unwrap_or(-10.0),
|
||
boxes: bv.as_chunks::<4>().0.to_vec(),
|
||
text,
|
||
})
|
||
}
|
||
}
|
||
|
||
/// The instances of `out` scoring above `threshold` (sigmoid(logit) × sigmoid(presence)),
|
||
/// merged into one mask (the per-pixel maximum of their probabilities, as logits).
|
||
pub fn merge(out: &Output, threshold: f32) -> Option<Prediction> {
|
||
let s = |x: f32| 1.0 / (1.0 + (-x).exp());
|
||
let presence = s(out.presence);
|
||
let keep: Vec<(usize, f32)> = out.logits.iter().enumerate().map(|(i, l)| (i, s(*l) * presence)).filter(|(_, p)| *p > threshold).collect();
|
||
let best = keep.iter().map(|k| k.1).fold(0f32, f32::max);
|
||
let mut merged: Option<Vec<f32>> = None;
|
||
for (i, _) in &keep {
|
||
let m = out.masks.i(*i).ok().and_then(|t| to_vec(&t).ok())?;
|
||
merged = Some(match merged {
|
||
None => m,
|
||
Some(acc) => acc.iter().zip(&m).map(|(a, b)| a.max(*b)).collect(),
|
||
});
|
||
}
|
||
merged.map(|logits| Prediction { logits, score: best })
|
||
}
|