Merge pull request #162 from GuintherKovalski/feat/sam3-object-masks
M12: SAM 3 Object and Describe masks (pure Rust, Metal)
This commit is contained in:
commit
b35487add2
|
|
@ -8,3 +8,6 @@ plan/
|
|||
.DS_Store
|
||||
.claude/
|
||||
/craft-fonts
|
||||
# model weights are never committed (SAM 3: docs/ai-masks.md)
|
||||
*.safetensors
|
||||
*.safetensors.part
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -29,6 +29,7 @@ lightcraft-merge = { path = "crates/merge" }
|
|||
lightcraft-engine = { path = "crates/engine" }
|
||||
lightcraft-ui-egui = { path = "crates/ui-egui" }
|
||||
lightcraft-mcp = { path = "crates/mcp" }
|
||||
lightcraft-segment = { path = "crates/segment" }
|
||||
|
||||
# foundation
|
||||
thiserror = "2"
|
||||
|
|
@ -62,6 +63,10 @@ jxl-oxide = { version = "0.12", default-features = false }
|
|||
ravif = { version = "0.13", default-features = false, features = ["threading"] } # no "asm": it needs nasm (non-Rust)
|
||||
|
||||
# gpu (crates/gpu; native backends only, no GL: everything stays pure Rust)
|
||||
# SAM 3 inference (lightcraft-segment). 0.9.2: later releases pull `tokenizers` with the
|
||||
# Oniguruma C library into candle-core, and the product is pure Rust.
|
||||
candle-core = { version = "=0.9.2", default-features = false }
|
||||
candle-nn = { version = "=0.9.2", default-features = false }
|
||||
wgpu = { version = "30", default-features = false, features = ["std", "parking_lot", "metal", "vulkan", "dx12", "wgsl"] }
|
||||
pollster = "1"
|
||||
|
||||
|
|
|
|||
12
NOTICE
12
NOTICE
|
|
@ -15,6 +15,18 @@ list (author, source, license) is in assets/ATTRIBUTION.md.
|
|||
Inter (assets/fonts/Inter-*.ttf) — SIL Open Font License 1.1. See assets/fonts/OFL-Inter.txt.
|
||||
Showcase photographs by Ansel Adams, Dorothea Lange and NASA (docs/images/) — public domain. See assets/ATTRIBUTION.md.
|
||||
|
||||
crates/segment (lightcraft-segment) ports the SAM 3 model code of Hugging Face Transformers
|
||||
(https://github.com/huggingface/transformers, src/transformers/models/sam3 and sam3_tracker),
|
||||
Copyright The HuggingFace Team and Meta Platforms, Inc., licensed under the Apache License 2.0.
|
||||
The port was modified by the LightCraft contributors (2026): translated from Python/PyTorch to
|
||||
Rust on candle and restructured; each ported file says so. Because of this, crates/segment is
|
||||
licensed under the Apache License 2.0 only (LICENSE-APACHE), not under the MIT option.
|
||||
|
||||
The SAM 3 weights are not part of LightCraft and are never distributed with it. A user downloads
|
||||
the facebook/sam3 checkpoint only when they choose to (in the app, or by hand), under Meta's SAM
|
||||
License (https://github.com/facebookresearch/sam3/blob/main/LICENSE), a separate, non-OSI
|
||||
licence. See docs/ai-masks.md.
|
||||
|
||||
Builds made with craft-fonts (the optional CRAFT_FONTS_DIR build input; all official releases)
|
||||
also embed the fonts listed in https://github.com/storytold/craft-fonts/blob/main/ATTRIBUTION.md
|
||||
(the Japanese fonts BIZ UDPGothic, BIZ UDMincho and Shippori Mincho; the web build embeds only BIZ
|
||||
|
|
|
|||
|
|
@ -39,7 +39,7 @@ duplicate Local entries #22, black GPU exports on an Intel iGPU #78).
|
|||
| Nikon / Sony / older-Canon raw shooter | ~65% | Camera colour fidelity and coverage (ARW and NEF preview estimates are only a starting point) |
|
||||
| Canon CR3 / Fujifilm / Olympus shooter | ~35% | Their raws open as embedded previews only (CR3: the full-size JPEG with full metadata) |
|
||||
| Lightroom Classic power user | ~45% | Print, Book, Map, publish, tethering |
|
||||
| Relies on AI (masks, denoise) | ~25% | No segmentation or AI denoise models |
|
||||
| Relies on AI (masks, denoise) | ~25% | Object / Describe masks via optional SAM 3 (download not yet hosted); no AI denoise |
|
||||
|
||||
## Where we're going
|
||||
|
||||
|
|
@ -57,7 +57,9 @@ Priorities, in order. Each points at tracker rows in [`docs/parity.md`](docs/par
|
|||
(references stay in the local `plan/`), then tune against the numbers.
|
||||
5. **Lens profiles of our own** (LR-EDIT-OPTICS-PROFILE, P1).
|
||||
6. **AI model strategy** (maintainer decision): which permissively licensed models (or our own training) for
|
||||
segmentation masks and denoise; then pure-Rust inference. Unblocks M12 and Enhance.
|
||||
segmentation masks and denoise; then pure-Rust inference. Unblocks M12 and Enhance. Object / Describe masks now
|
||||
run Meta's SAM 3 in pure Rust (optional, user-consented download under the non-OSI SAM License; CDN mirrors still
|
||||
to be set up, docs/ai-masks.md); Subject / Sky / People and Enhance are still open.
|
||||
7. **Then:** HDR (Q), the Classic output modules (Print first, then Map view, Book, Slideshow), video (R), localisation
|
||||
and accessibility.
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ default-run = "lightcraft"
|
|||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
lightcraft-engine = { workspace = true }
|
||||
lightcraft-engine = { workspace = true, features = ["sam"] }
|
||||
lightcraft-sysmem = { workspace = true }
|
||||
lightcraft-ui-egui = { workspace = true }
|
||||
lightcraft-codecs = { workspace = true }
|
||||
|
|
|
|||
|
|
@ -406,10 +406,42 @@ OPTIONS:
|
|||
ENVIRONMENT:
|
||||
LIGHTCRAFT_GPU_BACKEND=dx12|vulkan|metal|auto|off graphics backend (default: DX12 on Windows, Metal on macOS,
|
||||
Vulkan on Linux; off = render on the CPU); else WGPU_BACKEND. LIGHTCRAFT_GPU=0: CPU rendering.
|
||||
LIGHTCRAFT_SAM3_DIR=DIR the SAM 3 model for Object / Describe masks (default: <settings folder>/models/sam3;
|
||||
optional: LightCraft offers to download it when first needed)
|
||||
LIGHTCRAFT_SAM3_MIRRORS=URL,… where to download the SAM 3 model from (base URLs, tried in order)
|
||||
";
|
||||
|
||||
/// Warnings and errors (failed commands, AI mask analysis) on stderr; `LIGHTCRAFT_LOG=info`
|
||||
/// (or `debug`) shows more.
|
||||
struct StderrLog(log::LevelFilter);
|
||||
|
||||
impl log::Log for StderrLog {
|
||||
fn enabled(&self, m: &log::Metadata) -> bool {
|
||||
m.level() <= self.0 && (m.level() <= log::Level::Warn || m.target().starts_with("lightcraft"))
|
||||
}
|
||||
fn log(&self, r: &log::Record) {
|
||||
if self.enabled(r.metadata()) {
|
||||
eprintln!("lightcraft: {} {}: {}", r.level(), r.target(), r.args());
|
||||
}
|
||||
}
|
||||
fn flush(&self) {}
|
||||
}
|
||||
|
||||
fn install_log() {
|
||||
let level = match std::env::var("LIGHTCRAFT_LOG").unwrap_or_default().as_str() {
|
||||
"debug" => log::LevelFilter::Debug,
|
||||
"info" => log::LevelFilter::Info,
|
||||
_ => log::LevelFilter::Warn,
|
||||
};
|
||||
static LOGGER: std::sync::OnceLock<StderrLog> = std::sync::OnceLock::new();
|
||||
if log::set_logger(LOGGER.get_or_init(|| StderrLog(level))).is_ok() {
|
||||
log::set_max_level(level);
|
||||
}
|
||||
}
|
||||
|
||||
fn main() -> eframe::Result {
|
||||
lightcraft_engine::guard::install_hook(std::env::temp_dir().join("lightcraft-panics.log"));
|
||||
install_log();
|
||||
alloc_release::install();
|
||||
let mut control_port: Option<u16> = std::env::var("LIGHTCRAFT_CONTROL_PORT").ok().and_then(|p| p.parse().ok());
|
||||
let mut files = Vec::new();
|
||||
|
|
@ -472,7 +504,13 @@ fn main() -> eframe::Result {
|
|||
"LightCraft",
|
||||
options,
|
||||
Box::new(move |cc| {
|
||||
let (session, problem) = open_session(in_memory, library_dir, seed_demo && files.is_empty());
|
||||
let (mut session, problem) = open_session(in_memory, library_dir, seed_demo && files.is_empty());
|
||||
// AI masks: the SAM 3 checkpoint (facebook/sam3) in <config>/models/sam3, or LIGHTCRAFT_SAM3_DIR
|
||||
// (never required: without it, AI masks offer to download it; see docs/ai-masks.md)
|
||||
session.segmenter.dir =
|
||||
std::env::var_os("LIGHTCRAFT_SAM3_DIR").map(std::path::PathBuf::from).or_else(|| config_dir().map(|d| d.join("models").join("sam3")));
|
||||
// the user's own download locations, one base URL per line (LIGHTCRAFT_SAM3_MIRRORS too)
|
||||
session.segmenter.mirrors_file = config_dir().map(|d| d.join("models").join("sam3-mirrors.txt"));
|
||||
let mut app = LightcraftApp::new(session, services());
|
||||
if let Some(ui) = prefs {
|
||||
app.ui = ui;
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ lightcraft-geom = { workspace = true }
|
|||
lightcraft-color = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
miniz_oxide = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
proptest = { workspace = true }
|
||||
|
|
|
|||
|
|
@ -8,10 +8,12 @@
|
|||
|
||||
pub mod controls;
|
||||
pub mod presets;
|
||||
pub mod segmask;
|
||||
pub mod settings;
|
||||
|
||||
pub use controls::{CONTROLS, ControlSpec, Section, Track};
|
||||
pub use presets::{Preset, SettingsGroup, apply_partial, extract_groups};
|
||||
pub use segmask::SegMask;
|
||||
pub use settings::*;
|
||||
|
||||
use serde_json::Value;
|
||||
|
|
|
|||
|
|
@ -0,0 +1,183 @@
|
|||
//! Model segmentations stored with a mask ([`SegMask`]), so rendering never needs the model.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Logits are stored in steps of 1/`SCALE`, clamped to ±127/`SCALE` (≈ ±15.9: a sigmoid of
|
||||
/// 1 − 1.2e-7, i.e. fully in or out).
|
||||
const SCALE: f32 = 8.0;
|
||||
/// Largest grid a stored segmentation may claim (hostile input caps the allocation).
|
||||
pub const MAX_SIDE: usize = 1024;
|
||||
|
||||
/// A segmentation computed by a model (SAM 3 in `lightcraft-segment`): `side × side` logits
|
||||
/// over the uncropped, oriented image stretched to a square (row `y`, column `x` covers
|
||||
/// normalized `((x + 0.5) / side, (y + 0.5) / side)`), quantized to `i8`, deflated and base64
|
||||
/// encoded. Coverage at a point is the sigmoid of the bilinearly sampled logit, so edges stay
|
||||
/// smooth at any output size.
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct SegMask {
|
||||
pub side: u32,
|
||||
pub data: String,
|
||||
/// The part of the image the grid covers, normalized `[x0, y0, x1, y1]` (`None` = all of
|
||||
/// it). A zoomed-in pass over one object covers just its surroundings, at a higher
|
||||
/// resolution than a pass over the whole photo.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub rect: Option<[f64; 4]>,
|
||||
}
|
||||
|
||||
impl SegMask {
|
||||
/// From `side × side` logits over the whole image.
|
||||
pub fn from_logits(side: usize, logits: &[f32]) -> SegMask {
|
||||
let q: Vec<u8> = logits.iter().take(side * side).map(|l| ((l * SCALE).round().clamp(-127.0, 127.0) as i8) as u8).collect();
|
||||
let z = miniz_oxide::deflate::compress_to_vec(&q, 8);
|
||||
SegMask { side: side as u32, data: base64_encode(&z), rect: None }
|
||||
}
|
||||
|
||||
/// From `side × side` logits over the part `rect` of the image.
|
||||
pub fn from_logits_in(side: usize, logits: &[f32], rect: [f64; 4]) -> SegMask {
|
||||
SegMask { rect: Some(rect), ..SegMask::from_logits(side, logits) }
|
||||
}
|
||||
|
||||
/// The covered part, normalized (the whole image without a `rect`); `None` when damaged.
|
||||
pub fn bounds(&self) -> Option<[f64; 4]> {
|
||||
let r = self.rect.unwrap_or([0.0, 0.0, 1.0, 1.0]);
|
||||
(r.iter().all(|v| v.is_finite()) && r[2] > r[0] && r[3] > r[1]).then_some(r)
|
||||
}
|
||||
|
||||
/// The logits, or `None` when the data is damaged or claims an absurd size.
|
||||
pub fn logits(&self) -> Option<Vec<f32>> {
|
||||
let side = self.side as usize;
|
||||
if side == 0 || side > MAX_SIDE {
|
||||
return None;
|
||||
}
|
||||
let z = base64_decode(&self.data)?;
|
||||
let q = miniz_oxide::inflate::decompress_to_vec_with_limit(&z, side * side).ok()?;
|
||||
(q.len() == side * side).then(|| q.iter().map(|b| f32::from(*b as i8) / SCALE).collect())
|
||||
}
|
||||
}
|
||||
|
||||
/// Most zoomed-in detail patches one AI mask component keeps (the app makes one; each costs a
|
||||
/// decode per render, so a hostile document can't make rendering explode).
|
||||
pub const MAX_DETAIL: usize = 4;
|
||||
|
||||
/// Deserialize a component's detail patches, keeping at most [`MAX_DETAIL`] (the rest are
|
||||
/// skipped without being decoded).
|
||||
pub(crate) fn de_detail<'de, D: serde::Deserializer<'de>>(d: D) -> Result<Vec<SegMask>, D::Error> {
|
||||
struct Capped;
|
||||
impl<'de> serde::de::Visitor<'de> for Capped {
|
||||
type Value = Vec<SegMask>;
|
||||
fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
||||
f.write_str("a list of segmentations")
|
||||
}
|
||||
fn visit_seq<A: serde::de::SeqAccess<'de>>(self, mut seq: A) -> Result<Vec<SegMask>, A::Error> {
|
||||
let mut out = Vec::new();
|
||||
while out.len() < MAX_DETAIL {
|
||||
match seq.next_element::<SegMask>()? {
|
||||
Some(m) => out.push(m),
|
||||
None => return Ok(out),
|
||||
}
|
||||
}
|
||||
while seq.next_element::<serde::de::IgnoredAny>()?.is_some() {}
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
d.deserialize_seq(Capped)
|
||||
}
|
||||
|
||||
const B64: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
|
||||
|
||||
fn base64_encode(bytes: &[u8]) -> String {
|
||||
let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4);
|
||||
for chunk in bytes.chunks(3) {
|
||||
let b = [chunk.first().copied().unwrap_or(0), chunk.get(1).copied().unwrap_or(0), chunk.get(2).copied().unwrap_or(0)];
|
||||
let n = (u32::from(b[0]) << 16) | (u32::from(b[1]) << 8) | u32::from(b[2]);
|
||||
for i in 0..4 {
|
||||
if i <= chunk.len() {
|
||||
out.push(char::from(B64[((n >> (18 - 6 * i)) & 63) as usize]));
|
||||
} else {
|
||||
out.push('=');
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn base64_decode(s: &str) -> Option<Vec<u8>> {
|
||||
let mut out = Vec::with_capacity(s.len() / 4 * 3);
|
||||
let (mut acc, mut bits) = (0u32, 0u32);
|
||||
for c in s.bytes().filter(|c| !c.is_ascii_whitespace()) {
|
||||
if c == b'=' {
|
||||
break;
|
||||
}
|
||||
let v = B64.iter().position(|b| *b == c)? as u32;
|
||||
acc = (acc << 6) | v;
|
||||
bits += 6;
|
||||
if bits >= 8 {
|
||||
bits -= 8;
|
||||
out.push((acc >> bits) as u8);
|
||||
acc &= (1 << bits) - 1;
|
||||
}
|
||||
}
|
||||
Some(out)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn base64_round_trips() {
|
||||
for n in 0..10 {
|
||||
let v: Vec<u8> = (0..n).map(|i| (i * 37 + 11) as u8).collect();
|
||||
assert_eq!(base64_decode(&base64_encode(&v)).unwrap(), v);
|
||||
}
|
||||
assert_eq!(base64_encode(b"Man"), "TWFu");
|
||||
assert_eq!(base64_encode(b"Ma"), "TWE=");
|
||||
assert!(base64_decode("TW!u").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn logits_round_trip_quantized() {
|
||||
let side = 16;
|
||||
let l: Vec<f32> = (0..side * side).map(|i| (i as f32 - 128.0) / 4.0).collect();
|
||||
let m = SegMask::from_logits(side, &l);
|
||||
let back = m.logits().unwrap();
|
||||
for (a, b) in l.iter().zip(&back) {
|
||||
assert!((a.clamp(-127.0 / SCALE, 127.0 / SCALE) - b).abs() <= 0.5 / SCALE + 1e-6);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_hostile_document_keeps_at_most_max_detail_patches() {
|
||||
let patch = serde_json::to_value(SegMask::from_logits(4, &[1.0; 16])).unwrap();
|
||||
let shape = serde_json::json!({"kind": "object", "hint": [{"x": 0.5, "y": 0.5}], "detail": vec![patch; 10_000]});
|
||||
let s: crate::MaskShape = serde_json::from_value(shape).unwrap();
|
||||
let crate::MaskShape::Object { detail, .. } = s else { panic!("an Object") };
|
||||
assert_eq!(detail.len(), MAX_DETAIL);
|
||||
// and a whole develop document with them still loads
|
||||
let mut d = crate::DevelopSettings::default();
|
||||
d.masks.push(crate::Mask {
|
||||
id: 1,
|
||||
components: vec![crate::MaskComponent {
|
||||
name: None,
|
||||
op: crate::MaskOp::Add,
|
||||
invert: false,
|
||||
shape: crate::MaskShape::Prompt { text: "x".into(), seg: None, detail: vec![SegMask::from_logits(4, &[1.0; 16]); 50], edge: 0.0 },
|
||||
}],
|
||||
..Default::default()
|
||||
});
|
||||
let back: crate::DevelopSettings = serde_json::from_str(&serde_json::to_string(&d).unwrap()).unwrap();
|
||||
let crate::MaskShape::Prompt { detail, .. } = &back.masks[0].components[0].shape else { panic!("a Prompt") };
|
||||
assert_eq!(detail.len(), MAX_DETAIL);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn damaged_data_is_none() {
|
||||
let mut m = SegMask::from_logits(8, &[1.0; 64]);
|
||||
m.side = 9;
|
||||
assert!(m.logits().is_none(), "size mismatch");
|
||||
m.side = 100_000;
|
||||
assert!(m.logits().is_none(), "absurd size");
|
||||
assert!(SegMask { side: 8, data: "not base64 !".into(), rect: None }.logits().is_none());
|
||||
assert!(SegMask { rect: Some([0.5, 0.0, 0.2, 1.0]), ..SegMask::from_logits(2, &[0.0; 4]) }.bounds().is_none());
|
||||
}
|
||||
}
|
||||
|
|
@ -772,8 +772,33 @@ pub enum MaskShape {
|
|||
Subject,
|
||||
Sky,
|
||||
Background,
|
||||
/// One object picked by clicks (SAM 3 point prompts): `hint` holds the clicks that include,
|
||||
/// `exclude` the ones that exclude; `seg` the segmentation computed from them.
|
||||
Object {
|
||||
hint: Vec<Point>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
exclude: Vec<Point>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
seg: Option<crate::SegMask>,
|
||||
/// Zoomed-in passes over parts of the image (higher resolution than `seg`), used
|
||||
/// inside their rectangles.
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty", deserialize_with = "crate::segmask::de_detail")]
|
||||
detail: Vec<crate::SegMask>,
|
||||
/// Edge −100..100: below 0 harder (a steeper transition), above 0 softer (feathered).
|
||||
#[serde(default, skip_serializing_if = "is_zero")]
|
||||
edge: f64,
|
||||
},
|
||||
/// Everything a description names ("sky", "the red car": SAM 3 concept prompts); `seg` is
|
||||
/// the segmentation computed for it.
|
||||
Prompt {
|
||||
text: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
seg: Option<crate::SegMask>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty", deserialize_with = "crate::segmask::de_detail")]
|
||||
detail: Vec<crate::SegMask>,
|
||||
/// Edge −100..100, as for `Object`.
|
||||
#[serde(default, skip_serializing_if = "is_zero")]
|
||||
edge: f64,
|
||||
},
|
||||
People {
|
||||
person: u32,
|
||||
|
|
|
|||
|
|
@ -30,6 +30,11 @@ serde_json = { workspace = true }
|
|||
thiserror = { workspace = true }
|
||||
log = { workspace = true }
|
||||
ab_glyph = { workspace = true }
|
||||
lightcraft-segment = { workspace = true, optional = true }
|
||||
|
||||
[features]
|
||||
# AI masks with SAM 3 (Object / Describe). Off for the web and CLI builds.
|
||||
sam = ["dep:lightcraft-segment"]
|
||||
|
||||
[dev-dependencies]
|
||||
libc = { workspace = true }
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ use lightcraft_geom::Point;
|
|||
use serde_json::{Value, json};
|
||||
|
||||
use super::{CommandSpec, bad, bool_or, cmd, f64_or, has_active, ok, point, str_param};
|
||||
use crate::segment::{Click, MAX_CLICKS};
|
||||
use crate::{Result, Session};
|
||||
|
||||
fn shape_from(kind: &str, p: &Value, c: &str) -> Result<MaskShape> {
|
||||
|
|
@ -38,10 +39,100 @@ fn shape_from(kind: &str, p: &Value, c: &str) -> Result<MaskShape> {
|
|||
.unwrap_or_default();
|
||||
MaskShape::ColorRange { samples, refine: f64_or(p, "refine", 50.0) }
|
||||
}
|
||||
other => return Err(bad(c, format!("unknown mask kind `{other}` (brush|linear|radial|sky|subject|background|luminanceRange|colorRange)"))),
|
||||
"object" => MaskShape::Object { hint: points(p, "points"), exclude: points(p, "exclude"), seg: seg_param(p), detail: vec![], edge: 0.0 },
|
||||
"prompt" => {
|
||||
MaskShape::Prompt { text: str_param(p, "text").unwrap_or_default().trim().to_string(), seg: seg_param(p), detail: vec![], edge: 0.0 }
|
||||
}
|
||||
other => {
|
||||
return Err(bad(
|
||||
c,
|
||||
format!("unknown mask kind `{other}` (brush|linear|radial|sky|subject|background|luminanceRange|colorRange|object|prompt)"),
|
||||
));
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// `[[x, y], …]` normalized points under `key`.
|
||||
fn points(p: &Value, key: &str) -> Vec<Point> {
|
||||
p.get(key)
|
||||
.and_then(Value::as_array)
|
||||
.map(|a| a.iter().filter_map(|q| Some(Point::new(q.get(0)?.as_f64()?, q.get(1)?.as_f64()?))).collect())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// A segmentation passed in (`seg`: a stored or replayed AI mask; the model isn't needed).
|
||||
fn seg_param(p: &Value) -> Option<lightcraft_develop::SegMask> {
|
||||
p.get("seg").and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
}
|
||||
|
||||
/// What a new AI shape still needs once it is added.
|
||||
enum Later {
|
||||
Nothing,
|
||||
/// Background mode: these Object clicks are segmented on the worker.
|
||||
Clicks(Vec<Point>, Vec<Point>),
|
||||
/// Background mode: the Describe selection is computed first; the mask is added when found.
|
||||
Text(String),
|
||||
}
|
||||
|
||||
/// Compute the segmentation of a new AI shape (Object with clicks, Prompt) — or, in background
|
||||
/// mode, say what to queue. An Object without clicks stays empty until the photo is clicked
|
||||
/// (the photo is analyzed meanwhile); a Prompt that matches nothing is an error. Shapes given
|
||||
/// with their `seg` need no model.
|
||||
fn resolve_ai(s: &mut Session, shape: &mut MaskShape, c: &str) -> Result<Later> {
|
||||
let id = s.active().ok_or_else(|| bad(c, "no active photo"))?;
|
||||
let background = s.segmenter.background;
|
||||
match shape {
|
||||
MaskShape::Object { seg: Some(_), .. } | MaskShape::Prompt { seg: Some(_), .. } => {}
|
||||
MaskShape::Object { hint, exclude, .. } if !hint.is_empty() && background => {
|
||||
s.segmenter.model_dir().map_err(ai_err)?;
|
||||
return Ok(Later::Clicks(std::mem::take(hint), std::mem::take(exclude)));
|
||||
}
|
||||
MaskShape::Object { hint, exclude, seg, .. } if !hint.is_empty() => {
|
||||
*seg = Some(s.segment_clicks(id, &clicks_of(hint, exclude)).map_err(ai_err)?);
|
||||
}
|
||||
MaskShape::Object { .. } => {
|
||||
// the clicks come later: refuse an Object mask that could never be computed (no
|
||||
// model) rather than leave an empty one; in the app, analyze the photo meanwhile
|
||||
s.segmenter.model_dir().map_err(ai_err)?;
|
||||
if background {
|
||||
s.segment_prepare(id).map_err(ai_err)?;
|
||||
}
|
||||
}
|
||||
MaskShape::Prompt { text, .. } if background => {
|
||||
s.segmenter.model_dir().map_err(ai_err)?;
|
||||
if text.trim().is_empty() {
|
||||
return Err(ai_err("describe what to select"));
|
||||
}
|
||||
return Ok(Later::Text(text.clone()));
|
||||
}
|
||||
MaskShape::Prompt { text, seg, .. } => {
|
||||
let found = s.segment_text(id, text).map_err(ai_err)?;
|
||||
*seg = Some(found.ok_or_else(|| ai_err(format!("nothing matching “{text}” was found in this photo")))?);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
Ok(Later::Nothing)
|
||||
}
|
||||
|
||||
/// Queue what a new AI component (component `comp` of mask `mask`) still needs.
|
||||
fn queue_later(s: &mut Session, later: Later, mask: u32, comp: usize, mut out: Value) -> Result<Value> {
|
||||
if let Later::Clicks(hint, exclude) = later {
|
||||
let id = s.active().ok_or_else(|| ai_err("no active photo"))?;
|
||||
s.segment_clicks_later(id, mask, comp, hint, exclude).map_err(ai_err)?;
|
||||
out["pending"] = json!(true);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// AI mask failures are shown as they are (not as "invalid parameters").
|
||||
fn ai_err(msg: impl Into<String>) -> crate::EngineError {
|
||||
crate::EngineError::Other(msg.into())
|
||||
}
|
||||
|
||||
pub(crate) fn clicks_of(include: &[Point], exclude: &[Point]) -> Vec<Click> {
|
||||
include.iter().map(|p| Click { at: *p, include: true }).chain(exclude.iter().map(|p| Click { at: *p, include: false })).collect()
|
||||
}
|
||||
|
||||
fn masks_edit(s: &mut Session, c: &str, label: &str, f: impl FnOnce(&mut Vec<Mask>, &mut Option<u32>) -> Result<()>) -> Result<Value> {
|
||||
let id = s.active().ok_or_else(|| bad(c, "no active photo"))?;
|
||||
let mut d = (*s.develop_of(id).unwrap_or_default()).clone();
|
||||
|
|
@ -67,14 +158,20 @@ pub fn specs() -> Vec<CommandSpec> {
|
|||
"Create New Mask",
|
||||
[],
|
||||
None,
|
||||
"{kind: brush|linear|radial|sky|subject|background|luminanceRange|colorRange, ...shape params (start/end, center/rx/ry/angle/feather, lo/hi…), name?}",
|
||||
"{kind: brush|linear|radial|sky|subject|background|luminanceRange|colorRange|object|prompt, ...shape params (start/end, center/rx/ry/angle/feather, lo/hi…; object: points/exclude [[x,y],…]; prompt: text; object/prompt: seg? a stored segmentation, no model needed), name?} — AI kinds need the SAM 3 model; in the app a prompt returns {pending} and the mask appears when found",
|
||||
has_active,
|
||||
|s, p| {
|
||||
let kind = str_param(p, "kind").unwrap_or("radial").to_string();
|
||||
let shape = shape_from(&kind, p, "mask.add")?;
|
||||
let mut shape = shape_from(&kind, p, "mask.add")?;
|
||||
let later = resolve_ai(s, &mut shape, "mask.add")?;
|
||||
let name = str_param(p, "name").map(str::to_string);
|
||||
if let Later::Text(text) = &later {
|
||||
let id = s.active().ok_or_else(|| bad("mask.add", "no active photo"))?;
|
||||
s.segment_text_later(id, None, "add", name, text).map_err(ai_err)?;
|
||||
return Ok(json!({"pending": true}));
|
||||
}
|
||||
let next = s.active().and_then(|id| s.develop_of(id)).map(|d| d.next_mask_id()).unwrap_or(1);
|
||||
masks_edit(s, "mask.add", "Add Mask", |masks, active| {
|
||||
let out = masks_edit(s, "mask.add", "Add Mask", |masks, active| {
|
||||
masks.push(Mask {
|
||||
id: next,
|
||||
name: name.unwrap_or_else(|| format!("Mask {next}")),
|
||||
|
|
@ -83,7 +180,8 @@ pub fn specs() -> Vec<CommandSpec> {
|
|||
});
|
||||
*active = Some(next);
|
||||
Ok(())
|
||||
})
|
||||
})?;
|
||||
queue_later(s, later, next, 0, out)
|
||||
}
|
||||
),
|
||||
cmd!(
|
||||
|
|
@ -95,16 +193,29 @@ pub fn specs() -> Vec<CommandSpec> {
|
|||
has_active,
|
||||
|s, p| {
|
||||
let kind = str_param(p, "kind").unwrap_or("brush").to_string();
|
||||
let shape = shape_from(&kind, p, "mask.addComponent")?;
|
||||
let op: MaskOp =
|
||||
serde_json::from_value(p.get("op").cloned().unwrap_or(json!("add"))).map_err(|e| bad("mask.addComponent", e.to_string()))?;
|
||||
let active = s.active_mask;
|
||||
let mid = mask_id(p, active, "mask.addComponent")?;
|
||||
masks_edit(s, "mask.addComponent", "Edit Mask", |masks, _| {
|
||||
let id = s.active().ok_or_else(|| bad("mask.addComponent", "no active photo"))?;
|
||||
if !s.develop_of(id).is_some_and(|d| d.masks.iter().any(|m| m.id == mid)) {
|
||||
return Err(bad("mask.addComponent", format!("no mask {mid}")));
|
||||
}
|
||||
let mut shape = shape_from(&kind, p, "mask.addComponent")?;
|
||||
let later = resolve_ai(s, &mut shape, "mask.addComponent")?;
|
||||
if let Later::Text(text) = &later {
|
||||
let op = p.get("op").and_then(Value::as_str).unwrap_or("add");
|
||||
s.segment_text_later(id, Some(mid), op, None, text).map_err(ai_err)?;
|
||||
return Ok(json!({"pending": true}));
|
||||
}
|
||||
let mut comp = 0;
|
||||
let out = masks_edit(s, "mask.addComponent", "Edit Mask", |masks, _| {
|
||||
let i = find(masks, mid, "mask.addComponent")?;
|
||||
masks[i].components.push(MaskComponent { name: None, op, invert: bool_or(p, "invert", false), shape });
|
||||
comp = masks[i].components.len() - 1;
|
||||
Ok(())
|
||||
})
|
||||
})?;
|
||||
queue_later(s, later, mid, comp, out)
|
||||
}
|
||||
),
|
||||
cmd!(
|
||||
|
|
@ -213,6 +324,158 @@ pub fn specs() -> Vec<CommandSpec> {
|
|||
Ok(json!({"samples": out}))
|
||||
}
|
||||
),
|
||||
cmd!(
|
||||
"mask.objectPoint",
|
||||
"Add Object Click",
|
||||
[],
|
||||
None,
|
||||
"{x, y: normalized image point, exclude?: bool (⌥-click: leave this part out), id?: maskId} — a click on the selected mask's Object selection (its last Object component); SAM 3 re-segments the object → {include, exclude}",
|
||||
has_active,
|
||||
|s, p| {
|
||||
let c = "mask.objectPoint";
|
||||
let at = Point::new(super::f64_req(p, "x", c)?, super::f64_req(p, "y", c)?);
|
||||
if !(0.0..=1.0).contains(&at.x) || !(0.0..=1.0).contains(&at.y) {
|
||||
return Err(bad(c, "the point is outside the photo"));
|
||||
}
|
||||
let exclude = bool_or(p, "exclude", false);
|
||||
let id = s.active().ok_or_else(|| bad(c, "no active photo"))?;
|
||||
let mid = mask_id(p, s.active_mask, c)?;
|
||||
let d = s.develop_of(id).unwrap_or_default();
|
||||
let m = d.masks.iter().find(|m| m.id == mid).ok_or_else(|| bad(c, format!("no mask {mid}")))?;
|
||||
let k = m
|
||||
.components
|
||||
.iter()
|
||||
.rposition(|x| matches!(x.shape, MaskShape::Object { .. }))
|
||||
.ok_or_else(|| bad(c, "the mask has no Object selection (add one first)"))?;
|
||||
let Some(MaskShape::Object { hint, exclude: ex, edge, .. }) = m.components.get(k).map(|x| &x.shape) else {
|
||||
return Err(bad(c, "not an Object component"));
|
||||
};
|
||||
let (mut hint, mut ex, edge) = (hint.clone(), ex.clone(), *edge);
|
||||
// clicks still on their way to the model count too
|
||||
if let Some(pending) = s.segmenter.pending_clicks().filter(|q| (q.photo, q.mask, q.comp) == (id, mid, k)) {
|
||||
(hint, ex) = (pending.hint.clone(), pending.exclude.clone());
|
||||
}
|
||||
if hint.len() + ex.len() >= MAX_CLICKS {
|
||||
return Err(ai_err(format!("an Object selection takes at most {MAX_CLICKS} clicks; start a new one")));
|
||||
}
|
||||
if exclude {
|
||||
ex.push(at)
|
||||
} else {
|
||||
hint.push(at)
|
||||
}
|
||||
if s.segmenter.background && !hint.is_empty() {
|
||||
let out = json!({"include": hint.len(), "exclude": ex.len(), "pending": true});
|
||||
s.segment_clicks_later(id, mid, k, hint, ex).map_err(ai_err)?;
|
||||
return Ok(out);
|
||||
}
|
||||
let seg = if hint.is_empty() { None } else { Some(s.segment_clicks(id, &clicks_of(&hint, &ex)).map_err(ai_err)?) };
|
||||
let out = json!({"include": hint.len(), "exclude": ex.len()});
|
||||
masks_edit(s, c, "Object Mask", |masks, _| {
|
||||
let i = find(masks, mid, c)?;
|
||||
if let Some(comp) = masks[i].components.get_mut(k) {
|
||||
comp.shape = MaskShape::Object { hint, exclude: ex, seg, detail: vec![], edge };
|
||||
}
|
||||
Ok(())
|
||||
})?;
|
||||
Ok(out)
|
||||
}
|
||||
),
|
||||
cmd!(
|
||||
"mask.refineDetail",
|
||||
"Refine AI Mask Detail",
|
||||
[],
|
||||
None,
|
||||
"{id?: maskId, component?} — a zoomed-in SAM 3 pass over an Object/Describe selection (default: the mask's last one): the photo around it is analyzed again at a higher resolution, in the background; the mask updates when it's done → {started}",
|
||||
has_active,
|
||||
|s, p| {
|
||||
let c = "mask.refineDetail";
|
||||
let id = s.active().ok_or_else(|| bad(c, "no active photo"))?;
|
||||
let mid = mask_id(p, s.active_mask, c)?;
|
||||
let d = s.develop_of(id).unwrap_or_default();
|
||||
let m = d.masks.iter().find(|m| m.id == mid).ok_or_else(|| bad(c, format!("no mask {mid}")))?;
|
||||
let k = match p.get("component").and_then(Value::as_u64) {
|
||||
Some(k) => k as usize,
|
||||
None => m
|
||||
.components
|
||||
.iter()
|
||||
.rposition(|x| matches!(x.shape, MaskShape::Object { .. } | MaskShape::Prompt { .. }))
|
||||
.ok_or_else(|| bad(c, "the mask has no Object or Describe selection"))?,
|
||||
};
|
||||
let started = s.segment_detail(id, mid, k).map_err(ai_err)?;
|
||||
Ok(json!({"started": started}))
|
||||
}
|
||||
),
|
||||
cmd!(
|
||||
query "segment.prepare",
|
||||
"Prepare AI Masks",
|
||||
[],
|
||||
None,
|
||||
"{} — load SAM 3 and analyze the active photo in the background, so Object clicks are instant → {busy}",
|
||||
has_active,
|
||||
|s, _| {
|
||||
let c = "segment.prepare";
|
||||
let id = s.active().ok_or_else(|| bad(c, "no active photo"))?;
|
||||
s.segment_prepare(id).map_err(ai_err)?;
|
||||
Ok(json!({"busy": s.segmenter.busy()}))
|
||||
}
|
||||
),
|
||||
cmd!(
|
||||
query "segment.model.status",
|
||||
"AI Mask Model Status",
|
||||
[],
|
||||
None,
|
||||
"{} — whether this build has AI masks, whether the SAM 3 model is installed (and where), loaded, busy, and the download's progress → {available, installed, dir, loaded, busy, analyzing, sizeBytes, license, licenseUrl, mirrors, download: {running, done, total, file, error, finished}}",
|
||||
super::always,
|
||||
|s, _| {
|
||||
let g = &s.segmenter;
|
||||
Ok(json!({
|
||||
"available": crate::segment::Segmenter::AVAILABLE,
|
||||
"installed": g.installed(),
|
||||
"dir": g.dir.as_ref().map(|d| d.display().to_string()),
|
||||
"loaded": g.loaded(),
|
||||
"busy": g.busy(),
|
||||
"analyzing": g.analyzing(),
|
||||
"sizeBytes": crate::segment::MODEL_BYTES,
|
||||
"license": crate::segment::LICENSE_NAME,
|
||||
"licenseUrl": crate::segment::LICENSE_URL,
|
||||
"mirrors": g.mirrors().len(),
|
||||
"download": g.download_status(),
|
||||
}))
|
||||
}
|
||||
),
|
||||
cmd!(
|
||||
query "segment.model.download",
|
||||
"Download AI Mask Model",
|
||||
[],
|
||||
None,
|
||||
"{acknowledged: true} — download the SAM 3 model (about 3.4 GB, Meta's SAM License, not LightCraft's) in the background, from the configured mirrors; only after the user agreed to it. Watch segment.model.status; segment.model.cancel stops it (it resumes later) → {started, installed, downloading}",
|
||||
super::always,
|
||||
|s, p| {
|
||||
let c = "segment.model.download";
|
||||
if p.get("acknowledged").and_then(Value::as_bool) != Some(true) {
|
||||
return Err(bad(
|
||||
c,
|
||||
format!(
|
||||
"the SAM 3 model is a {:.1} GB download under Meta's {} ({}): pass `acknowledged: true` once the user has agreed to download it",
|
||||
crate::segment::MODEL_BYTES as f64 / 1e9,
|
||||
crate::segment::LICENSE_NAME,
|
||||
crate::segment::LICENSE_URL
|
||||
),
|
||||
));
|
||||
}
|
||||
let started = s.segmenter.start_download().map_err(ai_err)?;
|
||||
Ok(json!({"started": started, "installed": s.segmenter.installed(), "downloading": s.segmenter.download_status().running}))
|
||||
}
|
||||
),
|
||||
cmd!(
|
||||
query "segment.model.cancel",
|
||||
"Cancel AI Mask Model Download",
|
||||
[],
|
||||
None,
|
||||
"{} — stop the SAM 3 download (what has arrived is kept, and a new download resumes from it) → {cancelled}",
|
||||
super::always,
|
||||
|s, _| Ok(json!({"cancelled": s.segmenter.cancel_download()}))
|
||||
),
|
||||
cmd!(
|
||||
"mask.refine",
|
||||
"Refine Mask Edges",
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ pub mod preset_import;
|
|||
pub mod preset_luminar;
|
||||
pub mod presets;
|
||||
pub mod rename;
|
||||
pub mod segment;
|
||||
pub mod sidecar;
|
||||
pub mod smart;
|
||||
mod view;
|
||||
|
|
@ -154,6 +155,8 @@ pub struct Session {
|
|||
depth: u32,
|
||||
/// Selected mask (Masking panel), by mask id.
|
||||
pub active_mask: Option<u32>,
|
||||
/// AI masks (SAM 3): the model and the last photo prepared for it.
|
||||
pub segmenter: segment::Segmenter,
|
||||
/// Selected spot (Remove panel), by index into the active photo's spots.
|
||||
pub active_spot: Option<usize>,
|
||||
/// The persistent library this session writes to (`None` = in-memory only).
|
||||
|
|
@ -242,6 +245,7 @@ impl Session {
|
|||
clock: Box::new(|| "2026-09-30T12:00:00".to_string()),
|
||||
depth: 0,
|
||||
active_mask: None,
|
||||
segmenter: segment::Segmenter::default(),
|
||||
active_spot: None,
|
||||
library: None,
|
||||
xmp: sidecar::XmpPrefs::default(),
|
||||
|
|
@ -723,6 +727,8 @@ mod tests_persist;
|
|||
#[cfg(test)]
|
||||
mod tests_prefs;
|
||||
#[cfg(test)]
|
||||
mod tests_segment;
|
||||
#[cfg(test)]
|
||||
mod tests_settings_files;
|
||||
#[cfg(test)]
|
||||
mod tests_spots;
|
||||
|
|
|
|||
|
|
@ -0,0 +1,146 @@
|
|||
//! The model download as a background task the session starts, watches and cancels without
|
||||
//! ever blocking: progress lives in atomics, messages behind a mutex only ever `try_lock`ed
|
||||
//! by the session (the download thread holds it for a few instructions at a time).
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use lightcraft_segment::fetch::{self, Options, Progress};
|
||||
|
||||
/// What the session sees of a download.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize)]
|
||||
pub struct Status {
|
||||
pub running: bool,
|
||||
pub done: u64,
|
||||
pub total: u64,
|
||||
/// The file being fetched.
|
||||
pub file: String,
|
||||
/// The last download's error (cleared when a new one starts).
|
||||
pub error: Option<String>,
|
||||
/// The last download finished and the files are in place.
|
||||
pub finished: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Text {
|
||||
file: String,
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Shared {
|
||||
running: AtomicBool,
|
||||
cancel: AtomicBool,
|
||||
finished: AtomicBool,
|
||||
done: AtomicU64,
|
||||
total: AtomicU64,
|
||||
text: Mutex<Text>,
|
||||
}
|
||||
|
||||
/// Clears `running` when the download thread ends, however it ends.
|
||||
struct Running(Arc<Shared>);
|
||||
|
||||
impl Drop for Running {
|
||||
fn drop(&mut self) {
|
||||
self.0.running.store(false, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct Downloader {
|
||||
shared: Arc<Shared>,
|
||||
/// The last text seen (when the download thread held the lock at that moment).
|
||||
last: Mutex<Text>,
|
||||
}
|
||||
|
||||
fn lock(m: &Mutex<Text>) -> std::sync::MutexGuard<'_, Text> {
|
||||
m.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
|
||||
impl Downloader {
|
||||
pub fn status(&self) -> Status {
|
||||
let s = &self.shared;
|
||||
if let Ok(t) = s.text.try_lock()
|
||||
&& let Ok(mut last) = self.last.try_lock()
|
||||
{
|
||||
last.file.clone_from(&t.file);
|
||||
last.error.clone_from(&t.error);
|
||||
}
|
||||
// (only the session's thread locks `last`)
|
||||
let (file, error) = self.last.try_lock().map(|l| (l.file.clone(), l.error.clone())).unwrap_or_default();
|
||||
Status {
|
||||
running: s.running.load(Ordering::SeqCst),
|
||||
done: s.done.load(Ordering::SeqCst),
|
||||
total: s.total.load(Ordering::SeqCst),
|
||||
file,
|
||||
error,
|
||||
finished: s.finished.load(Ordering::SeqCst),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn running(&self) -> bool {
|
||||
self.shared.running.load(Ordering::SeqCst)
|
||||
}
|
||||
|
||||
/// Ask a running download to stop (it keeps its partial files to resume later).
|
||||
pub fn cancel(&self) -> bool {
|
||||
let running = self.running();
|
||||
if running {
|
||||
self.shared.cancel.store(true, Ordering::SeqCst);
|
||||
}
|
||||
running
|
||||
}
|
||||
|
||||
/// Start downloading `files` from `mirrors` into `dir` on a background thread. False when
|
||||
/// one is already running.
|
||||
pub fn start(&self, files: &'static [fetch::FileSpec], mirrors: Vec<String>, dir: PathBuf, opts: Options) -> Result<bool, String> {
|
||||
let s = self.shared.clone();
|
||||
if s.running.swap(true, Ordering::SeqCst) {
|
||||
return Ok(false);
|
||||
}
|
||||
s.cancel.store(false, Ordering::SeqCst);
|
||||
s.finished.store(false, Ordering::SeqCst);
|
||||
s.done.store(0, Ordering::SeqCst);
|
||||
s.total.store(files.iter().filter_map(|f| f.size).sum(), Ordering::SeqCst);
|
||||
// nothing else runs while `running` was false
|
||||
*lock(&s.text) = Text::default();
|
||||
*lock(&self.last) = Text::default();
|
||||
let guard = Running(s.clone());
|
||||
let spawned = std::thread::Builder::new().name("sam3-download".into()).spawn(move || {
|
||||
let s = guard.0.clone();
|
||||
let mut last_file = String::new();
|
||||
let mut on_progress = |p: &Progress| {
|
||||
s.done.store(p.done, Ordering::SeqCst);
|
||||
s.total.store(p.total, Ordering::SeqCst);
|
||||
if p.file != last_file {
|
||||
last_file.clone_from(&p.file);
|
||||
lock(&s.text).file.clone_from(&p.file);
|
||||
}
|
||||
};
|
||||
let r =
|
||||
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| fetch::download(files, &mirrors, &dir, &opts, &s.cancel, &mut on_progress)));
|
||||
let error = match r {
|
||||
Ok(Ok(())) => {
|
||||
s.finished.store(true, Ordering::SeqCst);
|
||||
log::info!("SAM 3 model downloaded to {}", dir.display());
|
||||
None
|
||||
}
|
||||
Ok(Err(e)) => Some(e.to_string()),
|
||||
Err(_) => Some("the download failed unexpectedly".to_string()),
|
||||
};
|
||||
if let Some(e) = &error {
|
||||
log::warn!("SAM 3 download: {e}");
|
||||
}
|
||||
lock(&s.text).error = error;
|
||||
drop(guard);
|
||||
});
|
||||
match spawned {
|
||||
Ok(_) => Ok(true),
|
||||
Err(e) => {
|
||||
// the closure (and its guard) was dropped: `running` is false again
|
||||
Err(format!("could not start the download: {e}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,715 @@
|
|||
//! AI masks: Object (clicks) and Describe (text) selections computed with SAM 3
|
||||
//! (`lightcraft-segment`, cargo feature `sam`; the desktop app enables it).
|
||||
//!
|
||||
//! **Nothing requires the model.** It is not part of LightCraft (SAM License); the user
|
||||
//! downloads it when they first use an AI mask and agree to (`segment.model.download`), or
|
||||
//! puts the files in the model folder themselves. Without it, AI mask requests fail with a
|
||||
//! clear "not installed" error; renders, exports and every other feature never touch it
|
||||
//! (stored segmentations are plain data in the develop settings).
|
||||
//!
|
||||
//! The model sees the photo as developed, uncropped and without masks, at 1008 px — the frame
|
||||
//! mask coordinates live in — so a stored segmentation lines up with the photo at any size.
|
||||
//! Everything that touches the model (loading it, rendering and encoding the photo, clicks,
|
||||
//! descriptions, detail passes) runs on one worker thread ([`worker`]); the session talks to it
|
||||
//! over channels and never waits on a lock. With [`Segmenter::background`] (the desktop app)
|
||||
//! commands return at once and [`Session::segment_poll`] applies the results; without it (CLI,
|
||||
//! MCP, tests) they wait for the worker.
|
||||
|
||||
#[cfg(feature = "sam")]
|
||||
mod download;
|
||||
#[cfg(feature = "sam")]
|
||||
mod worker;
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use lightcraft_catalog::PhotoId;
|
||||
use lightcraft_develop::{MaskShape, SegMask};
|
||||
use lightcraft_geom::Point;
|
||||
|
||||
use crate::Session;
|
||||
|
||||
/// Long edge of the image the model sees.
|
||||
pub const INPUT_EDGE: usize = 1008;
|
||||
/// Download size of the model (`model.safetensors`; the tokenizer files add ~2 MB).
|
||||
pub const MODEL_BYTES: u64 = 3_439_938_512;
|
||||
/// The model's licence (not LightCraft's): shown before downloading.
|
||||
pub const LICENSE_NAME: &str = "SAM License (Meta)";
|
||||
pub const LICENSE_URL: &str = "https://github.com/facebookresearch/sam3/blob/main/LICENSE";
|
||||
/// How errors about a missing model start (the UI offers the download on it).
|
||||
pub const NOT_INSTALLED: &str = "The SAM 3 model is not installed";
|
||||
/// Longest a waiting command (CLI, MCP) waits for the model (loading on a CPU can take minutes).
|
||||
#[cfg_attr(not(feature = "sam"), allow(dead_code))]
|
||||
const WAIT: std::time::Duration = std::time::Duration::from_secs(15 * 60);
|
||||
/// Most clicks one Object selection takes (the model's prompt grows with each).
|
||||
pub const MAX_CLICKS: usize = 64;
|
||||
|
||||
/// One click of an Object selection (normalized image coordinates).
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct Click {
|
||||
pub at: Point,
|
||||
pub include: bool,
|
||||
}
|
||||
|
||||
/// What a request's result is for (how [`Session::segment_poll`] applies it).
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
#[cfg_attr(not(feature = "sam"), allow(dead_code))]
|
||||
pub(crate) enum Tag {
|
||||
/// A command waiting for the result itself.
|
||||
Wait,
|
||||
/// Load the model and encode the photo ahead of the first click.
|
||||
Prepare,
|
||||
/// Object clicks on component `comp` of mask `mask`.
|
||||
Clicks { photo: PhotoId, mask: u32, comp: usize, hint: Vec<Point>, exclude: Vec<Point> },
|
||||
/// A Describe selection: a new mask (`mask: None`) or a component combined by `op`.
|
||||
NewPrompt { photo: PhotoId, mask: Option<u32>, op: String, name: Option<String>, text: String },
|
||||
/// A zoomed-in pass for a component, computed from `shape` (without its detail).
|
||||
Detail { photo: PhotoId, mask: u32, comp: usize, shape: MaskShape },
|
||||
}
|
||||
|
||||
/// The model input: 8-bit RGB.
|
||||
#[cfg_attr(not(feature = "sam"), allow(dead_code))]
|
||||
pub(crate) struct Input {
|
||||
rgb: Vec<u8>,
|
||||
w: usize,
|
||||
h: usize,
|
||||
}
|
||||
|
||||
/// What a detail pass re-runs.
|
||||
#[cfg_attr(not(feature = "sam"), allow(dead_code))]
|
||||
pub(crate) enum Prompt {
|
||||
Clicks(Vec<Click>),
|
||||
Text(String),
|
||||
}
|
||||
|
||||
/// Object clicks sent to the worker and not applied yet (background mode): further clicks add
|
||||
/// to these, and the loupe shows them.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct PendingClicks {
|
||||
pub photo: PhotoId,
|
||||
pub mask: u32,
|
||||
pub comp: usize,
|
||||
pub hint: Vec<Point>,
|
||||
pub exclude: Vec<Point>,
|
||||
}
|
||||
|
||||
/// What [`Session::segment_poll`] did.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct Polled {
|
||||
/// A mask changed.
|
||||
pub changed: bool,
|
||||
/// Errors and notices from background work, for the user.
|
||||
pub messages: Vec<String>,
|
||||
/// A mask that got a new selection (worth a detail pass).
|
||||
pub refine: Option<u32>,
|
||||
}
|
||||
|
||||
/// The model download (`segment.model.*`).
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize)]
|
||||
pub struct DownloadStatus {
|
||||
pub running: bool,
|
||||
pub done: u64,
|
||||
pub total: u64,
|
||||
pub file: String,
|
||||
pub error: Option<String>,
|
||||
pub finished: bool,
|
||||
}
|
||||
|
||||
/// The model (on its worker thread), the download, and AI mask requests in flight.
|
||||
pub struct Segmenter {
|
||||
/// Where the `facebook/sam3` files live (set by the app; `None`: no AI masks).
|
||||
pub dir: Option<PathBuf>,
|
||||
/// The user's list of download mirrors (one base URL per line), if any.
|
||||
pub mirrors_file: Option<PathBuf>,
|
||||
/// Run requests in the background and apply them in [`Session::segment_poll`] (the
|
||||
/// desktop app); otherwise commands wait for their result (CLI, MCP, tests).
|
||||
pub background: bool,
|
||||
pending: Option<PendingClicks>,
|
||||
messages: Vec<String>,
|
||||
#[cfg(feature = "sam")]
|
||||
worker: worker::Worker,
|
||||
#[cfg(feature = "sam")]
|
||||
download: download::Downloader,
|
||||
#[cfg(feature = "sam")]
|
||||
results: (std::sync::mpsc::Sender<worker::Outcome>, std::sync::mpsc::Receiver<worker::Outcome>),
|
||||
}
|
||||
|
||||
// (derivable only without the `sam` feature: the results channel has no Default)
|
||||
#[allow(clippy::derivable_impls)]
|
||||
impl Default for Segmenter {
|
||||
fn default() -> Self {
|
||||
Segmenter {
|
||||
dir: None,
|
||||
mirrors_file: None,
|
||||
background: false,
|
||||
pending: None,
|
||||
messages: Vec::new(),
|
||||
#[cfg(feature = "sam")]
|
||||
worker: worker::Worker::default(),
|
||||
#[cfg(feature = "sam")]
|
||||
download: download::Downloader::default(),
|
||||
#[cfg(feature = "sam")]
|
||||
results: std::sync::mpsc::channel(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Segmenter {
|
||||
/// Whether this build can compute AI masks at all.
|
||||
pub const AVAILABLE: bool = cfg!(feature = "sam");
|
||||
|
||||
/// Whether the model's files are in place.
|
||||
pub fn installed(&self) -> bool {
|
||||
#[cfg(feature = "sam")]
|
||||
if let Some(dir) = &self.dir {
|
||||
return lightcraft_segment::is_model_dir(dir);
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// Whether requests are queued or running (loading, analyzing the photo, a click…).
|
||||
pub fn busy(&self) -> bool {
|
||||
#[cfg(feature = "sam")]
|
||||
return self.worker.pending() > 0;
|
||||
#[cfg(not(feature = "sam"))]
|
||||
false
|
||||
}
|
||||
|
||||
/// Whether the slow part runs now: loading the model or analyzing a photo.
|
||||
pub fn analyzing(&self) -> bool {
|
||||
#[cfg(feature = "sam")]
|
||||
return self.worker.encoding();
|
||||
#[cfg(not(feature = "sam"))]
|
||||
false
|
||||
}
|
||||
|
||||
/// Whether a zoomed-in detail pass is queued or running.
|
||||
pub fn detail_busy(&self) -> bool {
|
||||
#[cfg(feature = "sam")]
|
||||
return self.worker.detail() > 0;
|
||||
#[cfg(not(feature = "sam"))]
|
||||
false
|
||||
}
|
||||
|
||||
/// Whether the model is in memory (it is unloaded after a while without use).
|
||||
pub fn loaded(&self) -> bool {
|
||||
#[cfg(feature = "sam")]
|
||||
return self.worker.loaded();
|
||||
#[cfg(not(feature = "sam"))]
|
||||
false
|
||||
}
|
||||
|
||||
/// Clicks waiting for their segmentation.
|
||||
pub fn pending_clicks(&self) -> Option<&PendingClicks> {
|
||||
self.pending.as_ref()
|
||||
}
|
||||
|
||||
/// The model download's state.
|
||||
pub fn download_status(&self) -> DownloadStatus {
|
||||
#[cfg(feature = "sam")]
|
||||
{
|
||||
let s = self.download.status();
|
||||
DownloadStatus { running: s.running, done: s.done, total: s.total, file: s.file, error: s.error, finished: s.finished }
|
||||
}
|
||||
#[cfg(not(feature = "sam"))]
|
||||
DownloadStatus::default()
|
||||
}
|
||||
|
||||
/// Where the model is downloaded from, in order (the user's mirrors first).
|
||||
pub fn mirrors(&self) -> Vec<String> {
|
||||
#[cfg(feature = "sam")]
|
||||
{
|
||||
let env = std::env::var(lightcraft_segment::fetch::MIRRORS_ENV).ok();
|
||||
lightcraft_segment::fetch::mirrors(env.as_deref(), self.mirrors_file.as_deref())
|
||||
}
|
||||
#[cfg(not(feature = "sam"))]
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// The model folder, if the model is there; otherwise an error that says what to do.
|
||||
pub fn model_dir(&self) -> Result<PathBuf, String> {
|
||||
if !Self::AVAILABLE {
|
||||
return Err("AI masks are not available in this build".into());
|
||||
}
|
||||
let dir = self.dir.clone().ok_or("no folder is set for the SAM 3 model")?;
|
||||
if self.installed() {
|
||||
return Ok(dir);
|
||||
}
|
||||
let d = self.download_status();
|
||||
if d.running {
|
||||
let pct = d.done.saturating_mul(100).checked_div(d.total).unwrap_or(0);
|
||||
return Err(format!("{NOT_INSTALLED} yet: it is downloading ({pct} %)."));
|
||||
}
|
||||
Err(format!(
|
||||
"{NOT_INSTALLED}. Download it (about {:.1} GB, {LICENSE_NAME}) when LightCraft offers it, with `segment.model.download {{\"acknowledged\": true}}`, or put model.safetensors, vocab.json and merges.txt in {}.",
|
||||
MODEL_BYTES as f64 / 1e9,
|
||||
dir.display()
|
||||
))
|
||||
}
|
||||
|
||||
/// Start downloading the model (on a background thread; see `segment.model.download`).
|
||||
/// `Ok(false)` when it is installed or downloading already.
|
||||
pub fn start_download(&self) -> Result<bool, String> {
|
||||
if !Self::AVAILABLE {
|
||||
return Err("AI masks are not available in this build".into());
|
||||
}
|
||||
let dir = self.dir.clone().ok_or("no folder is set for the SAM 3 model")?;
|
||||
if self.installed() {
|
||||
return Ok(false);
|
||||
}
|
||||
#[cfg(feature = "sam")]
|
||||
{
|
||||
let mirrors = self.mirrors();
|
||||
if mirrors.is_empty() {
|
||||
return Err(lightcraft_segment::fetch::DownloadError::NoMirrors.to_string());
|
||||
}
|
||||
self.download.start(lightcraft_segment::fetch::SAM3_FILES, mirrors, dir, lightcraft_segment::fetch::Options::default())
|
||||
}
|
||||
#[cfg(not(feature = "sam"))]
|
||||
{
|
||||
let _ = dir;
|
||||
Ok(false)
|
||||
}
|
||||
}
|
||||
|
||||
/// Stop a running download (its partial files stay, to resume). False when none runs.
|
||||
pub fn cancel_download(&self) -> bool {
|
||||
#[cfg(feature = "sam")]
|
||||
return self.download.cancel();
|
||||
#[cfg(not(feature = "sam"))]
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Mix `v` into a 64-bit hash.
|
||||
#[cfg_attr(not(feature = "sam"), allow(dead_code))]
|
||||
fn mix(h: u64, v: u64) -> u64 {
|
||||
(h ^ v).wrapping_mul(0x0100_0000_01b3).rotate_left(29)
|
||||
}
|
||||
|
||||
/// Objects covering more than this much of the photo gain little from a zoomed pass.
|
||||
const DETAIL_MAX_AREA: f64 = 0.4;
|
||||
|
||||
/// The part of the photo to zoom into for `seg`'s selection: its bounding box with a margin of
|
||||
/// a quarter of its size, at least 12 % of the long edge, at most 2:1, inside the photo
|
||||
/// (`aspect` = width / height). `None` when nothing is selected or it is most of the photo.
|
||||
pub fn detail_region(seg: &SegMask, aspect: f64) -> Option<[f64; 4]> {
|
||||
let logits = seg.logits()?;
|
||||
let r = seg.bounds()?;
|
||||
let side = seg.side as usize;
|
||||
let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0, 0);
|
||||
for (i, l) in logits.iter().enumerate() {
|
||||
if *l > 0.0 {
|
||||
let (x, y) = (i % side, i / side);
|
||||
(x0, y0, x1, y1) = (x0.min(x), y0.min(y), x1.max(x + 1), y1.max(y + 1));
|
||||
}
|
||||
}
|
||||
if x0 == usize::MAX {
|
||||
return None;
|
||||
}
|
||||
let (rw, rh) = (r[2] - r[0], r[3] - r[1]);
|
||||
let nx = |c: usize| r[0] + c as f64 / side as f64 * rw;
|
||||
let ny = |c: usize| r[1] + c as f64 / side as f64 * rh;
|
||||
let (bx0, by0, bx1, by1) = (nx(x0), ny(y0), nx(x1), ny(y1));
|
||||
if (bx1 - bx0) * (by1 - by0) > DETAIL_MAX_AREA {
|
||||
return None;
|
||||
}
|
||||
// in units of the photo's height: x × aspect
|
||||
let a = if aspect.is_finite() { aspect.clamp(0.1, 10.0) } else { 1.0 };
|
||||
let (cx, cy) = ((bx0 + bx1) / 2.0 * a, (by0 + by1) / 2.0);
|
||||
let size = ((bx1 - bx0) * a).max(by1 - by0);
|
||||
let m = 0.25 * size;
|
||||
let mut w = ((bx1 - bx0) * a + 2.0 * m).max(0.12 * a.max(1.0));
|
||||
let mut h = ((by1 - by0) + 2.0 * m).max(0.12 * a.max(1.0));
|
||||
w = w.max(h / 2.0).min(a);
|
||||
h = h.max(w / 2.0).min(1.0);
|
||||
let x = (cx - w / 2.0).clamp(0.0, (a - w).max(0.0));
|
||||
let y = (cy - h / 2.0).clamp(0.0, (1.0 - h).max(0.0));
|
||||
let out = [x / a, y, (x + w) / a, y + h];
|
||||
// a region nearly the whole photo isn't worth a second pass
|
||||
((out[2] - out[0]) * (out[3] - out[1]) < 0.8).then_some(out)
|
||||
}
|
||||
|
||||
/// `shape` without its detail patches (what a detail pass was computed from)
|
||||
#[cfg_attr(not(feature = "sam"), allow(dead_code))]
|
||||
fn strip_detail(mut shape: MaskShape) -> MaskShape {
|
||||
// (nor its edge setting: moving the slider meanwhile keeps the pass)
|
||||
if let MaskShape::Object { detail, edge, .. } | MaskShape::Prompt { detail, edge, .. } = &mut shape {
|
||||
detail.clear();
|
||||
*edge = 0.0;
|
||||
}
|
||||
shape
|
||||
}
|
||||
|
||||
fn clicks_of(include: &[Point], exclude: &[Point]) -> Vec<Click> {
|
||||
include.iter().map(|p| Click { at: *p, include: true }).chain(exclude.iter().map(|p| Click { at: *p, include: false })).collect()
|
||||
}
|
||||
|
||||
impl Session {
|
||||
/// The key and settings of photo `id` as the model sees it: its look without masks, uncropped.
|
||||
#[cfg_attr(not(feature = "sam"), allow(dead_code))]
|
||||
fn segment_key(&self, id: PhotoId) -> Option<(u64, lightcraft_develop::DevelopSettings)> {
|
||||
let p = self.catalog.photo(id)?;
|
||||
let mut d = (*p.develop).clone();
|
||||
d.masks.clear();
|
||||
let content = crate::media::content_key(p);
|
||||
let mut h = mix(0xcbf2_9ce4_8422_2325, d.hash64());
|
||||
h = mix(h, id.0);
|
||||
for b in content.bytes() {
|
||||
h = mix(h, u64::from(b));
|
||||
}
|
||||
Some((h, d))
|
||||
}
|
||||
|
||||
/// Clicks (in normalized photo coordinates) for the model.
|
||||
#[cfg(feature = "sam")]
|
||||
fn model_clicks(clicks: &[Click]) -> Vec<lightcraft_segment::Click> {
|
||||
clicks.iter().map(|c| lightcraft_segment::Click { x: c.at.x as f32, y: c.at.y as f32, positive: c.include }).collect()
|
||||
}
|
||||
|
||||
/// A worker request about photo `id` (the render job for its model input goes along; the
|
||||
/// worker runs it only when its cache doesn't hold the photo). `region` renders that part
|
||||
/// of the photo, from a source sharp enough for it (detail passes).
|
||||
#[cfg(feature = "sam")]
|
||||
fn segment_job(
|
||||
&mut self,
|
||||
id: PhotoId,
|
||||
kind: worker::Kind,
|
||||
tag: Tag,
|
||||
reply: std::sync::mpsc::Sender<worker::Outcome>,
|
||||
region: Option<[f64; 4]>,
|
||||
) -> Result<worker::Job, String> {
|
||||
let dir = self.segmenter.model_dir()?;
|
||||
let (key, mut d) = self.segment_key(id).ok_or("no such photo")?;
|
||||
let p = self.catalog.photo(id).ok_or("no such photo")?.clone();
|
||||
let missing = match &p.source {
|
||||
lightcraft_catalog::Source::File { path } if !std::path::Path::new(path).exists() => Some(path.clone()),
|
||||
_ => None,
|
||||
};
|
||||
let render = match region {
|
||||
None => self.preview_job(id, INPUT_EDGE, INPUT_EDGE, false, &d).ok_or("no such photo")?,
|
||||
Some(region) => {
|
||||
d.crop.geometry =
|
||||
lightcraft_geom::CropGeometry { rect: lightcraft_geom::Rect::new(region[0], region[1], region[2], region[3]), angle: 0.0 };
|
||||
d.crop.flip_h = false;
|
||||
d.crop.flip_v = false;
|
||||
d.geometry.constrain_crop = false;
|
||||
let mut job = self.preview_job(id, INPUT_EDGE, INPUT_EDGE, true, &d).ok_or("no such photo")?;
|
||||
let frac = (region[2] - region[0]).max(region[3] - region[1]).max(0.01);
|
||||
job.level = crate::media::SourceLevel::for_size((INPUT_EDGE as f64 / frac) as usize);
|
||||
job.source = self.media.source_ref(&p, job.level);
|
||||
job
|
||||
}
|
||||
};
|
||||
Ok(worker::Job { dir, key: if region.is_some() { 0 } else { key }, render: Some(render), missing, kind, tag, reply })
|
||||
}
|
||||
|
||||
/// Send a request and wait for its result (CLI, MCP, tests; or when not in background mode).
|
||||
#[cfg(feature = "sam")]
|
||||
fn segment_wait(&mut self, id: PhotoId, kind: worker::Kind, region: Option<[f64; 4]>) -> Result<Option<SegMask>, String> {
|
||||
let (tx, rx) = std::sync::mpsc::channel();
|
||||
let job = self.segment_job(id, kind, Tag::Wait, tx, region)?;
|
||||
self.segmenter.worker.submit(job)?;
|
||||
match rx.recv_timeout(WAIT) {
|
||||
Ok(o) => o.result,
|
||||
Err(_) => Err("the AI model took too long; try again".into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a request whose result [`Session::segment_poll`] applies.
|
||||
#[cfg(feature = "sam")]
|
||||
fn segment_queue(&mut self, id: PhotoId, kind: worker::Kind, tag: Tag, region: Option<[f64; 4]>) -> Result<(), String> {
|
||||
let reply = self.segmenter.results.0.clone();
|
||||
let job = self.segment_job(id, kind, tag, reply, region)?;
|
||||
self.segmenter.worker.submit(job)
|
||||
}
|
||||
|
||||
/// Start preparing photo `id` for AI masks (load the model, encode the image), so the first
|
||||
/// click is fast. In background mode it returns at once.
|
||||
pub fn segment_prepare(&mut self, id: PhotoId) -> Result<(), String> {
|
||||
self.segmenter.model_dir()?;
|
||||
#[cfg(feature = "sam")]
|
||||
{
|
||||
if self.segmenter.background {
|
||||
return self.segment_queue(id, worker::Kind::Prepare, Tag::Prepare, None);
|
||||
}
|
||||
self.segment_wait(id, worker::Kind::Prepare, None)?;
|
||||
}
|
||||
#[cfg(not(feature = "sam"))]
|
||||
let _ = id;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Segment the object `clicks` pick on photo `id`, waiting for the model.
|
||||
pub fn segment_clicks(&mut self, id: PhotoId, clicks: &[Click]) -> Result<SegMask, String> {
|
||||
self.segmenter.model_dir()?;
|
||||
if !clicks.iter().any(|c| c.include) {
|
||||
return Err("click the object to include it first".into());
|
||||
}
|
||||
#[cfg(feature = "sam")]
|
||||
{
|
||||
self.segment_wait(id, worker::Kind::Clicks(Self::model_clicks(clicks)), None)?.ok_or_else(|| "nothing was selected".to_string())
|
||||
}
|
||||
#[cfg(not(feature = "sam"))]
|
||||
{
|
||||
let _ = id;
|
||||
Err("AI masks are not available in this build".into())
|
||||
}
|
||||
}
|
||||
|
||||
/// Segment everything `text` describes on photo `id`, waiting for the model; `Ok(None)`
|
||||
/// when nothing matches.
|
||||
pub fn segment_text(&mut self, id: PhotoId, text: &str) -> Result<Option<SegMask>, String> {
|
||||
self.segmenter.model_dir()?;
|
||||
let text = text.trim();
|
||||
if text.is_empty() {
|
||||
return Err("describe what to select".into());
|
||||
}
|
||||
#[cfg(feature = "sam")]
|
||||
{
|
||||
self.segment_wait(id, worker::Kind::Text(text.to_string()), None)
|
||||
}
|
||||
#[cfg(not(feature = "sam"))]
|
||||
{
|
||||
let _ = id;
|
||||
Err("AI masks are not available in this build".into())
|
||||
}
|
||||
}
|
||||
|
||||
/// Background mode: queue Object clicks for component `comp` of mask `mask`; the result
|
||||
/// (clicks and segmentation, one undo step) is applied by [`Session::segment_poll`].
|
||||
pub(crate) fn segment_clicks_later(&mut self, id: PhotoId, mask: u32, comp: usize, hint: Vec<Point>, exclude: Vec<Point>) -> Result<(), String> {
|
||||
self.segmenter.model_dir()?;
|
||||
if hint.is_empty() {
|
||||
return Err("click the object to include it first".into());
|
||||
}
|
||||
#[cfg(feature = "sam")]
|
||||
{
|
||||
let clicks = Self::model_clicks(&clicks_of(&hint, &exclude));
|
||||
let tag = Tag::Clicks { photo: id, mask, comp, hint: hint.clone(), exclude: exclude.clone() };
|
||||
self.segment_queue(id, worker::Kind::Clicks(clicks), tag, None)?;
|
||||
self.segmenter.pending = Some(PendingClicks { photo: id, mask, comp, hint, exclude });
|
||||
Ok(())
|
||||
}
|
||||
#[cfg(not(feature = "sam"))]
|
||||
{
|
||||
let _ = (id, mask, comp, exclude);
|
||||
Err("AI masks are not available in this build".into())
|
||||
}
|
||||
}
|
||||
|
||||
/// Background mode: queue a Describe selection; the mask (or component, combined by `op`
|
||||
/// into mask `mask`) is created by [`Session::segment_poll`] when something matches.
|
||||
pub(crate) fn segment_text_later(&mut self, id: PhotoId, mask: Option<u32>, op: &str, name: Option<String>, text: &str) -> Result<(), String> {
|
||||
self.segmenter.model_dir()?;
|
||||
let text = text.trim();
|
||||
if text.is_empty() {
|
||||
return Err("describe what to select".into());
|
||||
}
|
||||
#[cfg(feature = "sam")]
|
||||
{
|
||||
let tag = Tag::NewPrompt { photo: id, mask, op: op.to_string(), name, text: text.to_string() };
|
||||
self.segment_queue(id, worker::Kind::Text(text.to_string()), tag, None)
|
||||
}
|
||||
#[cfg(not(feature = "sam"))]
|
||||
{
|
||||
let _ = (id, mask, op, name);
|
||||
Err("AI masks are not available in this build".into())
|
||||
}
|
||||
}
|
||||
|
||||
/// Start a zoomed-in pass for component `comp` of mask `mask` on photo `id` (an Object or
|
||||
/// Describe selection): the photo around the selection is rendered at a higher resolution
|
||||
/// and segmented again. In background mode [`Session::segment_poll`] applies the result;
|
||||
/// otherwise it is applied before this returns. `Ok(false)` when there is nothing to
|
||||
/// refine (no selection yet, or it is most of the photo) or a pass is running already.
|
||||
pub fn segment_detail(&mut self, id: PhotoId, mask: u32, comp: usize) -> Result<bool, String> {
|
||||
self.segmenter.model_dir()?;
|
||||
let d = self.develop_of(id).ok_or("no such photo")?;
|
||||
let shape = d.masks.iter().find(|m| m.id == mask).and_then(|m| m.components.get(comp)).map(|c| c.shape.clone()).ok_or("no such mask")?;
|
||||
let (seg, prompt) = match &shape {
|
||||
// (a document with absurdly many clicks isn't run through the model)
|
||||
MaskShape::Object { hint, exclude, seg: Some(seg), .. } if !hint.is_empty() && hint.len() + exclude.len() <= MAX_CLICKS => {
|
||||
(seg.clone(), Prompt::Clicks(clicks_of(hint, exclude)))
|
||||
}
|
||||
MaskShape::Prompt { text, seg: Some(seg), .. } => (seg.clone(), Prompt::Text(text.clone())),
|
||||
_ => return Ok(false),
|
||||
};
|
||||
let p = self.catalog.photo(id).ok_or("no such photo")?;
|
||||
let (pw, ph) = (f64::from(p.width.max(1)), f64::from(p.height.max(1)));
|
||||
let aspect = if d.orientation.swaps_axes() { ph / pw } else { pw / ph };
|
||||
let Some(region) = detail_region(&seg, aspect) else { return Ok(false) };
|
||||
#[cfg(feature = "sam")]
|
||||
{
|
||||
let kind = worker::Kind::Detail { prompt, region };
|
||||
let tag = Tag::Detail { photo: id, mask, comp, shape: strip_detail(shape) };
|
||||
if self.segmenter.background {
|
||||
if self.segmenter.detail_busy() {
|
||||
return Ok(false);
|
||||
}
|
||||
self.segment_queue(id, kind, tag, Some(region))?;
|
||||
} else {
|
||||
let patch = self.segment_wait(id, kind, Some(region))?;
|
||||
let outcome = worker::Outcome { tag, result: Ok(patch), superseded: false };
|
||||
let mut polled = Polled::default();
|
||||
self.apply_outcome(outcome, &mut polled);
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
#[cfg(not(feature = "sam"))]
|
||||
{
|
||||
let _ = (prompt, region, mask);
|
||||
Ok(false)
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply finished background requests (to masks that haven't changed since) and collect
|
||||
/// their messages. Cheap when nothing is pending: call it every frame.
|
||||
pub fn segment_poll(&mut self) -> Polled {
|
||||
#[cfg_attr(not(feature = "sam"), allow(unused_mut))]
|
||||
let mut polled = Polled { messages: std::mem::take(&mut self.segmenter.messages), ..Default::default() };
|
||||
#[cfg(feature = "sam")]
|
||||
{
|
||||
// (before draining: a worker that is idle now has sent everything)
|
||||
let idle = !self.segmenter.busy();
|
||||
let done: Vec<worker::Outcome> = self.segmenter.results.1.try_iter().collect();
|
||||
for o in done {
|
||||
self.apply_outcome(o, &mut polled);
|
||||
}
|
||||
// clicks whose request was lost (the worker restarted) aren't pending forever
|
||||
if idle {
|
||||
self.segmenter.pending = None;
|
||||
}
|
||||
}
|
||||
polled
|
||||
}
|
||||
|
||||
#[cfg(feature = "sam")]
|
||||
fn apply_outcome(&mut self, o: worker::Outcome, polled: &mut Polled) {
|
||||
if o.superseded {
|
||||
return;
|
||||
}
|
||||
match o.tag {
|
||||
Tag::Wait => {}
|
||||
Tag::Prepare => {
|
||||
if let Err(e) = o.result {
|
||||
polled.messages.push(e);
|
||||
}
|
||||
}
|
||||
Tag::Clicks { photo, mask, comp, hint, exclude } => {
|
||||
let latest = PendingClicks { photo, mask, comp, hint, exclude };
|
||||
if self.segmenter.pending.as_ref() != Some(&latest) {
|
||||
return; // newer clicks are on their way
|
||||
}
|
||||
self.segmenter.pending = None;
|
||||
let seg = match o.result {
|
||||
Ok(Some(seg)) => seg,
|
||||
Ok(None) => return,
|
||||
Err(e) => return polled.messages.push(e),
|
||||
};
|
||||
let Some(d) = self.develop_of(photo) else { return };
|
||||
let mut d = (*d).clone();
|
||||
let Some(c) = d.masks.iter_mut().find(|m| m.id == mask).and_then(|m| m.components.get_mut(comp)) else { return };
|
||||
let MaskShape::Object { edge, .. } = c.shape else { return };
|
||||
c.shape = MaskShape::Object { hint: latest.hint, exclude: latest.exclude, seg: Some(seg), detail: vec![], edge };
|
||||
match self.set_develop(photo, d, "Object Mask") {
|
||||
Ok(()) => {
|
||||
polled.changed = true;
|
||||
polled.refine = Some(mask);
|
||||
}
|
||||
Err(e) => polled.messages.push(e.to_string()),
|
||||
}
|
||||
}
|
||||
Tag::NewPrompt { photo, mask, op, name, text } => {
|
||||
let seg = match o.result {
|
||||
Ok(Some(seg)) => seg,
|
||||
Ok(None) => return polled.messages.push(format!("Nothing matching “{text}” was found in this photo.")),
|
||||
Err(e) => return polled.messages.push(e),
|
||||
};
|
||||
if self.active() != Some(photo) {
|
||||
return polled.messages.push(format!("The photo changed before “{text}” was found; try again."));
|
||||
}
|
||||
let Ok(seg) = serde_json::to_value(seg) else { return };
|
||||
let mut p = serde_json::json!({"kind": "prompt", "text": text, "seg": seg});
|
||||
if let Some(n) = name {
|
||||
p["name"] = serde_json::json!(n);
|
||||
}
|
||||
let r = match mask {
|
||||
None => self.execute("mask.add", &p),
|
||||
Some(m) => {
|
||||
if !self.develop_of(photo).is_some_and(|d| d.masks.iter().any(|x| x.id == m)) {
|
||||
return polled.messages.push("The mask was removed before the selection was found.".into());
|
||||
}
|
||||
self.active_mask = Some(m);
|
||||
p["op"] = serde_json::json!(op);
|
||||
self.execute("mask.addComponent", &p)
|
||||
}
|
||||
};
|
||||
match r {
|
||||
Ok(_) => {
|
||||
polled.changed = true;
|
||||
polled.refine = self.active_mask;
|
||||
}
|
||||
Err(e) => polled.messages.push(e.to_string()),
|
||||
}
|
||||
}
|
||||
Tag::Detail { photo, mask, comp, shape } => {
|
||||
let patch = match o.result {
|
||||
Ok(Some(p)) => p,
|
||||
Ok(None) => return,
|
||||
Err(e) => return log::warn!("SAM 3 detail pass: {e}"),
|
||||
};
|
||||
let Some(d) = self.develop_of(photo) else { return };
|
||||
let mut d = (*d).clone();
|
||||
let Some(c) = d.masks.iter_mut().find(|m| m.id == mask).and_then(|m| m.components.get_mut(comp)) else { return };
|
||||
if strip_detail(c.shape.clone()) != shape {
|
||||
return; // clicked again meanwhile: a newer pass will follow
|
||||
}
|
||||
match &mut c.shape {
|
||||
MaskShape::Object { detail, .. } | MaskShape::Prompt { detail, .. } => *detail = vec![patch],
|
||||
_ => return,
|
||||
}
|
||||
match self.set_develop(photo, d, "AI Mask Detail") {
|
||||
Ok(()) => polled.changed = true,
|
||||
Err(e) => log::warn!("AI mask detail: {e}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn seg_box(side: usize, x0: usize, y0: usize, x1: usize, y1: usize) -> SegMask {
|
||||
let l: Vec<f32> =
|
||||
(0..side * side).map(|i| if (x0..x1).contains(&(i % side)) && (y0..y1).contains(&(i / side)) { 8.0 } else { -8.0 }).collect();
|
||||
SegMask::from_logits(side, &l)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detail_region_frames_the_selection_with_a_margin() {
|
||||
// a small object in the middle of a square photo
|
||||
let r = detail_region(&seg_box(100, 40, 45, 50, 55), 1.0).unwrap();
|
||||
assert!(r[0] < 0.4 && r[2] > 0.5 && r[1] < 0.45 && r[3] > 0.55, "contains it: {r:?}");
|
||||
assert!(r[2] - r[0] < 0.25 && r[3] - r[1] < 0.25, "zoomed in: {r:?}");
|
||||
assert!(r.iter().all(|v| (0.0..=1.0).contains(v)));
|
||||
// at the corner: shifted inside the photo, not cut
|
||||
let r = detail_region(&seg_box(100, 0, 0, 5, 5), 1.5).unwrap();
|
||||
assert!(r[0] >= 0.0 && r[1] >= 0.0 && r[0] < 1e-9 && r[1] < 1e-9);
|
||||
// a thin object stays at most 2:1 (in photo pixels, 1.5:1 photo)
|
||||
let r = detail_region(&seg_box(100, 10, 50, 90, 52), 1.5).unwrap();
|
||||
let (w, h) = ((r[2] - r[0]) * 1.5, r[3] - r[1]);
|
||||
assert!(w / h <= 2.0 + 1e-9, "{w} × {h}");
|
||||
// a hostile aspect ratio stays sane
|
||||
assert!(detail_region(&seg_box(100, 40, 45, 50, 55), f64::NAN).is_some_and(|r| r.iter().all(|v| v.is_finite())));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_region_for_nothing_or_most_of_the_photo() {
|
||||
assert!(detail_region(&seg_box(50, 0, 0, 0, 0), 1.0).is_none());
|
||||
assert!(detail_region(&seg_box(50, 2, 2, 48, 48), 1.0).is_none());
|
||||
let damaged = SegMask { side: 50, data: "?".into(), rect: None };
|
||||
assert!(detail_region(&damaged, 1.0).is_none());
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,426 @@
|
|||
//! The SAM 3 worker thread: it owns the model and the last encoded photo, and runs every
|
||||
//! request (load, render the model input, encode, clicks, descriptions, detail passes) in
|
||||
//! turn. The session only sends jobs and receives outcomes over channels, so no thread ever
|
||||
//! waits on a lock the model holds, and the UI thread never runs the model.
|
||||
//!
|
||||
//! Every job runs under `catch_unwind`: a panic inside candle becomes an error outcome, drops
|
||||
//! the (possibly inconsistent) model and leaves the worker ready for the next job. The busy
|
||||
//! counters are decremented by drop guards, so they can't stay stuck. The model is unloaded
|
||||
//! after [`IDLE_UNLOAD`] without requests (it holds several GB).
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::panic::AssertUnwindSafe;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
use std::sync::mpsc::{self, RecvTimeoutError};
|
||||
use std::time::Duration;
|
||||
|
||||
use lightcraft_develop::SegMask;
|
||||
use lightcraft_segment::{Encoded, Sam3};
|
||||
|
||||
use super::{Input, Prompt, Tag};
|
||||
use crate::media::RenderJob;
|
||||
|
||||
/// Unload the model after this long without a request.
|
||||
pub const IDLE_UNLOAD: Duration = Duration::from_secs(10 * 60);
|
||||
|
||||
/// What to compute.
|
||||
pub enum Kind {
|
||||
/// Load the model and encode the photo (so the first click is fast).
|
||||
Prepare,
|
||||
Clicks(Vec<lightcraft_segment::Click>),
|
||||
Text(String),
|
||||
/// A zoomed-in pass over `region` (normalized) of the photo, rendered by the job's `render`.
|
||||
Detail {
|
||||
prompt: Prompt,
|
||||
region: [f64; 4],
|
||||
},
|
||||
/// Panics (tests of the worker's recovery).
|
||||
#[cfg(test)]
|
||||
Panic,
|
||||
}
|
||||
|
||||
/// One request.
|
||||
pub struct Job {
|
||||
/// The model folder.
|
||||
pub dir: PathBuf,
|
||||
/// Which photo and look the request is about (the encode cache's key).
|
||||
pub key: u64,
|
||||
/// Renders the model input when the cache doesn't hold `key` (always for a detail pass).
|
||||
pub render: Option<RenderJob>,
|
||||
/// The photo's file, when it is missing on disk (for a clearer error if rendering fails).
|
||||
pub missing: Option<String>,
|
||||
pub kind: Kind,
|
||||
pub tag: Tag,
|
||||
pub reply: mpsc::Sender<Outcome>,
|
||||
}
|
||||
|
||||
/// A finished request: the segmentation (`None`: nothing selected / nothing to report).
|
||||
pub struct Outcome {
|
||||
pub tag: Tag,
|
||||
pub result: Result<Option<SegMask>, String>,
|
||||
/// A newer request for the same selection came in; this one was skipped.
|
||||
pub superseded: bool,
|
||||
}
|
||||
|
||||
/// Counters the session reads without locking.
|
||||
#[derive(Default)]
|
||||
pub struct Shared {
|
||||
/// Requests other than detail passes, queued or running.
|
||||
pub pending: AtomicUsize,
|
||||
/// Detail passes queued or running.
|
||||
pub detail: AtomicUsize,
|
||||
/// The model is in memory.
|
||||
pub loaded: AtomicBool,
|
||||
/// Loading the model or encoding a photo right now (the slow parts).
|
||||
pub encoding: AtomicBool,
|
||||
}
|
||||
|
||||
/// Decrements a counter when dropped (also when unwinding).
|
||||
struct Done<'a>(&'a AtomicUsize);
|
||||
|
||||
impl Drop for Done<'_> {
|
||||
fn drop(&mut self) {
|
||||
let _ = self.0.try_update(Ordering::SeqCst, Ordering::SeqCst, |n| Some(n.saturating_sub(1)));
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets a flag while alive.
|
||||
struct Flag<'a>(&'a AtomicBool);
|
||||
|
||||
impl<'a> Flag<'a> {
|
||||
fn raise(f: &'a AtomicBool) -> Self {
|
||||
f.store(true, Ordering::SeqCst);
|
||||
Flag(f)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Flag<'_> {
|
||||
fn drop(&mut self) {
|
||||
self.0.store(false, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
|
||||
/// The session's handle on the worker thread (started on the first request, restarted if it
|
||||
/// ever ends).
|
||||
#[derive(Default)]
|
||||
pub struct Worker {
|
||||
tx: Option<mpsc::Sender<Job>>,
|
||||
thread: Option<std::thread::JoinHandle<()>>,
|
||||
pub shared: Arc<Shared>,
|
||||
}
|
||||
|
||||
impl Worker {
|
||||
fn counter(&self, job: &Job) -> &AtomicUsize {
|
||||
if matches!(job.kind, Kind::Detail { .. }) { &self.shared.detail } else { &self.shared.pending }
|
||||
}
|
||||
|
||||
/// Whether the thread has ended (it never should); its counters are then meaningless.
|
||||
fn dead(&self) -> bool {
|
||||
self.thread.as_ref().is_some_and(|t| t.is_finished())
|
||||
}
|
||||
|
||||
pub fn pending(&self) -> usize {
|
||||
if self.dead() { 0 } else { self.shared.pending.load(Ordering::SeqCst) }
|
||||
}
|
||||
|
||||
pub fn detail(&self) -> usize {
|
||||
if self.dead() { 0 } else { self.shared.detail.load(Ordering::SeqCst) }
|
||||
}
|
||||
|
||||
pub fn encoding(&self) -> bool {
|
||||
!self.dead() && self.shared.encoding.load(Ordering::SeqCst)
|
||||
}
|
||||
|
||||
pub fn loaded(&self) -> bool {
|
||||
!self.dead() && self.shared.loaded.load(Ordering::SeqCst)
|
||||
}
|
||||
|
||||
/// Queue `job`; returns at once.
|
||||
pub fn submit(&mut self, job: Job) -> Result<(), String> {
|
||||
if self.dead() {
|
||||
// a thread that ended took its counters' meaning with it
|
||||
self.tx = None;
|
||||
self.thread = None;
|
||||
self.shared = Arc::new(Shared::default());
|
||||
}
|
||||
self.counter(&job).fetch_add(1, Ordering::SeqCst);
|
||||
let job = match &self.tx {
|
||||
Some(tx) => match tx.send(job) {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(mpsc::SendError(job)) => job,
|
||||
},
|
||||
None => job,
|
||||
};
|
||||
// (re)start the thread
|
||||
let (tx, rx) = mpsc::channel();
|
||||
let shared = self.shared.clone();
|
||||
match std::thread::Builder::new().name("sam3".into()).spawn(move || run(&rx, &shared, IDLE_UNLOAD)) {
|
||||
Ok(t) => {
|
||||
self.thread = Some(t);
|
||||
self.tx = Some(tx.clone());
|
||||
tx.send(job).map_err(|_| "the AI mask worker stopped".to_string())
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = self.counter(&job).try_update(Ordering::SeqCst, Ordering::SeqCst, |n| Some(n.saturating_sub(1)));
|
||||
Err(format!("could not start the AI mask worker: {e}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The model and the last encoded photo.
|
||||
#[derive(Default)]
|
||||
struct State {
|
||||
model: Option<(PathBuf, Sam3)>,
|
||||
cache: Option<(u64, Encoded)>,
|
||||
}
|
||||
|
||||
/// Whether `later` makes `job` pointless (it asks for the same selection again).
|
||||
fn supersedes(later: &Job, job: &Job) -> bool {
|
||||
match (&later.tag, &job.tag) {
|
||||
(Tag::Prepare, Tag::Prepare) => later.key == job.key,
|
||||
(Tag::Clicks { photo: a, mask: am, comp: ac, .. }, Tag::Clicks { photo: b, mask: bm, comp: bc, .. }) => (a, am, ac) == (b, bm, bc),
|
||||
(Tag::Detail { photo: a, mask: am, comp: ac, .. }, Tag::Detail { photo: b, mask: bm, comp: bc, .. }) => (a, am, ac) == (b, bm, bc),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn run(rx: &mpsc::Receiver<Job>, shared: &Shared, idle: Duration) {
|
||||
let mut state = State::default();
|
||||
let mut queue: VecDeque<Job> = VecDeque::new();
|
||||
loop {
|
||||
if queue.is_empty() {
|
||||
match rx.recv_timeout(idle) {
|
||||
Ok(j) => queue.push_back(j),
|
||||
Err(RecvTimeoutError::Timeout) => {
|
||||
if state.model.is_some() {
|
||||
state = State::default();
|
||||
shared.loaded.store(false, Ordering::SeqCst);
|
||||
log::info!("SAM 3 unloaded after {} minutes without use", idle.as_secs() / 60);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
Err(RecvTimeoutError::Disconnected) => return,
|
||||
}
|
||||
}
|
||||
while let Ok(j) = rx.try_recv() {
|
||||
queue.push_back(j);
|
||||
}
|
||||
let Some(mut job) = queue.pop_front() else { continue };
|
||||
let is_detail = matches!(job.kind, Kind::Detail { .. });
|
||||
let _done = Done(if is_detail { &shared.detail } else { &shared.pending });
|
||||
if queue.iter().any(|later| supersedes(later, &job)) {
|
||||
let _ = job.reply.send(Outcome { tag: job.tag, result: Ok(None), superseded: true });
|
||||
continue;
|
||||
}
|
||||
let result = match std::panic::catch_unwind(AssertUnwindSafe(|| run_job(&mut state, shared, &mut job))) {
|
||||
Ok(r) => r,
|
||||
Err(p) => {
|
||||
let why = p.downcast_ref::<&str>().map(|s| s.to_string()).or_else(|| p.downcast_ref::<String>().cloned()).unwrap_or_default();
|
||||
log::error!("SAM 3 panicked: {why}");
|
||||
// don't trust what the model kept: load it again next time
|
||||
state = State::default();
|
||||
Err("the AI model failed unexpectedly (it will be reloaded on the next try)".to_string())
|
||||
}
|
||||
};
|
||||
if result.is_err() && !is_detail {
|
||||
state.cache = None;
|
||||
}
|
||||
shared.loaded.store(state.model.is_some(), Ordering::SeqCst);
|
||||
let _ = job.reply.send(Outcome { tag: job.tag, result, superseded: false });
|
||||
}
|
||||
}
|
||||
|
||||
/// The model for `dir`, loading it when needed.
|
||||
fn model<'a>(state: &'a mut State, shared: &Shared, dir: &std::path::Path) -> Result<&'a mut Sam3, String> {
|
||||
if state.model.as_ref().is_none_or(|(d, _)| d != dir) {
|
||||
state.model = None;
|
||||
state.cache = None;
|
||||
let _busy = Flag::raise(&shared.encoding);
|
||||
let t = web_time::Instant::now();
|
||||
let m = Sam3::load(dir).map_err(|e| e.to_string())?;
|
||||
log::info!("SAM 3 loaded in {:?}", t.elapsed());
|
||||
state.model = Some((dir.to_path_buf(), m));
|
||||
shared.loaded.store(true, Ordering::SeqCst);
|
||||
}
|
||||
state.model.as_mut().map(|(_, m)| m).ok_or_else(|| "the model did not load".to_string())
|
||||
}
|
||||
|
||||
/// Render the model input.
|
||||
fn render(job: &mut Job) -> Result<Input, String> {
|
||||
let r = job.render.take().ok_or("the photo is not available")?.run();
|
||||
let img = match r.rendered {
|
||||
Ok(r) => r.image,
|
||||
Err(e) => {
|
||||
return Err(match &job.missing {
|
||||
Some(path) => format!(
|
||||
"This photo's file is missing (moved or deleted), so it can't be analyzed for AI masks: {path}. Library ▸ Find Missing Photos can relink it."
|
||||
),
|
||||
None => format!("couldn't render the photo for AI masks: {e}"),
|
||||
});
|
||||
}
|
||||
};
|
||||
let rgb = img.data.iter().flat_map(|p| [p[0], p[1], p[2]]).collect();
|
||||
Ok(Input { rgb, w: img.width, h: img.height })
|
||||
}
|
||||
|
||||
/// The model and the encoded photo `job.key` (rendering and encoding it when needed).
|
||||
fn encoded<'a>(state: &'a mut State, shared: &Shared, job: &mut Job) -> Result<(&'a mut Sam3, &'a mut Encoded), String> {
|
||||
model(state, shared, &job.dir)?;
|
||||
if state.cache.as_ref().is_none_or(|(k, _)| *k != job.key) {
|
||||
state.cache = None;
|
||||
let input = render(job)?;
|
||||
let _busy = Flag::raise(&shared.encoding);
|
||||
let m = state.model.as_mut().map(|(_, m)| m).ok_or("the model did not load")?;
|
||||
let t = web_time::Instant::now();
|
||||
let mut enc = m.encode(&input.rgb, input.w, input.h).map_err(|e| e.to_string())?;
|
||||
m.prepare_clicks(&mut enc).map_err(|e| e.to_string())?;
|
||||
log::info!("SAM 3 encoded {}×{} in {:?}", input.w, input.h, t.elapsed());
|
||||
state.cache = Some((job.key, enc));
|
||||
}
|
||||
let State { model, cache } = state;
|
||||
match (model.as_mut(), cache.as_mut()) {
|
||||
(Some((_, m)), Some((_, e))) => Ok((m, e)),
|
||||
_ => Err("the model is not ready".into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Every phrase of `text` ("car, road"), merged (the per-pixel maximum); `None` when nothing
|
||||
/// matches.
|
||||
fn text_logits(model: &mut Sam3, enc: &mut Encoded, text: &str) -> Result<Option<Vec<f32>>, String> {
|
||||
let mut merged: Option<Vec<f32>> = None;
|
||||
for phrase in text.split([',', ';']).map(str::trim).filter(|p| !p.is_empty()) {
|
||||
if let Some(p) = model.segment_text(enc, phrase, 0.5).map_err(|e| e.to_string())? {
|
||||
merged = Some(match merged {
|
||||
None => p.logits,
|
||||
Some(m) => m.iter().zip(&p.logits).map(|(a, b)| a.max(*b)).collect(),
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(merged)
|
||||
}
|
||||
|
||||
fn run_job(state: &mut State, shared: &Shared, job: &mut Job) -> Result<Option<SegMask>, String> {
|
||||
let side = lightcraft_segment::MASK_SIDE;
|
||||
match std::mem::replace(&mut job.kind, Kind::Prepare) {
|
||||
Kind::Prepare => {
|
||||
encoded(state, shared, job)?;
|
||||
Ok(None)
|
||||
}
|
||||
Kind::Clicks(clicks) => {
|
||||
let (m, enc) = encoded(state, shared, job)?;
|
||||
let p = m.segment_clicks(enc, &clicks).map_err(|e| e.to_string())?;
|
||||
Ok(Some(SegMask::from_logits(side, &p.logits)))
|
||||
}
|
||||
Kind::Text(text) => {
|
||||
let (m, enc) = encoded(state, shared, job)?;
|
||||
Ok(text_logits(m, enc, &text)?.map(|l| SegMask::from_logits(side, &l)))
|
||||
}
|
||||
#[cfg(test)]
|
||||
#[allow(clippy::panic)]
|
||||
Kind::Panic => panic!("a test panic"),
|
||||
Kind::Detail { prompt, region } => {
|
||||
let input = render(job)?;
|
||||
let m = model(state, shared, &job.dir)?;
|
||||
let t = web_time::Instant::now();
|
||||
let mut enc = m.encode(&input.rgb, input.w, input.h).map_err(|e| e.to_string())?;
|
||||
let (rw, rh) = (region[2] - region[0], region[3] - region[1]);
|
||||
if !(rw > 0.0 && rh > 0.0) {
|
||||
return Ok(None);
|
||||
}
|
||||
let logits = match prompt {
|
||||
Prompt::Clicks(clicks) => {
|
||||
let inside: Vec<lightcraft_segment::Click> = clicks
|
||||
.iter()
|
||||
.map(|c| ((c.at.x - region[0]) / rw, (c.at.y - region[1]) / rh, c.include))
|
||||
.filter(|(x, y, _)| (0.0..=1.0).contains(x) && (0.0..=1.0).contains(y))
|
||||
.map(|(x, y, positive)| lightcraft_segment::Click { x: x as f32, y: y as f32, positive })
|
||||
.collect();
|
||||
if !inside.iter().any(|c| c.positive) {
|
||||
return Ok(None);
|
||||
}
|
||||
Some(m.segment_clicks(&mut enc, &inside).map_err(|e| e.to_string())?.logits)
|
||||
}
|
||||
Prompt::Text(text) => text_logits(m, &mut enc, &text)?,
|
||||
};
|
||||
log::info!("SAM 3 detail pass ({}×{}) in {:?}", input.w, input.h, t.elapsed());
|
||||
Ok(logits.map(|l| SegMask::from_logits_in(side, &l, region)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn job(tag: Tag, reply: &mpsc::Sender<Outcome>) -> Job {
|
||||
// an empty folder: loading fails with an error (no model files)
|
||||
let dir = std::env::temp_dir().join(format!("lc-sam3-worker-none-{}", std::process::id()));
|
||||
Job { dir, key: 1, render: None, missing: None, kind: Kind::Prepare, tag, reply: reply.clone() }
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn errors_come_back_and_the_counters_return_to_zero() {
|
||||
let mut w = Worker::default();
|
||||
let (tx, rx) = mpsc::channel();
|
||||
for _ in 0..3 {
|
||||
w.submit(job(Tag::Wait, &tx)).unwrap();
|
||||
}
|
||||
for _ in 0..3 {
|
||||
let o = rx.recv_timeout(Duration::from_secs(30)).unwrap();
|
||||
assert!(o.result.unwrap_err().contains("not found"));
|
||||
}
|
||||
let t = web_time::Instant::now();
|
||||
while w.pending() > 0 && t.elapsed() < Duration::from_secs(5) {
|
||||
std::thread::yield_now();
|
||||
}
|
||||
assert_eq!(w.pending(), 0);
|
||||
assert!(!w.encoding() && !w.loaded());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_panic_in_a_job_is_an_error_and_the_worker_carries_on() {
|
||||
let shared = Shared::default();
|
||||
let (tx, rx) = mpsc::channel();
|
||||
let (jtx, jrx) = mpsc::channel();
|
||||
let mut j = job(Tag::Wait, &tx);
|
||||
j.kind = Kind::Panic;
|
||||
jtx.send(j).unwrap();
|
||||
jtx.send(job(Tag::Wait, &tx)).unwrap();
|
||||
shared.pending.store(2, Ordering::SeqCst);
|
||||
drop(jtx);
|
||||
run(&jrx, &shared, Duration::from_secs(60));
|
||||
let first = rx.recv().unwrap();
|
||||
assert!(first.result.unwrap_err().contains("failed unexpectedly"));
|
||||
let second = rx.recv().unwrap();
|
||||
assert!(second.result.is_err(), "the next job still runs");
|
||||
assert_eq!(shared.pending.load(Ordering::SeqCst), 0, "counters reset");
|
||||
assert!(!shared.encoding.load(Ordering::SeqCst));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_newer_request_for_the_same_selection_supersedes_a_queued_one() {
|
||||
let (tx, rx) = mpsc::channel();
|
||||
let (jtx, jrx) = mpsc::channel();
|
||||
let clicks = |n: usize| Tag::Clicks {
|
||||
photo: lightcraft_catalog::PhotoId(1),
|
||||
mask: 1,
|
||||
comp: 0,
|
||||
hint: vec![lightcraft_geom::Point::new(0.5, 0.5); n],
|
||||
exclude: vec![],
|
||||
};
|
||||
for n in 1..=3 {
|
||||
jtx.send(job(clicks(n), &tx)).unwrap();
|
||||
}
|
||||
drop(jtx);
|
||||
let shared = Shared::default();
|
||||
shared.pending.store(3, Ordering::SeqCst);
|
||||
run(&jrx, &shared, Duration::from_secs(60));
|
||||
let out: Vec<Outcome> = rx.try_iter().collect();
|
||||
assert_eq!(out.len(), 3);
|
||||
assert!(out[0].superseded && out[1].superseded && !out[2].superseded);
|
||||
assert_eq!(shared.pending.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
}
|
||||
|
|
@ -728,6 +728,40 @@ fn curve_reset_by_channel_and_whole() {
|
|||
assert!(s.execute("curve.reset", &json!({"channel": "alpha"})).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ai_masks_without_the_model_explain_themselves() {
|
||||
let mut s = demo();
|
||||
let id = s.active().unwrap();
|
||||
// an Object mask that could never be computed is refused, not left empty
|
||||
let e = s.execute("mask.add", &json!({"kind": "object"})).unwrap_err().to_string();
|
||||
assert!(e.contains("SAM 3") || e.contains("not available"), "{e}");
|
||||
assert!(!e.contains("invalid parameters"), "shown as is: {e}");
|
||||
assert!(s.develop_of(id).unwrap().masks.is_empty());
|
||||
// an empty Object component (as an older document may hold) still takes clicks
|
||||
let shape = lightcraft_develop::MaskShape::Object { hint: vec![], exclude: vec![], seg: None, detail: vec![], edge: 0.0 };
|
||||
let mut d = (*s.develop_of(id).unwrap()).clone();
|
||||
d.masks.push(lightcraft_develop::Mask {
|
||||
id: 1,
|
||||
components: vec![lightcraft_develop::MaskComponent { name: None, op: lightcraft_develop::MaskOp::Add, invert: false, shape }],
|
||||
..Default::default()
|
||||
});
|
||||
s.set_develop(id, d, "test").unwrap();
|
||||
s.active_mask = Some(1);
|
||||
// clicks and descriptions need the model: an error that says what's missing, nothing changed
|
||||
let e = s.execute("mask.objectPoint", &json!({"x": 0.5, "y": 0.5})).unwrap_err().to_string();
|
||||
assert!(e.contains("SAM 3") || e.contains("not available"), "{e}");
|
||||
let e = s.execute("mask.add", &json!({"kind": "prompt", "text": "sky"})).unwrap_err().to_string();
|
||||
assert!(e.contains("SAM 3") || e.contains("not available"), "{e}");
|
||||
let empty = std::env::temp_dir().join(format!("lc-no-sam3-{}", std::process::id()));
|
||||
s.segmenter.dir = Some(empty.clone());
|
||||
let e = s.execute("mask.objectPoint", &json!({"x": 0.5, "y": 0.5})).unwrap_err().to_string();
|
||||
assert!(e.contains(&empty.display().to_string()) || e.contains("not available"), "names the folder: {e}");
|
||||
assert_eq!(s.develop_of(id).unwrap().masks.len(), 1);
|
||||
// bad input is an error, not a crash
|
||||
assert!(s.execute("mask.objectPoint", &json!({"x": 1.5, "y": 0.5})).is_err());
|
||||
assert!(s.execute("mask.add", &json!({"kind": "prompt", "text": " "})).is_err());
|
||||
}
|
||||
|
||||
/// A People card's close-up is a square (in pixels) inside the photo, around the face, whatever the
|
||||
/// face rectangle says: corners, oversized and degenerate rectangles never leave the frame or panic.
|
||||
#[test]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,232 @@
|
|||
//! AI masks without the SAM 3 model (it is never required): every request is a clear error that
|
||||
//! comes back at once, nothing waits on the network or the model, a damaged model is an error
|
||||
//! (not a crash or a stuck "busy"), and the download command needs consent and fails cleanly.
|
||||
//! No test touches the internet: downloads go to a local server.
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use crate::Session;
|
||||
use crate::segment::{NOT_INSTALLED, Segmenter};
|
||||
|
||||
fn tmp(name: &str) -> std::path::PathBuf {
|
||||
let d = std::env::temp_dir().join(format!("lc-sam3-engine-{name}-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
d
|
||||
}
|
||||
|
||||
/// A folder that looks installed but holds no usable model.
|
||||
fn damaged_model(name: &str) -> std::path::PathBuf {
|
||||
let d = tmp(name);
|
||||
std::fs::create_dir_all(&d).unwrap();
|
||||
std::fs::write(d.join("model.safetensors"), b"\x10\x00\x00\x00\x00\x00\x00\x00{not json at all}").unwrap();
|
||||
std::fs::write(d.join("vocab.json"), b"{}").unwrap();
|
||||
std::fs::write(d.join("merges.txt"), b"#version: 0.2\n").unwrap();
|
||||
d
|
||||
}
|
||||
|
||||
/// Poll until the worker is idle (or `limit` passes), collecting messages.
|
||||
fn settle(s: &mut Session, limit: Duration) -> Vec<String> {
|
||||
let t = Instant::now();
|
||||
let mut messages = Vec::new();
|
||||
loop {
|
||||
let p = s.segment_poll();
|
||||
messages.extend(p.messages);
|
||||
if !s.segmenter.busy() && !s.segmenter.detail_busy() {
|
||||
messages.extend(s.segment_poll().messages);
|
||||
return messages;
|
||||
}
|
||||
assert!(t.elapsed() < limit, "the worker never finished");
|
||||
std::thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn status_and_model_commands_without_the_model() {
|
||||
let mut s = Session::with_demo();
|
||||
let st = s.execute("segment.model.status", &json!({})).unwrap();
|
||||
assert_eq!(st["available"], Segmenter::AVAILABLE);
|
||||
assert_eq!(st["installed"], false);
|
||||
assert_eq!(st["download"]["running"], false);
|
||||
assert!(st["license"].as_str().unwrap().contains("SAM License"));
|
||||
// downloading needs the user's consent, and names the size and the licence
|
||||
let e = s.execute("segment.model.download", &json!({})).unwrap_err().to_string();
|
||||
assert!(e.contains("acknowledged") && e.contains("SAM License") && e.contains("GB"), "{e}");
|
||||
let e = s.execute("segment.model.download", &json!({"acknowledged": "yes"})).unwrap_err().to_string();
|
||||
assert!(e.contains("acknowledged"), "{e}");
|
||||
assert_eq!(s.execute("segment.model.cancel", &json!({})).unwrap()["cancelled"], false);
|
||||
if !Segmenter::AVAILABLE {
|
||||
assert!(s.execute("segment.model.download", &json!({"acknowledged": true})).is_err());
|
||||
return;
|
||||
}
|
||||
// no folder, then no download location: clear errors, nothing started
|
||||
let e = s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap_err().to_string();
|
||||
assert!(e.contains("no folder"), "{e}");
|
||||
let dir = tmp("nomirrors");
|
||||
s.segmenter.dir = Some(dir.clone());
|
||||
s.segmenter.mirrors_file = Some(dir.join("none.txt"));
|
||||
if s.segmenter.mirrors().is_empty() {
|
||||
let e = s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap_err().to_string();
|
||||
assert!(e.contains("LIGHTCRAFT_SAM3_MIRRORS"), "{e}");
|
||||
}
|
||||
assert!(!s.segmenter.download_status().running);
|
||||
}
|
||||
|
||||
#[cfg(feature = "sam")]
|
||||
#[test]
|
||||
fn a_failing_download_ends_with_an_error_and_never_blocks() {
|
||||
use std::io::{BufRead, BufReader, Write};
|
||||
// a local mirror that has nothing
|
||||
let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let base = format!("http://{}/sam3", l.local_addr().unwrap());
|
||||
std::thread::spawn(move || {
|
||||
for mut c in l.incoming().flatten() {
|
||||
let mut r = BufReader::new(c.try_clone().unwrap());
|
||||
let mut line = String::new();
|
||||
while r.read_line(&mut line).unwrap_or(0) > 2 {
|
||||
line.clear();
|
||||
}
|
||||
let _ = c.write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 13\r\nConnection: close\r\n\r\n<h1>404</h1>\n");
|
||||
}
|
||||
});
|
||||
let dir = tmp("download404");
|
||||
let mirrors = dir.with_extension("mirrors.txt");
|
||||
std::fs::write(&mirrors, format!("# test mirror\n{base}\n")).unwrap();
|
||||
let mut s = Session::with_demo();
|
||||
s.segmenter.dir = Some(dir.clone());
|
||||
s.segmenter.mirrors_file = Some(mirrors.clone());
|
||||
let t = Instant::now();
|
||||
let r = s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap();
|
||||
assert_eq!(r["started"], true);
|
||||
assert!(t.elapsed() < Duration::from_secs(1), "starting a download returns at once");
|
||||
// a second request while it runs starts nothing
|
||||
let again = s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap();
|
||||
assert!(again["started"] == false || again["downloading"] == false);
|
||||
let t = Instant::now();
|
||||
while s.segmenter.download_status().running {
|
||||
assert!(t.elapsed() < Duration::from_secs(30), "the download never ended");
|
||||
// AI mask requests meanwhile: an error at once, never a wait
|
||||
let q = Instant::now();
|
||||
let e = s.execute("mask.add", &json!({"kind": "prompt", "text": "sky"})).unwrap_err().to_string();
|
||||
assert!(e.contains(NOT_INSTALLED), "{e}");
|
||||
assert!(q.elapsed() < Duration::from_millis(500));
|
||||
std::thread::sleep(Duration::from_millis(20));
|
||||
}
|
||||
let st = s.segmenter.download_status();
|
||||
assert!(!st.finished);
|
||||
let e = st.error.unwrap();
|
||||
assert!(e.contains("not found"), "{e}");
|
||||
assert!(!s.segmenter.installed());
|
||||
// nothing was left behind under a real name, and no error page was saved
|
||||
for f in ["model.safetensors", "vocab.json", "merges.txt", "vocab.json.part"] {
|
||||
assert!(!dir.join(f).exists(), "{f}");
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
let _ = std::fs::remove_file(&mirrors);
|
||||
}
|
||||
|
||||
/// The app (background mode) never waits: with the model missing every request fails at once;
|
||||
/// with a damaged model the requests are queued, fail on the worker, and come back as messages.
|
||||
#[test]
|
||||
fn background_requests_never_wait_and_failures_come_back_as_messages() {
|
||||
let mut s = Session::with_demo();
|
||||
s.segmenter.background = true;
|
||||
let id = s.active().unwrap();
|
||||
s.segmenter.dir = Some(tmp("missing"));
|
||||
for (cmd, p) in [("mask.add", json!({"kind": "object"})), ("mask.add", json!({"kind": "prompt", "text": "sky"})), ("segment.prepare", json!({}))]
|
||||
{
|
||||
let t = Instant::now();
|
||||
let e = s.execute(cmd, &p).unwrap_err().to_string();
|
||||
assert!(t.elapsed() < Duration::from_millis(500), "{cmd}");
|
||||
if Segmenter::AVAILABLE {
|
||||
assert!(e.starts_with(NOT_INSTALLED), "{cmd}: {e}");
|
||||
} else {
|
||||
assert!(e.contains("not available"), "{cmd}: {e}");
|
||||
}
|
||||
}
|
||||
assert!(s.develop_of(id).unwrap().masks.is_empty());
|
||||
if !Segmenter::AVAILABLE {
|
||||
return;
|
||||
}
|
||||
// a model folder whose weights are damaged
|
||||
let dir = damaged_model("damaged");
|
||||
s.segmenter.dir = Some(dir.clone());
|
||||
assert!(s.segmenter.installed());
|
||||
let t = Instant::now();
|
||||
s.execute("mask.add", &json!({"kind": "object"})).unwrap();
|
||||
let r = s.execute("mask.objectPoint", &json!({"x": 0.5, "y": 0.5})).unwrap();
|
||||
assert_eq!(r["pending"], true);
|
||||
assert!(s.segmenter.pending_clicks().is_some_and(|p| p.hint.len() == 1));
|
||||
// a second click while the first is on its way builds on it
|
||||
let r = s.execute("mask.objectPoint", &json!({"x": 0.6, "y": 0.5, "exclude": true})).unwrap();
|
||||
assert_eq!((r["include"].as_u64(), r["exclude"].as_u64()), (Some(1), Some(1)));
|
||||
let r = s.execute("mask.add", &json!({"kind": "prompt", "text": "the sky"})).unwrap();
|
||||
assert_eq!(r["pending"], true);
|
||||
assert!(t.elapsed() < Duration::from_secs(1), "nothing waited for the model: {:?}", t.elapsed());
|
||||
let messages = settle(&mut s, Duration::from_secs(60));
|
||||
assert!(!messages.is_empty() && messages.iter().all(|m| !m.is_empty()), "{messages:?}");
|
||||
assert!(messages.iter().any(|m| m.contains("model.safetensors") || m.contains("SAM 3")), "{messages:?}");
|
||||
assert!(s.segmenter.pending_clicks().is_none());
|
||||
// the Object mask is there, without a selection; no Describe mask was made from nothing
|
||||
let d = s.develop_of(id).unwrap();
|
||||
assert_eq!(d.masks.len(), 1);
|
||||
assert!(matches!(&d.masks[0].components[0].shape, lightcraft_develop::MaskShape::Object { seg: None, .. }));
|
||||
assert!(!s.segmenter.busy() && !s.segmenter.loaded());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
/// CLI / MCP (waiting mode): a damaged model is an error from the command, quickly.
|
||||
#[test]
|
||||
fn waiting_requests_with_a_damaged_model_are_errors() {
|
||||
if !Segmenter::AVAILABLE {
|
||||
return;
|
||||
}
|
||||
let mut s = Session::with_demo();
|
||||
let id = s.active().unwrap();
|
||||
let dir = damaged_model("damaged-wait");
|
||||
s.segmenter.dir = Some(dir.clone());
|
||||
let e = s.execute("mask.add", &json!({"kind": "prompt", "text": "sky"})).unwrap_err().to_string();
|
||||
assert!(e.contains("model.safetensors") || e.contains("SAM 3"), "{e}");
|
||||
let e = s.execute("mask.add", &json!({"kind": "object", "points": [[0.5, 0.5]]})).unwrap_err().to_string();
|
||||
assert!(!e.is_empty());
|
||||
assert!(s.develop_of(id).unwrap().masks.is_empty());
|
||||
assert!(!s.segmenter.busy());
|
||||
// a stored segmentation needs no model at all (replayed or pasted masks)
|
||||
let seg = serde_json::to_value(lightcraft_develop::SegMask::from_logits(4, &[5.0; 16])).unwrap();
|
||||
s.segmenter.dir = None;
|
||||
s.execute("mask.add", &json!({"kind": "prompt", "text": "sky", "seg": seg})).unwrap();
|
||||
assert_eq!(s.develop_of(id).unwrap().masks.len(), 1);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
/// Too many clicks on one Object selection are refused, not sent to the model.
|
||||
#[test]
|
||||
fn object_clicks_are_capped() {
|
||||
if !Segmenter::AVAILABLE {
|
||||
return;
|
||||
}
|
||||
let mut s = Session::with_demo();
|
||||
s.segmenter.background = true;
|
||||
let dir = damaged_model("clicks");
|
||||
s.segmenter.dir = Some(dir.clone());
|
||||
let id = s.active().unwrap();
|
||||
let mut d = (*s.develop_of(id).unwrap()).clone();
|
||||
let many = vec![lightcraft_geom::Point::new(0.5, 0.5); crate::segment::MAX_CLICKS];
|
||||
d.masks.push(lightcraft_develop::Mask {
|
||||
id: 1,
|
||||
components: vec![lightcraft_develop::MaskComponent {
|
||||
name: None,
|
||||
op: lightcraft_develop::MaskOp::Add,
|
||||
invert: false,
|
||||
shape: lightcraft_develop::MaskShape::Object { hint: many, exclude: vec![], seg: None, detail: vec![], edge: 0.0 },
|
||||
}],
|
||||
..Default::default()
|
||||
});
|
||||
s.set_develop(id, d, "test").unwrap();
|
||||
s.active_mask = Some(1);
|
||||
let e = s.execute("mask.objectPoint", &json!({"x": 0.5, "y": 0.5})).unwrap_err().to_string();
|
||||
assert!(e.contains("at most"), "{e}");
|
||||
assert!(!s.segmenter.busy());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
|
@ -3,7 +3,7 @@
|
|||
//! Shapes are defined in normalized oriented-image coordinates; sizes are in "long-edge units".
|
||||
//! Components combine with Add (max), Subtract (a·(1−c)) and Intersect (a·c).
|
||||
|
||||
use lightcraft_develop::{BrushStroke, LocalAdjustments, Mask, MaskOp, MaskShape};
|
||||
use lightcraft_develop::{BrushStroke, LocalAdjustments, Mask, MaskOp, MaskShape, SegMask};
|
||||
use lightcraft_geom::Point;
|
||||
use lightcraft_raster::{Plane, Rgb32f};
|
||||
|
||||
|
|
@ -153,6 +153,19 @@ pub fn shape_alpha(shape: &MaskShape, frame: &Frame, w: usize, h: usize, img: &R
|
|||
}
|
||||
smooth_plane(&mut out, 0.01 * frame_px(frame, w));
|
||||
}
|
||||
MaskShape::Object { seg, detail, edge, .. } | MaskShape::Prompt { seg, detail, edge, .. } if seg.is_some() || !detail.is_empty() => {
|
||||
// Edge: below 0 a steeper transition (up to 8× the logits), above 0 feathered (up to
|
||||
// 2 % of the long edge, so previews and exports match)
|
||||
let e = (*edge / 100.0).clamp(-1.0, 1.0) as f32;
|
||||
let gain = if e < 0.0 { 1.0 - 7.0 * e } else { 1.0 };
|
||||
sample_seg(seg.as_ref(), detail, gain, frame, w, h, &mut out);
|
||||
if e > 0.0 {
|
||||
smooth_plane(&mut out, e * 0.02 * frame_px(frame, w));
|
||||
}
|
||||
}
|
||||
// not computed yet (no clicks, or a build without the model): nothing selected
|
||||
MaskShape::Object { hint, .. } if hint.is_empty() => {}
|
||||
MaskShape::Prompt { .. } => {}
|
||||
MaskShape::Subject | MaskShape::Object { .. } | MaskShape::People { .. } => {
|
||||
// Saliency heuristic: centre-weighted local contrast (replaced by the segmenter in M12).
|
||||
let m = frame.out_to_norm(w, h);
|
||||
|
|
@ -176,6 +189,61 @@ pub fn shape_alpha(shape: &MaskShape, frame: &Frame, w: usize, h: usize, img: &R
|
|||
out
|
||||
}
|
||||
|
||||
/// A decoded [`SegMask`]: its logits and the normalized rectangle they cover.
|
||||
struct SegGrid {
|
||||
logits: Vec<f32>,
|
||||
side: usize,
|
||||
r: [f64; 4],
|
||||
}
|
||||
|
||||
impl SegGrid {
|
||||
fn new(seg: &SegMask) -> Option<SegGrid> {
|
||||
Some(SegGrid { logits: seg.logits()?, side: seg.side as usize, r: seg.bounds()? })
|
||||
}
|
||||
|
||||
fn contains(&self, n: Point) -> bool {
|
||||
n.x >= self.r[0] && n.x <= self.r[2] && n.y >= self.r[1] && n.y <= self.r[3]
|
||||
}
|
||||
|
||||
/// The bilinearly sampled logit at normalized image point `n` (cell centres at
|
||||
/// `(i + 0.5) / side` of the rectangle).
|
||||
fn logit(&self, n: Point) -> f32 {
|
||||
let side = self.side;
|
||||
let at = |x: usize, y: usize| self.logits.get(y.min(side - 1) * side + x.min(side - 1)).copied().unwrap_or(-16.0);
|
||||
let u = (n.x - self.r[0]) / (self.r[2] - self.r[0]);
|
||||
let v = (n.y - self.r[1]) / (self.r[3] - self.r[1]);
|
||||
let (fx, fy) = ((u * side as f64 - 0.5) as f32, (v * side as f64 - 0.5) as f32);
|
||||
let (fx, fy) = (fx.clamp(0.0, (side - 1) as f32), fy.clamp(0.0, (side - 1) as f32));
|
||||
let (x0, y0) = (fx.floor() as usize, fy.floor() as usize);
|
||||
let (tx, ty) = (fx - x0 as f32, fy - y0 as f32);
|
||||
let top = at(x0, y0) * (1.0 - tx) + at(x0 + 1, y0) * tx;
|
||||
let bot = at(x0, y0 + 1) * (1.0 - tx) + at(x0 + 1, y0 + 1) * tx;
|
||||
top * (1.0 - ty) + bot * ty
|
||||
}
|
||||
}
|
||||
|
||||
/// Stored model segmentations at each output pixel: the zoomed-in `detail` passes inside
|
||||
/// their rectangles (the highest of them where they overlap), the whole-photo `seg`
|
||||
/// elsewhere; the logit (times `gain`: > 1 is a harder edge) through a sigmoid. Damaged
|
||||
/// data selects nothing.
|
||||
fn sample_seg(seg: Option<&SegMask>, detail: &[SegMask], gain: f32, frame: &Frame, w: usize, h: usize, out: &mut Plane) {
|
||||
let coarse = seg.and_then(SegGrid::new);
|
||||
// (capped: a hostile document can hold many)
|
||||
let patches: Vec<SegGrid> = detail.iter().take(lightcraft_develop::segmask::MAX_DETAIL).filter_map(SegGrid::new).collect();
|
||||
if coarse.is_none() && patches.is_empty() {
|
||||
return;
|
||||
}
|
||||
let m = frame.out_to_norm(w, h);
|
||||
for_rows(&mut out.data, w, |y, row| {
|
||||
for (x, v) in row.iter_mut().enumerate() {
|
||||
let n = m.apply(Point::new(x as f64 + 0.5, y as f64 + 0.5));
|
||||
let fine = patches.iter().filter(|p| p.contains(n)).map(|p| p.logit(n)).reduce(f32::max);
|
||||
let l = fine.or_else(|| coarse.as_ref().map(|c| c.logit(n))).unwrap_or(-16.0);
|
||||
*v = 1.0 / (1.0 + (-l * gain).exp());
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn frame_px(frame: &Frame, w: usize) -> f32 {
|
||||
frame.px_per_long(w) as f32
|
||||
}
|
||||
|
|
@ -356,6 +424,93 @@ mod tests {
|
|||
assert!((a.get(50, 25) - 0.5).abs() < 0.05);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stored_segmentations_render_at_any_size_and_follow_the_crop() {
|
||||
// the left half selected, on a 16² logit grid
|
||||
let side = 16;
|
||||
let logits: Vec<f32> = (0..side * side).map(|i| if i % side < side / 2 { 12.0 } else { -12.0 }).collect();
|
||||
let seg = SegMask::from_logits(side, &logits);
|
||||
for shape in [
|
||||
MaskShape::Object { hint: vec![Point::new(0.25, 0.5)], exclude: vec![], seg: Some(seg.clone()), detail: vec![], edge: 0.0 },
|
||||
MaskShape::Prompt { text: "left".into(), seg: Some(seg.clone()), detail: vec![], edge: 0.0 },
|
||||
] {
|
||||
for (w, h) in [(40, 20), (400, 200)] {
|
||||
let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0);
|
||||
assert!(a.get(1, h / 2) > 0.99 && a.get(w - 2, h / 2) < 0.01, "{w}×{h}");
|
||||
if w >= 400 {
|
||||
let soft = (0..w).filter(|x| (0.05..0.95).contains(&a.get(*x, h / 2))).count();
|
||||
assert!((2..30).contains(&soft), "an anti-aliased edge, not a ramp: {soft} px");
|
||||
}
|
||||
}
|
||||
}
|
||||
// cropped to the right half: nothing left selected
|
||||
let mut d = DevelopSettings::default();
|
||||
d.crop.geometry.rect = lightcraft_geom::Rect::from_center(Point::new(0.75, 0.5), 0.5, 1.0);
|
||||
let f = Frame::new(40, 20, &d, true);
|
||||
let shape = MaskShape::Prompt { text: "left".into(), seg: Some(seg), detail: vec![], edge: 0.0 };
|
||||
let a = shape_alpha(&shape, &f, 20, 20, &Rgb32f::new(20, 20), &Plane::new(20, 20), 0.0);
|
||||
assert!(a.data.iter().all(|v| *v < 0.05));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detail_patches_replace_the_coarse_mask_inside_their_rectangle() {
|
||||
// coarse: nothing; a patch over the right half selects its own left half (x 0.5..0.75)
|
||||
let coarse = SegMask::from_logits(4, &[-12.0; 16]);
|
||||
let side = 16;
|
||||
let l: Vec<f32> = (0..side * side).map(|i| if i % side < side / 2 { 12.0 } else { -12.0 }).collect();
|
||||
let patch = SegMask::from_logits_in(side, &l, [0.5, 0.0, 1.0, 1.0]);
|
||||
let shape = MaskShape::Object { hint: vec![Point::new(0.6, 0.5)], exclude: vec![], seg: Some(coarse), detail: vec![patch], edge: 0.0 };
|
||||
let (w, h) = (200, 100);
|
||||
let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0);
|
||||
assert!(a.get(20, 50) < 0.01, "outside the patch: the coarse mask");
|
||||
assert!(a.get(120, 50) > 0.99, "inside the patch, its selected half");
|
||||
assert!(a.get(190, 50) < 0.01, "inside the patch, its unselected half");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_hostile_number_of_detail_patches_is_capped() {
|
||||
// built in memory (a deserialized document is capped already): only the first
|
||||
// MAX_DETAIL are decoded and sampled, so this needs 4 × 4 MB, not 20 000 × 4 MB
|
||||
let side = 1024;
|
||||
let l: Vec<f32> = vec![12.0; side * side];
|
||||
let patch = SegMask::from_logits_in(side, &l, [0.0, 0.0, 0.5, 1.0]);
|
||||
let shape = MaskShape::Object { hint: vec![Point::new(0.2, 0.5)], exclude: vec![], seg: None, detail: vec![patch; 20_000], edge: 0.0 };
|
||||
let (w, h) = (64, 32);
|
||||
let t = std::time::Instant::now();
|
||||
let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0);
|
||||
assert!(a.get(5, 16) > 0.99 && a.get(60, 16) < 0.01);
|
||||
assert!(t.elapsed() < std::time::Duration::from_secs(20), "{:?}", t.elapsed());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_edge_setting_hardens_or_feathers_the_transition() {
|
||||
let side = 16;
|
||||
let l: Vec<f32> = (0..side * side).map(|i| ((i % side) as f32 - 7.5) * -0.6).collect(); // a gentle ramp
|
||||
let seg = SegMask::from_logits(side, &l);
|
||||
let (w, h) = (400, 40);
|
||||
let soft_px = |edge: f64| {
|
||||
let shape = MaskShape::Prompt { text: "x".into(), seg: Some(seg.clone()), detail: vec![], edge };
|
||||
let a = shape_alpha(&shape, &frame(w, h), w, h, &Rgb32f::new(w, h), &Plane::new(w, h), 0.0);
|
||||
assert!(a.get(2, h / 2) > 0.9 && a.get(w - 3, h / 2) < 0.1, "still the same selection at {edge}");
|
||||
(0..w).filter(|x| (0.1..0.9).contains(&a.get(*x, h / 2))).count()
|
||||
};
|
||||
let (hard, plain, soft) = (soft_px(-100.0), soft_px(0.0), soft_px(100.0));
|
||||
assert!(hard < plain && plain < soft, "transition widths {hard} < {plain} < {soft}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ai_shapes_without_a_segmentation_select_nothing() {
|
||||
let f = frame(30, 30);
|
||||
let (img, l) = (Rgb32f::new(30, 30), Plane::new(30, 30));
|
||||
for shape in [
|
||||
MaskShape::Object { hint: vec![], exclude: vec![], seg: None, detail: vec![], edge: 0.0 },
|
||||
MaskShape::Prompt { text: "sky".into(), seg: None, detail: vec![], edge: 0.0 },
|
||||
MaskShape::Prompt { text: "sky".into(), seg: Some(SegMask { side: 4, data: "damaged!".into(), rect: None }), detail: vec![], edge: 0.0 },
|
||||
] {
|
||||
assert!(shape_alpha(&shape, &f, 30, 30, &img, &l, 0.0).data.iter().all(|v| *v == 0.0));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn radial_inside_and_invert() {
|
||||
let f = frame(100, 100);
|
||||
|
|
|
|||
|
|
@ -0,0 +1,28 @@
|
|||
[package]
|
||||
name = "lightcraft-segment"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license = "Apache-2.0"
|
||||
description = "LightCraft segmentation: SAM 3 (Segment Anything with Concepts) inference in pure Rust (L3)"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
log = { workspace = true }
|
||||
|
||||
# the model runs natively only (the web build has no AI masks)
|
||||
[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
|
||||
candle-core = { workspace = true }
|
||||
candle-nn = { workspace = true }
|
||||
rustls = { version = "0.23", default-features = false, features = ["std", "tls12"] }
|
||||
rustls-rustcrypto = "0.0.2-alpha"
|
||||
webpki-roots = "1"
|
||||
sha2 = { version = "0.10", default-features = false, features = ["std"] }
|
||||
|
||||
[target.'cfg(target_os = "macos")'.dependencies]
|
||||
candle-core = { workspace = true, features = ["metal"] }
|
||||
candle-nn = { workspace = true, features = ["metal"] }
|
||||
|
|
@ -0,0 +1,25 @@
|
|||
//! Times SAM 3 on this machine (a synthetic gradient image; the first run includes compiling
|
||||
//! the GPU kernels): `cargo run -p lightcraft-segment --release --example sam3_bench -- <model dir>`.
|
||||
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let mut args = std::env::args().skip(1);
|
||||
let dir = std::path::PathBuf::from(args.next().ok_or("usage: sam3_bench <model dir>")?);
|
||||
let t = std::time::Instant::now();
|
||||
let mut model = lightcraft_segment::Sam3::load(&dir)?;
|
||||
eprintln!("load: {:?} on {:?}", t.elapsed(), model.device());
|
||||
let (w, h) = (1200, 800);
|
||||
let rgb: Vec<u8> = (0..w * h * 3).map(|i| ((i / 3 % w) * 255 / w) as u8).collect();
|
||||
for run in 0..3 {
|
||||
let t = std::time::Instant::now();
|
||||
let mut enc = model.encode(&rgb, w, h)?;
|
||||
lightcraft_segment::sync(&model)?;
|
||||
let e = t.elapsed();
|
||||
let t = std::time::Instant::now();
|
||||
let p = model.segment_clicks(&mut enc, &[lightcraft_segment::Click { x: 0.5, y: 0.5, positive: true }])?;
|
||||
let c = t.elapsed();
|
||||
let t = std::time::Instant::now();
|
||||
let d = model.segment_text(&mut enc, "sky", 0.5)?;
|
||||
eprintln!("run {run}: encode {e:?}, click {c:?} (score {:.2}), text {:?} ({})", p.score, t.elapsed(), d.is_some());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -0,0 +1,81 @@
|
|||
//! The CLIP text encoder (24 layers, 1024 wide, causal) that turns a prompt into the token
|
||||
//! features SAM 3's detector attends to.
|
||||
//!
|
||||
//! 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`, CLIP text model), Copyright The HuggingFace
|
||||
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
|
||||
|
||||
use candle_core::{Device, Module, Tensor};
|
||||
use candle_nn::{Embedding, LayerNorm, VarBuilder};
|
||||
|
||||
use crate::Result;
|
||||
use crate::nn::{Attention, Mlp, ln};
|
||||
use crate::tokenizer::CONTEXT;
|
||||
|
||||
const HIDDEN: usize = 1024;
|
||||
const LAYERS: usize = 24;
|
||||
const HEADS: usize = 16;
|
||||
const EPS: f64 = 1e-5;
|
||||
|
||||
struct Layer {
|
||||
ln1: LayerNorm,
|
||||
attn: Attention,
|
||||
ln2: LayerNorm,
|
||||
mlp: Mlp,
|
||||
}
|
||||
|
||||
pub struct TextEncoder {
|
||||
tokens: Embedding,
|
||||
positions: Tensor,
|
||||
layers: Vec<Layer>,
|
||||
final_ln: LayerNorm,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl TextEncoder {
|
||||
/// `vb` at `detector_model.text_encoder.text_model`.
|
||||
pub fn new(vb: VarBuilder) -> Result<Self> {
|
||||
let e = vb.pp("embeddings");
|
||||
let mut layers = Vec::with_capacity(LAYERS);
|
||||
for i in 0..LAYERS {
|
||||
let l = vb.pp(format!("encoder.layers.{i}"));
|
||||
layers.push(Layer {
|
||||
ln1: ln(HIDDEN, EPS, l.pp("layer_norm1"))?,
|
||||
attn: Attention::named(HIDDEN, HIDDEN, HEADS, "out_proj", l.pp("self_attn"))?,
|
||||
ln2: ln(HIDDEN, EPS, l.pp("layer_norm2"))?,
|
||||
mlp: Mlp::new(HIDDEN, 4 * HIDDEN, true, l.pp("mlp"))?,
|
||||
});
|
||||
}
|
||||
Ok(Self {
|
||||
tokens: candle_nn::embedding(49408, HIDDEN, e.pp("token_embedding"))?,
|
||||
positions: e.get((CONTEXT, HIDDEN), "position_embedding.weight")?.unsqueeze(0)?,
|
||||
layers,
|
||||
final_ln: ln(HIDDEN, EPS, vb.pp("final_layer_norm"))?,
|
||||
device: vb.device().clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// `[1, 32, 1024]` features for token `ids` (`valid` marks real tokens).
|
||||
pub fn forward(&self, ids: &[u32], valid: &[bool]) -> Result<Tensor> {
|
||||
let ids = Tensor::from_slice(ids, (1, ids.len()), &self.device)?;
|
||||
let mut x = self.tokens.forward(&ids)?.broadcast_add(&self.positions)?;
|
||||
// causal, and padding keys masked
|
||||
let n = valid.len();
|
||||
let mut bias = vec![0f32; n * n];
|
||||
for (q, row) in bias.chunks_mut(n).enumerate() {
|
||||
for (k, v) in row.iter_mut().enumerate() {
|
||||
if k > q || !valid.get(k).copied().unwrap_or(false) {
|
||||
*v = f32::MIN;
|
||||
}
|
||||
}
|
||||
}
|
||||
let bias = Tensor::from_vec(bias, (1, 1, n, n), &self.device)?;
|
||||
for l in &self.layers {
|
||||
let h = l.ln1.forward(&x)?;
|
||||
x = (&x + l.attn.forward(&h, &h, &h, Some(&bias))?)?;
|
||||
let h = l.ln2.forward(&x)?;
|
||||
x = (&x + l.mlp.forward(&h)?)?;
|
||||
}
|
||||
Ok(self.final_ln.forward(&x)?)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,308 @@
|
|||
//! 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 })
|
||||
}
|
||||
|
|
@ -0,0 +1,541 @@
|
|||
//! A small HTTP/1.1 GET client for model downloads, in pure Rust: `std::net` sockets and, for
|
||||
//! `https://`, rustls with the RustCrypto provider (`rustls-rustcrypto`; no `ring` or
|
||||
//! `aws-lc-rs`, so nothing is compiled from C or assembly) and the Mozilla root certificates
|
||||
//! (`webpki-roots`).
|
||||
//!
|
||||
//! Only what a large-file download needs: one request per connection (`Connection: close`),
|
||||
//! `Range` requests for resuming, redirects (followed by the caller), `Content-Length`,
|
||||
//! chunked and read-to-close bodies. Every read waits at most `stall` for data and checks the
|
||||
//! cancel flag at least twice a second, so a dead server or a cancelled download never hangs
|
||||
//! the downloading thread. No proxies (see docs/ai-masks.md).
|
||||
|
||||
use std::io::{ErrorKind, Read, Write};
|
||||
use std::net::{TcpStream, ToSocketAddrs};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
/// Why a request failed.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum HttpError {
|
||||
/// The URL can't be used (scheme, host, characters).
|
||||
BadUrl(String),
|
||||
/// Name resolution or connecting failed (or took longer than the connect timeout).
|
||||
Connect(String),
|
||||
/// No data arrived for the stall timeout.
|
||||
Stalled,
|
||||
/// The cancel flag was raised.
|
||||
Cancelled,
|
||||
/// TLS failed (certificate, handshake).
|
||||
Tls(String),
|
||||
/// The server's response could not be understood.
|
||||
Protocol(String),
|
||||
/// Any other I/O error.
|
||||
Io(String),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for HttpError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
HttpError::BadUrl(e) => write!(f, "unusable URL: {e}"),
|
||||
HttpError::Connect(e) => write!(f, "could not connect: {e}"),
|
||||
HttpError::Stalled => write!(f, "the server stopped sending data"),
|
||||
HttpError::Cancelled => write!(f, "cancelled"),
|
||||
HttpError::Tls(e) => write!(f, "secure connection failed: {e}"),
|
||||
HttpError::Protocol(e) => write!(f, "unexpected response: {e}"),
|
||||
HttpError::Io(e) => write!(f, "{e}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A parsed `http://` or `https://` URL.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct Url {
|
||||
pub tls: bool,
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
/// Path and query, starting with `/`.
|
||||
pub path: String,
|
||||
}
|
||||
|
||||
impl Url {
|
||||
pub fn parse(s: &str) -> Result<Url, HttpError> {
|
||||
let bad = |why: &str| HttpError::BadUrl(format!("{s}: {why}"));
|
||||
let s = s.trim();
|
||||
if s.chars().any(|c| c.is_control() || c == ' ') {
|
||||
return Err(bad("contains spaces or control characters"));
|
||||
}
|
||||
let (tls, rest) = if let Some(r) = s.strip_prefix("https://") {
|
||||
(true, r)
|
||||
} else if let Some(r) = s.strip_prefix("http://") {
|
||||
(false, r)
|
||||
} else {
|
||||
return Err(bad("only http:// and https:// are supported"));
|
||||
};
|
||||
let rest = rest.split('#').next().unwrap_or_default();
|
||||
let (authority, path) = match rest.find(['/', '?']) {
|
||||
Some(i) => (rest.get(..i).unwrap_or_default(), rest.get(i..).unwrap_or_default()),
|
||||
None => (rest, ""),
|
||||
};
|
||||
let path = if path.starts_with('/') { path.to_string() } else { format!("/{path}") };
|
||||
if authority.contains('@') {
|
||||
return Err(bad("user names and passwords in URLs are not supported"));
|
||||
}
|
||||
let (host, port) = if let Some(v6) = authority.strip_prefix('[') {
|
||||
let end = v6.find(']').ok_or_else(|| bad("unterminated IPv6 address"))?;
|
||||
let host = v6.get(..end).unwrap_or_default();
|
||||
let after = v6.get(end + 1..).unwrap_or_default();
|
||||
(host.to_string(), after.strip_prefix(':'))
|
||||
} else {
|
||||
match authority.rsplit_once(':') {
|
||||
Some((h, p)) => (h.to_string(), Some(p)),
|
||||
None => (authority.to_string(), None),
|
||||
}
|
||||
};
|
||||
let port = match port {
|
||||
Some(p) => p.parse::<u16>().ok().filter(|p| *p != 0).ok_or_else(|| bad("bad port"))?,
|
||||
None if tls => 443,
|
||||
None => 80,
|
||||
};
|
||||
if host.is_empty() || !host.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '.' | ':' | '_')) {
|
||||
return Err(bad("bad host name"));
|
||||
}
|
||||
Ok(Url { tls, host: host.to_ascii_lowercase(), port, path })
|
||||
}
|
||||
|
||||
/// `Location` of a redirect, resolved against this URL.
|
||||
pub fn join(&self, location: &str) -> Result<Url, HttpError> {
|
||||
let l = location.trim();
|
||||
if l.starts_with("http://") || l.starts_with("https://") {
|
||||
return Url::parse(l);
|
||||
}
|
||||
// another scheme (`ftp:`, `file:`, `data:`…)
|
||||
if let Some((scheme, _)) = l.split_once(':')
|
||||
&& !scheme.is_empty()
|
||||
&& scheme.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '-' | '.'))
|
||||
{
|
||||
return Err(HttpError::BadUrl(format!("redirect to an unsupported location `{scheme}:`")));
|
||||
}
|
||||
let scheme = if self.tls { "https" } else { "http" };
|
||||
let host = if self.host.contains(':') { format!("[{}]", self.host) } else { self.host.clone() };
|
||||
let base = format!("{scheme}://{host}:{}", self.port);
|
||||
if let Some(rest) = l.strip_prefix("//") {
|
||||
return Url::parse(&format!("{scheme}://{rest}"));
|
||||
}
|
||||
if l.starts_with('/') {
|
||||
return Url::parse(&format!("{base}{l}"));
|
||||
}
|
||||
// relative to the current directory
|
||||
let dir = self.path.split('?').next().unwrap_or("/");
|
||||
let dir = dir.rsplit_once('/').map_or("", |(d, _)| d);
|
||||
Url::parse(&format!("{base}{dir}/{l}"))
|
||||
}
|
||||
|
||||
fn host_header(&self) -> String {
|
||||
let host = if self.host.contains(':') { format!("[{}]", self.host) } else { self.host.clone() };
|
||||
let default = if self.tls { 443 } else { 80 };
|
||||
if self.port == default { host } else { format!("{host}:{}", self.port) }
|
||||
}
|
||||
|
||||
/// The URL without its query (for messages: a query may hold a signed token).
|
||||
pub fn display(&self) -> String {
|
||||
let scheme = if self.tls { "https" } else { "http" };
|
||||
format!("{scheme}://{}{}", self.host_header(), self.path.split('?').next().unwrap_or("/"))
|
||||
}
|
||||
}
|
||||
|
||||
/// Timeouts and the cancel flag for one request.
|
||||
pub struct Limits<'a> {
|
||||
pub connect: Duration,
|
||||
pub stall: Duration,
|
||||
pub cancel: &'a AtomicBool,
|
||||
}
|
||||
|
||||
/// How often a blocked read wakes up to check the cancel flag.
|
||||
const TICK: Duration = Duration::from_millis(250);
|
||||
/// Longest status line + headers accepted.
|
||||
const MAX_HEAD: usize = 64 * 1024;
|
||||
|
||||
enum Stream {
|
||||
Plain(TcpStream),
|
||||
Tls(Box<rustls::ClientConnection>, TcpStream),
|
||||
}
|
||||
|
||||
fn is_timeout(e: &std::io::Error) -> bool {
|
||||
matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut | ErrorKind::Interrupted)
|
||||
}
|
||||
|
||||
fn tls_config() -> Result<Arc<rustls::ClientConfig>, HttpError> {
|
||||
let roots = rustls::RootCertStore { roots: webpki_roots::TLS_SERVER_ROOTS.to_vec() };
|
||||
let config = rustls::ClientConfig::builder_with_provider(Arc::new(rustls_rustcrypto::provider()))
|
||||
.with_safe_default_protocol_versions()
|
||||
.map_err(|e| HttpError::Tls(e.to_string()))?
|
||||
.with_root_certificates(roots)
|
||||
.with_no_client_auth();
|
||||
Ok(Arc::new(config))
|
||||
}
|
||||
|
||||
impl Stream {
|
||||
fn connect(url: &Url, limits: &Limits) -> Result<Stream, HttpError> {
|
||||
let addrs: Vec<_> = (url.host.as_str(), url.port).to_socket_addrs().map_err(|e| HttpError::Connect(format!("{}: {e}", url.host)))?.collect();
|
||||
if addrs.is_empty() {
|
||||
return Err(HttpError::Connect(format!("{}: no address", url.host)));
|
||||
}
|
||||
let mut last = String::new();
|
||||
let mut sock = None;
|
||||
for a in addrs {
|
||||
if limits.cancel.load(Ordering::Relaxed) {
|
||||
return Err(HttpError::Cancelled);
|
||||
}
|
||||
match TcpStream::connect_timeout(&a, limits.connect) {
|
||||
Ok(s) => {
|
||||
sock = Some(s);
|
||||
break;
|
||||
}
|
||||
Err(e) => last = format!("{}: {e}", url.host),
|
||||
}
|
||||
}
|
||||
let sock = sock.ok_or(HttpError::Connect(last))?;
|
||||
sock.set_read_timeout(Some(TICK)).map_err(|e| HttpError::Io(e.to_string()))?;
|
||||
sock.set_write_timeout(Some(TICK)).map_err(|e| HttpError::Io(e.to_string()))?;
|
||||
let _ = sock.set_nodelay(true);
|
||||
if !url.tls {
|
||||
return Ok(Stream::Plain(sock));
|
||||
}
|
||||
let name = rustls::pki_types::ServerName::try_from(url.host.clone()).map_err(|e| HttpError::Tls(format!("{}: {e}", url.host)))?;
|
||||
let conn = rustls::ClientConnection::new(tls_config()?, name).map_err(|e| HttpError::Tls(e.to_string()))?;
|
||||
let mut s = Stream::Tls(Box::new(conn), sock);
|
||||
s.handshake(limits)?;
|
||||
Ok(s)
|
||||
}
|
||||
|
||||
/// Wait for `f` to make progress, retrying timeouts until `stall` passes without any.
|
||||
fn retry<T>(limits: &Limits, mut f: impl FnMut() -> std::io::Result<T>) -> Result<T, HttpError> {
|
||||
let start = Instant::now();
|
||||
loop {
|
||||
if limits.cancel.load(Ordering::Relaxed) {
|
||||
return Err(HttpError::Cancelled);
|
||||
}
|
||||
match f() {
|
||||
Ok(v) => return Ok(v),
|
||||
Err(e) if is_timeout(&e) => {
|
||||
if start.elapsed() >= limits.stall {
|
||||
return Err(HttpError::Stalled);
|
||||
}
|
||||
}
|
||||
Err(e) if e.kind() == ErrorKind::InvalidData => return Err(HttpError::Tls(e.to_string())),
|
||||
Err(e) => return Err(HttpError::Io(e.to_string())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn handshake(&mut self, limits: &Limits) -> Result<(), HttpError> {
|
||||
let Stream::Tls(conn, sock) = self else { return Ok(()) };
|
||||
while conn.is_handshaking() {
|
||||
Self::retry(limits, || conn.complete_io(sock))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_all(&mut self, data: &[u8], limits: &Limits) -> Result<(), HttpError> {
|
||||
match self {
|
||||
Stream::Plain(sock) => {
|
||||
let mut rest = data;
|
||||
while !rest.is_empty() {
|
||||
let n = Self::retry(limits, || sock.write(rest))?;
|
||||
if n == 0 {
|
||||
return Err(HttpError::Io("connection closed while sending the request".into()));
|
||||
}
|
||||
rest = rest.get(n..).unwrap_or_default();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Stream::Tls(conn, sock) => {
|
||||
// buffered by rustls, then sent as TLS records
|
||||
conn.writer().write_all(data).map_err(|e| HttpError::Tls(e.to_string()))?;
|
||||
while conn.wants_write() {
|
||||
Self::retry(limits, || conn.write_tls(sock))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Read some bytes (0 at the end of the stream), waiting at most `stall` for them.
|
||||
fn read(&mut self, buf: &mut [u8], limits: &Limits) -> Result<usize, HttpError> {
|
||||
match self {
|
||||
Stream::Plain(sock) => Self::retry(limits, || sock.read(buf)),
|
||||
Stream::Tls(conn, sock) => {
|
||||
let start = Instant::now();
|
||||
loop {
|
||||
if limits.cancel.load(Ordering::Relaxed) {
|
||||
return Err(HttpError::Cancelled);
|
||||
}
|
||||
match conn.reader().read(buf) {
|
||||
Ok(n) => return Ok(n),
|
||||
// closed without close_notify: the caller checks the length it got
|
||||
Err(e) if e.kind() == ErrorKind::UnexpectedEof => return Ok(0),
|
||||
Err(e) if e.kind() == ErrorKind::WouldBlock => {}
|
||||
Err(e) => return Err(HttpError::Tls(e.to_string())),
|
||||
}
|
||||
// no plaintext buffered: more TLS records from the socket
|
||||
match conn.read_tls(sock) {
|
||||
Ok(0) => return Ok(0),
|
||||
Ok(_) => {
|
||||
conn.process_new_packets().map_err(|e| HttpError::Tls(e.to_string()))?;
|
||||
while conn.wants_write() {
|
||||
Self::retry(limits, || conn.write_tls(sock))?;
|
||||
}
|
||||
}
|
||||
Err(e) if is_timeout(&e) => {
|
||||
if start.elapsed() >= limits.stall {
|
||||
return Err(HttpError::Stalled);
|
||||
}
|
||||
}
|
||||
Err(e) => return Err(HttpError::Io(e.to_string())),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The connection with a read buffer (headers and chunk sizes are read line by line).
|
||||
struct Conn {
|
||||
stream: Stream,
|
||||
buf: Vec<u8>,
|
||||
pos: usize,
|
||||
}
|
||||
|
||||
impl Conn {
|
||||
fn fill(&mut self, limits: &Limits) -> Result<usize, HttpError> {
|
||||
if self.pos >= self.buf.len() {
|
||||
self.buf.clear();
|
||||
self.pos = 0;
|
||||
}
|
||||
let mut tmp = [0u8; 64 * 1024];
|
||||
let n = self.stream.read(&mut tmp, limits)?;
|
||||
self.buf.extend_from_slice(tmp.get(..n).unwrap_or_default());
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
/// One line without its CRLF; at most `max` bytes.
|
||||
fn line(&mut self, max: usize, limits: &Limits) -> Result<String, HttpError> {
|
||||
loop {
|
||||
let avail = self.buf.get(self.pos..).unwrap_or_default();
|
||||
if let Some(i) = avail.iter().position(|b| *b == b'\n') {
|
||||
let line = avail.get(..i).unwrap_or_default();
|
||||
let line = line.strip_suffix(b"\r").unwrap_or(line);
|
||||
let s = String::from_utf8_lossy(line).into_owned();
|
||||
self.pos = self.pos.saturating_add(i + 1);
|
||||
return Ok(s);
|
||||
}
|
||||
if avail.len() > max {
|
||||
return Err(HttpError::Protocol("header line too long".into()));
|
||||
}
|
||||
// keep the partial line, read more after it
|
||||
if self.pos > 0 {
|
||||
self.buf.drain(..self.pos.min(self.buf.len()));
|
||||
self.pos = 0;
|
||||
}
|
||||
let mut tmp = [0u8; 16 * 1024];
|
||||
let n = self.stream.read(&mut tmp, limits)?;
|
||||
if n == 0 {
|
||||
return Err(HttpError::Protocol("connection closed in the middle of the response".into()));
|
||||
}
|
||||
self.buf.extend_from_slice(tmp.get(..n).unwrap_or_default());
|
||||
}
|
||||
}
|
||||
|
||||
/// Up to `out.len()` body bytes (0 at the end of the stream).
|
||||
fn read(&mut self, out: &mut [u8], limits: &Limits) -> Result<usize, HttpError> {
|
||||
if self.pos >= self.buf.len() && self.fill(limits)? == 0 {
|
||||
return Ok(0);
|
||||
}
|
||||
let avail = self.buf.get(self.pos..).unwrap_or_default();
|
||||
let n = avail.len().min(out.len());
|
||||
if let (Some(dst), Some(src)) = (out.get_mut(..n), avail.get(..n)) {
|
||||
dst.copy_from_slice(src);
|
||||
}
|
||||
self.pos = self.pos.saturating_add(n);
|
||||
Ok(n)
|
||||
}
|
||||
}
|
||||
|
||||
/// A response's status, headers and body.
|
||||
pub struct Response {
|
||||
pub status: u16,
|
||||
headers: Vec<(String, String)>,
|
||||
conn: Conn,
|
||||
body: BodyMode,
|
||||
}
|
||||
|
||||
enum BodyMode {
|
||||
Length(u64),
|
||||
Chunked { left: u64, done: bool },
|
||||
Close,
|
||||
Done,
|
||||
}
|
||||
|
||||
impl Response {
|
||||
/// The first header named `name` (case-insensitive).
|
||||
pub fn header(&self, name: &str) -> Option<&str> {
|
||||
self.headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str())
|
||||
}
|
||||
|
||||
/// The body's length when the server announced it.
|
||||
pub fn content_length(&self) -> Option<u64> {
|
||||
match self.body {
|
||||
BodyMode::Length(n) => Some(n),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Read body bytes into `out` (0 at the end of the body).
|
||||
pub fn read(&mut self, out: &mut [u8], limits: &Limits) -> Result<usize, HttpError> {
|
||||
match &mut self.body {
|
||||
BodyMode::Done => Ok(0),
|
||||
BodyMode::Length(left) => {
|
||||
if *left == 0 {
|
||||
self.body = BodyMode::Done;
|
||||
return Ok(0);
|
||||
}
|
||||
let want = usize::try_from(*left).unwrap_or(usize::MAX).min(out.len());
|
||||
let n = self.conn.read(out.get_mut(..want).unwrap_or_default(), limits)?;
|
||||
if n == 0 {
|
||||
return Err(HttpError::Protocol("the connection closed before the end of the file".into()));
|
||||
}
|
||||
*left = left.saturating_sub(n as u64);
|
||||
Ok(n)
|
||||
}
|
||||
BodyMode::Close => {
|
||||
let n = self.conn.read(out, limits)?;
|
||||
if n == 0 {
|
||||
self.body = BodyMode::Done;
|
||||
}
|
||||
Ok(n)
|
||||
}
|
||||
BodyMode::Chunked { left, done } => {
|
||||
if *done {
|
||||
return Ok(0);
|
||||
}
|
||||
if *left == 0 {
|
||||
let line = self.conn.line(1024, limits)?;
|
||||
let hex = line.split(';').next().unwrap_or_default().trim();
|
||||
let size = u64::from_str_radix(hex, 16).map_err(|_| HttpError::Protocol(format!("bad chunk size `{hex}`")))?;
|
||||
if size == 0 {
|
||||
// trailers, up to the empty line
|
||||
let mut n = 0;
|
||||
while !self.conn.line(MAX_HEAD, limits)?.is_empty() {
|
||||
n += 1;
|
||||
if n > 100 {
|
||||
return Err(HttpError::Protocol("too many trailers".into()));
|
||||
}
|
||||
}
|
||||
*done = true;
|
||||
return Ok(0);
|
||||
}
|
||||
*left = size;
|
||||
}
|
||||
let want = usize::try_from(*left).unwrap_or(usize::MAX).min(out.len());
|
||||
let n = self.conn.read(out.get_mut(..want).unwrap_or_default(), limits)?;
|
||||
if n == 0 {
|
||||
return Err(HttpError::Protocol("the connection closed in the middle of a chunk".into()));
|
||||
}
|
||||
*left = left.saturating_sub(n as u64);
|
||||
if *left == 0 && !self.conn.line(16, limits)?.is_empty() {
|
||||
return Err(HttpError::Protocol("missing CRLF after a chunk".into()));
|
||||
}
|
||||
Ok(n)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Send `GET url` with extra `headers` and read the status line and headers.
|
||||
pub fn get(url: &Url, headers: &[(&str, String)], limits: &Limits) -> Result<Response, HttpError> {
|
||||
let mut req = format!(
|
||||
"GET {} HTTP/1.1\r\nHost: {}\r\nUser-Agent: LightCraft/{}\r\nAccept: */*\r\nAccept-Encoding: identity\r\nConnection: close\r\n",
|
||||
url.path,
|
||||
url.host_header(),
|
||||
env!("CARGO_PKG_VERSION")
|
||||
);
|
||||
for (k, v) in headers {
|
||||
if v.chars().any(|c| c == '\r' || c == '\n') || k.chars().any(|c| c == '\r' || c == '\n' || c == ':') {
|
||||
return Err(HttpError::BadUrl("header with a line break".into()));
|
||||
}
|
||||
req.push_str(&format!("{k}: {v}\r\n"));
|
||||
}
|
||||
req.push_str("\r\n");
|
||||
let mut stream = Stream::connect(url, limits)?;
|
||||
stream.write_all(req.as_bytes(), limits)?;
|
||||
let mut conn = Conn { stream, buf: Vec::new(), pos: 0 };
|
||||
let status_line = conn.line(MAX_HEAD, limits)?;
|
||||
let mut parts = status_line.split_whitespace();
|
||||
let version = parts.next().unwrap_or_default();
|
||||
if !version.starts_with("HTTP/1.") {
|
||||
return Err(HttpError::Protocol(format!("not an HTTP/1.x response: {}", status_line.chars().take(60).collect::<String>())));
|
||||
}
|
||||
let status = parts.next().and_then(|s| s.parse::<u16>().ok()).ok_or_else(|| HttpError::Protocol("bad status line".into()))?;
|
||||
let mut headers = Vec::new();
|
||||
let mut total = status_line.len();
|
||||
loop {
|
||||
let line = conn.line(MAX_HEAD, limits)?;
|
||||
if line.is_empty() {
|
||||
break;
|
||||
}
|
||||
total = total.saturating_add(line.len());
|
||||
if total > MAX_HEAD || headers.len() > 200 {
|
||||
return Err(HttpError::Protocol("response headers too large".into()));
|
||||
}
|
||||
if let Some((k, v)) = line.split_once(':') {
|
||||
headers.push((k.trim().to_string(), v.trim().to_string()));
|
||||
}
|
||||
}
|
||||
let find = |name: &str| headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str());
|
||||
let body = if status == 204 || status == 304 || (100..200).contains(&status) {
|
||||
BodyMode::Done
|
||||
} else if find("transfer-encoding").is_some_and(|v| v.to_ascii_lowercase().contains("chunked")) {
|
||||
BodyMode::Chunked { left: 0, done: false }
|
||||
} else if let Some(n) = find("content-length") {
|
||||
BodyMode::Length(n.trim().parse::<u64>().map_err(|_| HttpError::Protocol(format!("bad Content-Length `{n}`")))?)
|
||||
} else {
|
||||
BodyMode::Close
|
||||
};
|
||||
Ok(Response { status, headers, conn, body })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parses_urls() {
|
||||
let u = Url::parse("https://cdn.example.com/models/sam3/model.safetensors?sig=1#x").unwrap();
|
||||
assert_eq!((u.tls, u.host.as_str(), u.port, u.path.as_str()), (true, "cdn.example.com", 443, "/models/sam3/model.safetensors?sig=1"));
|
||||
assert_eq!(u.display(), "https://cdn.example.com/models/sam3/model.safetensors");
|
||||
let u = Url::parse("http://127.0.0.1:8080").unwrap();
|
||||
assert_eq!((u.tls, u.port, u.path.as_str()), (false, 8080, "/"));
|
||||
let u = Url::parse("http://[::1]:9/a").unwrap();
|
||||
assert_eq!((u.host.as_str(), u.port), ("::1", 9));
|
||||
for bad in ["ftp://x/y", "https://", "https://a b/", "https://u:p@h/", "http://h:0/", "http://h:99999/", "https://h/\r\nX: y", "file:///etc"]
|
||||
{
|
||||
assert!(Url::parse(bad).is_err(), "{bad}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_redirects() {
|
||||
let u = Url::parse("https://a.example/x/y/file?q").unwrap();
|
||||
assert_eq!(u.join("https://b.example/z").unwrap().host, "b.example");
|
||||
assert_eq!(u.join("/root").unwrap().path, "/root");
|
||||
assert_eq!(u.join("other").unwrap().path, "/x/y/other");
|
||||
assert_eq!(u.join("//c.example/p").unwrap().host, "c.example");
|
||||
assert!(u.join("ftp://nope").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_tls_config_builds_with_the_pure_rust_provider() {
|
||||
assert!(tls_config().is_ok());
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,424 @@
|
|||
//! Downloading the SAM 3 model, only when the user asked for it: from an ordered list of mirrors
|
||||
//! (the next one when a mirror fails), resuming partial files, with connect and stall timeouts,
|
||||
//! a size cap, progress and cancellation, and checked (exact size, SHA-256 where pinned) before
|
||||
//! each file is moved into the model folder. Runs on whatever thread calls [`download`] (the
|
||||
//! engine uses a background thread; the UI thread never waits for the network).
|
||||
//!
|
||||
//! Files are written to `<name>.part` next to their final place and renamed when verified, so a
|
||||
//! model folder never holds a half-written file under a real name, and an interrupted download
|
||||
//! resumes from where it stopped. A file that fails its check is deleted, never kept.
|
||||
//!
|
||||
//! The weights are never part of LightCraft (SAM License, see docs/ai-masks.md); where they are
|
||||
//! downloaded from is configured by [`DEFAULT_MIRRORS`] and the user's own list.
|
||||
|
||||
pub mod http;
|
||||
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::io::{Read, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use http::{HttpError, Limits, Url};
|
||||
|
||||
/// One file of the model.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct FileSpec {
|
||||
pub name: &'static str,
|
||||
/// Exact size in bytes, when pinned (then the download resumes and is checked against it).
|
||||
pub size: Option<u64>,
|
||||
/// SHA-256 (lower-case hex), when pinned.
|
||||
pub sha256: Option<&'static str>,
|
||||
/// Largest size accepted when `size` is not pinned.
|
||||
pub max: u64,
|
||||
}
|
||||
|
||||
/// The `facebook/sam3` checkpoint files [`crate::Sam3::load`] needs.
|
||||
///
|
||||
/// `model.safetensors` is pinned to the official checkpoint (Hugging Face LFS SHA-256). The two
|
||||
/// tokenizer files are small text files without a pinned hash yet: they are only parsed by the
|
||||
/// tokenizer (never executed) and capped in size.
|
||||
// TODO(maintainer): pin `vocab.json` and `merges.txt` (size + SHA-256) when the files are
|
||||
// uploaded to LightCraft's CDN.
|
||||
pub const SAM3_FILES: &[FileSpec] = &[
|
||||
FileSpec { name: "vocab.json", size: None, sha256: None, max: 16 << 20 },
|
||||
FileSpec { name: "merges.txt", size: None, sha256: None, max: 16 << 20 },
|
||||
FileSpec {
|
||||
name: crate::WEIGHTS_FILE,
|
||||
size: Some(SAM3_WEIGHTS_SIZE),
|
||||
sha256: Some("6d06f0a5f84e435071fe6603e61d0b4cc7b40e0d39d487cfd4d67d8cc11cc14a"),
|
||||
max: SAM3_WEIGHTS_SIZE,
|
||||
},
|
||||
];
|
||||
|
||||
/// Size of the official `model.safetensors` (3.44 GB).
|
||||
pub const SAM3_WEIGHTS_SIZE: u64 = 3_439_938_512;
|
||||
|
||||
/// Where LightCraft downloads the model from, in order: each is a base URL, and file `f` is at
|
||||
/// `<base>/<f>`. **Empty for now**: LightCraft has no download location of its own yet, so the
|
||||
/// in-app download needs the user's list (`LIGHTCRAFT_SAM3_MIRRORS` or the mirrors file, see
|
||||
/// [`mirrors`] and docs/ai-masks.md) until this is filled in.
|
||||
///
|
||||
/// Hugging Face's `facebook/sam3` can't be a default: it is gated (each user must accept the
|
||||
/// SAM License there and download with their own token).
|
||||
// TODO(maintainer): add LightCraft's CDN locations (https, primary first), e.g.
|
||||
// "https://<cdn-host>/models/sam3/<revision>" — the files there must match `SAM3_FILES`.
|
||||
pub const DEFAULT_MIRRORS: &[&str] = &[];
|
||||
|
||||
/// Environment variable with extra mirrors (base URLs separated by commas, spaces or newlines),
|
||||
/// tried before the defaults.
|
||||
pub const MIRRORS_ENV: &str = "LIGHTCRAFT_SAM3_MIRRORS";
|
||||
|
||||
/// Most mirrors used (a hostile list can't make a download loop forever).
|
||||
const MAX_MIRRORS: usize = 16;
|
||||
/// Redirects followed per request.
|
||||
const MAX_REDIRECTS: usize = 8;
|
||||
|
||||
/// The mirrors to try, in order: the environment variable's, then the mirrors file's (one base
|
||||
/// URL per line, `#` comments), then [`DEFAULT_MIRRORS`]; duplicates and unusable URLs dropped.
|
||||
pub fn mirrors(env: Option<&str>, file: Option<&Path>) -> Vec<String> {
|
||||
let mut out: Vec<String> = Vec::new();
|
||||
let mut add = |s: &str| {
|
||||
let s = s.trim().trim_end_matches('/');
|
||||
if !s.is_empty() && Url::parse(s).is_ok() && !out.iter().any(|m| m == s) && out.len() < MAX_MIRRORS {
|
||||
out.push(s.to_string());
|
||||
}
|
||||
};
|
||||
if let Some(env) = env {
|
||||
env.split([',', ' ', '\n', '\t']).for_each(&mut add);
|
||||
}
|
||||
// a small text file; anything larger isn't a mirror list
|
||||
if let Some(text) = file.and_then(|f| std::fs::metadata(f).ok().filter(|m| m.len() < 64 * 1024).and_then(|_| std::fs::read_to_string(f).ok())) {
|
||||
text.lines().map(|l| l.split('#').next().unwrap_or_default()).for_each(&mut add);
|
||||
}
|
||||
DEFAULT_MIRRORS.iter().for_each(|m| add(m));
|
||||
out
|
||||
}
|
||||
|
||||
/// Download settings.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Options {
|
||||
/// Longest wait to connect to a mirror.
|
||||
pub connect_timeout: Duration,
|
||||
/// Longest wait for the next bytes once connected.
|
||||
pub stall_timeout: Duration,
|
||||
/// Tries per mirror and file (each resumes where the last one stopped).
|
||||
pub attempts: u32,
|
||||
}
|
||||
|
||||
impl Default for Options {
|
||||
fn default() -> Self {
|
||||
Options { connect_timeout: Duration::from_secs(15), stall_timeout: Duration::from_secs(30), attempts: 3 }
|
||||
}
|
||||
}
|
||||
|
||||
/// Where a download is.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct Progress {
|
||||
pub file: String,
|
||||
/// Bytes of all files on disk so far, and the total expected (pinned sizes + announced ones).
|
||||
pub done: u64,
|
||||
pub total: u64,
|
||||
/// The mirror in use (without query).
|
||||
pub mirror: String,
|
||||
}
|
||||
|
||||
/// Why a download failed.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum DownloadError {
|
||||
/// No mirror is configured.
|
||||
NoMirrors,
|
||||
Cancelled,
|
||||
/// A local file problem (folder not writable, disk full…).
|
||||
Disk(String),
|
||||
/// Every mirror failed for `file`; the reason for each.
|
||||
AllMirrorsFailed {
|
||||
file: String,
|
||||
errors: Vec<String>,
|
||||
},
|
||||
}
|
||||
|
||||
impl std::fmt::Display for DownloadError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
DownloadError::NoMirrors => write!(
|
||||
f,
|
||||
"no download location is configured for the SAM 3 model in this build (set {MIRRORS_ENV}, or put the files in the model folder yourself; see docs/ai-masks.md)"
|
||||
),
|
||||
DownloadError::Cancelled => write!(f, "the download was cancelled"),
|
||||
DownloadError::Disk(e) => write!(f, "{e}"),
|
||||
DownloadError::AllMirrorsFailed { file, errors } => write!(f, "could not download {file}: {}", errors.join("; ")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One mirror's failure for one file; `Retry` resumes on the same mirror.
|
||||
enum Fail {
|
||||
Retry(String),
|
||||
NextMirror(String),
|
||||
Cancelled,
|
||||
Disk(String),
|
||||
}
|
||||
|
||||
/// Download `files` into `dir` from `mirrors` (in order; see [`mirrors`]). Files already there
|
||||
/// and verified are kept. `progress` is called as bytes arrive (often: throttle in it).
|
||||
pub fn download(
|
||||
files: &[FileSpec],
|
||||
mirrors: &[String],
|
||||
dir: &Path,
|
||||
opts: &Options,
|
||||
cancel: &AtomicBool,
|
||||
progress: &mut dyn FnMut(&Progress),
|
||||
) -> Result<(), DownloadError> {
|
||||
std::fs::create_dir_all(dir).map_err(|e| DownloadError::Disk(format!("can't create {}: {e}", dir.display())))?;
|
||||
// what's left to fetch
|
||||
let mut todo = Vec::new();
|
||||
for f in files {
|
||||
if !verified(f, &dir.join(f.name))? {
|
||||
todo.push(f);
|
||||
}
|
||||
}
|
||||
if todo.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
if mirrors.is_empty() {
|
||||
return Err(DownloadError::NoMirrors);
|
||||
}
|
||||
let total: u64 = files.iter().map(|f| f.size.unwrap_or(0)).sum();
|
||||
let mut before: u64 = files.iter().filter(|f| !todo.contains(f)).filter_map(|f| f.size).sum();
|
||||
for f in todo {
|
||||
let mut errors = Vec::new();
|
||||
let mut ok = false;
|
||||
'mirrors: for m in mirrors {
|
||||
let url = format!("{m}/{}", f.name);
|
||||
for _ in 0..opts.attempts.max(1) {
|
||||
if cancel.load(Ordering::Relaxed) {
|
||||
return Err(DownloadError::Cancelled);
|
||||
}
|
||||
let mut report = |n: u64, announced: u64| {
|
||||
let mirror = Url::parse(m).map(|u| u.display()).unwrap_or_default();
|
||||
progress(&Progress {
|
||||
file: f.name.to_string(),
|
||||
done: before.saturating_add(n),
|
||||
total: total.max(before.saturating_add(announced)),
|
||||
mirror,
|
||||
})
|
||||
};
|
||||
match fetch_file(f, &url, dir, opts, cancel, &mut report) {
|
||||
Ok(()) => {
|
||||
ok = true;
|
||||
break 'mirrors;
|
||||
}
|
||||
Err(Fail::Cancelled) => return Err(DownloadError::Cancelled),
|
||||
Err(Fail::Disk(e)) => return Err(DownloadError::Disk(e)),
|
||||
Err(Fail::Retry(e)) => {
|
||||
log::warn!("SAM 3 download of {url}: {e} (retrying)");
|
||||
errors.push(format!("{}: {e}", short(m)));
|
||||
}
|
||||
Err(Fail::NextMirror(e)) => {
|
||||
log::warn!("SAM 3 download of {url}: {e}");
|
||||
errors.push(format!("{}: {e}", short(m)));
|
||||
continue 'mirrors;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if !ok {
|
||||
return Err(DownloadError::AllMirrorsFailed { file: f.name.to_string(), errors });
|
||||
}
|
||||
before = before.saturating_add(std::fs::metadata(dir.join(f.name)).map(|m| m.len()).unwrap_or(0));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// A mirror for messages (no query string: it may hold a token).
|
||||
fn short(m: &str) -> String {
|
||||
Url::parse(m).map(|u| u.display()).unwrap_or_else(|_| "mirror".into())
|
||||
}
|
||||
|
||||
/// Whether `path` is already the right file (exact size and hash when pinned).
|
||||
fn verified(f: &FileSpec, path: &Path) -> Result<bool, DownloadError> {
|
||||
let Ok(meta) = std::fs::metadata(path) else { return Ok(false) };
|
||||
if !meta.is_file() || meta.len() == 0 {
|
||||
return Ok(false);
|
||||
}
|
||||
if f.size.is_some_and(|s| s != meta.len()) || meta.len() > f.max {
|
||||
return Ok(false);
|
||||
}
|
||||
match f.sha256 {
|
||||
Some(want) => Ok(sha256_file(path).map_err(DownloadError::Disk)? == want),
|
||||
None => Ok(true),
|
||||
}
|
||||
}
|
||||
|
||||
fn sha256_file(path: &Path) -> Result<String, String> {
|
||||
let mut file = File::open(path).map_err(|e| format!("{}: {e}", path.display()))?;
|
||||
let mut h = Sha256::new();
|
||||
let mut buf = vec![0u8; 1 << 20];
|
||||
loop {
|
||||
let n = file.read(&mut buf).map_err(|e| format!("{}: {e}", path.display()))?;
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
h.update(buf.get(..n).unwrap_or_default());
|
||||
}
|
||||
Ok(h.finalize().iter().map(|b| format!("{b:02x}")).collect())
|
||||
}
|
||||
|
||||
fn part_path(dir: &Path, name: &str) -> PathBuf {
|
||||
dir.join(format!("{name}.part"))
|
||||
}
|
||||
|
||||
/// Fetch one file from one mirror into `<name>.part` (resuming it when its size is pinned),
|
||||
/// check it and move it into place.
|
||||
fn fetch_file(f: &FileSpec, url: &str, dir: &Path, opts: &Options, cancel: &AtomicBool, report: &mut dyn FnMut(u64, u64)) -> Result<(), Fail> {
|
||||
let part = part_path(dir, f.name);
|
||||
let limit = f.size.unwrap_or(f.max);
|
||||
let disk = |e: std::io::Error| Fail::Disk(format!("{}: {e}", part.display()));
|
||||
// resume only a pinned file (an unpinned one can't be checked after resuming)
|
||||
let mut have = match std::fs::metadata(&part) {
|
||||
Ok(m) if f.size.is_some() && m.len() <= limit => m.len(),
|
||||
Ok(_) => {
|
||||
std::fs::remove_file(&part).map_err(disk)?;
|
||||
0
|
||||
}
|
||||
Err(_) => 0,
|
||||
};
|
||||
let fetched = f.size != Some(have) || have == 0;
|
||||
if fetched {
|
||||
let limits = Limits { connect: opts.connect_timeout, stall: opts.stall_timeout, cancel };
|
||||
let mut url = Url::parse(url).map_err(|e| Fail::NextMirror(e.to_string()))?;
|
||||
let origin = url.host.clone();
|
||||
let mut redirects = 0;
|
||||
let mut resp = loop {
|
||||
let mut headers = vec![];
|
||||
if have > 0 {
|
||||
headers.push(("Range", format!("bytes={have}-")));
|
||||
}
|
||||
// a token for the configured host only (never sent on to a redirect's host)
|
||||
if let Ok(t) = std::env::var("LIGHTCRAFT_SAM3_TOKEN")
|
||||
&& !t.trim().is_empty()
|
||||
&& url.tls
|
||||
&& url.host == origin
|
||||
{
|
||||
headers.push(("Authorization", format!("Bearer {}", t.trim())));
|
||||
}
|
||||
let r = http::get(&url, &headers, &limits).map_err(http_fail)?;
|
||||
match r.status {
|
||||
301 | 302 | 303 | 307 | 308 => {
|
||||
redirects += 1;
|
||||
if redirects > MAX_REDIRECTS {
|
||||
return Err(Fail::NextMirror("too many redirects".into()));
|
||||
}
|
||||
let loc = r.header("location").ok_or_else(|| Fail::NextMirror("redirect without a location".into()))?;
|
||||
url = url.join(loc).map_err(|e| Fail::NextMirror(e.to_string()))?;
|
||||
}
|
||||
_ => break r,
|
||||
}
|
||||
};
|
||||
let mut file = match resp.status {
|
||||
200 => {
|
||||
// the whole file (from the start, also when a range was asked for)
|
||||
if resp.content_length().is_some_and(|n| n > limit || f.size.is_some_and(|s| s != n)) {
|
||||
return Err(Fail::NextMirror(format!("the server's file has the wrong size ({:?} bytes)", resp.content_length())));
|
||||
}
|
||||
have = 0;
|
||||
File::create(&part).map_err(disk)?
|
||||
}
|
||||
206 if have > 0 => {
|
||||
let range = resp.header("content-range").unwrap_or_default();
|
||||
let (start, total) = parse_content_range(range).ok_or_else(|| Fail::NextMirror(format!("bad Content-Range `{range}`")))?;
|
||||
if start != have || total.is_some_and(|t| Some(t) != f.size) {
|
||||
return Err(Fail::NextMirror(format!("the server resumed at the wrong place ({range})")));
|
||||
}
|
||||
OpenOptions::new().append(true).open(&part).map_err(disk)?
|
||||
}
|
||||
416 => {
|
||||
// the partial file doesn't fit what the server has: start over
|
||||
std::fs::remove_file(&part).map_err(disk)?;
|
||||
return Err(Fail::Retry("the partial download didn't match; starting over".into()));
|
||||
}
|
||||
401 | 403 => return Err(Fail::NextMirror(format!("access denied (HTTP {})", resp.status))),
|
||||
404 | 410 => return Err(Fail::NextMirror(format!("not found (HTTP {})", resp.status))),
|
||||
s if s >= 500 => return Err(Fail::Retry(format!("server error (HTTP {s})"))),
|
||||
s => return Err(Fail::NextMirror(format!("HTTP {s}"))),
|
||||
};
|
||||
let announced = f.size.unwrap_or_else(|| resp.content_length().unwrap_or(0));
|
||||
report(have, announced);
|
||||
let mut buf = vec![0u8; 256 * 1024];
|
||||
loop {
|
||||
let n = match resp.read(&mut buf, &limits) {
|
||||
Ok(n) => n,
|
||||
Err(e) => {
|
||||
let _ = file.flush();
|
||||
if f.size.is_none() {
|
||||
drop(file);
|
||||
let _ = std::fs::remove_file(&part);
|
||||
}
|
||||
return Err(http_fail(e));
|
||||
}
|
||||
};
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
if have.saturating_add(n as u64) > limit {
|
||||
drop(file);
|
||||
let _ = std::fs::remove_file(&part);
|
||||
return Err(Fail::NextMirror(format!("the server sent more than the expected {limit} bytes")));
|
||||
}
|
||||
file.write_all(buf.get(..n).unwrap_or_default()).map_err(disk)?;
|
||||
have = have.saturating_add(n as u64);
|
||||
report(have, announced);
|
||||
}
|
||||
file.sync_all().map_err(disk)?;
|
||||
}
|
||||
// checks: exact size, hash
|
||||
let len = std::fs::metadata(&part).map_err(disk)?.len();
|
||||
if f.size.is_some_and(|s| s != len) {
|
||||
return Err(Fail::Retry(format!("incomplete ({len} of {} bytes)", f.size.unwrap_or(0))));
|
||||
}
|
||||
if len == 0 {
|
||||
let _ = std::fs::remove_file(&part);
|
||||
return Err(Fail::NextMirror("the server sent an empty file".into()));
|
||||
}
|
||||
if let Some(want) = f.sha256 {
|
||||
let got = sha256_file(&part).map_err(Fail::Disk)?;
|
||||
if got != want {
|
||||
std::fs::remove_file(&part).map_err(disk)?;
|
||||
let why = format!("the file is damaged or not the official one (SHA-256 {got})");
|
||||
// a damaged partial file from an earlier run isn't this mirror's fault
|
||||
return Err(if fetched { Fail::NextMirror(why) } else { Fail::Retry(why) });
|
||||
}
|
||||
}
|
||||
std::fs::rename(&part, dir.join(f.name)).map_err(disk)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn http_fail(e: HttpError) -> Fail {
|
||||
match e {
|
||||
HttpError::Cancelled => Fail::Cancelled,
|
||||
HttpError::Stalled | HttpError::Io(_) | HttpError::Protocol(_) => Fail::Retry(e.to_string()),
|
||||
HttpError::BadUrl(_) | HttpError::Connect(_) | HttpError::Tls(_) => Fail::NextMirror(e.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
/// `bytes a-b/total` → (a, total) (`*` total → None).
|
||||
fn parse_content_range(v: &str) -> Option<(u64, Option<u64>)> {
|
||||
let v = v.trim().strip_prefix("bytes")?.trim();
|
||||
let (range, total) = v.split_once('/')?;
|
||||
let (a, b) = range.split_once('-')?;
|
||||
let (a, b) = (a.trim().parse::<u64>().ok()?, b.trim().parse::<u64>().ok()?);
|
||||
if b < a {
|
||||
return None;
|
||||
}
|
||||
let total = match total.trim() {
|
||||
"*" => None,
|
||||
t => Some(t.parse::<u64>().ok()?),
|
||||
};
|
||||
Some((a, total))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
|
@ -0,0 +1,314 @@
|
|||
//! The downloader against a local HTTP server (std `TcpListener`; never the internet): mirror
|
||||
//! fallback, redirects, resuming, chunked bodies, stalls, hash mismatches, size caps, cancel.
|
||||
|
||||
use std::io::{BufRead, BufReader, Write};
|
||||
use std::net::{TcpListener, TcpStream};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use super::*;
|
||||
|
||||
/// The served file: 300 kB of a deterministic pattern.
|
||||
fn payload() -> Vec<u8> {
|
||||
(0..300_000u32).map(|i| (i.wrapping_mul(2_654_435_761) >> 13) as u8).collect()
|
||||
}
|
||||
|
||||
fn sha_hex(b: &[u8]) -> String {
|
||||
Sha256::digest(b).iter().map(|b| format!("{b:02x}")).collect()
|
||||
}
|
||||
|
||||
fn leak(s: String) -> &'static str {
|
||||
Box::leak(s.into_boxed_str())
|
||||
}
|
||||
|
||||
fn pinned() -> FileSpec {
|
||||
let p = payload();
|
||||
FileSpec { name: "model.bin", size: Some(p.len() as u64), sha256: Some(leak(sha_hex(&p))), max: p.len() as u64 }
|
||||
}
|
||||
|
||||
struct Server {
|
||||
base: String,
|
||||
hits: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
/// Serve on 127.0.0.1; the first path segment picks the behaviour.
|
||||
fn server() -> Server {
|
||||
let l = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let base = format!("http://{}", l.local_addr().unwrap());
|
||||
let hits = Arc::new(AtomicUsize::new(0));
|
||||
let h = hits.clone();
|
||||
std::thread::spawn(move || {
|
||||
for s in l.incoming().flatten() {
|
||||
let h = h.clone();
|
||||
std::thread::spawn(move || handle(s, &h));
|
||||
}
|
||||
});
|
||||
Server { base, hits }
|
||||
}
|
||||
|
||||
fn handle(mut s: TcpStream, hits: &AtomicUsize) {
|
||||
let n = hits.fetch_add(1, Ordering::SeqCst);
|
||||
let mut r = BufReader::new(s.try_clone().unwrap());
|
||||
let mut first = String::new();
|
||||
if r.read_line(&mut first).is_err() {
|
||||
return;
|
||||
}
|
||||
let mut range: Option<u64> = None;
|
||||
loop {
|
||||
let mut l = String::new();
|
||||
if r.read_line(&mut l).unwrap_or(0) == 0 || l.trim().is_empty() {
|
||||
break;
|
||||
}
|
||||
if let Some(v) = l.to_ascii_lowercase().strip_prefix("range: bytes=") {
|
||||
range = v.trim().trim_end_matches('-').parse().ok();
|
||||
}
|
||||
}
|
||||
let path = first.split_whitespace().nth(1).unwrap_or("/").to_string();
|
||||
let mode = path.trim_start_matches('/').split('/').next().unwrap_or("").to_string();
|
||||
let body = payload();
|
||||
let send = |s: &mut TcpStream, status: &str, headers: &str, data: &[u8]| {
|
||||
let _ = s.write_all(format!("HTTP/1.1 {status}\r\n{headers}Connection: close\r\n\r\n").as_bytes());
|
||||
let _ = s.write_all(data);
|
||||
};
|
||||
match mode.as_str() {
|
||||
"missing" => send(&mut s, "404 Not Found", "Content-Length: 9\r\n", b"not found"),
|
||||
"error" => send(&mut s, "500 Oops", "Content-Length: 25\r\n", b"<html>server error</html>"),
|
||||
"redirect" => send(&mut s, "302 Found", &format!("Location: {}\r\nContent-Length: 0\r\n", path.replacen("/redirect", "/good", 1)), b""),
|
||||
"wrong" => {
|
||||
let mut b = body.clone();
|
||||
b[1000] ^= 0xff;
|
||||
send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", b.len()), &b)
|
||||
}
|
||||
"toolong" => {
|
||||
// no length announced, more bytes than the file can have
|
||||
let mut b = body.clone();
|
||||
b.extend_from_slice(&[7u8; 5000]);
|
||||
send(&mut s, "200 OK", "", &b)
|
||||
}
|
||||
"chunked" => {
|
||||
let _ = s.write_all(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n");
|
||||
for c in body.chunks(70_001) {
|
||||
let _ = s.write_all(format!("{:x};ext=1\r\n", c.len()).as_bytes());
|
||||
let _ = s.write_all(c);
|
||||
let _ = s.write_all(b"\r\n");
|
||||
}
|
||||
let _ = s.write_all(b"0\r\nX-Trailer: 1\r\n\r\n");
|
||||
}
|
||||
"stall" => {
|
||||
// half the file, then nothing
|
||||
let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes());
|
||||
let _ = s.write_all(&body[..body.len() / 2]);
|
||||
std::thread::sleep(Duration::from_secs(5));
|
||||
}
|
||||
"slow" => {
|
||||
let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes());
|
||||
for c in body.chunks(1000) {
|
||||
if s.write_all(c).is_err() {
|
||||
return;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(20));
|
||||
}
|
||||
}
|
||||
"flaky" if n.is_multiple_of(2) && range.is_none() => {
|
||||
// the first try breaks off after a third
|
||||
let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes());
|
||||
let _ = s.write_all(&body[..body.len() / 3]);
|
||||
}
|
||||
"range416" if range.is_some() => send(&mut s, "416 Range Not Satisfiable", "Content-Length: 0\r\n", b""),
|
||||
"norange" | "range416" => send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", body.len()), &body),
|
||||
_ => match range {
|
||||
Some(from) if from as usize >= body.len() => send(&mut s, "416 Range Not Satisfiable", "Content-Length: 0\r\n", b""),
|
||||
Some(from) => send(
|
||||
&mut s,
|
||||
"206 Partial Content",
|
||||
&format!("Content-Range: bytes {from}-{}/{}\r\nContent-Length: {}\r\n", body.len() - 1, body.len(), body.len() - from as usize),
|
||||
&body[from as usize..],
|
||||
),
|
||||
None => send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", body.len()), &body),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn tmp(name: &str) -> PathBuf {
|
||||
let d = std::env::temp_dir().join(format!("lc-sam3-fetch-{name}-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
d
|
||||
}
|
||||
|
||||
fn quick() -> Options {
|
||||
Options { connect_timeout: Duration::from_secs(2), stall_timeout: Duration::from_millis(600), attempts: 2 }
|
||||
}
|
||||
|
||||
fn run(files: &[FileSpec], mirrors: &[String], dir: &Path, cancel: &AtomicBool) -> (Result<(), DownloadError>, Vec<Progress>) {
|
||||
let mut seen = Vec::new();
|
||||
let r = download(files, mirrors, dir, &quick(), cancel, &mut |p| seen.push(p.clone()));
|
||||
(r, seen)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn falls_back_to_the_next_mirror_and_verifies() {
|
||||
let srv = server();
|
||||
let dir = tmp("fallback");
|
||||
let mirrors = vec![
|
||||
format!("{}/missing", srv.base),
|
||||
format!("{}/error", srv.base),
|
||||
"http://127.0.0.1:1/unreachable".to_string(),
|
||||
format!("{}/redirect", srv.base),
|
||||
];
|
||||
let (r, seen) = run(&[pinned()], &mirrors, &dir, &AtomicBool::new(false));
|
||||
r.unwrap();
|
||||
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
|
||||
assert!(!part_path(&dir, "model.bin").exists());
|
||||
let last = seen.last().unwrap();
|
||||
assert_eq!((last.done, last.total), (300_000, 300_000));
|
||||
// an error page never ends up in the file: nothing left over from the failed mirrors
|
||||
assert_eq!(std::fs::read_dir(&dir).unwrap().count(), 1);
|
||||
// installed: nothing is fetched again
|
||||
let hits = srv.hits.load(Ordering::SeqCst);
|
||||
run(&[pinned()], &mirrors, &dir, &AtomicBool::new(false)).0.unwrap();
|
||||
assert_eq!(srv.hits.load(Ordering::SeqCst), hits);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_hash_mismatch_is_deleted_and_the_next_mirror_used() {
|
||||
let srv = server();
|
||||
let dir = tmp("hash");
|
||||
let (r, _) = run(&[pinned()], &[format!("{}/wrong", srv.base)], &dir, &AtomicBool::new(false));
|
||||
let e = r.unwrap_err().to_string();
|
||||
assert!(e.contains("SHA-256"), "{e}");
|
||||
assert!(!dir.join("model.bin").exists() && !part_path(&dir, "model.bin").exists(), "a damaged file is never kept");
|
||||
run(&[pinned()], &[format!("{}/wrong", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
|
||||
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resumes_partial_files_and_restarts_bad_ones() {
|
||||
let srv = server();
|
||||
let dir = tmp("resume");
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
// a third already there
|
||||
std::fs::write(part_path(&dir, "model.bin"), &payload()[..100_000]).unwrap();
|
||||
run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
|
||||
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
|
||||
// a partial file with the wrong bytes: the hash fails, it's deleted, the next try starts over
|
||||
std::fs::remove_file(dir.join("model.bin")).unwrap();
|
||||
std::fs::write(part_path(&dir, "model.bin"), vec![0u8; 100_000]).unwrap();
|
||||
run(&[pinned()], &[format!("{}/good", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
|
||||
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
|
||||
// a server ignoring Range: the whole file again, not appended
|
||||
std::fs::remove_file(dir.join("model.bin")).unwrap();
|
||||
std::fs::write(part_path(&dir, "model.bin"), &payload()[..5]).unwrap();
|
||||
run(&[pinned()], &[format!("{}/norange", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
|
||||
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
|
||||
// a connection breaking off is resumed on the same mirror
|
||||
std::fs::remove_file(dir.join("model.bin")).unwrap();
|
||||
run(&[pinned()], &[format!("{}/flaky", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
|
||||
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_416_or_a_bad_complete_part_starts_over() {
|
||||
let srv = server();
|
||||
let dir = tmp("416");
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
// the server refuses the range: the part is dropped and the file fetched whole (no loop)
|
||||
std::fs::write(part_path(&dir, "model.bin"), &payload()[..100]).unwrap();
|
||||
run(&[pinned()], &[format!("{}/range416", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
|
||||
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
|
||||
// as long as the file but wrong: hash fails → deleted → downloaded again from the same mirror
|
||||
std::fs::remove_file(dir.join("model.bin")).unwrap();
|
||||
std::fs::write(part_path(&dir, "model.bin"), vec![1u8; 300_000]).unwrap();
|
||||
run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
|
||||
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chunked_bodies_and_unpinned_files() {
|
||||
let srv = server();
|
||||
let dir = tmp("chunked");
|
||||
let f = FileSpec { name: "vocab.json", size: None, sha256: None, max: 1 << 20 };
|
||||
run(std::slice::from_ref(&f), &[format!("{}/chunked", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
|
||||
assert_eq!(std::fs::read(dir.join("vocab.json")).unwrap(), payload());
|
||||
// more than the cap: refused, nothing kept
|
||||
let small = FileSpec { max: 1000, ..f };
|
||||
let dir2 = tmp("cap");
|
||||
let e = run(&[small], &[format!("{}/good", srv.base)], &dir2, &AtomicBool::new(false)).0.unwrap_err().to_string();
|
||||
assert!(e.contains("wrong size") || e.contains("more than"), "{e}");
|
||||
assert!(!dir2.join("vocab.json").exists() && !part_path(&dir2, "vocab.json").exists());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
let _ = std::fs::remove_dir_all(&dir2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_server_sending_too_much_is_cut_off() {
|
||||
let srv = server();
|
||||
let dir = tmp("toolong");
|
||||
let e = run(&[pinned()], &[format!("{}/toolong", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap_err().to_string();
|
||||
assert!(e.contains("more than"), "{e}");
|
||||
assert!(!dir.join("model.bin").exists() && !part_path(&dir, "model.bin").exists());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_stalled_server_times_out_and_the_next_mirror_finishes() {
|
||||
let srv = server();
|
||||
let dir = tmp("stall");
|
||||
let t = Instant::now();
|
||||
let (r, _) = run(&[pinned()], &[format!("{}/stall", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false));
|
||||
r.unwrap();
|
||||
assert!(t.elapsed() < Duration::from_secs(4), "{:?}", t.elapsed());
|
||||
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancel_stops_quickly_and_keeps_the_part_for_later() {
|
||||
let srv = server();
|
||||
let dir = tmp("cancel");
|
||||
let cancel = Arc::new(AtomicBool::new(false));
|
||||
let c = cancel.clone();
|
||||
std::thread::spawn(move || {
|
||||
std::thread::sleep(Duration::from_millis(400));
|
||||
c.store(true, Ordering::SeqCst);
|
||||
});
|
||||
let t = Instant::now();
|
||||
let (r, _) = run(&[pinned()], &[format!("{}/slow", srv.base)], &dir, &cancel);
|
||||
assert_eq!(r, Err(DownloadError::Cancelled));
|
||||
assert!(t.elapsed() < Duration::from_secs(3), "{:?}", t.elapsed());
|
||||
assert!(!dir.join("model.bin").exists());
|
||||
let kept = std::fs::metadata(part_path(&dir, "model.bin")).map(|m| m.len()).unwrap_or(0);
|
||||
assert!(kept > 0 && kept < 300_000, "{kept}");
|
||||
// and the next run resumes it
|
||||
run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
|
||||
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_mirrors_is_a_clear_error_and_mirror_lists_are_parsed() {
|
||||
let dir = tmp("nomirror");
|
||||
let e = run(&[pinned()], &[], &dir, &AtomicBool::new(false)).0.unwrap_err();
|
||||
assert_eq!(e, DownloadError::NoMirrors);
|
||||
assert!(e.to_string().contains(MIRRORS_ENV));
|
||||
let file = std::env::temp_dir().join(format!("lc-sam3-mirrors-{}.txt", std::process::id()));
|
||||
std::fs::write(&file, "# mine\nhttps://b.example/m/ # second\n\nnot a url\nhttps://a.example/x\n").unwrap();
|
||||
let m = mirrors(Some("https://a.example/x/, ftp://no"), Some(&file));
|
||||
assert_eq!(m[..2], ["https://a.example/x".to_string(), "https://b.example/m".to_string()]);
|
||||
assert_eq!(m.len(), 2 + DEFAULT_MIRRORS.len());
|
||||
assert_eq!(mirrors(None, Some(Path::new("/nonexistent/mirrors.txt"))).len(), DEFAULT_MIRRORS.len());
|
||||
let _ = std::fs::remove_file(&file);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn content_ranges() {
|
||||
assert_eq!(parse_content_range("bytes 100-299/300"), Some((100, Some(300))));
|
||||
assert_eq!(parse_content_range("bytes 0-0/*"), Some((0, None)));
|
||||
assert_eq!(parse_content_range("bytes 9-1/300"), None);
|
||||
assert_eq!(parse_content_range("items 1-2/3"), None);
|
||||
}
|
||||
|
|
@ -0,0 +1,263 @@
|
|||
//! SAM 3 (Segment Anything with Concepts, Meta, 2025) inference in pure Rust on candle.
|
||||
//!
|
||||
//! - **Point prompts** ([`Sam3::segment_clicks`]): positive and negative clicks select one
|
||||
//! object, via SAM 3's interactive (SAM 2 style) prompt encoder and mask decoder.
|
||||
//! - **Text prompts** ([`Sam3::segment_text`]): a noun phrase ("the red car", "sky") selects
|
||||
//! every instance of that concept, via the CLIP text encoder and the DETR detector.
|
||||
//!
|
||||
//! The expensive part, the 32-layer ViT over the 1008 × 1008 input, runs once per image
|
||||
//! ([`Sam3::encode`]); each click or prompt then reuses the [`Encoded`] features.
|
||||
//!
|
||||
//! On macOS the model runs on the GPU through Metal; elsewhere (for now) on the CPU. The
|
||||
//! weights are not part of LightCraft: they are the `facebook/sam3` checkpoint (SAM License),
|
||||
//! read from a directory the user downloads them to ([`Sam3::load`]).
|
||||
//!
|
||||
//! Ported from the Hugging Face `transformers` implementation (Apache-2.0); see `NOTICE`.
|
||||
//! This crate is therefore licensed under the Apache License 2.0 only (not the MIT option of
|
||||
//! the rest of LightCraft). No UI dependencies (L3).
|
||||
//!
|
||||
//! [`fetch`] downloads the model, only when the user asks for it (from configurable mirrors,
|
||||
//! verified before use). Nothing in LightCraft requires the model: without it, AI masks report
|
||||
//! that it isn't installed and everything else works.
|
||||
//!
|
||||
//! 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 (`src/transformers/models/sam3` and `sam3_tracker`), Copyright The HuggingFace
|
||||
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
|
||||
#![forbid(unsafe_code)]
|
||||
#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)]
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
mod clip;
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
mod detector;
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
pub mod fetch;
|
||||
pub mod mask;
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
mod neck;
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
mod nn;
|
||||
pub mod tokenizer;
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
mod tracker;
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
mod vit;
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
mod weights;
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
use std::sync::Arc;
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
use candle_nn::VarBuilder;
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
pub use tracker::Click;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("SAM 3 model: {0}")]
|
||||
Model(String),
|
||||
#[error("SAM 3 model files not found in {0} (expected model.safetensors, vocab.json, merges.txt)")]
|
||||
Missing(PathBuf),
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
#[error(transparent)]
|
||||
Candle(#[from] candle_core::Error),
|
||||
}
|
||||
|
||||
pub type Result<T> = std::result::Result<T, Error>;
|
||||
|
||||
/// Side of the square model input.
|
||||
pub const INPUT: usize = 1008;
|
||||
/// Side of the low-resolution masks the model predicts.
|
||||
pub const MASK_SIDE: usize = 288;
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
const _: () = assert!(MASK_SIDE == tracker::LOW_RES && INPUT == vit::IMAGE);
|
||||
|
||||
/// The files of a model directory.
|
||||
pub const WEIGHTS_FILE: &str = "model.safetensors";
|
||||
|
||||
/// Whether `dir` holds a usable checkpoint.
|
||||
pub fn is_model_dir(dir: &Path) -> bool {
|
||||
dir.join(WEIGHTS_FILE).is_file() && dir.join("vocab.json").is_file() && dir.join("merges.txt").is_file()
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
/// The best device here: Metal on macOS (when a GPU is available), else the CPU.
|
||||
pub fn best_device() -> Device {
|
||||
#[cfg(target_os = "macos")]
|
||||
if let Ok(d) = Device::new_metal(0) {
|
||||
return d;
|
||||
}
|
||||
Device::Cpu
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
/// A loaded model. The vision backbone loads with it; the point and text heads load on first use.
|
||||
pub struct Sam3 {
|
||||
dir: PathBuf,
|
||||
device: Device,
|
||||
weights: Arc<weights::Weights>,
|
||||
backbone: vit::Vit,
|
||||
tracker_neck: Option<neck::Neck>,
|
||||
tracker: Option<tracker::Tracker>,
|
||||
detector: Option<detector::Detector>,
|
||||
tokenizer: Option<tokenizer::Tokenizer>,
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
/// An image run through the backbone, with the head-specific features computed on demand.
|
||||
pub struct Encoded {
|
||||
backbone: Tensor,
|
||||
tracker: Option<tracker::Features>,
|
||||
detector: Option<detector::Features>,
|
||||
}
|
||||
|
||||
/// A predicted mask: `MASK_SIDE²` logits over the model's square input (which the image was
|
||||
/// stretched to), and a confidence 0..1.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Prediction {
|
||||
pub logits: Vec<f32>,
|
||||
pub score: f32,
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
impl Sam3 {
|
||||
/// Load from a directory holding `model.safetensors` (the `facebook/sam3` checkpoint in
|
||||
/// Hugging Face format), `vocab.json` and `merges.txt`.
|
||||
pub fn load(dir: &Path) -> Result<Self> {
|
||||
Self::load_on(dir, best_device())
|
||||
}
|
||||
|
||||
pub fn load_on(dir: &Path, device: Device) -> Result<Self> {
|
||||
if !is_model_dir(dir) {
|
||||
return Err(Error::Missing(dir.to_path_buf()));
|
||||
}
|
||||
let weights = weights::Weights::open(&dir.join(WEIGHTS_FILE))?;
|
||||
let backbone = vit::Vit::new(Self::vb(&weights, &device).pp("detector_model.vision_encoder.backbone"))?;
|
||||
Ok(Self { dir: dir.to_path_buf(), device, weights, backbone, tracker_neck: None, tracker: None, detector: None, tokenizer: None })
|
||||
}
|
||||
|
||||
fn vb<'a>(w: &Arc<weights::Weights>, device: &Device) -> VarBuilder<'a> {
|
||||
VarBuilder::from_backend(Box::new(weights::Backend(w.clone())), DType::F32, device.clone())
|
||||
}
|
||||
|
||||
pub fn device(&self) -> &Device {
|
||||
&self.device
|
||||
}
|
||||
|
||||
/// Run the backbone on an RGB image (8-bit, row-major, `w × h`), stretched to the square
|
||||
/// input like the reference preprocessing.
|
||||
pub fn encode(&self, rgb: &[u8], w: usize, h: usize) -> Result<Encoded> {
|
||||
let pixels = mask::preprocess(rgb, w, h, INPUT)?;
|
||||
let pixels = Tensor::from_vec(pixels, (1, 3, INPUT, INPUT), &self.device)?;
|
||||
self.encode_pixels(&pixels)
|
||||
}
|
||||
|
||||
/// Run the backbone on preprocessed pixels `[1, 3, 1008, 1008]` (−1..1).
|
||||
pub fn encode_pixels(&self, pixels: &Tensor) -> Result<Encoded> {
|
||||
Ok(Encoded { backbone: self.backbone.forward(&pixels.to_device(&self.device)?)?, tracker: None, detector: None })
|
||||
}
|
||||
|
||||
fn tracker_parts(&mut self) -> Result<(&neck::Neck, &tracker::Tracker)> {
|
||||
if self.tracker_neck.is_none() {
|
||||
self.tracker_neck = Some(neck::Neck::new(Self::vb(&self.weights, &self.device).pp("tracker_neck"))?);
|
||||
}
|
||||
if self.tracker.is_none() {
|
||||
self.tracker = Some(tracker::Tracker::new(Self::vb(&self.weights, &self.device).pp("tracker_model"))?);
|
||||
}
|
||||
match (&self.tracker_neck, &self.tracker) {
|
||||
(Some(n), Some(t)) => Ok((n, t)),
|
||||
_ => Err(Error::Model("tracker failed to load".into())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute the click decoder's features for `enc` now (so the first click is fast).
|
||||
pub fn prepare_clicks(&mut self, enc: &mut Encoded) -> Result<()> {
|
||||
let (neck, tracker) = self.tracker_parts()?;
|
||||
if enc.tracker.is_none() {
|
||||
let levels = neck.forward(&enc.backbone, 3)?;
|
||||
enc.tracker = Some(tracker.features(&levels)?);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Segment one object from clicks (positions normalized to the image, 0..1).
|
||||
pub fn segment_clicks(&mut self, enc: &mut Encoded, clicks: &[Click]) -> Result<Prediction> {
|
||||
if clicks.is_empty() {
|
||||
return Err(Error::Model("no clicks".into()));
|
||||
}
|
||||
let (neck, tracker) = self.tracker_parts()?;
|
||||
if enc.tracker.is_none() {
|
||||
let levels = neck.forward(&enc.backbone, 3)?;
|
||||
enc.tracker = Some(tracker.features(&levels)?);
|
||||
}
|
||||
let f = enc.tracker.as_ref().ok_or_else(|| Error::Model("tracker features".into()))?;
|
||||
// one click: SAM's three alternatives, best by predicted IoU; more: the single mask
|
||||
let p = tracker.predict(f, clicks, clicks.len() == 1)?;
|
||||
Ok(Prediction { logits: p.logits, score: p.iou * sigmoid(p.object_score) })
|
||||
}
|
||||
|
||||
fn detector_parts(&mut self) -> Result<(&detector::Detector, &tokenizer::Tokenizer)> {
|
||||
if self.detector.is_none() {
|
||||
self.detector = Some(detector::Detector::new(Self::vb(&self.weights, &self.device).pp("detector_model"))?);
|
||||
}
|
||||
if self.tokenizer.is_none() {
|
||||
self.tokenizer = Some(tokenizer::Tokenizer::load(&self.dir)?);
|
||||
}
|
||||
match (&self.detector, &self.tokenizer) {
|
||||
(Some(d), Some(t)) => Ok((d, t)),
|
||||
_ => Err(Error::Model("detector failed to load".into())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Segment every instance of the concept `text` names. Instances scoring above
|
||||
/// `threshold` (0..1, 0.5 is the reference default) are merged into one mask; `None` when
|
||||
/// nothing matches.
|
||||
pub fn segment_text(&mut self, enc: &mut Encoded, text: &str, threshold: f32) -> Result<Option<Prediction>> {
|
||||
let (det, tok) = self.detector_parts()?;
|
||||
if enc.detector.is_none() {
|
||||
enc.detector = Some(det.features(&enc.backbone)?);
|
||||
}
|
||||
let f = enc.detector.as_ref().ok_or_else(|| Error::Model("detector features".into()))?;
|
||||
let (ids, valid) = tok.encode(text);
|
||||
let out = det.predict(f, &ids, &valid)?;
|
||||
Ok(detector::merge(&out, threshold))
|
||||
}
|
||||
|
||||
/// Raw detector outputs for `text` (tests and tools).
|
||||
pub fn detect(&mut self, enc: &mut Encoded, text: &str) -> Result<detector::Output> {
|
||||
let (det, tok) = self.detector_parts()?;
|
||||
if enc.detector.is_none() {
|
||||
enc.detector = Some(det.features(&enc.backbone)?);
|
||||
}
|
||||
let f = enc.detector.as_ref().ok_or_else(|| Error::Model("detector features".into()))?;
|
||||
let (ids, valid) = tok.encode(text);
|
||||
det.predict(f, &ids, &valid)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
impl Encoded {
|
||||
/// The backbone output `[1, 1024, 72, 72]` (tests).
|
||||
pub fn backbone(&self) -> &Tensor {
|
||||
&self.backbone
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
fn sigmoid(x: f32) -> f32 {
|
||||
1.0 / (1.0 + (-x).exp())
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
pub use detector::Output as DetectorOutput;
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
/// Wait for queued GPU work to finish (benchmarks).
|
||||
pub fn sync(model: &Sam3) -> Result<()> {
|
||||
Ok(model.device.synchronize()?)
|
||||
}
|
||||
|
|
@ -0,0 +1,130 @@
|
|||
//! Image preprocessing and mask post-processing (CPU, no model needed).
|
||||
//!
|
||||
//! 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/image_processing_sam3*.py`), Copyright The HuggingFace
|
||||
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
|
||||
|
||||
use crate::{Error, Result};
|
||||
|
||||
/// Separable resampling weights for `src → dst` samples with an antialiased triangle filter
|
||||
/// (Pillow's BILINEAR: the support widens with the downscale factor).
|
||||
fn weights(src: usize, dst: usize) -> Vec<(usize, Vec<f32>)> {
|
||||
let scale = src as f64 / dst as f64;
|
||||
let ss = scale.max(1.0);
|
||||
let support = ss;
|
||||
(0..dst)
|
||||
.map(|o| {
|
||||
let center = (o as f64 + 0.5) * scale;
|
||||
let lo = ((center - support + 0.5).floor().max(0.0)) as usize;
|
||||
let hi = ((center + support + 0.5).floor() as usize).min(src);
|
||||
let mut w: Vec<f32> = (lo..hi).map(|i| (1.0 - ((i as f64 - center + 0.5) / ss).abs()).max(0.0) as f32).collect();
|
||||
let sum: f32 = w.iter().sum();
|
||||
if sum > 0.0 {
|
||||
w.iter_mut().for_each(|v| *v /= sum);
|
||||
}
|
||||
(lo, w)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// An 8-bit RGB image (`w × h`, row-major) stretched to `side × side`, scaled to −1..1 and laid
|
||||
/// out channel-first (`[3, side, side]`), as SAM 3's processor does.
|
||||
pub fn preprocess(rgb: &[u8], w: usize, h: usize, side: usize) -> Result<Vec<f32>> {
|
||||
let n = w.checked_mul(h).and_then(|n| n.checked_mul(3)).ok_or_else(|| Error::Model("image too large".into()))?;
|
||||
if w == 0 || h == 0 || rgb.len() < n {
|
||||
return Err(Error::Model(format!("image buffer of {} bytes for {w} × {h}", rgb.len())));
|
||||
}
|
||||
let (wx, wy) = (weights(w, side), weights(h, side));
|
||||
// horizontal pass: h rows × side columns × 3
|
||||
let mut tmp = vec![0f32; h * side * 3];
|
||||
for y in 0..h {
|
||||
for (ox, (lo, ws)) in wx.iter().enumerate() {
|
||||
let mut acc = [0f32; 3];
|
||||
for (k, wt) in ws.iter().enumerate() {
|
||||
let i = (y * w + lo + k) * 3;
|
||||
for (c, a) in acc.iter_mut().enumerate() {
|
||||
*a += wt * f32::from(rgb.get(i + c).copied().unwrap_or(0));
|
||||
}
|
||||
}
|
||||
if let Some(dst) = tmp.get_mut((y * side + ox) * 3..(y * side + ox) * 3 + 3) {
|
||||
dst.copy_from_slice(&acc);
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut out = vec![0f32; 3 * side * side];
|
||||
for (oy, (lo, ws)) in wy.iter().enumerate() {
|
||||
for ox in 0..side {
|
||||
let mut acc = [0f32; 3];
|
||||
for (k, wt) in ws.iter().enumerate() {
|
||||
let i = ((lo + k) * side + ox) * 3;
|
||||
for (c, a) in acc.iter_mut().enumerate() {
|
||||
*a += wt * tmp.get(i + c).copied().unwrap_or(0.0);
|
||||
}
|
||||
}
|
||||
for (c, a) in acc.iter().enumerate() {
|
||||
// Pillow rounds to 8 bits; then /255 and (x − 0.5) / 0.5
|
||||
let v = a.round().clamp(0.0, 255.0) / 255.0;
|
||||
if let Some(d) = out.get_mut((c * side + oy) * side + ox) {
|
||||
*d = (v - 0.5) / 0.5;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Mask `logits` (`side × side`) resampled bilinearly (half-pixel centres) to `w × h` and
|
||||
/// mapped to 0..1 coverage (sigmoid).
|
||||
pub fn alpha(logits: &[f32], side: usize, w: usize, h: usize) -> Vec<f32> {
|
||||
let mut out = vec![0f32; w * h];
|
||||
if side == 0 || logits.len() < side * side {
|
||||
return out;
|
||||
}
|
||||
let sample = |x: usize, y: usize| logits.get(y.min(side - 1) * side + x.min(side - 1)).copied().unwrap_or(-20.0);
|
||||
let (sx, sy) = (side as f32 / w.max(1) as f32, side as f32 / h.max(1) as f32);
|
||||
for y in 0..h {
|
||||
let fy = ((y as f32 + 0.5) * sy - 0.5).max(0.0);
|
||||
let (y0, ty) = (fy.floor() as usize, fy.fract());
|
||||
for x in 0..w {
|
||||
let fx = ((x as f32 + 0.5) * sx - 0.5).max(0.0);
|
||||
let (x0, tx) = (fx.floor() as usize, fx.fract());
|
||||
let top = sample(x0, y0) * (1.0 - tx) + sample(x0 + 1, y0) * tx;
|
||||
let bot = sample(x0, y0 + 1) * (1.0 - tx) + sample(x0 + 1, y0 + 1) * tx;
|
||||
let l = top * (1.0 - ty) + bot * ty;
|
||||
if let Some(o) = out.get_mut(y * w + x) {
|
||||
*o = 1.0 / (1.0 + (-l).exp());
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn preprocess_is_bounded_and_channel_first() {
|
||||
let rgb: Vec<u8> = (0..20 * 10 * 3).map(|i| (i % 256) as u8).collect();
|
||||
let p = preprocess(&rgb, 20, 10, 8).unwrap();
|
||||
assert_eq!(p.len(), 3 * 64);
|
||||
assert!(p.iter().all(|v| (-1.0..=1.0).contains(v)));
|
||||
// a flat image stays flat
|
||||
let flat = vec![255u8; 30 * 30 * 3];
|
||||
assert!(preprocess(&flat, 30, 30, 7).unwrap().iter().all(|v| (*v - 1.0).abs() < 1e-6));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn short_buffers_are_errors() {
|
||||
assert!(preprocess(&[0; 5], 2, 2, 4).is_err());
|
||||
assert!(preprocess(&[], 0, 0, 4).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn alpha_follows_the_logits() {
|
||||
let logits: Vec<f32> = (0..16).map(|i| if i % 4 < 2 { 10.0 } else { -10.0 }).collect();
|
||||
let a = alpha(&logits, 4, 8, 2);
|
||||
assert!(a[0] > 0.99 && a[7] < 0.01);
|
||||
assert!(alpha(&[], 4, 3, 3).iter().all(|v| *v == 0.0));
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
//! The feature pyramid (`Sam3VisionNeck`): the backbone's 72 × 72 map at four scales
|
||||
//! (288, 144, 72 and 36 per side), each projected to 256 channels.
|
||||
//!
|
||||
//! 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::{Module, Tensor};
|
||||
use candle_nn::{Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, VarBuilder};
|
||||
|
||||
use crate::Result;
|
||||
use crate::vit::HIDDEN;
|
||||
|
||||
pub const FPN: usize = 256;
|
||||
|
||||
enum Scale {
|
||||
/// Two 2× transposed convolutions with GELU between (4×).
|
||||
Up4(ConvTranspose2d, ConvTranspose2d),
|
||||
Up2(ConvTranspose2d),
|
||||
Same,
|
||||
/// 2 × 2 max pooling (½×).
|
||||
Down2,
|
||||
}
|
||||
|
||||
struct Level {
|
||||
scale: Scale,
|
||||
proj1: Conv2d,
|
||||
proj2: Conv2d,
|
||||
}
|
||||
|
||||
pub struct Neck {
|
||||
levels: Vec<Level>,
|
||||
}
|
||||
|
||||
impl Neck {
|
||||
/// `vb` at `…neck` (detector) or `tracker_neck`.
|
||||
pub fn new(vb: VarBuilder) -> Result<Self> {
|
||||
let up = ConvTranspose2dConfig { stride: 2, ..Default::default() };
|
||||
let mut levels = Vec::new();
|
||||
for (i, factor) in [4.0f32, 2.0, 1.0, 0.5].into_iter().enumerate() {
|
||||
let l = vb.pp(format!("fpn_layers.{i}"));
|
||||
let s = l.pp("scale_layers");
|
||||
let (scale, inter) = if factor == 4.0 {
|
||||
let a = candle_nn::conv_transpose2d(HIDDEN, HIDDEN / 2, 2, up, s.pp("0"))?;
|
||||
let b = candle_nn::conv_transpose2d(HIDDEN / 2, HIDDEN / 4, 2, up, s.pp("2"))?;
|
||||
(Scale::Up4(a, b), HIDDEN / 4)
|
||||
} else if factor == 2.0 {
|
||||
(Scale::Up2(candle_nn::conv_transpose2d(HIDDEN, HIDDEN / 2, 2, up, s.pp("0"))?), HIDDEN / 2)
|
||||
} else if factor == 1.0 {
|
||||
(Scale::Same, HIDDEN)
|
||||
} else {
|
||||
(Scale::Down2, HIDDEN)
|
||||
};
|
||||
levels.push(Level {
|
||||
scale,
|
||||
proj1: candle_nn::conv2d(inter, FPN, 1, Conv2dConfig::default(), l.pp("proj1"))?,
|
||||
proj2: candle_nn::conv2d(FPN, FPN, 3, Conv2dConfig { padding: 1, ..Default::default() }, l.pp("proj2"))?,
|
||||
});
|
||||
}
|
||||
Ok(Self { levels })
|
||||
}
|
||||
|
||||
/// `x`: the backbone's `[1, 1024, 72, 72]`; `n`: how many levels (from the finest) to compute.
|
||||
pub fn forward(&self, x: &Tensor, n: usize) -> Result<Vec<Tensor>> {
|
||||
let mut out = Vec::with_capacity(n);
|
||||
for l in self.levels.iter().take(n) {
|
||||
let y = match &l.scale {
|
||||
Scale::Up4(a, b) => b.forward(&a.forward(x)?.gelu_erf()?)?,
|
||||
Scale::Up2(a) => a.forward(x)?,
|
||||
Scale::Same => x.clone(),
|
||||
Scale::Down2 => x.max_pool2d(2)?,
|
||||
};
|
||||
out.push(l.proj2.forward(&l.proj1.forward(&y)?)?);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,242 @@
|
|||
//! Shared building blocks: attention, MLPs, layer norms, sine position encodings.
|
||||
//!
|
||||
//! 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`, `models/sam3_tracker`), Copyright The HuggingFace
|
||||
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
|
||||
|
||||
use candle_core::{D, DType, Device, Module, Tensor};
|
||||
use candle_nn::{LayerNorm, Linear, VarBuilder};
|
||||
|
||||
use crate::Result;
|
||||
|
||||
/// `nn.LayerNorm(d)` (PyTorch's default eps).
|
||||
pub const TORCH_LN_EPS: f64 = 1e-5;
|
||||
|
||||
pub fn ln(d: usize, eps: f64, vb: VarBuilder) -> Result<LayerNorm> {
|
||||
Ok(candle_nn::layer_norm(d, eps, vb)?)
|
||||
}
|
||||
|
||||
pub fn linear(i: usize, o: usize, vb: VarBuilder) -> Result<Linear> {
|
||||
Ok(candle_nn::linear(i, o, vb)?)
|
||||
}
|
||||
|
||||
/// Largest attention-score block computed at once (elements): the 5184-token global layers
|
||||
/// would otherwise need ~1.7 GB for all 16 heads.
|
||||
const SCORE_BLOCK: usize = 64 << 20;
|
||||
|
||||
/// Scaled dot-product attention on `[B, H, L, d]` tensors, with an optional additive bias
|
||||
/// broadcastable to `[B, H, Lq, Lk]`. Heads are processed in groups to bound memory.
|
||||
pub fn sdpa(q: &Tensor, k: &Tensor, v: &Tensor, bias: Option<&Tensor>, scale: f64) -> Result<Tensor> {
|
||||
let (b, h, lq, _) = q.dims4()?;
|
||||
let lk = k.dim(2)?;
|
||||
let per_head = b.saturating_mul(lq).saturating_mul(lk).max(1);
|
||||
let group = (SCORE_BLOCK / per_head).clamp(1, h);
|
||||
let bias_heads = match bias {
|
||||
Some(t) => t.dim(1)?,
|
||||
None => 1,
|
||||
};
|
||||
let mut outs = Vec::with_capacity(h.div_ceil(group));
|
||||
let mut start = 0;
|
||||
while start < h {
|
||||
let n = group.min(h - start);
|
||||
let qs = q.narrow(1, start, n)?;
|
||||
let ks = k.narrow(1, start, n)?;
|
||||
let vs = v.narrow(1, start, n)?;
|
||||
let mut att = (qs.matmul(&ks.t()?)? * scale)?;
|
||||
if let Some(bias) = bias {
|
||||
let bs = if bias_heads > 1 { bias.narrow(1, start, n)? } else { bias.clone() };
|
||||
att = att.broadcast_add(&bs)?;
|
||||
}
|
||||
let att = candle_nn::ops::softmax_last_dim(&att)?;
|
||||
outs.push(att.matmul(&vs)?);
|
||||
start += n;
|
||||
}
|
||||
Ok(Tensor::cat(&outs, 1)?)
|
||||
}
|
||||
|
||||
/// Multi-head attention with q/k/v/o projections (`Sam3Attention`, `Sam3TrackerAttention`
|
||||
/// with `internal` = hidden / downsample rate).
|
||||
pub struct Attention {
|
||||
q: Linear,
|
||||
k: Linear,
|
||||
v: Linear,
|
||||
o: Linear,
|
||||
heads: usize,
|
||||
head_dim: usize,
|
||||
}
|
||||
|
||||
impl Attention {
|
||||
pub fn new(hidden: usize, internal: usize, heads: usize, vb: VarBuilder) -> Result<Self> {
|
||||
Self::named(hidden, internal, heads, "o_proj", vb)
|
||||
}
|
||||
|
||||
/// With the output projection called `out` (CLIP's `out_proj`).
|
||||
pub fn named(hidden: usize, internal: usize, heads: usize, out: &str, vb: VarBuilder) -> Result<Self> {
|
||||
Ok(Self {
|
||||
q: linear(hidden, internal, vb.pp("q_proj"))?,
|
||||
k: linear(hidden, internal, vb.pp("k_proj"))?,
|
||||
v: linear(hidden, internal, vb.pp("v_proj"))?,
|
||||
o: linear(internal, hidden, vb.pp(out))?,
|
||||
heads,
|
||||
head_dim: internal / heads,
|
||||
})
|
||||
}
|
||||
|
||||
fn split(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let (b, l, _) = x.dims3()?;
|
||||
Ok(x.reshape((b, l, self.heads, self.head_dim))?.transpose(1, 2)?.contiguous()?)
|
||||
}
|
||||
|
||||
/// The q, k and v projections of `x` (for attention variants that transform them).
|
||||
pub fn qkv(&self, x: &Tensor) -> Result<(Tensor, Tensor, Tensor)> {
|
||||
Ok((self.q.forward(x)?, self.k.forward(x)?, self.v.forward(x)?))
|
||||
}
|
||||
|
||||
/// The output projection.
|
||||
pub fn out(&self, x: &Tensor) -> Result<Tensor> {
|
||||
Ok(self.o.forward(x)?)
|
||||
}
|
||||
|
||||
/// `q`, `k`, `v`: `[B, L, hidden]`; `bias` broadcastable to `[B, heads, Lq, Lk]`.
|
||||
pub fn forward(&self, q: &Tensor, k: &Tensor, v: &Tensor, bias: Option<&Tensor>) -> Result<Tensor> {
|
||||
let (b, lq, _) = q.dims3()?;
|
||||
let q = self.split(&self.q.forward(q)?)?;
|
||||
let k = self.split(&self.k.forward(k)?)?;
|
||||
let v = self.split(&self.v.forward(v)?)?;
|
||||
let o = sdpa(&q, &k, &v, bias, (self.head_dim as f64).powf(-0.5))?;
|
||||
let o = o.transpose(1, 2)?.reshape((b, lq, self.heads * self.head_dim))?;
|
||||
Ok(self.o.forward(&o)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// An additive key-padding bias `[1, 1, 1, L]` (0 for valid tokens, the f32 minimum for padding).
|
||||
pub fn padding_bias(valid: &[bool], device: &Device) -> Result<Tensor> {
|
||||
let v: Vec<f32> = valid.iter().map(|&ok| if ok { 0.0 } else { f32::MIN }).collect();
|
||||
Ok(Tensor::from_vec(v, (1, 1, 1, valid.len()), device)?)
|
||||
}
|
||||
|
||||
/// `Sam3MLP`: fc1 → activation → fc2.
|
||||
pub struct Mlp {
|
||||
fc1: Linear,
|
||||
fc2: Linear,
|
||||
gelu: bool,
|
||||
}
|
||||
|
||||
impl Mlp {
|
||||
pub fn new(hidden: usize, inter: usize, gelu: bool, vb: VarBuilder) -> Result<Self> {
|
||||
Ok(Self { fc1: linear(hidden, inter, vb.pp("fc1"))?, fc2: linear(inter, hidden, vb.pp("fc2"))?, gelu })
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let x = self.fc1.forward(x)?;
|
||||
let x = if self.gelu { x.gelu_erf()? } else { x.relu()? };
|
||||
Ok(self.fc2.forward(&x)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// `Sam3DecoderMLP`: `layer1 … layerN` with ReLU between (2 or 3 layers).
|
||||
pub struct DecoderMlp {
|
||||
layers: Vec<Linear>,
|
||||
}
|
||||
|
||||
impl DecoderMlp {
|
||||
pub fn new(input: usize, hidden: usize, output: usize, n: usize, vb: VarBuilder) -> Result<Self> {
|
||||
let mut layers = Vec::with_capacity(n);
|
||||
for i in 0..n {
|
||||
let (a, b) = (if i == 0 { input } else { hidden }, if i + 1 == n { output } else { hidden });
|
||||
layers.push(linear(a, b, vb.pp(format!("layer{}", i + 1)))?);
|
||||
}
|
||||
Ok(Self { layers })
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let mut x = x.clone();
|
||||
for (i, l) in self.layers.iter().enumerate() {
|
||||
x = l.forward(&x)?;
|
||||
if i + 1 < self.layers.len() {
|
||||
x = x.relu()?;
|
||||
}
|
||||
}
|
||||
Ok(x)
|
||||
}
|
||||
}
|
||||
|
||||
/// `Sam3TrackerFeedForward`: proj_in → act → (layers.i → act)* → proj_out (→ sigmoid).
|
||||
pub struct FeedForward {
|
||||
proj_in: Linear,
|
||||
layers: Vec<Linear>,
|
||||
proj_out: Linear,
|
||||
sigmoid: bool,
|
||||
}
|
||||
|
||||
impl FeedForward {
|
||||
pub fn new(input: usize, hidden: usize, output: usize, n: usize, sigmoid: bool, vb: VarBuilder) -> Result<Self> {
|
||||
let mut layers = Vec::new();
|
||||
for i in 0..n.saturating_sub(2) {
|
||||
layers.push(linear(hidden, hidden, vb.pp(format!("layers.{i}")))?);
|
||||
}
|
||||
Ok(Self { proj_in: linear(input, hidden, vb.pp("proj_in"))?, layers, proj_out: linear(hidden, output, vb.pp("proj_out"))?, sigmoid })
|
||||
}
|
||||
|
||||
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 { candle_nn::ops::sigmoid(&x)? } else { x })
|
||||
}
|
||||
}
|
||||
|
||||
/// Layer norm over the channel axis of an NCHW tensor (`Sam3TrackerLayerNorm`, channels_first).
|
||||
pub fn ln_channels(x: &Tensor, norm: &LayerNorm) -> Result<Tensor> {
|
||||
let x = x.permute((0, 2, 3, 1))?;
|
||||
Ok(norm.forward(&x)?.permute((0, 3, 1, 2))?.contiguous()?)
|
||||
}
|
||||
|
||||
/// The DETR sine embedding of an `h × w` grid, normalized (`Sam3SinePositionEmbedding` with
|
||||
/// `normalize=True`, scale 2π): `[1, 2·feats, h, w]`, y features first.
|
||||
pub fn sine_grid(h: usize, w: usize, feats: usize, device: &Device) -> Result<Tensor> {
|
||||
let scale = 2.0 * std::f64::consts::PI;
|
||||
let dim_t: Vec<f64> = (0..feats).map(|i| 10000f64.powf(2.0 * (i / 2) as f64 / feats as f64)).collect();
|
||||
let mut out = vec![0f32; 2 * feats * h * w];
|
||||
for y in 0..h {
|
||||
let ye = (y + 1) as f64 / (h as f64 + 1e-6) * scale;
|
||||
for x in 0..w {
|
||||
let xe = (x + 1) as f64 / (w as f64 + 1e-6) * scale;
|
||||
for (i, d) in dim_t.iter().enumerate() {
|
||||
let (py, px) = (ye / d, xe / d);
|
||||
let (vy, vx) = if i % 2 == 0 { (py.sin(), px.sin()) } else { (py.cos(), px.cos()) };
|
||||
if let Some(v) = out.get_mut((i * h + y) * w + x) {
|
||||
*v = vy as f32;
|
||||
}
|
||||
if let Some(v) = out.get_mut(((feats + i) * h + y) * w + x) {
|
||||
*v = vx as f32;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Tensor::from_vec(out, (1, 2 * feats, h, w), device)?)
|
||||
}
|
||||
|
||||
/// `inverse_sigmoid` with the reference's clamping (eps 1e-3).
|
||||
pub fn inverse_sigmoid(x: &Tensor) -> Result<Tensor> {
|
||||
let x = x.clamp(0f32, 1f32)?;
|
||||
let x1 = x.clamp(1e-3f32, f32::MAX)?;
|
||||
let x2 = (1.0 - &x)?.clamp(1e-3f32, f32::MAX)?;
|
||||
Ok((x1 / x2)?.log()?)
|
||||
}
|
||||
|
||||
/// `[..., 4]` boxes (cx, cy, w, h) → (x0, y0, x1, y1).
|
||||
pub fn cxcywh_to_xyxy(b: &Tensor) -> Result<Tensor> {
|
||||
let cx = b.narrow(D::Minus1, 0, 1)?;
|
||||
let cy = b.narrow(D::Minus1, 1, 1)?;
|
||||
let w = (b.narrow(D::Minus1, 2, 1)? * 0.5)?;
|
||||
let h = (b.narrow(D::Minus1, 3, 1)? * 0.5)?;
|
||||
Ok(Tensor::cat(&[(&cx - &w)?, (&cy - &h)?, (&cx + &w)?, (&cy + &h)?], D::Minus1)?)
|
||||
}
|
||||
|
||||
/// A tensor's values as `f32` on the CPU.
|
||||
pub fn to_vec(t: &Tensor) -> Result<Vec<f32>> {
|
||||
Ok(t.to_dtype(DType::F32)?.flatten_all()?.to_vec1::<f32>()?)
|
||||
}
|
||||
|
|
@ -0,0 +1,177 @@
|
|||
//! The CLIP byte-pair-encoding tokenizer (lower-cased, `</w>` word ends), padded to the
|
||||
//! 32-token context SAM 3's text encoder takes.
|
||||
//!
|
||||
//! 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 (the CLIP tokenizer used by `models/sam3`), Copyright The HuggingFace
|
||||
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use crate::{Error, Result};
|
||||
|
||||
pub const CONTEXT: usize = 32;
|
||||
const BOS: u32 = 49406;
|
||||
const EOS: u32 = 49407;
|
||||
|
||||
pub struct Tokenizer {
|
||||
vocab: HashMap<String, u32>,
|
||||
ranks: HashMap<(String, String), usize>,
|
||||
byte_char: [char; 256],
|
||||
}
|
||||
|
||||
/// GPT-2's reversible byte → printable character map.
|
||||
fn bytes_to_unicode() -> [char; 256] {
|
||||
let mut map = ['\0'; 256];
|
||||
let mut n = 0u32;
|
||||
for b in 0..=255u8 {
|
||||
let printable = (b'!'..=b'~').contains(&b) || (0xA1..=0xAC).contains(&b) || (0xAE..=0xFF).contains(&b);
|
||||
let c = if printable {
|
||||
u32::from(b)
|
||||
} else {
|
||||
n += 1;
|
||||
255 + n
|
||||
};
|
||||
if let Some(slot) = map.get_mut(usize::from(b)) {
|
||||
*slot = char::from_u32(c).unwrap_or('?');
|
||||
}
|
||||
}
|
||||
map
|
||||
}
|
||||
|
||||
impl Tokenizer {
|
||||
/// From `vocab.json` and `merges.txt` in `dir`.
|
||||
pub fn load(dir: &Path) -> Result<Self> {
|
||||
let read = |name: &str| std::fs::read_to_string(dir.join(name)).map_err(|e| Error::Model(format!("{name}: {e}")));
|
||||
let vocab: HashMap<String, u32> = serde_json::from_str(&read("vocab.json")?).map_err(|e| Error::Model(format!("vocab.json: {e}")))?;
|
||||
let merges = read("merges.txt")?;
|
||||
Ok(Self::new(vocab, &merges))
|
||||
}
|
||||
|
||||
pub fn new(vocab: HashMap<String, u32>, merges: &str) -> Self {
|
||||
let ranks = merges
|
||||
.lines()
|
||||
.filter(|l| !l.starts_with("#version") && !l.trim().is_empty())
|
||||
.filter_map(|l| l.split_once(' ').map(|(a, b)| (a.to_string(), b.to_string())))
|
||||
.enumerate()
|
||||
.map(|(i, p)| (p, i))
|
||||
.collect();
|
||||
Self { vocab, ranks, byte_char: bytes_to_unicode() }
|
||||
}
|
||||
|
||||
/// Token ids (start, words, end, padding) and which positions are real tokens.
|
||||
pub fn encode(&self, text: &str) -> (Vec<u32>, Vec<bool>) {
|
||||
let mut ids = vec![BOS];
|
||||
for word in pre_tokenize(&text.split_whitespace().collect::<Vec<_>>().join(" ").to_lowercase()) {
|
||||
let mapped: String = word.bytes().map(|b| self.byte_char.get(usize::from(b)).copied().unwrap_or('?')).collect();
|
||||
for piece in self.bpe(&mapped) {
|
||||
ids.push(self.vocab.get(&piece).copied().unwrap_or(EOS));
|
||||
}
|
||||
}
|
||||
ids.truncate(CONTEXT - 1);
|
||||
ids.push(EOS);
|
||||
let n = ids.len();
|
||||
ids.resize(CONTEXT, EOS);
|
||||
let valid = (0..CONTEXT).map(|i| i < n).collect();
|
||||
(ids, valid)
|
||||
}
|
||||
|
||||
fn bpe(&self, word: &str) -> Vec<String> {
|
||||
let mut syms: Vec<String> = word.chars().map(String::from).collect();
|
||||
if let Some(last) = syms.last_mut() {
|
||||
last.push_str("</w>");
|
||||
}
|
||||
loop {
|
||||
let best = syms
|
||||
.windows(2)
|
||||
.enumerate()
|
||||
.filter_map(|(i, w)| match w {
|
||||
[a, b] => self.ranks.get(&(a.clone(), b.clone())).map(|r| (*r, i)),
|
||||
_ => None,
|
||||
})
|
||||
.min();
|
||||
let Some((_, i)) = best else { break };
|
||||
let (Some(a), Some(b)) = (syms.get(i).cloned(), syms.get(i + 1).cloned()) else { break };
|
||||
let merged = a.clone() + &b;
|
||||
// merge every occurrence of the pair, left to right
|
||||
let mut out = Vec::with_capacity(syms.len());
|
||||
let mut j = 0;
|
||||
while j < syms.len() {
|
||||
if j + 1 < syms.len() && syms.get(j) == Some(&a) && syms.get(j + 1) == Some(&b) {
|
||||
out.push(merged.clone());
|
||||
j += 2;
|
||||
} else {
|
||||
out.extend(syms.get(j).cloned());
|
||||
j += 1;
|
||||
}
|
||||
}
|
||||
syms = out;
|
||||
if syms.len() == 1 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
syms
|
||||
}
|
||||
}
|
||||
|
||||
/// CLIP's pre-tokenizer: contractions, letter runs, single digits, and runs of other
|
||||
/// non-space characters.
|
||||
fn pre_tokenize(text: &str) -> Vec<String> {
|
||||
let chars: Vec<char> = text.chars().collect();
|
||||
let mut out = Vec::new();
|
||||
let mut i = 0;
|
||||
while let Some(&c) = chars.get(i) {
|
||||
if c.is_whitespace() {
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
if c == '\'' {
|
||||
let rest: String = chars.get(i + 1..(i + 3).min(chars.len())).map(|s| s.iter().collect()).unwrap_or_default();
|
||||
if let Some(k) = ["re", "ve", "ll"]
|
||||
.iter()
|
||||
.find(|k| rest.starts_with(**k))
|
||||
.map(|k| k.len())
|
||||
.or_else(|| rest.chars().next().filter(|c| matches!(c, 's' | 't' | 'm' | 'd')).map(|_| 1))
|
||||
{
|
||||
out.push(chars.get(i..i + 1 + k).map(|s| s.iter().collect()).unwrap_or_default());
|
||||
i += 1 + k;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
let start = i;
|
||||
if c.is_alphabetic() {
|
||||
while chars.get(i).is_some_and(|c| c.is_alphabetic()) {
|
||||
i += 1;
|
||||
}
|
||||
} else if c.is_numeric() {
|
||||
i += 1;
|
||||
} else {
|
||||
while chars.get(i).is_some_and(|c| !c.is_whitespace() && !c.is_alphabetic() && !c.is_numeric()) {
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
out.push(chars.get(start..i).map(|s| s.iter().collect()).unwrap_or_default());
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn splits_like_clip() {
|
||||
assert_eq!(pre_tokenize("a dog's toy, 25 cats!"), ["a", "dog", "'s", "toy", ",", "2", "5", "cats", "!"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merges_by_rank() {
|
||||
let vocab: HashMap<String, u32> =
|
||||
[("c", 1), ("a", 2), ("t</w>", 3), ("ca", 4), ("cat</w>", 5)].into_iter().map(|(k, v)| (k.to_string(), v)).collect();
|
||||
let t = Tokenizer::new(vocab, "#version: 0.2\nc a\nca t</w>\n");
|
||||
let (ids, valid) = t.encode(" CAT ");
|
||||
assert_eq!(&ids[..3], &[BOS, 5, EOS]);
|
||||
assert_eq!(ids.len(), CONTEXT);
|
||||
assert_eq!(valid.iter().filter(|v| **v).count(), 3);
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,253 @@
|
|||
//! Point prompts (SAM 3's interactive "tracker" head, SAM 2 style): a prompt encoder for
|
||||
//! positive/negative clicks and a two-way-transformer mask decoder over the image features.
|
||||
//!
|
||||
//! 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_tracker/modeling_sam3_tracker.py`), Copyright The HuggingFace
|
||||
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
|
||||
|
||||
use candle_core::{DType, Device, IndexOp, Module, Tensor};
|
||||
use candle_nn::{Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, LayerNorm, VarBuilder};
|
||||
|
||||
use crate::nn::{Attention, FeedForward, TORCH_LN_EPS, ln, ln_channels};
|
||||
use crate::vit::{GRID, IMAGE};
|
||||
use crate::{Result, nn::to_vec};
|
||||
|
||||
const HIDDEN: usize = 256;
|
||||
const MASK_TOKENS: usize = 4;
|
||||
/// Low-resolution mask side (4 × the 72 grid).
|
||||
pub const LOW_RES: usize = 288;
|
||||
|
||||
/// The image features the decoder needs, computed once per image.
|
||||
pub struct Features {
|
||||
/// Level 2 (72 × 72, 256 channels).
|
||||
pub image: Tensor,
|
||||
/// `conv_s0`(level 0): `[1, 32, 288, 288]`.
|
||||
pub s0: Tensor,
|
||||
/// `conv_s1`(level 1): `[1, 64, 144, 144]`.
|
||||
pub s1: Tensor,
|
||||
}
|
||||
|
||||
/// One click: position normalized to the image (0..1), and whether it includes or excludes.
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct Click {
|
||||
pub x: f32,
|
||||
pub y: f32,
|
||||
pub positive: bool,
|
||||
}
|
||||
|
||||
/// A predicted mask: 288 × 288 logits over the (squared) model input, and its quality.
|
||||
pub struct Prediction {
|
||||
pub logits: Vec<f32>,
|
||||
pub iou: f32,
|
||||
pub object_score: f32,
|
||||
}
|
||||
|
||||
/// Random-Fourier positional encoding of points normalized to 0..1 (`Sam3TrackerPositionalEmbedding`).
|
||||
struct Fourier {
|
||||
pe: Vec<f32>, // [2, 128]
|
||||
}
|
||||
|
||||
impl Fourier {
|
||||
fn encode(&self, x: f32, y: f32) -> Vec<f32> {
|
||||
let half = self.pe.len() / 2;
|
||||
let (cx, cy) = (2.0 * x - 1.0, 2.0 * y - 1.0);
|
||||
let proj: Vec<f32> = (0..half)
|
||||
.map(|i| {
|
||||
let a = cx * self.pe.get(i).copied().unwrap_or(0.0) + cy * self.pe.get(half + i).copied().unwrap_or(0.0);
|
||||
2.0 * std::f32::consts::PI * a
|
||||
})
|
||||
.collect();
|
||||
proj.iter().map(|a| a.sin()).chain(proj.iter().map(|a| a.cos())).collect()
|
||||
}
|
||||
}
|
||||
|
||||
struct TwoWayBlock {
|
||||
self_attn: Attention,
|
||||
ln1: LayerNorm,
|
||||
t2i: Attention,
|
||||
ln2: LayerNorm,
|
||||
mlp: FeedForward,
|
||||
ln3: LayerNorm,
|
||||
ln4: LayerNorm,
|
||||
i2t: Attention,
|
||||
skip_first_pe: bool,
|
||||
}
|
||||
|
||||
/// `Sam3TrackerAttention` on `[1, L, C]` (the point batch of one is folded away).
|
||||
fn attn(a: &Attention, q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Tensor> {
|
||||
a.forward(q, k, v, None)
|
||||
}
|
||||
|
||||
impl TwoWayBlock {
|
||||
fn forward(&self, queries: &Tensor, keys: &Tensor, qpe: &Tensor, kpe: &Tensor) -> Result<(Tensor, Tensor)> {
|
||||
let queries = if self.skip_first_pe {
|
||||
attn(&self.self_attn, queries, queries, queries)?
|
||||
} else {
|
||||
let q = (queries + qpe)?;
|
||||
(queries + attn(&self.self_attn, &q, &q, queries)?)?
|
||||
};
|
||||
let queries = self.ln1.forward(&queries)?;
|
||||
let (q, k) = ((&queries + qpe)?, (keys + kpe)?);
|
||||
let queries = self.ln2.forward(&(&queries + attn(&self.t2i, &q, &k, keys)?)?)?;
|
||||
let queries = self.ln3.forward(&(&queries + self.mlp.forward(&queries)?)?)?;
|
||||
let (q, k) = ((&queries + qpe)?, (keys + kpe)?);
|
||||
let keys = self.ln4.forward(&(keys + attn(&self.i2t, &k, &q, &queries)?)?)?;
|
||||
Ok((queries, keys))
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Tracker {
|
||||
device: Device,
|
||||
point_pe: Fourier,
|
||||
image_pe: Tensor,
|
||||
point_embed: Tensor,
|
||||
not_a_point: Tensor,
|
||||
no_mask: Tensor,
|
||||
conv_s0: Conv2d,
|
||||
conv_s1: Conv2d,
|
||||
output_tokens: Tensor,
|
||||
blocks: Vec<TwoWayBlock>,
|
||||
final_t2i: Attention,
|
||||
ln_final: LayerNorm,
|
||||
up1: ConvTranspose2d,
|
||||
up_ln: LayerNorm,
|
||||
up2: ConvTranspose2d,
|
||||
hyper: Vec<FeedForward>,
|
||||
iou_head: FeedForward,
|
||||
obj_head: FeedForward,
|
||||
}
|
||||
|
||||
fn read_vec(vb: &VarBuilder, shape: (usize, usize), name: &str) -> Result<Vec<f32>> {
|
||||
to_vec(&vb.get(shape, name)?)
|
||||
}
|
||||
|
||||
impl Tracker {
|
||||
/// `vb` at `tracker_model`.
|
||||
pub fn new(vb: VarBuilder) -> Result<Self> {
|
||||
let device = vb.device().clone();
|
||||
let pe = vb.pp("prompt_encoder");
|
||||
let md = vb.pp("mask_decoder");
|
||||
let image_fourier = Fourier { pe: read_vec(&vb, (2, HIDDEN / 2), "shared_image_embedding.positional_embedding")? };
|
||||
// image-wide positional embedding: pixel centres of the 72 × 72 grid
|
||||
let mut grid = Vec::with_capacity(GRID * GRID * HIDDEN);
|
||||
for r in 0..GRID {
|
||||
for c in 0..GRID {
|
||||
grid.extend(image_fourier.encode((c as f32 + 0.5) / GRID as f32, (r as f32 + 0.5) / GRID as f32));
|
||||
}
|
||||
}
|
||||
let image_pe = Tensor::from_vec(grid, (1, GRID * GRID, HIDDEN), &device)?;
|
||||
let t = md.pp("transformer");
|
||||
let mut blocks = Vec::new();
|
||||
for i in 0..2 {
|
||||
let b = t.pp(format!("layers.{i}"));
|
||||
blocks.push(TwoWayBlock {
|
||||
self_attn: Attention::new(HIDDEN, HIDDEN, 8, b.pp("self_attn"))?,
|
||||
ln1: ln(HIDDEN, TORCH_LN_EPS, b.pp("layer_norm1"))?,
|
||||
t2i: Attention::new(HIDDEN, HIDDEN / 2, 8, b.pp("cross_attn_token_to_image"))?,
|
||||
ln2: ln(HIDDEN, TORCH_LN_EPS, b.pp("layer_norm2"))?,
|
||||
mlp: FeedForward::new(HIDDEN, 2048, HIDDEN, 2, false, b.pp("mlp"))?,
|
||||
ln3: ln(HIDDEN, TORCH_LN_EPS, b.pp("layer_norm3"))?,
|
||||
ln4: ln(HIDDEN, TORCH_LN_EPS, b.pp("layer_norm4"))?,
|
||||
i2t: Attention::new(HIDDEN, HIDDEN / 2, 8, b.pp("cross_attn_image_to_token"))?,
|
||||
skip_first_pe: i == 0,
|
||||
});
|
||||
}
|
||||
let output_tokens = Tensor::cat(
|
||||
&[
|
||||
md.get((1, HIDDEN), "obj_score_token.weight")?,
|
||||
md.get((1, HIDDEN), "iou_token.weight")?,
|
||||
md.get((MASK_TOKENS, HIDDEN), "mask_tokens.weight")?,
|
||||
],
|
||||
0,
|
||||
)?;
|
||||
let up = ConvTranspose2dConfig { stride: 2, ..Default::default() };
|
||||
let one = Conv2dConfig::default();
|
||||
Ok(Self {
|
||||
point_pe: Fourier { pe: read_vec(&pe, (2, HIDDEN / 2), "shared_embedding.positional_embedding")? },
|
||||
image_pe,
|
||||
point_embed: pe.get((4, HIDDEN), "point_embed.weight")?.to_device(&Device::Cpu)?,
|
||||
not_a_point: pe.get((1, HIDDEN), "not_a_point_embed.weight")?.to_device(&Device::Cpu)?,
|
||||
no_mask: pe.get((1, HIDDEN), "no_mask_embed.weight")?.reshape((1, HIDDEN, 1, 1))?,
|
||||
conv_s0: candle_nn::conv2d(HIDDEN, HIDDEN / 8, 1, one, md.pp("conv_s0"))?,
|
||||
conv_s1: candle_nn::conv2d(HIDDEN, HIDDEN / 4, 1, one, md.pp("conv_s1"))?,
|
||||
output_tokens,
|
||||
blocks,
|
||||
final_t2i: Attention::new(HIDDEN, HIDDEN / 2, 8, t.pp("final_attn_token_to_image"))?,
|
||||
ln_final: ln(HIDDEN, TORCH_LN_EPS, t.pp("layer_norm_final_attn"))?,
|
||||
up1: candle_nn::conv_transpose2d(HIDDEN, HIDDEN / 4, 2, up, md.pp("upscale_conv1"))?,
|
||||
up_ln: ln(HIDDEN / 4, 1e-6, md.pp("upscale_layer_norm"))?,
|
||||
up2: candle_nn::conv_transpose2d(HIDDEN / 4, HIDDEN / 8, 2, up, md.pp("upscale_conv2"))?,
|
||||
hyper: (0..MASK_TOKENS)
|
||||
.map(|i| FeedForward::new(HIDDEN, HIDDEN, HIDDEN / 8, 3, false, md.pp(format!("output_hypernetworks_mlps.{i}"))))
|
||||
.collect::<Result<_>>()?,
|
||||
iou_head: FeedForward::new(HIDDEN, HIDDEN, MASK_TOKENS, 3, true, md.pp("iou_prediction_head"))?,
|
||||
obj_head: FeedForward::new(HIDDEN, HIDDEN, 1, 3, false, md.pp("pred_obj_score_head"))?,
|
||||
device,
|
||||
})
|
||||
}
|
||||
|
||||
/// The decoder's image features from the tracker neck's levels (288, 144, 72 per side).
|
||||
pub fn features(&self, levels: &[Tensor]) -> Result<Features> {
|
||||
let [l0, l1, l2, ..] = levels else { return Err(crate::Error::Model("tracker needs three feature levels".into())) };
|
||||
Ok(Features { image: l2.clone(), s0: self.conv_s0.forward(l0)?, s1: self.conv_s1.forward(l1)? })
|
||||
}
|
||||
|
||||
/// Sparse prompt embeddings `[1, n + 1, 256]` (a padding point closes the list).
|
||||
fn embed_clicks(&self, clicks: &[Click]) -> Result<Tensor> {
|
||||
let point_embed = to_vec(&self.point_embed)?;
|
||||
let not_a_point = to_vec(&self.not_a_point)?;
|
||||
let mut rows = Vec::with_capacity((clicks.len() + 1) * HIDDEN);
|
||||
for c in clicks {
|
||||
// the reference shifts pixel coordinates (in the 1008² input) by half a pixel
|
||||
let (x, y) = ((c.x * IMAGE as f32 + 0.5) / IMAGE as f32, (c.y * IMAGE as f32 + 0.5) / IMAGE as f32);
|
||||
let e = self.point_pe.encode(x, y);
|
||||
let label = usize::from(c.positive);
|
||||
rows.extend(e.iter().enumerate().map(|(i, v)| v + point_embed.get(label * HIDDEN + i).copied().unwrap_or(0.0)));
|
||||
}
|
||||
rows.extend_from_slice(¬_a_point);
|
||||
Ok(Tensor::from_vec(rows, (1, clicks.len() + 1, HIDDEN), &self.device)?)
|
||||
}
|
||||
|
||||
/// Predict masks for `clicks`. `multimask`: the three alternatives (best by predicted IoU is
|
||||
/// returned), else the single-mask output with SAM 2's stability fallback.
|
||||
pub fn predict(&self, f: &Features, clicks: &[Click], multimask: bool) -> Result<Prediction> {
|
||||
let sparse = self.embed_clicks(clicks)?;
|
||||
let tokens = Tensor::cat(&[self.output_tokens.unsqueeze(0)?, sparse], 1)?; // [1, T, 256]
|
||||
let image = f.image.broadcast_add(&self.no_mask)?; // dense prompt: "no mask"
|
||||
let (b, c, h, w) = image.dims4()?;
|
||||
let keys = image.flatten_from(2)?.transpose(1, 2)?.contiguous()?; // [1, HW, 256]
|
||||
let (mut q, mut k) = (tokens.clone(), keys);
|
||||
for blk in &self.blocks {
|
||||
(q, k) = blk.forward(&q, &k, &tokens, &self.image_pe)?;
|
||||
}
|
||||
let qf = (&q + &tokens)?;
|
||||
let kf = (&k + &self.image_pe)?;
|
||||
let q = self.ln_final.forward(&(&q + attn(&self.final_t2i, &qf, &kf, &k)?)?)?;
|
||||
let iou_token = q.i((.., 1, ..))?;
|
||||
let image = k.transpose(1, 2)?.reshape((b, c, h, w))?;
|
||||
let up = (self.up1.forward(&image)? + &f.s1)?;
|
||||
let up = ln_channels(&up, &self.up_ln)?.gelu_erf()?;
|
||||
let up = (self.up2.forward(&up)? + &f.s0)?.gelu_erf()?; // [1, 32, 288, 288]
|
||||
let (_, uc, uh, uw) = up.dims4()?;
|
||||
let hyper = (0..MASK_TOKENS)
|
||||
.map(|i| self.hyper.get(i).map_or(Err(crate::Error::Model("hypernetwork".into())), |m| m.forward(&q.i((.., 2 + i, ..))?)))
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
let hyper = Tensor::stack(&hyper, 1)?; // [1, 4, 32]
|
||||
let masks = hyper.matmul(&up.reshape((b, uc, uh * uw))?)?.reshape((b, MASK_TOKENS, uh, uw))?;
|
||||
let iou = to_vec(&self.iou_head.forward(&iou_token)?)?;
|
||||
let object_score = to_vec(&self.obj_head.forward(&q.i((.., 0, ..))?)?)?.first().copied().unwrap_or(0.0);
|
||||
let pick = |i: usize| -> Result<Prediction> {
|
||||
Ok(Prediction { logits: to_vec(&masks.i((0, i))?)?, iou: iou.get(i).copied().unwrap_or(0.0), object_score })
|
||||
};
|
||||
let best_multi = || (1..MASK_TOKENS).max_by(|a, b| iou.get(*a).partial_cmp(&iou.get(*b)).unwrap_or(std::cmp::Ordering::Equal)).unwrap_or(1);
|
||||
if multimask {
|
||||
return pick(best_multi());
|
||||
}
|
||||
// single-mask output unless it is unstable (area changes > 2 % when the threshold moves ±0.05)
|
||||
let single = masks.i((0, 0))?;
|
||||
let inter = single.gt(0.05)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::<f32>()?;
|
||||
let union = single.gt(-0.05)?.to_dtype(DType::F32)?.sum_all()?.to_scalar::<f32>()?;
|
||||
let stability = if union > 0.0 { inter / union } else { 1.0 };
|
||||
if stability >= 0.98 { pick(0) } else { pick(best_multi()) }
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,170 @@
|
|||
//! The SAM 3 vision backbone: a 32-layer ViT (Perception Encoder) with 2-D axial rotary
|
||||
//! position encoding, windowed attention (24 × 24) in most layers and global attention in four.
|
||||
//!
|
||||
//! 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, Module, Tensor};
|
||||
use candle_nn::{Conv2d, Conv2dConfig, LayerNorm, VarBuilder};
|
||||
|
||||
use crate::nn::{Attention, Mlp, ln, sdpa};
|
||||
use crate::{Error, Result};
|
||||
|
||||
pub const IMAGE: usize = 1008;
|
||||
pub const PATCH: usize = 14;
|
||||
/// Patches per side (72).
|
||||
pub const GRID: usize = IMAGE / PATCH;
|
||||
pub const HIDDEN: usize = 1024;
|
||||
const HEADS: usize = 16;
|
||||
const HEAD_DIM: usize = HIDDEN / HEADS;
|
||||
const LAYERS: usize = 32;
|
||||
const WINDOW: usize = 24;
|
||||
const GLOBAL: [usize; 4] = [7, 15, 23, 31];
|
||||
const PRETRAIN_GRID: usize = 336 / PATCH;
|
||||
const EPS: f64 = 1e-6;
|
||||
|
||||
/// Cosine and sine tables `[N, head_dim]` for a `side × side` grid whose coordinates are scaled
|
||||
/// by `scale` (`Sam3ViTRotaryEmbedding`: x frequencies then y, each duplicated pairwise).
|
||||
fn rope_tables(side: usize, scale: f64, device: &Device) -> Result<(Tensor, Tensor)> {
|
||||
let spatial = HEAD_DIM / 2;
|
||||
let inv: Vec<f64> = (0..spatial).step_by(2).map(|i| 1.0 / 10000f64.powf(i as f64 / spatial as f64)).collect();
|
||||
let n = side * side;
|
||||
let (mut cos, mut sin) = (Vec::with_capacity(n * HEAD_DIM), Vec::with_capacity(n * HEAD_DIM));
|
||||
for r in 0..side {
|
||||
for c in 0..side {
|
||||
let (x, y) = (c as f64 * scale, r as f64 * scale);
|
||||
for pos in [x, y] {
|
||||
for f in &inv {
|
||||
let a = (pos * f) as f32 as f64;
|
||||
let (s, co) = (a.sin() as f32, a.cos() as f32);
|
||||
cos.extend([co, co]);
|
||||
sin.extend([s, s]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok((Tensor::from_vec(cos, (n, HEAD_DIM), device)?, Tensor::from_vec(sin, (n, HEAD_DIM), device)?))
|
||||
}
|
||||
|
||||
/// `x·cos + rotate_pairwise(x)·sin`, where rotate_pairwise maps (a, b) pairs to (−b, a).
|
||||
fn apply_rope(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result<Tensor> {
|
||||
let dims = x.dims().to_vec();
|
||||
let last = *dims.last().ok_or_else(|| Error::Model("rope on a scalar".into()))?;
|
||||
let mut paired = dims.clone();
|
||||
paired.pop();
|
||||
paired.extend([last / 2, 2]);
|
||||
let p = x.reshape(paired.as_slice())?;
|
||||
let a = p.narrow(D::Minus1, 0, 1)?;
|
||||
let b = p.narrow(D::Minus1, 1, 1)?;
|
||||
let rot = Tensor::cat(&[b.neg()?, a], D::Minus1)?.reshape(dims.as_slice())?;
|
||||
Ok((x.broadcast_mul(cos)? + rot.broadcast_mul(sin)?)?)
|
||||
}
|
||||
|
||||
struct Layer {
|
||||
ln1: LayerNorm,
|
||||
attn: Attention,
|
||||
ln2: LayerNorm,
|
||||
mlp: Mlp,
|
||||
window: usize,
|
||||
}
|
||||
|
||||
impl Layer {
|
||||
/// `x`: `[B, H, W, C]`.
|
||||
fn forward(&self, x: &Tensor, rope: &(Tensor, Tensor)) -> Result<Tensor> {
|
||||
let (b, h, w, c) = x.dims4()?;
|
||||
let y = self.ln1.forward(x)?;
|
||||
let y = match (h.checked_div(self.window), w.checked_div(self.window)) {
|
||||
(Some(nh), Some(nw)) => {
|
||||
// non-overlapping windows (72 is a multiple of 24: no padding)
|
||||
let ws = self.window;
|
||||
let y = y.reshape((b, nh, ws, nw, ws, c))?.permute((0, 1, 3, 2, 4, 5))?.reshape((b * nh * nw, ws, ws, c))?;
|
||||
let y = self.attend(&y, rope)?;
|
||||
y.reshape((b, nh, nw, ws, ws, c))?.permute((0, 1, 3, 2, 4, 5))?.reshape((b, h, w, c))?
|
||||
}
|
||||
// global attention
|
||||
_ => self.attend(&y, rope)?,
|
||||
};
|
||||
let x = (x + y)?;
|
||||
let y = self.mlp.forward(&self.ln2.forward(&x)?)?;
|
||||
Ok((x + y)?)
|
||||
}
|
||||
|
||||
/// RoPE self-attention over `[B, H, W, C]`.
|
||||
fn attend(&self, x: &Tensor, rope: &(Tensor, Tensor)) -> Result<Tensor> {
|
||||
let (b, h, w, c) = x.dims4()?;
|
||||
let seq = x.reshape((b, h * w, c))?;
|
||||
Ok(self.attn_rope(&seq, rope)?.reshape((b, h, w, c))?)
|
||||
}
|
||||
|
||||
fn attn_rope(&self, x: &Tensor, rope: &(Tensor, Tensor)) -> Result<Tensor> {
|
||||
let (b, n, _) = x.dims3()?;
|
||||
let split = |t: Tensor| -> Result<Tensor> { Ok(t.reshape((b, n, HEADS, HEAD_DIM))?.transpose(1, 2)?.contiguous()?) };
|
||||
let (q, k, v) = self.attn.qkv(x)?;
|
||||
let q = apply_rope(&split(q)?, &rope.0, &rope.1)?;
|
||||
let k = apply_rope(&split(k)?, &rope.0, &rope.1)?;
|
||||
let v = split(v)?;
|
||||
let o = sdpa(&q, &k, &v, None, (HEAD_DIM as f64).powf(-0.5))?;
|
||||
let o = o.transpose(1, 2)?.reshape((b, n, HIDDEN))?;
|
||||
self.attn.out(&o)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Vit {
|
||||
patch: Conv2d,
|
||||
pos: Tensor,
|
||||
ln_pre: LayerNorm,
|
||||
layers: Vec<Layer>,
|
||||
rope_window: (Tensor, Tensor),
|
||||
rope_global: (Tensor, Tensor),
|
||||
}
|
||||
|
||||
impl Vit {
|
||||
/// `vb` at `detector_model.vision_encoder.backbone`.
|
||||
pub fn new(vb: VarBuilder) -> Result<Self> {
|
||||
let device = vb.device().clone();
|
||||
let cfg = Conv2dConfig { stride: PATCH, ..Default::default() };
|
||||
let patch = candle_nn::conv2d_no_bias(3, HIDDEN, PATCH, cfg, vb.pp("embeddings.patch_embeddings.projection"))?;
|
||||
let pos = vb.get((1, PRETRAIN_GRID * PRETRAIN_GRID, HIDDEN), "embeddings.position_embeddings")?;
|
||||
// tile (not interpolate) the 24 × 24 pretraining table over the 72 × 72 grid
|
||||
let reps = GRID.div_ceil(PRETRAIN_GRID);
|
||||
let pos = pos
|
||||
.reshape((1, PRETRAIN_GRID, PRETRAIN_GRID, HIDDEN))?
|
||||
.repeat((1, reps, reps, 1))?
|
||||
.narrow(1, 0, GRID)?
|
||||
.narrow(2, 0, GRID)?
|
||||
.reshape((1, GRID * GRID, HIDDEN))?;
|
||||
let mut layers = Vec::with_capacity(LAYERS);
|
||||
for i in 0..LAYERS {
|
||||
let l = vb.pp(format!("layers.{i}"));
|
||||
layers.push(Layer {
|
||||
ln1: ln(HIDDEN, EPS, l.pp("layer_norm1"))?,
|
||||
attn: Attention::new(HIDDEN, HIDDEN, HEADS, l.pp("attention"))?,
|
||||
ln2: ln(HIDDEN, EPS, l.pp("layer_norm2"))?,
|
||||
mlp: Mlp::new(HIDDEN, 4736, true, l.pp("mlp"))?,
|
||||
window: if GLOBAL.contains(&i) { 0 } else { WINDOW },
|
||||
});
|
||||
}
|
||||
Ok(Self {
|
||||
patch,
|
||||
pos,
|
||||
ln_pre: ln(HIDDEN, EPS, vb.pp("layer_norm"))?,
|
||||
layers,
|
||||
rope_window: rope_tables(WINDOW, 1.0, &device)?,
|
||||
rope_global: rope_tables(GRID, WINDOW as f64 / GRID as f64, &device)?,
|
||||
})
|
||||
}
|
||||
|
||||
/// `pixels`: `[1, 3, 1008, 1008]` normalized to −1..1. Returns `[1, 1024, 72, 72]`.
|
||||
pub fn forward(&self, pixels: &Tensor) -> Result<Tensor> {
|
||||
let x = self.patch.forward(pixels)?; // [1, C, 72, 72]
|
||||
let b = x.dim(0)?;
|
||||
let x = x.flatten_from(2)?.transpose(1, 2)?.broadcast_add(&self.pos)?;
|
||||
let mut x = self.ln_pre.forward(&x.reshape((b, GRID, GRID, HIDDEN))?)?;
|
||||
for l in &self.layers {
|
||||
let rope = if l.window > 0 { &self.rope_window } else { &self.rope_global };
|
||||
x = l.forward(&x, rope)?;
|
||||
}
|
||||
Ok(x.permute((0, 3, 1, 2))?.contiguous()?)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,151 @@
|
|||
//! Safetensors weights read on demand.
|
||||
//!
|
||||
//! Tensors are read from the file when a module asks for them (positional reads, no memory
|
||||
//! map: `unsafe` stays in `lightcraft-sysmem`), so building only the point-prompt model never
|
||||
//! touches the text encoder's bytes.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::fs::File;
|
||||
use std::io::{Read, Seek, SeekFrom};
|
||||
use std::path::Path;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use candle_core::{DType, Device, Shape, Tensor};
|
||||
use candle_nn::var_builder::SimpleBackend;
|
||||
|
||||
use crate::{Error, Result};
|
||||
|
||||
struct Entry {
|
||||
shape: Vec<usize>,
|
||||
start: u64,
|
||||
end: u64,
|
||||
}
|
||||
|
||||
/// An open safetensors file: its header, and the file to read tensor bytes from.
|
||||
pub struct Weights {
|
||||
file: Mutex<File>,
|
||||
entries: HashMap<String, Entry>,
|
||||
data_start: u64,
|
||||
}
|
||||
|
||||
impl Weights {
|
||||
pub fn open(path: &Path) -> Result<Arc<Self>> {
|
||||
let mut file = File::open(path).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?;
|
||||
let file_len = file.metadata().map_err(|e| Error::Model(format!("{}: {e}", path.display())))?.len();
|
||||
let mut len = [0u8; 8];
|
||||
file.read_exact(&mut len).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?;
|
||||
let n = u64::from_le_bytes(len);
|
||||
if n > 100 << 20 || n.saturating_add(8) > file_len {
|
||||
return Err(Error::Model(format!("{}: not a safetensors file (header of {n} bytes)", path.display())));
|
||||
}
|
||||
let mut header = vec![0u8; n as usize];
|
||||
file.read_exact(&mut header).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?;
|
||||
let header: serde_json::Map<String, serde_json::Value> =
|
||||
serde_json::from_slice(&header).map_err(|e| Error::Model(format!("{}: bad header: {e}", path.display())))?;
|
||||
let mut entries = HashMap::new();
|
||||
for (name, v) in header {
|
||||
if name == "__metadata__" {
|
||||
continue;
|
||||
}
|
||||
let dtype = match v.get("dtype").and_then(|d| d.as_str()) {
|
||||
Some("F32") => DType::F32,
|
||||
other => return Err(Error::Model(format!("{name}: unsupported dtype {other:?} (the SAM 3 checkpoint is F32)"))),
|
||||
};
|
||||
let shape: Vec<usize> = v
|
||||
.get("shape")
|
||||
.and_then(|s| s.as_array())
|
||||
.map(|a| a.iter().filter_map(|x| x.as_u64().map(|x| x as usize)).collect())
|
||||
.unwrap_or_default();
|
||||
let offsets: Vec<u64> =
|
||||
v.get("data_offsets").and_then(|s| s.as_array()).map(|a| a.iter().filter_map(|x| x.as_u64()).collect()).unwrap_or_default();
|
||||
let [start, end] = offsets[..] else { return Err(Error::Model(format!("{name}: bad data_offsets"))) };
|
||||
let want = shape.iter().try_fold(dtype.size_in_bytes(), |a, d| a.checked_mul(*d));
|
||||
if end.checked_sub(start) != want.map(|w| w as u64) {
|
||||
return Err(Error::Model(format!("{name}: size does not match its shape")));
|
||||
}
|
||||
// the bytes must be in the file (a damaged or truncated download fails here, not later)
|
||||
if (8 + n).checked_add(end).is_none_or(|e| e > file_len) {
|
||||
return Err(Error::Model(format!("{}: tensor {name} lies outside the file (incomplete or damaged download?)", path.display())));
|
||||
}
|
||||
entries.insert(name, Entry { shape, start, end });
|
||||
}
|
||||
Ok(Arc::new(Self { file: Mutex::new(file), entries, data_start: 8 + n }))
|
||||
}
|
||||
|
||||
pub fn contains(&self, name: &str) -> bool {
|
||||
self.entries.contains_key(name)
|
||||
}
|
||||
|
||||
/// Tensor `name` as stored, on the CPU.
|
||||
pub fn read(&self, name: &str) -> Result<Tensor> {
|
||||
let e = self.entries.get(name).ok_or_else(|| Error::Model(format!("missing tensor `{name}`")))?;
|
||||
// checked against the file's length in `open`
|
||||
let len = usize::try_from(e.end.saturating_sub(e.start)).map_err(|_| Error::Model(format!("{name}: too large")))?;
|
||||
let mut bytes = vec![0u8; len];
|
||||
{
|
||||
let mut f = self.file.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
f.seek(SeekFrom::Start(self.data_start + e.start)).map_err(|err| Error::Model(format!("{name}: {err}")))?;
|
||||
f.read_exact(&mut bytes).map_err(|err| Error::Model(format!("{name}: {err}")))?;
|
||||
}
|
||||
let v: Vec<f32> = bytes.as_chunks::<4>().0.iter().map(|c| f32::from_le_bytes(*c)).collect();
|
||||
Ok(Tensor::from_vec(v, e.shape.as_slice(), &Device::Cpu)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// [`Weights`] as a candle `VarBuilder` backend.
|
||||
pub struct Backend(pub Arc<Weights>);
|
||||
|
||||
impl SimpleBackend for Backend {
|
||||
fn get(&self, s: Shape, name: &str, _h: candle_nn::Init, dtype: DType, dev: &Device) -> candle_core::Result<Tensor> {
|
||||
let t = self.get_unchecked(name, dtype, dev)?;
|
||||
if t.shape() != &s {
|
||||
candle_core::bail!("shape mismatch for {name}: expected {s:?}, got {:?}", t.shape())
|
||||
}
|
||||
Ok(t)
|
||||
}
|
||||
|
||||
fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> candle_core::Result<Tensor> {
|
||||
let t = self.0.read(name).map_err(|e| candle_core::Error::Msg(e.to_string()))?;
|
||||
t.to_dtype(dtype)?.to_device(dev)
|
||||
}
|
||||
|
||||
fn contains_tensor(&self, name: &str) -> bool {
|
||||
self.0.contains(name)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn write(name: &str, header: &str, data_len: usize) -> std::path::PathBuf {
|
||||
let p = std::env::temp_dir().join(format!("lc-weights-{name}-{}.safetensors", std::process::id()));
|
||||
let mut b = (header.len() as u64).to_le_bytes().to_vec();
|
||||
b.extend_from_slice(header.as_bytes());
|
||||
b.extend(std::iter::repeat_n(0u8, data_len));
|
||||
std::fs::write(&p, b).unwrap();
|
||||
p
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tensors_must_lie_inside_the_file() {
|
||||
let ok = write("ok", r#"{"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}}"#, 16);
|
||||
let w = Weights::open(&ok).unwrap();
|
||||
assert_eq!(w.read("a").unwrap().dims(), &[2, 2]);
|
||||
// a truncated download: the header promises bytes that aren't there
|
||||
let short = write("short", r#"{"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}}"#, 8);
|
||||
assert!(Weights::open(&short).err().unwrap().to_string().contains("outside the file"));
|
||||
// hostile offsets: huge, or beyond u64
|
||||
let huge = write("huge", r#"{"a":{"dtype":"F32","shape":[1073741824,1073741824],"data_offsets":[0,4611686018427387904]}}"#, 16);
|
||||
assert!(Weights::open(&huge).is_err());
|
||||
let wrap = write("wrap", r#"{"a":{"dtype":"F32","shape":[1],"data_offsets":[18446744073709551612,18446744073709551616]}}"#, 16);
|
||||
assert!(Weights::open(&wrap).is_err());
|
||||
// a header longer than the file
|
||||
let p = std::env::temp_dir().join(format!("lc-weights-hdr-{}.safetensors", std::process::id()));
|
||||
std::fs::write(&p, 1000u64.to_le_bytes()).unwrap();
|
||||
assert!(Weights::open(&p).is_err());
|
||||
for f in [ok, short, huge, wrap, p] {
|
||||
let _ = std::fs::remove_file(f);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,97 @@
|
|||
//! The port against the reference implementation (Hugging Face `transformers`, fp32).
|
||||
//!
|
||||
//! Needs the checkpoint and a reference file made by `tools/sam3_reference.py`:
|
||||
//! `LIGHTCRAFT_SAM3_DIR=<model dir> LIGHTCRAFT_SAM3_REF=<ref.safetensors> cargo test -p lightcraft-segment --release -- --nocapture`.
|
||||
//! Skipped (passes) when either is missing.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use candle_core::{Device, IndexOp, Tensor};
|
||||
use lightcraft_segment::{Click, Sam3, tokenizer::Tokenizer};
|
||||
|
||||
fn setup() -> Option<(PathBuf, HashMap<String, Tensor>)> {
|
||||
let dir = PathBuf::from(std::env::var_os("LIGHTCRAFT_SAM3_DIR")?);
|
||||
let reference = PathBuf::from(std::env::var_os("LIGHTCRAFT_SAM3_REF")?);
|
||||
let r = candle_core::safetensors::load(&reference, &Device::Cpu).ok()?;
|
||||
Some((dir, r))
|
||||
}
|
||||
|
||||
fn vec(t: &Tensor) -> Vec<f32> {
|
||||
t.to_device(&Device::Cpu).unwrap().flatten_all().unwrap().to_vec1::<f32>().unwrap()
|
||||
}
|
||||
|
||||
/// Max |a − b| relative to the reference's max |b|.
|
||||
fn rel_err(name: &str, a: &[f32], b: &[f32]) -> f32 {
|
||||
assert_eq!(a.len(), b.len(), "{name}: length");
|
||||
let scale = b.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
|
||||
let err = a.iter().zip(b).fold(0f32, |m, (x, y)| m.max((x - y).abs()));
|
||||
eprintln!("{name:>18}: max abs err {err:.2e}, rel {:.2e} (scale {scale:.2e})", err / scale);
|
||||
err / scale
|
||||
}
|
||||
|
||||
/// Fraction of pixels where the two logit maps disagree on the mask (> 0).
|
||||
fn mask_disagreement(a: &[f32], b: &[f32]) -> f32 {
|
||||
a.iter().zip(b).filter(|(x, y)| (**x > 0.0) != (**y > 0.0)).count() as f32 / a.len() as f32
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_reference() {
|
||||
let Some((dir, r)) = setup() else {
|
||||
eprintln!("skipped: set LIGHTCRAFT_SAM3_DIR and LIGHTCRAFT_SAM3_REF");
|
||||
return;
|
||||
};
|
||||
let mut model = Sam3::load(&dir).unwrap();
|
||||
eprintln!("device: {:?}", model.device());
|
||||
let t = std::time::Instant::now();
|
||||
let mut enc = model.encode_pixels(&r["pixel_values"]).unwrap();
|
||||
let backbone = enc.backbone().permute((0, 2, 3, 1)).unwrap();
|
||||
assert!(rel_err("backbone", &vec(&backbone), &vec(&r["backbone"])) < 1e-3);
|
||||
eprintln!("backbone: {:?}", t.elapsed());
|
||||
|
||||
// points: two clicks (single-mask output), one click (best of three)
|
||||
let clicks = [Click { x: 0.5, y: 0.55, positive: true }, Click { x: 0.2, y: 0.2, positive: false }];
|
||||
let t = std::time::Instant::now();
|
||||
let p = model.segment_clicks(&mut enc, &clicks).unwrap();
|
||||
eprintln!("clicks (incl. tracker load + features): {:?}", t.elapsed());
|
||||
let want = vec(&r["trk_masks"]);
|
||||
assert!(rel_err("2-click logits", &p.logits, &want) < 1e-2);
|
||||
assert!(mask_disagreement(&p.logits, &want) < 1e-3);
|
||||
let t = std::time::Instant::now();
|
||||
let p = model.segment_clicks(&mut enc, &clicks[..1]).unwrap();
|
||||
eprintln!("one more click: {:?}", t.elapsed());
|
||||
let iou = vec(&r["trk_multi_iou"]);
|
||||
let best = (0..3).max_by(|a, b| iou[*a].partial_cmp(&iou[*b]).unwrap()).unwrap();
|
||||
let want = vec(&r["trk_multi_masks"].i((0, 0, best)).unwrap());
|
||||
assert!(rel_err("1-click logits", &p.logits, &want) < 1e-2);
|
||||
assert!(mask_disagreement(&p.logits, &want) < 1e-3);
|
||||
|
||||
// tokenizer
|
||||
let tok = Tokenizer::load(&dir).unwrap();
|
||||
let cases = ["person", "Red Car", "a dog's toy", "sky, clouds!", "café 2 people", "tree-line"];
|
||||
for (i, c) in cases.iter().enumerate() {
|
||||
let (ids, _) = tok.encode(c);
|
||||
let want: Vec<i64> = r[&format!("tok{i}")].flatten_all().unwrap().to_vec1().unwrap();
|
||||
assert_eq!(ids.iter().map(|v| *v as i64).collect::<Vec<_>>(), want, "tokens of {c:?}");
|
||||
}
|
||||
|
||||
// text prompts
|
||||
for (i, text) in ["moon", "sky", "bird"].iter().enumerate() {
|
||||
let t = std::time::Instant::now();
|
||||
let out = model.detect(&mut enc, text).unwrap();
|
||||
eprintln!("detect {text:?}: {:?}", t.elapsed());
|
||||
assert!(rel_err("text features", &vec(&out.text), &vec(&r[&format!("det{i}_text")])) < 1e-3);
|
||||
assert!(rel_err("query logits", &out.logits, &vec(&r[&format!("det{i}_logits")])) < 1e-2);
|
||||
let presence = vec(&r[&format!("det{i}_presence")])[0];
|
||||
assert!((out.presence - presence).abs() < 1e-2, "presence {} vs {presence}", out.presence);
|
||||
let boxes: Vec<f32> = out.boxes.iter().flatten().copied().collect();
|
||||
assert!(rel_err("boxes", &boxes, &vec(&r[&format!("det{i}_boxes")])) < 1e-2);
|
||||
// the best instance's mask
|
||||
let logits = vec(&r[&format!("det{i}_logits")]);
|
||||
let q = (0..logits.len()).max_by(|a, b| logits[*a].partial_cmp(&logits[*b]).unwrap()).unwrap();
|
||||
let ours = vec(&out.masks.i(q).unwrap());
|
||||
let want = vec(&r[&format!("det{i}_masks")].i((0, q)).unwrap());
|
||||
assert!(rel_err("best mask", &ours, &want) < 2e-2);
|
||||
assert!(mask_disagreement(&ours, &want) < 2e-3);
|
||||
}
|
||||
}
|
||||
|
|
@ -278,7 +278,9 @@ pub fn handle(app: &mut LightcraftApp, ctx: &egui::Context, req: &ControlRequest
|
|||
Some(d) => {
|
||||
let r = crate::panels::dialogs::confirm_dialog(app, &d);
|
||||
// the import review stays open on an error, as with its button
|
||||
if r.is_err() && matches!(d, crate::state::Dialog::Import { .. }) {
|
||||
if (r.is_err() && matches!(d, crate::state::Dialog::Import { .. } | crate::state::Dialog::SamModel { .. }))
|
||||
|| (r.is_ok() && crate::panels::dialogs::keeps_open(app, &d))
|
||||
{
|
||||
app.ui.dialog = Some(d);
|
||||
}
|
||||
wrap(r)
|
||||
|
|
|
|||
|
|
@ -203,7 +203,9 @@ pub struct LightcraftApp {
|
|||
}
|
||||
|
||||
impl LightcraftApp {
|
||||
pub fn new(session: Session, services: Services) -> Self {
|
||||
pub fn new(mut session: Session, services: Services) -> Self {
|
||||
// AI mask requests run on the model's worker; frames apply their results (never wait)
|
||||
session.segmenter.background = true;
|
||||
Self {
|
||||
session,
|
||||
ui: UiState::default(),
|
||||
|
|
@ -264,6 +266,7 @@ impl LightcraftApp {
|
|||
}
|
||||
let r = self.session.execute(id, ¶ms).map_err(|e| e.to_string());
|
||||
if let Err(e) = &r {
|
||||
log::warn!("{id}: {e}");
|
||||
self.ui.status = e.clone();
|
||||
}
|
||||
r
|
||||
|
|
@ -357,6 +360,84 @@ impl LightcraftApp {
|
|||
self.ui.toast = Some((text.into(), t + secs));
|
||||
}
|
||||
|
||||
/// AI masks: apply finished background requests (clicks, descriptions, detail passes) and
|
||||
/// show their errors; start a zoomed-in detail pass once the clicking stops
|
||||
/// (`ui.detail_due`); watch the model download. Never waits for the model.
|
||||
fn ai_mask_detail(&mut self, ctx: &egui::Context) {
|
||||
let polled = self.session.segment_poll();
|
||||
let now = ctx.input(|i| i.time);
|
||||
if polled.changed {
|
||||
ctx.request_repaint();
|
||||
}
|
||||
if let Some(mask) = polled.refine {
|
||||
self.ui.detail_due = Some((now + 0.3, mask));
|
||||
}
|
||||
if let Some(e) = polled.messages.into_iter().last() {
|
||||
self.ai_error(ctx, e, None);
|
||||
}
|
||||
let seg = &self.session.segmenter;
|
||||
if seg.busy() || seg.pending_clicks().is_some() {
|
||||
ctx.request_repaint_after(std::time::Duration::from_millis(50));
|
||||
}
|
||||
let download = seg.download_status();
|
||||
if download.running {
|
||||
ctx.request_repaint_after(std::time::Duration::from_millis(250));
|
||||
}
|
||||
// a download finishing while its dialog is closed: say so once
|
||||
if self.ui.sam_downloading && !download.running {
|
||||
self.ui.sam_downloading = false;
|
||||
let open = matches!(self.ui.dialog, Some(state::Dialog::SamModel { .. }));
|
||||
match (&download.error, open) {
|
||||
(_, true) => {}
|
||||
(Some(e), false) if e.contains("cancelled") => {
|
||||
self.toast(ctx, crate::i18n::tr("SAM 3 download stopped: it resumes where it left off next time."))
|
||||
}
|
||||
(Some(e), false) => self.toast_error(ctx, format!("The SAM 3 download failed: {e}")),
|
||||
(None, false) if download.finished => {
|
||||
self.toast_error(ctx, crate::i18n::tr("The SAM 3 model is installed: Object and Describe masks are ready."))
|
||||
}
|
||||
(None, false) => {}
|
||||
}
|
||||
}
|
||||
if let Some((due, mask)) = self.ui.detail_due {
|
||||
if now >= due && !self.session.segmenter.busy() && !self.session.segmenter.detail_busy() {
|
||||
self.ui.detail_due = None;
|
||||
if let Err(e) = self.run("mask.refineDetail", serde_json::json!({"id": mask})) {
|
||||
log::warn!("detail pass: {e}");
|
||||
}
|
||||
} else {
|
||||
ctx.request_repaint_after(std::time::Duration::from_millis(150));
|
||||
}
|
||||
}
|
||||
if self.session.segmenter.detail_busy() {
|
||||
ctx.request_repaint_after(std::time::Duration::from_millis(200));
|
||||
}
|
||||
}
|
||||
|
||||
/// An AI mask error: when the model isn't installed, the dialog that offers to download it
|
||||
/// (`then`: the AI mask to start afterwards); otherwise a toast.
|
||||
pub fn ai_error(&mut self, ctx: &egui::Context, e: impl Into<String>, then: Option<(&str, &str)>) {
|
||||
let e = e.into();
|
||||
if lightcraft_engine::segment::Segmenter::AVAILABLE && e.starts_with(lightcraft_engine::segment::NOT_INSTALLED) {
|
||||
self.offer_sam_download(then);
|
||||
} else {
|
||||
self.toast_error(ctx, e);
|
||||
}
|
||||
}
|
||||
|
||||
/// Open the dialog offering the SAM 3 download (never downloads by itself).
|
||||
pub fn offer_sam_download(&mut self, then: Option<(&str, &str)>) {
|
||||
if self.ui.dialog.is_none() || matches!(self.ui.dialog, Some(state::Dialog::SamModel { .. })) {
|
||||
self.ui.dialog = Some(state::Dialog::SamModel { then: then.map(|(k, o)| (k.to_string(), o.to_string())), error: None });
|
||||
}
|
||||
}
|
||||
|
||||
/// A toast for an error the user has to read and act on (stays 6 s).
|
||||
pub fn toast_error(&mut self, ctx: &egui::Context, text: impl Into<String>) {
|
||||
let t = ctx.input(|i| i.time);
|
||||
self.ui.toast = Some((text.into(), t + 6.0));
|
||||
}
|
||||
|
||||
fn drain_control(&mut self, ctx: &egui::Context) {
|
||||
let Some(rx) = self.control_rx.take() else { return };
|
||||
while let Ok(req) = rx.try_recv() {
|
||||
|
|
@ -688,6 +769,7 @@ impl LightcraftApp {
|
|||
}
|
||||
// panels set it again this frame while the pointer rests on a preset or profile
|
||||
self.hover_preview = None;
|
||||
self.ai_mask_detail(&ctx);
|
||||
if self.ui.fullscreen {
|
||||
// full-screen preview: the photo alone on black
|
||||
egui::CentralPanel::default().frame(egui::Frame::NONE.fill(egui::Color32::BLACK)).show(ui, |ui| panels::detail::show(self, ui));
|
||||
|
|
|
|||
|
|
@ -1172,6 +1172,47 @@ fn mask_overlay(app: &mut LightcraftApp, ui: &mut egui::Ui, resp: &egui::Respons
|
|||
}
|
||||
}
|
||||
}
|
||||
// Object selections (SAM 3) of the selected mask: each click, green to include, red to exclude
|
||||
// (clicks still on their way to the model included)
|
||||
let pending = app.session.segmenter.pending_clicks().cloned();
|
||||
for m in d.masks.iter().filter(|m| Some(m.id) == active) {
|
||||
for (k, c) in m.components.iter().enumerate() {
|
||||
if let MaskShape::Object { hint, exclude, .. } = &c.shape {
|
||||
let (hint, exclude) = match pending.as_ref().filter(|q| (q.mask, q.comp) == (m.id, k)) {
|
||||
Some(q) => (&q.hint, &q.exclude),
|
||||
None => (hint, exclude),
|
||||
};
|
||||
for (pts, col) in [(hint, Color32::from_rgb(40, 200, 90)), (exclude, Color32::from_rgb(230, 60, 60))] {
|
||||
for q in pts {
|
||||
let q = map.screen(*q);
|
||||
p.circle_filled(q, 5.0, col);
|
||||
p.circle_stroke(q, 5.0, Stroke::new(1.5, Color32::WHITE));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Object tool: a click includes what's under it, ⌥-click leaves it out
|
||||
if app.ui.tool == "object" {
|
||||
if resp.hovered() {
|
||||
ui.ctx().set_cursor_icon(egui::CursorIcon::Crosshair);
|
||||
}
|
||||
if resp.clicked()
|
||||
&& let Some(q) = resp.interact_pointer_pos()
|
||||
{
|
||||
let n = map.norm(q);
|
||||
let exclude = ui.input(|i| i.modifiers.alt);
|
||||
match app.run("mask.objectPoint", json!({"x": n.x, "y": n.y, "exclude": exclude})) {
|
||||
Ok(_) => {
|
||||
if let Some(mid) = app.session.active_mask {
|
||||
app.ui.detail_due = Some((ui.input(|i| i.time) + 1.0, mid));
|
||||
}
|
||||
}
|
||||
Err(e) => app.ai_error(ui.ctx(), e, Some(("object", "new"))),
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
// brush tool
|
||||
if app.ui.tool == "brush" {
|
||||
let r = (app.ui.brush_size as f64 * long) as f32;
|
||||
|
|
|
|||
|
|
@ -94,6 +94,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
|||
Dialog::Merge { opts } => opts.title(),
|
||||
Dialog::Settings { .. } => "Settings",
|
||||
Dialog::ConfirmDelete { .. } => "Delete Photos",
|
||||
Dialog::SamModel { .. } => "Download the SAM 3 Model?",
|
||||
Dialog::About => "About LightCraft",
|
||||
Dialog::Shortcuts => "Keyboard Shortcuts",
|
||||
}
|
||||
|
|
@ -700,6 +701,7 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
|||
Dialog::Merge { opts } => crate::merge::body(app, ui, opts),
|
||||
Dialog::Import { opts } => crate::import::body(app, ui, opts),
|
||||
Dialog::Settings { tab } => crate::panels::settings::body(app, ui, tab),
|
||||
Dialog::SamModel { error, .. } => sam_model_body(app, ui, error.as_deref()),
|
||||
Dialog::ConfirmDelete { count } => {
|
||||
let what = if *count == 1 { crate::i18n::tr("this photo").to_string() } else { crate::i18n::tr_format!("these {count} photos", count = count) };
|
||||
ui.label(crate::i18n::tr_format!("Move {what} to Recently Deleted?", what = what));
|
||||
|
|
@ -767,8 +769,21 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
|||
ui.add_space(4.0);
|
||||
ui.horizontal(|ui| {
|
||||
let informational = matches!(dlg, Dialog::About | Dialog::Shortcuts | Dialog::Settings { .. });
|
||||
if !informational && ui.button(crate::i18n::tr("Cancel")).clicked() {
|
||||
close = true;
|
||||
let sam = &app.session.segmenter;
|
||||
let (sam_installed, sam_running, sam_failed) = (sam.installed(), sam.download_status().running, sam.download_status().error.is_some());
|
||||
// no download location in this build: nothing to offer but the manual install
|
||||
let sam_nowhere = !sam_installed && !sam_running && sam.mirrors().is_empty();
|
||||
let cancel = match &dlg {
|
||||
Dialog::SamModel { .. } if sam_running || sam_installed || sam_nowhere => "Close",
|
||||
Dialog::SamModel { .. } => "Not Now",
|
||||
_ => "Cancel",
|
||||
};
|
||||
if !informational {
|
||||
let r = ui.button(crate::i18n::tr(cancel));
|
||||
crate::widgets::register(ui.ctx(), "button:dialogCancel", r.rect);
|
||||
if r.clicked() {
|
||||
close = true;
|
||||
}
|
||||
}
|
||||
let add_label;
|
||||
let ok = match &dlg {
|
||||
|
|
@ -780,10 +795,18 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
|||
}
|
||||
Dialog::Merge { .. } => "Merge",
|
||||
Dialog::ConfirmDelete { .. } => "Delete",
|
||||
Dialog::SamModel { then: Some(_), .. } if sam_installed => "Continue",
|
||||
Dialog::SamModel { .. } if sam_installed || sam_running || sam_nowhere => "",
|
||||
Dialog::SamModel { error, .. } if error.is_some() || sam_failed => "Try Again",
|
||||
Dialog::SamModel { .. } => "Download",
|
||||
_ if informational => "Close",
|
||||
_ => "OK",
|
||||
};
|
||||
if ui.button(crate::i18n::tr(ok)).clicked() {
|
||||
let r = (!ok.is_empty()).then(|| ui.button(crate::i18n::tr(ok)));
|
||||
if let Some(r) = &r {
|
||||
crate::widgets::register(ui.ctx(), "button:dialogOk", r.rect);
|
||||
}
|
||||
if r.is_some_and(|r| r.clicked()) {
|
||||
if informational {
|
||||
close = true;
|
||||
} else {
|
||||
|
|
@ -803,6 +826,17 @@ pub fn show(app: &mut LightcraftApp, ctx: &egui::Context) {
|
|||
match confirm_dialog(app, &dlg) {
|
||||
// the import review stays open on an error (e.g. an unusable folder template)
|
||||
Err(e) if matches!(dlg, Dialog::Import { .. }) => app.toast(ctx, e),
|
||||
// the SAM 3 dialog stays open to show the download (or why it can't start)
|
||||
Err(e) if matches!(dlg, Dialog::SamModel { .. }) => {
|
||||
if let Dialog::SamModel { error, .. } = &mut dlg {
|
||||
*error = Some(e);
|
||||
}
|
||||
}
|
||||
Ok(_) if keeps_open(app, &dlg) => {
|
||||
if let Dialog::SamModel { error, .. } = &mut dlg {
|
||||
*error = None;
|
||||
}
|
||||
}
|
||||
_ => close = true,
|
||||
}
|
||||
}
|
||||
|
|
@ -823,8 +857,86 @@ pub fn fmt_gap(v: f64) -> String {
|
|||
}
|
||||
}
|
||||
|
||||
/// Whether a dialog stays open after its action succeeded (the SAM 3 dialog while the model
|
||||
/// downloads).
|
||||
pub fn keeps_open(app: &LightcraftApp, dlg: &Dialog) -> bool {
|
||||
matches!(dlg, Dialog::SamModel { .. }) && !app.session.segmenter.installed()
|
||||
}
|
||||
|
||||
/// The SAM 3 dialog: what the model is, its size and licence, and the download's progress.
|
||||
fn sam_model_body(app: &mut LightcraftApp, ui: &mut egui::Ui, error: Option<&str>) {
|
||||
use lightcraft_engine::segment::{LICENSE_NAME, LICENSE_URL, MODEL_BYTES};
|
||||
let t = Tokens::get(ui.ctx());
|
||||
let seg = &app.session.segmenter;
|
||||
let d = seg.download_status();
|
||||
if seg.installed() {
|
||||
ui.label(crate::i18n::tr("The SAM 3 model is installed: Object and Describe masks are ready."));
|
||||
return;
|
||||
}
|
||||
let gb = |b: u64| b as f64 / 1e9;
|
||||
ui.label(crate::i18n::tr(
|
||||
"Object and Describe masks use SAM 3, Meta's segmentation model. It isn't part of LightCraft, and everything else works without it.",
|
||||
));
|
||||
let dir = seg.dir.as_ref().map(|d| d.display().to_string()).unwrap_or_default();
|
||||
ui.label(format!("{} {:.1} GB, {} {dir}", crate::i18n::tr("A one-time download of about"), gb(MODEL_BYTES), crate::i18n::tr("saved in")));
|
||||
ui.label(
|
||||
egui::RichText::new(format!(
|
||||
"{} {LICENSE_NAME} — {}",
|
||||
crate::i18n::tr("Licence:"),
|
||||
crate::i18n::tr("Meta's terms, not LightCraft's. Downloading it means accepting them.")
|
||||
))
|
||||
.color(t.text_label),
|
||||
);
|
||||
let r = ui.link(crate::i18n::tr("Read the SAM License")).on_hover_text(LICENSE_URL);
|
||||
crate::widgets::register(ui.ctx(), "link:samLicense", r.rect);
|
||||
if r.clicked() {
|
||||
let _ = crate::links::open(app, LICENSE_URL);
|
||||
}
|
||||
let seg = &app.session.segmenter;
|
||||
if d.running {
|
||||
ui.add_space(4.0);
|
||||
let frac = if d.total > 0 { d.done as f64 / d.total as f64 } else { 0.0 };
|
||||
let text = format!("{:.2} / {:.2} GB · {}", gb(d.done), gb(d.total), d.file);
|
||||
let r = ui.add(egui::ProgressBar::new(frac as f32).text(text));
|
||||
crate::widgets::register(ui.ctx(), "progress:samDownload", r.rect);
|
||||
let r = ui.button(crate::i18n::tr("Cancel Download"));
|
||||
crate::widgets::register(ui.ctx(), "button:samCancel", r.rect);
|
||||
if r.clicked() {
|
||||
app.session.segmenter.cancel_download();
|
||||
}
|
||||
ui.label(egui::RichText::new(crate::i18n::tr("You can close this: the download continues, and resumes if interrupted.")).color(t.text_dim));
|
||||
ui.ctx().request_repaint_after(std::time::Duration::from_millis(250));
|
||||
return;
|
||||
}
|
||||
if seg.mirrors().is_empty() {
|
||||
ui.label(
|
||||
egui::RichText::new(crate::i18n::tr(
|
||||
"This build has no download location for the model yet: put the files in the folder above yourself (see docs/ai-masks.md).",
|
||||
))
|
||||
.color(t.text_dim),
|
||||
);
|
||||
}
|
||||
if let Some(e) = error.map(str::to_string).or(d.error) {
|
||||
ui.label(egui::RichText::new(format!("{} {e}", crate::i18n::tr("The download didn't work:"))).color(egui::Color32::from_rgb(230, 90, 80)));
|
||||
}
|
||||
}
|
||||
|
||||
pub fn confirm_dialog(app: &mut LightcraftApp, dlg: &Dialog) -> Result<serde_json::Value, String> {
|
||||
match dlg {
|
||||
Dialog::SamModel { then, .. } => {
|
||||
if app.session.segmenter.installed() {
|
||||
// installed: start what the user was doing
|
||||
if let Some((kind, op)) = then {
|
||||
crate::panels::masking::begin_ai(app, kind, op)?;
|
||||
}
|
||||
return Ok(serde_json::Value::Null);
|
||||
}
|
||||
let r = app.run("segment.model.download", json!({"acknowledged": true}));
|
||||
if r.is_ok() {
|
||||
app.ui.sam_downloading = true;
|
||||
}
|
||||
r
|
||||
}
|
||||
Dialog::NewAlbum { name, folder } => app.run("album.create", json!({"name": name, "folder": folder, "addSelected": !folder})),
|
||||
Dialog::RenameAlbum { id, name } => app.run("album.rename", json!({"id": id, "name": name})),
|
||||
Dialog::TextPrompt { value, command, params, key, .. } => {
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ fn kind_label(s: &MaskShape) -> (&'static str, Icon) {
|
|||
MaskShape::Sky => ("Sky", Icon::Sky),
|
||||
MaskShape::Background => ("Background", Icon::Subject),
|
||||
MaskShape::Object { .. } => ("Object", Icon::Subject),
|
||||
MaskShape::Prompt { .. } => ("Describe", Icon::Subject),
|
||||
MaskShape::People { .. } => ("People", Icon::Subject),
|
||||
MaskShape::Landscape { .. } => ("Landscape", Icon::Sky),
|
||||
}
|
||||
|
|
@ -75,7 +76,9 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) {
|
|||
egui::Frame::NONE.inner_margin(egui::Margin { left: 24, right: 22, top: 0, bottom: 10 }).show(ui, |ui| {
|
||||
ui.label(egui::RichText::new(crate::i18n::tr("Create New Mask")).color(t.text_dim));
|
||||
ui.add_space(6.0);
|
||||
let tiles: [(&str, &str, Icon); 8] = [
|
||||
let tiles: [(&str, &str, Icon); 10] = [
|
||||
("object", "Object", Icon::Subject),
|
||||
("prompt", "Describe", Icon::Subject),
|
||||
("subject", "Subject", Icon::Subject),
|
||||
("sky", "Sky", Icon::Sky),
|
||||
("background", "Background", Icon::Subject),
|
||||
|
|
@ -102,6 +105,8 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) {
|
|||
app.ui.tool = "colorRange".into();
|
||||
app.toast(ui.ctx(), "Click the photo to pick a colour · ⇧-click adds more");
|
||||
}
|
||||
"object" => start_object(app, ui.ctx(), "new"),
|
||||
"prompt" => start_describe(app, "new"),
|
||||
"brush" | "linear" | "radial" => {
|
||||
app.ui.tool = kind.to_string();
|
||||
if *kind != "brush" {
|
||||
|
|
@ -118,6 +123,37 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) {
|
|||
}
|
||||
}
|
||||
});
|
||||
describe_field(app, ui, true);
|
||||
let seg = &app.session.segmenter;
|
||||
let status = if seg.analyzing() {
|
||||
Some("Analyzing the photo for AI masks…")
|
||||
} else if seg.busy() {
|
||||
Some("Selecting…")
|
||||
} else if seg.detail_busy() || app.ui.detail_due.is_some() {
|
||||
Some("Refining the mask's detail…")
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if let Some(status) = status {
|
||||
ui.add_space(4.0);
|
||||
ui.label(egui::RichText::new(crate::i18n::tr(status)).color(t.text_dim));
|
||||
ui.ctx().request_repaint_after(std::time::Duration::from_millis(250));
|
||||
}
|
||||
// the model download, while its dialog is closed
|
||||
let download = seg.download_status();
|
||||
if download.running && app.ui.dialog.is_none() {
|
||||
ui.add_space(4.0);
|
||||
ui.horizontal(|ui| {
|
||||
let pct = if download.total > 0 { download.done as f64 / download.total as f64 } else { 0.0 };
|
||||
let r = ui.add(
|
||||
egui::ProgressBar::new(pct as f32).desired_width(ui.available_width() - 70.0).text(format!("SAM 3: {} %", (pct * 100.0).floor())),
|
||||
);
|
||||
register(ui.ctx(), "maskSamProgress", r.rect);
|
||||
if text_button(ui, "maskSamDetails", crate::i18n::tr("Details"), false).clicked() {
|
||||
app.offer_sam_download(None);
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
divider(ui);
|
||||
// mask list
|
||||
|
|
@ -201,6 +237,18 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) {
|
|||
app.ui.renaming_mask = Some((m.id, m.name.clone()));
|
||||
}
|
||||
resp.context_menu(|ui| mask_menu(app, ui, m.id, &m.name, m.visible, index, count));
|
||||
if sel {
|
||||
// right under the selected mask: + adds to it, − takes away from it (any mask
|
||||
// type; Describe… opens its field here)
|
||||
ui.horizontal(|ui| {
|
||||
ui.add_space(30.0);
|
||||
let plus = text_button(ui, &format!("maskPlus:{}", m.id), "+", false).on_hover_text(crate::i18n::tr("Add to this mask"));
|
||||
egui::Popup::menu(&plus).show(|ui| component_menu(app, ui, "add"));
|
||||
let minus = text_button(ui, &format!("maskMinus:{}", m.id), "−", false).on_hover_text(crate::i18n::tr("Subtract from this mask"));
|
||||
egui::Popup::menu(&minus).show(|ui| component_menu(app, ui, "subtract"));
|
||||
});
|
||||
describe_field(app, ui, false);
|
||||
}
|
||||
}
|
||||
if !d.masks.is_empty() {
|
||||
ui.add_space(6.0);
|
||||
|
|
@ -213,7 +261,10 @@ pub fn show(app: &mut LightcraftApp, ui: &mut egui::Ui, id: PhotoId) {
|
|||
egui::Frame::NONE.inner_margin(egui::Margin { left: 24, right: 22, top: 8, bottom: 8 }).show(ui, |ui| {
|
||||
for (i, c) in m.components.iter().enumerate() {
|
||||
let (kind, icon) = kind_label(&c.shape);
|
||||
let label = c.name.clone().unwrap_or_else(|| kind.to_string());
|
||||
let label = c.name.clone().unwrap_or_else(|| match &c.shape {
|
||||
MaskShape::Prompt { text, .. } => format!("“{text}”"),
|
||||
_ => kind.to_string(),
|
||||
});
|
||||
ui.horizontal(|ui| {
|
||||
let (r, _) = ui.allocate_exact_size(vec2(16.0, 16.0), Sense::hover());
|
||||
paint(ui.painter(), r, icon, t.text_label);
|
||||
|
|
@ -486,6 +537,30 @@ fn range_controls(app: &mut LightcraftApp, ui: &mut egui::Ui, comp: usize, shape
|
|||
}
|
||||
}
|
||||
}
|
||||
MaskShape::Object { edge, .. } | MaskShape::Prompt { edge, .. } => {
|
||||
// Edge: how crisp the selection's border is (−100 hard … 0 as computed … 100 soft)
|
||||
let spec = ControlSpec {
|
||||
id: "edge",
|
||||
label: "Edge",
|
||||
section: Section::Light,
|
||||
min: -100.0,
|
||||
max: 100.0,
|
||||
default: 0.0,
|
||||
step: 1.0,
|
||||
decimals: 0,
|
||||
track: Track::Centered,
|
||||
};
|
||||
let out = slider(ui, &spec, *edge, true, None);
|
||||
let base = shape.clone();
|
||||
apply_slider_out(app, &spec, out, |app, v| {
|
||||
let mut s = base.clone();
|
||||
if let MaskShape::Object { edge, .. } | MaskShape::Prompt { edge, .. } = &mut s {
|
||||
*edge = v.clamp(-100.0, 100.0);
|
||||
}
|
||||
app.run("mask.update", json!({"component": comp, "shape": s}))
|
||||
});
|
||||
ui.label(egui::RichText::new(crate::i18n::tr("− harder border · + softer border")).color(t.text_dim).size(11.0));
|
||||
}
|
||||
MaskShape::ColorRange { samples, refine } => {
|
||||
ui.horizontal(|ui| {
|
||||
ui.label(
|
||||
|
|
@ -560,7 +635,121 @@ fn component_row_menu(app: &mut LightcraftApp, ui: &mut egui::Ui, mask: u32, k:
|
|||
}
|
||||
}
|
||||
|
||||
/// Whether the SAM 3 model is missing in a build that could use it: then the download is
|
||||
/// offered (with `then` to start afterwards) instead of starting an AI mask.
|
||||
fn needs_model(app: &mut LightcraftApp, kind: &str, op: &str) -> bool {
|
||||
let seg = &app.session.segmenter;
|
||||
let missing = lightcraft_engine::segment::Segmenter::AVAILABLE && seg.dir.is_some() && !seg.installed();
|
||||
if missing {
|
||||
app.offer_sam_download(Some((kind, op)));
|
||||
}
|
||||
missing
|
||||
}
|
||||
|
||||
/// Start an Object selection (SAM 3 clicks): a new mask, or a component of the selected one
|
||||
/// combined by `op`; the photo is analyzed meanwhile (in the background).
|
||||
pub(crate) fn start_object(app: &mut LightcraftApp, ctx: &egui::Context, op: &str) {
|
||||
if needs_model(app, "object", op) {
|
||||
return;
|
||||
}
|
||||
let r =
|
||||
if op == "new" { app.run("mask.add", json!({"kind": "object"})) } else { app.run("mask.addComponent", json!({"op": op, "kind": "object"})) };
|
||||
match r {
|
||||
Ok(_) => {
|
||||
app.ui.tool = "object".into();
|
||||
app.toast(ctx, "Click the object to select it · ⌥-click leaves a part out");
|
||||
}
|
||||
Err(e) => app.ai_error(ctx, e, Some(("object", op))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Start an AI mask of `kind` (object|prompt) combined by `op` (new|add|subtract|intersect),
|
||||
/// after the model was installed.
|
||||
pub(crate) fn begin_ai(app: &mut LightcraftApp, kind: &str, op: &str) -> Result<serde_json::Value, String> {
|
||||
match kind {
|
||||
"object" => {
|
||||
let r = if op == "new" {
|
||||
app.run("mask.add", json!({"kind": "object"}))
|
||||
} else {
|
||||
app.run("mask.addComponent", json!({"op": op, "kind": "object"}))
|
||||
}?;
|
||||
app.ui.tool = "object".into();
|
||||
Ok(r)
|
||||
}
|
||||
_ => {
|
||||
app.ui.describe = Some((op.to_string(), String::new()));
|
||||
Ok(serde_json::Value::Null)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Open the Describe field (a new mask, or a component combined by `op`).
|
||||
pub(crate) fn start_describe(app: &mut LightcraftApp, op: &str) {
|
||||
if needs_model(app, "prompt", op) {
|
||||
return;
|
||||
}
|
||||
app.ui.describe = Some((op.to_string(), String::new()));
|
||||
}
|
||||
|
||||
/// The Describe field: type what to select ("sky", "the red car") and press Return.
|
||||
fn describe_field(app: &mut LightcraftApp, ui: &mut egui::Ui, new_mask: bool) {
|
||||
let Some((op, mut text)) = app.ui.describe.clone() else { return };
|
||||
// a new mask's field sits under the tiles; one that combines, under the selected mask
|
||||
if (op == "new") != new_mask {
|
||||
return;
|
||||
}
|
||||
ui.add_space(8.0);
|
||||
let prompt = match op.as_str() {
|
||||
"subtract" => "Describe what to leave out",
|
||||
"intersect" => "Describe what to keep",
|
||||
_ => "Describe what to select",
|
||||
};
|
||||
ui.label(crate::i18n::tr(prompt));
|
||||
let mut submit = false;
|
||||
ui.horizontal(|ui| {
|
||||
let r =
|
||||
ui.add(egui::TextEdit::singleline(&mut text).hint_text("e.g. sky · the red car · car, road").desired_width(ui.available_width() - 64.0));
|
||||
register(ui.ctx(), "maskDescribe", r.rect);
|
||||
if !r.has_focus() && !r.lost_focus() && text.is_empty() {
|
||||
r.request_focus();
|
||||
}
|
||||
if r.lost_focus() && ui.input(|i| i.key_pressed(egui::Key::Enter)) {
|
||||
submit = true;
|
||||
}
|
||||
if text_button(ui, "maskDescribeGo", crate::i18n::tr("Select"), false).clicked() {
|
||||
submit = true;
|
||||
}
|
||||
});
|
||||
if ui.input(|i| i.key_pressed(egui::Key::Escape)) {
|
||||
app.ui.describe = None;
|
||||
return;
|
||||
}
|
||||
app.ui.describe = Some((op.clone(), text.clone()));
|
||||
if submit && !text.trim().is_empty() {
|
||||
let r = if op == "new" {
|
||||
app.run("mask.add", json!({"kind": "prompt", "text": text}))
|
||||
} else {
|
||||
app.run("mask.addComponent", json!({"op": op, "kind": "prompt", "text": text}))
|
||||
};
|
||||
match r {
|
||||
// the mask appears when the model has found it (a detail pass follows)
|
||||
Ok(_) => app.ui.describe = None,
|
||||
Err(e) => app.ai_error(ui.ctx(), e, Some(("prompt", op.as_str()))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn component_menu(app: &mut LightcraftApp, ui: &mut egui::Ui, op: &str) {
|
||||
let b = ui.button(crate::i18n::tr("Object"));
|
||||
register(ui.ctx(), format!("maskComp:{op}:object"), b.rect);
|
||||
if b.clicked() {
|
||||
start_object(app, ui.ctx(), op);
|
||||
}
|
||||
let b = ui.button(crate::i18n::tr("Describe…"));
|
||||
register(ui.ctx(), format!("maskComp:{op}:prompt"), b.rect);
|
||||
if b.clicked() {
|
||||
start_describe(app, op);
|
||||
}
|
||||
for (kind, label) in [
|
||||
("brush", "Brush"),
|
||||
("linear", "Linear Gradient"),
|
||||
|
|
|
|||
|
|
@ -287,6 +287,17 @@ pub struct UiState {
|
|||
/// A mask being renamed in the Masks list: its id and the edited name.
|
||||
#[serde(skip)]
|
||||
pub renaming_mask: Option<(u32, String)>,
|
||||
/// The Describe field (AI mask from a text prompt) while open: how the selection combines
|
||||
/// (`new` mask, or `add`/`subtract`/`intersect` on the selected one) and the text typed.
|
||||
#[serde(skip)]
|
||||
pub describe: Option<(String, String)>,
|
||||
/// A SAM 3 download was started from the app (to report its end once).
|
||||
#[serde(skip)]
|
||||
pub sam_downloading: bool,
|
||||
/// When to start the zoomed-in detail pass of an AI mask (app time) and which mask: set by
|
||||
/// each click or description, so the pass runs once the clicking stops.
|
||||
#[serde(skip)]
|
||||
pub detail_due: Option<(f64, u32)>,
|
||||
/// A mask component being renamed inline: (mask id, component index, name).
|
||||
pub renaming_component: Option<(u32, usize, String)>,
|
||||
/// Close the window on the next frame (File → Quit).
|
||||
|
|
@ -506,6 +517,15 @@ pub enum Dialog {
|
|||
Settings {
|
||||
tab: String,
|
||||
},
|
||||
/// Object and Describe masks need the SAM 3 model, which isn't installed: offer to download
|
||||
/// it (size, licence, progress). `then`: the AI mask to start once it is there (`kind`
|
||||
/// object|prompt, `op` new|add|subtract|intersect).
|
||||
SamModel {
|
||||
then: Option<(String, String)>,
|
||||
/// Why the download couldn't start (shown in the dialog).
|
||||
#[serde(default)]
|
||||
error: Option<String>,
|
||||
},
|
||||
/// Confirm moving photos to Recently Deleted.
|
||||
ConfirmDelete {
|
||||
count: usize,
|
||||
|
|
@ -561,6 +581,9 @@ impl Default for UiState {
|
|||
search: String::new(),
|
||||
focus_search: false,
|
||||
renaming_mask: None,
|
||||
describe: None,
|
||||
detail_due: None,
|
||||
sam_downloading: false,
|
||||
renaming_component: None,
|
||||
quit: false,
|
||||
dragging_photos: None,
|
||||
|
|
|
|||
|
|
@ -620,3 +620,58 @@ fn soft_proofing_flags_out_of_gamut_colours_and_makes_proof_copies() {
|
|||
assert!(!h.app.ui.soft_proof);
|
||||
assert_ne!(h.app.session.catalog.stack_of(copy.id).unwrap().collapsed, stack.collapsed, "S toggled the stack");
|
||||
}
|
||||
|
||||
/// Object / Describe without the SAM 3 model: the download is offered (never started without a
|
||||
/// yes), the dialog says why it can't start when no mirror is configured, nothing freezes and no
|
||||
/// empty mask is left behind.
|
||||
#[test]
|
||||
fn ai_masks_without_the_model_offer_the_download() {
|
||||
use crate::state::Dialog;
|
||||
use lightcraft_engine::segment::Segmenter;
|
||||
let mut h = detail("panel.masking");
|
||||
let dir = std::env::temp_dir().join(format!("lc-ui-no-sam3-{}", std::process::id()));
|
||||
h.app.session.segmenter.dir = Some(dir.clone());
|
||||
h.app.session.segmenter.mirrors_file = Some(dir.join("none.txt"));
|
||||
let t = std::time::Instant::now();
|
||||
let r = h.request("ui.clickWidget", json!({"id": "maskNew:object"}), T);
|
||||
assert_eq!(r["ok"], true, "{r}");
|
||||
assert!(develop(&h).masks.is_empty());
|
||||
if !Segmenter::AVAILABLE {
|
||||
// a build without AI masks says so (a toast), no dialog
|
||||
assert_eq!(h.app.ui.dialog, None);
|
||||
return;
|
||||
}
|
||||
assert_eq!(h.app.ui.dialog, Some(Dialog::SamModel { then: Some(("object".into(), "new".into())), error: None }));
|
||||
assert!(!h.app.session.segmenter.download_status().running, "nothing downloads without a yes");
|
||||
if h.app.session.segmenter.mirrors().is_empty() {
|
||||
// no location configured in this build: no Download button (only Close), and an agent
|
||||
// confirming anyway gets the reason; the dialog stays
|
||||
h.step();
|
||||
let r = h.request("ui.clickWidget", json!({"id": "button:dialogOk"}), T);
|
||||
assert_eq!(r["ok"], false, "{r}");
|
||||
let r = h.request("ui.dialog.confirm", json!({}), T);
|
||||
assert_eq!(r["ok"], false, "{r}");
|
||||
assert!(r["error"].as_str().unwrap_or_default().contains("LIGHTCRAFT_SAM3_MIRRORS"), "{r}");
|
||||
assert!(matches!(h.app.ui.dialog, Some(Dialog::SamModel { .. })), "stays open");
|
||||
} else {
|
||||
// with a mirror: Download starts it in the background (here it fails: nothing listens)
|
||||
let r = h.request("ui.clickWidget", json!({"id": "button:dialogOk"}), T);
|
||||
assert_eq!(r["ok"], true, "{r}");
|
||||
}
|
||||
// (a frame with the message laid out, so the buttons are where they are drawn)
|
||||
h.step();
|
||||
h.step();
|
||||
let r = h.request("ui.clickWidget", json!({"id": "button:dialogCancel"}), T);
|
||||
assert_eq!(r["ok"], true, "{r}");
|
||||
assert_eq!(h.app.ui.dialog, None);
|
||||
// Describe asks too
|
||||
let r = h.request("ui.clickWidget", json!({"id": "maskNew:prompt"}), T);
|
||||
assert_eq!(r["ok"], true, "{r}");
|
||||
assert_eq!(h.app.ui.dialog, Some(Dialog::SamModel { then: Some(("prompt".into(), "new".into())), error: None }));
|
||||
assert_eq!(h.app.ui.describe, None);
|
||||
// a click command from an agent gets the not-installed error at once, and the app offers the model
|
||||
let e = h.app.run("mask.add", json!({"kind": "prompt", "text": "sky"})).unwrap_err();
|
||||
assert!(e.starts_with(lightcraft_engine::segment::NOT_INSTALLED), "{e}");
|
||||
assert!(develop(&h).masks.is_empty());
|
||||
assert!(t.elapsed() < SETTLE, "{:?}", t.elapsed());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,198 @@
|
|||
# AI masks (Object and Describe)
|
||||
|
||||
LightCraft selects objects with **SAM 3** (Segment Anything with Concepts, Meta 2025), run in
|
||||
pure Rust by `crates/segment` on [candle](https://github.com/huggingface/candle): on the GPU
|
||||
through Metal on macOS, on the CPU elsewhere (for now).
|
||||
|
||||
**The model is optional.** It is not part of LightCraft, and nothing requires it: without it,
|
||||
Object and Describe offer to download it, every other feature (and every existing mask, see
|
||||
below) works, and no render, export or command ever waits for it.
|
||||
|
||||
## Using them
|
||||
|
||||
In the Masking panel (M):
|
||||
|
||||
- **Object**: click the thing you want; every click refines the selection. ⌥-click (Alt) a part
|
||||
to leave it out. Clicks show as green (include) and red (leave out) dots.
|
||||
- **Describe**: type what to select — `sky`, `trees`, `the red car` — and press Return. Every
|
||||
instance the model finds (score above 0.5) joins the mask; if nothing matches, nothing is
|
||||
created and you are told so.
|
||||
- Several things at once: `car, road` (comma-separated) selects each and merges them.
|
||||
- **+ / −** under the selected mask (and the Add / Subtract / Intersect menus) combine
|
||||
selections: Describe `car`, click the mask, + ▸ Describe… `road`; − ▸ Describe… `people` takes
|
||||
the people out of a sky mask; − ▸ Object then clicks remove one object.
|
||||
- **Edge** (per selection, −100…100): below 0 a crisper border, above 0 a feathered one (up to 2 %
|
||||
of the long edge, so previews and exports match).
|
||||
- **Detail:** about a second after the last click (or right after a description), the photo
|
||||
around the selection is analyzed again, zoomed in, from a sharper source (~5 s on an M4 Pro,
|
||||
in the background); the result is kept as a high-resolution patch (at most 4 per selection),
|
||||
so small objects get 5–10× finer edges than one pass over the whole photo gives.
|
||||
|
||||
The first Object or Describe on a photo analyzes it (the image encoder runs once per photo and
|
||||
look: ~3 s on an Apple M4 Pro, plus a one-time ~6 s for compiling GPU kernels and loading the
|
||||
model after launch). That starts in the background as soon as you pick Object; after it, each
|
||||
click takes ~35 ms and each description ~0.6 s. Everything the model does runs on its own worker
|
||||
thread: the window stays responsive, a click shows its dot at once and the selection follows
|
||||
when it's computed ("Selecting…" / "Analyzing the photo…" under the mask tiles). The model is
|
||||
unloaded after 10 minutes without use (it takes several GB of memory) and loaded again when
|
||||
needed.
|
||||
|
||||
The result is stored with the mask: the model's 288 × 288 mask logits over the uncropped photo
|
||||
(quantized, compressed, ~10–30 KB). Renders and exports sample them at any resolution and never
|
||||
need the model, so masks survive moving the library to a machine without it, and the web build
|
||||
renders them. Editing the photo's look later does not recompute a mask (click again to update).
|
||||
|
||||
## Getting the model
|
||||
|
||||
### In the app (recommended)
|
||||
|
||||
The first time you pick Object or Describe without the model, LightCraft asks:
|
||||
|
||||
> **Download the SAM 3 model?** Object and Describe masks use SAM 3, Meta's segmentation model.
|
||||
> It isn't part of LightCraft, and everything else works without it. Download it once (about
|
||||
> 3.4 GB) · *folder*. Licence: SAM License (Meta) — Meta's terms, not LightCraft's. Downloading
|
||||
> it means accepting them. [Read the SAM License] — **Not Now** / **Download**
|
||||
|
||||
Nothing is downloaded until you choose **Download**. The download runs in the background (close
|
||||
the dialog and keep working; the Masking panel shows its progress), can be cancelled, resumes
|
||||
where it stopped if interrupted (also after quitting), and tries the next download location when
|
||||
one fails or stalls. Each file is written as `<name>.part` and moved into the model folder only
|
||||
when complete and verified — `model.safetensors` must have the exact size and SHA-256 of Meta's
|
||||
official checkpoint; a file that doesn't match is deleted, never used. If no location works, the
|
||||
dialog says why and LightCraft carries on without the model.
|
||||
|
||||
Agents and scripts use the same commands (control channel, `lightcraft-cli mcp --connect`):
|
||||
`segment.model.status`, `segment.model.download {"acknowledged": true}` (refused without the
|
||||
explicit `acknowledged`: pass it only after the user agreed to the download and the licence),
|
||||
`segment.model.cancel`.
|
||||
|
||||
### Download locations
|
||||
|
||||
The model is fetched from an ordered list of mirrors: base URLs where `<base>/model.safetensors`,
|
||||
`<base>/vocab.json` and `<base>/merges.txt` live. The list is, in order:
|
||||
|
||||
1. `LIGHTCRAFT_SAM3_MIRRORS` (environment variable; URLs separated by commas or spaces);
|
||||
2. the file `models/sam3-mirrors.txt` in LightCraft's settings folder (one URL per line, `#`
|
||||
comments) — `~/Library/Application Support/LightCraft/` (macOS), `%APPDATA%\LightCraft\`
|
||||
(Windows), `~/.config/lightcraft/` (Linux);
|
||||
3. the built-in list, `DEFAULT_MIRRORS` in `crates/segment/src/fetch/mod.rs`.
|
||||
|
||||
> **Maintainers:** the built-in list is **empty** until LightCraft's own CDN locations exist
|
||||
> (see the `TODO(maintainer)` there): add them in order of preference, host the three files
|
||||
> unchanged, and pin `vocab.json` / `merges.txt` (size + SHA-256) in `SAM3_FILES` at the same
|
||||
> time. Until then the in-app download needs a user-configured mirror, and the dialog says so.
|
||||
|
||||
Hugging Face's `facebook/sam3` can't be a default: it is **gated** (each person must accept the
|
||||
licence there, wait for approval and download with their own token). Someone with access can
|
||||
use it as their mirror:
|
||||
|
||||
```sh
|
||||
LIGHTCRAFT_SAM3_MIRRORS=https://huggingface.co/facebook/sam3/resolve/main \
|
||||
LIGHTCRAFT_SAM3_TOKEN=hf_... lightcraft
|
||||
```
|
||||
|
||||
(`LIGHTCRAFT_SAM3_TOKEN` is sent as a bearer token only over https, and only to the mirror's own
|
||||
host — not to the storage host it redirects to.)
|
||||
|
||||
The download is plain HTTP/1.1 over TLS in pure Rust (rustls with the RustCrypto provider and
|
||||
the Mozilla root certificates; no OpenSSL, `ring` or `aws-lc`, nothing compiled from C), with a
|
||||
15 s connect timeout and a 30 s stall timeout; it follows redirects but doesn't use HTTP proxies.
|
||||
|
||||
### By hand
|
||||
|
||||
Put `model.safetensors`, `vocab.json` and `merges.txt` from `facebook/sam3` in the model folder:
|
||||
`models/sam3/` in the settings folder above, or the folder `LIGHTCRAFT_SAM3_DIR` names. Developers
|
||||
can use the installer scripts, which download from Hugging Face with your token and verify the
|
||||
weights:
|
||||
|
||||
```sh
|
||||
HF_TOKEN=hf_... tools/install-sam3.sh # macOS, Linux
|
||||
tools/install-sam3.sh --check # verify an installation
|
||||
```
|
||||
|
||||
```powershell
|
||||
$env:HF_TOKEN = "hf_..."; .\tools\install-sam3.ps1 # Windows
|
||||
```
|
||||
|
||||
## Licence of the model
|
||||
|
||||
LightCraft's code is MIT OR Apache-2.0; `crates/segment` (a port of Apache-2.0 code) is
|
||||
Apache-2.0. **The SAM 3 weights are neither**: they are Meta's, released under the **SAM
|
||||
License** (<https://github.com/facebookresearch/sam3/blob/main/LICENSE>), a custom licence that
|
||||
is not an OSI open-source licence. It grants a non-exclusive, worldwide, non-transferable,
|
||||
royalty-free licence to use, reproduce, distribute, copy, modify and create derivative works of
|
||||
the "SAM Materials", subject to its terms, including: anyone redistributing them must provide a
|
||||
copy of the licence with them, and derivative works are distributed under its terms; uses
|
||||
subject to ITAR, military or warfare purposes, nuclear industries, espionage, guns or illegal
|
||||
weapons, and reverse engineering are prohibited, and trade controls (sanctions, export rules)
|
||||
apply; the materials come "as is", without warranty; Meta may terminate the licence on a breach
|
||||
(then the materials must be deleted); California law governs. Read the licence itself — this
|
||||
summary is not legal advice.
|
||||
|
||||
Therefore:
|
||||
|
||||
- **The weights are never part of this repository, its releases or installers**, and must not be
|
||||
committed.
|
||||
- They reach a user's computer only when that user asks for them, after the dialog names the
|
||||
licence (or when they install them by hand).
|
||||
- Whoever hosts a mirror redistributes the SAM Materials and must follow the licence (ship it
|
||||
next to the files).
|
||||
|
||||
## Commands (control channel, MCP)
|
||||
|
||||
| Command | Params | |
|
||||
|---|---|---|
|
||||
| `mask.add` | `{kind: "object", points?: [[x,y],…], exclude?: [[x,y],…], seg?}` | an Object mask (empty until clicked) |
|
||||
| `mask.add` | `{kind: "prompt", text, seg?}` | a Describe mask |
|
||||
| `mask.addComponent` | `{op: add\|subtract\|intersect, kind: "object"\|"prompt", …}` | the same as a component |
|
||||
| `mask.objectPoint` | `{x, y, exclude?: bool, id?}` | one click on the selected mask's Object selection (≤ 64 clicks) |
|
||||
| `mask.refineDetail` | `{id?, component?}` | the zoomed-in detail pass → `{started}` |
|
||||
| `segment.prepare` | `{}` | load the model and analyze the active photo → `{busy}` |
|
||||
| `segment.model.status` | `{}` | `{available, installed, dir, loaded, busy, analyzing, sizeBytes, license, licenseUrl, mirrors, download}` |
|
||||
| `segment.model.download` | `{acknowledged: true}` | start the download in the background → `{started}` |
|
||||
| `segment.model.cancel` | `{}` | stop it (it resumes next time) → `{cancelled}` |
|
||||
|
||||
Coordinates are normalized to the uncropped, oriented photo, like every mask shape. `seg` passes
|
||||
a stored segmentation (as `mask.list` / the develop settings hold it): the mask is made without
|
||||
the model. Without the model, the AI kinds fail with an error starting "The SAM 3 model is not
|
||||
installed" that says how to get it.
|
||||
|
||||
In the desktop app the commands don't wait for the model: `mask.objectPoint` returns
|
||||
`{pending: true}` and the selection is applied when it's computed (one undo step per settled
|
||||
selection); `mask.add {kind: "prompt"}` returns `{pending: true}` and the mask appears when
|
||||
something is found (or a message says nothing was). `lightcraft-cli` and the headless MCP server
|
||||
wait and return the result (when built with the `sam` feature; the default CLI has no AI masks).
|
||||
|
||||
## Implementation notes
|
||||
|
||||
- `crates/segment` ports the Hugging Face `transformers` implementation (Apache-2.0; see
|
||||
`NOTICE` — the ported files say so and that they were modified): the 32-layer ViT backbone
|
||||
with 2-D RoPE and windowed attention, the feature pyramids, the SAM 2-style prompt encoder and
|
||||
two-way mask decoder for clicks, and the CLIP text encoder, DETR encoder/decoder (box
|
||||
relative-position bias, presence token) and pixel decoder for text. The CLIP tokenizer is a
|
||||
small BPE in `tokenizer.rs`. `fetch/` is the downloader (original code).
|
||||
- candle is pinned at 0.9.2: later releases make `candle-core` depend on `tokenizers` with the
|
||||
Oniguruma C library, and the product is pure Rust. Weights are read with positional reads (no
|
||||
memory map: `unsafe` stays in `lightcraft-sysmem`), only the tensors a path needs, and every
|
||||
tensor's byte range is checked against the file's length when it is opened (a truncated or
|
||||
damaged file is an error, never a crash).
|
||||
- The engine (`crates/engine/src/segment/`) owns one worker thread with the model and the last
|
||||
encoded photo; the session sends it jobs and receives results over channels, never holding a
|
||||
lock the model holds. Each job runs under `catch_unwind`: a panic inside candle becomes an
|
||||
error and the model is reloaded on the next request; busy counters are reset by drop guards.
|
||||
- Hostile documents: a selection keeps at most 4 detail patches (when read and when rendered),
|
||||
each at most 1024² logits; more than 64 clicks are never sent to the model.
|
||||
- Accuracy: `crates/segment/tests/reference.rs` compares every stage against the Python
|
||||
reference (`tools/sam3_reference.py`); on Metal the backbone matches to 7e-5 relative error
|
||||
and masks to ~5e-5. f16 was tried and rejected (overflow, ~7 % faster).
|
||||
- The engine feature `sam` (on in the desktop app, off for web and CLI) gates the model; without
|
||||
it the commands report that the build has no AI masks.
|
||||
|
||||
**Developers:** the accuracy test against the Python reference needs `torch` and
|
||||
`transformers` (`tools/requirements-sam3-reference.txt`):
|
||||
|
||||
```sh
|
||||
python3 -m venv .venv-sam3 && .venv-sam3/bin/pip install -r tools/requirements-sam3-reference.txt
|
||||
.venv-sam3/bin/python tools/sam3_reference.py photo.jpg ref.safetensors
|
||||
LIGHTCRAFT_SAM3_DIR=<model dir> LIGHTCRAFT_SAM3_REF=ref.safetensors cargo test -p lightcraft-segment --release -- --nocapture
|
||||
```
|
||||
|
|
@ -31,7 +31,7 @@ and prints the summary; `cargo xtask parity --write` refreshes the summary table
|
|||
| H. Crop & rotate (CROP) | 9 | 0 | 0 | 1 | 6/6 (100%) | 3/3 (100%) |
|
||||
| I. Remove / healing (REM) | 7 | 1 | 2 | 2 | 4/4 (100%) | 2/3 (67%) |
|
||||
| J. Red eye (EYE) | 2 | 0 | 0 | 0 | — | 1/1 (100%) |
|
||||
| K. Masking (MASK) | 14 | 3 | 6 | 0 | 8/8 (100%) | 5/5 (100%) |
|
||||
| K. Masking (MASK) | 14 | 4 | 5 | 0 | 8/8 (100%) | 5/5 (100%) |
|
||||
| L. Presets (PRE) | 6 | 0 | 1 | 1 | 2/2 (100%) | 2/2 (100%) |
|
||||
| M. Versions & history (VER) | 5 | 0 | 0 | 0 | 1/1 (100%) | 3/3 (100%) |
|
||||
| N. Copy / paste / sync (SYNC) | 5 | 0 | 0 | 0 | 3/3 (100%) | 1/1 (100%) |
|
||||
|
|
@ -48,9 +48,9 @@ and prints the summary; `cargo xtask parity --write` refreshes the summary table
|
|||
| Y. Menus | 81 | 0 | 4 | 8 | 47/47 (100%) | 23/23 (100%) |
|
||||
| Z. Keyboard shortcuts (desktop) | 74 | 3 | 3 | 1 | 49/52 (94%) | 22/23 (96%) |
|
||||
| Lightroom Classic extras | 31 | 10 | 39 | 9 | — | 21/22 (95%) |
|
||||
| **Total** | 388 | 33 | 87 | 37 | 194/200 (97%) | 139/149 (93%) |
|
||||
| **Total** | 388 | 34 | 86 | 37 | 194/200 (97%) | 139/149 (93%) |
|
||||
|
||||
Weighted completion (✅ = 1, 🟡 = ½, 🚫 left out): **79.6%** of 508 in-scope rows — P0 98.5% of 200 · P1 95.6% of 149 · P2 40.9% of 159.
|
||||
Weighted completion (✅ = 1, 🟡 = ½, 🚫 left out): **79.7%** of 508 in-scope rows — P0 98.5% of 200 · P1 95.6% of 149 · P2 41.2% of 159.
|
||||
<!-- /parity:summary -->
|
||||
|
||||
## Top gaps
|
||||
|
|
@ -69,8 +69,10 @@ Take the first one nobody is working on.
|
|||
then tune tone, highlights, texture/clarity/dehaze, NR and sharpening against it.
|
||||
5. **LR-EDIT-OPTICS-PROFILE** (P1): a lens-profile database of our own (embedded DNG/maker corrections work today).
|
||||
6. **AI masks and Enhance** (LR-MASK-SUBJECT / SKY / PEOPLE / OBJECTS, LR-EDIT-DETAIL-DENOISE, SUPERRES, LENSBLUR):
|
||||
blocked on a model strategy (permissively licensed weights or our own training, pure-Rust inference). A maintainer
|
||||
decision, not just engineering.
|
||||
Object and Describe masks run SAM 3 in pure Rust (`crates/segment`); the weights (SAM License) are never bundled:
|
||||
the app offers a consented, verified download, but **LightCraft's CDN mirrors are not configured yet** (the
|
||||
built-in list is empty, so today users need their own mirror or a manual install). Subject / Sky / People could
|
||||
use the same model with fixed prompts; denoise / super-resolution models remain a maintainer decision.
|
||||
7. **HDR** (Q. HDR, LR-EXP-HDR), **video** (R. Video), **Classic output modules** (Map view, Book, Slideshow module,
|
||||
Print, publish): large, well understood, lower priority than 1–5.
|
||||
|
||||
|
|
@ -328,7 +330,7 @@ under the UI command that wraps it (a no-op error); the UI command now wins.
|
|||
| LR-MASK-SUBJECT | Select subject | P2 | 🟡 | `cmd:mask.add` (`subject`), `crates/pipeline/src/masks.rs` | saliency heuristic, no segmentation model |
|
||||
| LR-MASK-SKY | Select sky | P2 | 🟡 | `cmd:mask.add` (`sky`) | heuristic |
|
||||
| LR-MASK-BACKGROUND | Select background | P2 | 🟡 | `cmd:mask.add` (`background`) | inverse of the subject heuristic |
|
||||
| LR-MASK-OBJECTS | Object selection | P2 | ⬜ | | shape exists (falls back to the subject heuristic); no UI |
|
||||
| LR-MASK-OBJECTS | Object selection | P2 | 🟡 | `cmd:mask.add` (`object`, `prompt`), `cmd:mask.objectPoint`, `cmd:mask.refineDetail`, `cmd:segment.prepare`, `cmd:segment.model.status`, `cmd:segment.model.download`, `cmd:segment.model.cancel`, `crates/segment`, `crates/engine/src/segment/mod.rs` | SAM 3 in pure Rust (candle; Metal on macOS, CPU elsewhere): Object tile → click to include, ⌥-click to leave out; Describe tile → a text prompt selects every instance ("sky", "the red car"); both also as Add/Subtract/Intersect components; + / − under the selected mask; comma lists (`car, road`); a zoomed-in detail pass for 5–10× finer edges on small objects; per-selection Edge (hard ↔ soft). The model runs on its own worker thread (the UI never waits; panics become errors), is unloaded after 10 min idle, and is optional: the segmentation is stored with the mask (288² logits), so renders and exports never need it. Weights are never bundled (SAM License): without them the app offers a consented background download (mirrors, resume, timeouts, SHA-256) — **but no default download location is configured yet** (users set `LIGHTCRAFT_SAM3_MIRRORS` or install by hand; see docs/ai-masks.md). Not verified against Lightroom's Select Object; no brush/box object mode; the first click on a photo waits for its analysis (~4 s on an M4 Pro, much longer on CPU) |
|
||||
| LR-MASK-PEOPLE | People parts | P2 | ⬜ | | |
|
||||
| LR-MASK-LANDSCAPE | Landscape classes | P2 | ⬜ | | shape exists, evaluates empty |
|
||||
| LR-MASK-BRUSH | Brush mask | P0 | ✅ | `cmd:tool.brush`, `cmd:mask.brushStroke` (`autoMask`), `crates/pipeline/src/masks.rs` | size/feather/flow/density/erase; Auto Mask: dabs weighted by similarity to the colour under the dab centre, refined by a guided filter on luminance (CPU + GPU); no A/B brushes, no pressure |
|
||||
|
|
|
|||
|
|
@ -131,6 +131,12 @@
|
|||
lost everything at quit; a temporary session shows a banner the whole time and never writes to your library.
|
||||
|
||||
### Editing
|
||||
- AI masks with SAM 3 (Object and Describe in the Masking panel): click an object to select it (⌥-click leaves a
|
||||
part out), or type what to select ("sky", "the red car", "car, road"); both combine with other masks, have an
|
||||
Edge setting, and get a sharper zoomed-in pass in the background. The model runs inside LightCraft in pure Rust
|
||||
and never freezes the window. It is optional and not part of LightCraft (Meta's SAM License, about 3.4 GB): the
|
||||
first time you use an AI mask, LightCraft asks before downloading it, shows the progress, can cancel and resume,
|
||||
and checks the file before using it. Masks keep their selection, so they render and export without the model.
|
||||
- Auto Sync: edits apply to every selected photo. Auto B&W mix. Automatic versions.
|
||||
- Colour-range masks: click the photo to sample. Luminance ranges: range bar, smoothness, luminance map.
|
||||
- ⌘-drag to straighten, ⇧G Guided Upright, a grid while transforming.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,60 @@
|
|||
# Install the SAM 3 model for LightCraft's Object and Describe masks on Windows
|
||||
# (see docs/ai-masks.md; tools/install-sam3.sh is the macOS/Linux version). The desktop app
|
||||
# offers to download the model itself when an AI mask first needs it; this is for developers.
|
||||
#
|
||||
# facebook/sam3 is gated: accept the SAM License at https://huggingface.co/facebook/sam3, wait
|
||||
# for approval, then: $env:HF_TOKEN = "hf_..." ; .\tools\install-sam3.ps1
|
||||
param(
|
||||
[string]$Dir = $(if ($env:LIGHTCRAFT_SAM3_DIR) { $env:LIGHTCRAFT_SAM3_DIR } else { Join-Path $env:APPDATA "LightCraft\models\sam3" }),
|
||||
[string]$Repo = "facebook/sam3",
|
||||
[switch]$Check
|
||||
)
|
||||
$ErrorActionPreference = "Stop"
|
||||
$Sha256 = "6d06f0a5f84e435071fe6603e61d0b4cc7b40e0d39d487cfd4d67d8cc11cc14a"
|
||||
$Files = "config.json", "vocab.json", "merges.txt", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "processor_config.json"
|
||||
|
||||
function Test-Install {
|
||||
foreach ($f in "model.safetensors", "vocab.json", "merges.txt") {
|
||||
if (-not (Test-Path (Join-Path $Dir $f))) { throw "missing: $(Join-Path $Dir $f)" }
|
||||
}
|
||||
Write-Host "checking model.safetensors (SHA-256 of 3.4 GB)..."
|
||||
$got = (Get-FileHash (Join-Path $Dir "model.safetensors") -Algorithm SHA256).Hash.ToLower()
|
||||
if ($got -ne $Sha256) { throw "model.safetensors does not match the official checkpoint (got $got)" }
|
||||
Write-Host "SAM 3 is installed in $Dir"
|
||||
}
|
||||
|
||||
if ($Check) { Test-Install; exit 0 }
|
||||
|
||||
$Token = if ($env:HF_TOKEN) { $env:HF_TOKEN } else { $tf = Join-Path $HOME ".cache\huggingface\token"; if (Test-Path $tf) { (Get-Content $tf -Raw).Trim() } }
|
||||
if (-not $Token -and $Repo -eq "facebook/sam3") {
|
||||
throw "No Hugging Face token. Accept the license at https://huggingface.co/facebook/sam3, then set HF_TOKEN or run 'hf auth login'."
|
||||
}
|
||||
New-Item -ItemType Directory -Force -Path $Dir | Out-Null
|
||||
$Headers = @{}
|
||||
if ($Token) { $Headers["Authorization"] = "Bearer $Token" }
|
||||
$Base = "https://huggingface.co/$Repo/resolve/main"
|
||||
foreach ($f in $Files + "model.safetensors") {
|
||||
Write-Host "fetching $f"
|
||||
# into a .part file, moved into place only when complete (and, for the weights, verified):
|
||||
# a failed or damaged download never ends up under the real name
|
||||
$final = Join-Path $Dir $f
|
||||
$part = "$final.part"
|
||||
Remove-Item -Force -ErrorAction SilentlyContinue $part
|
||||
try {
|
||||
Invoke-WebRequest -Uri "$Base/$f" -Headers $Headers -OutFile $part
|
||||
} catch {
|
||||
Remove-Item -Force -ErrorAction SilentlyContinue $part
|
||||
throw "download of $Repo/$f failed: $($_.Exception.Message)"
|
||||
}
|
||||
if ($f -eq "model.safetensors") {
|
||||
Write-Host "checking model.safetensors (SHA-256 of 3.4 GB)..."
|
||||
$got = (Get-FileHash $part -Algorithm SHA256).Hash.ToLower()
|
||||
if ($got -ne $Sha256) {
|
||||
Remove-Item -Force $part
|
||||
throw "the download is damaged or not the official checkpoint (got $got); deleted it"
|
||||
}
|
||||
}
|
||||
Move-Item -Force $part $final
|
||||
}
|
||||
Test-Install
|
||||
Write-Host "Restart LightCraft; Object and Describe in the Masking panel now use it."
|
||||
|
|
@ -0,0 +1,130 @@
|
|||
#!/bin/sh
|
||||
# Install the SAM 3 model for LightCraft's Object and Describe masks (see docs/ai-masks.md).
|
||||
#
|
||||
# The desktop app offers to download the model itself when an AI mask first needs it; this script
|
||||
# is for developers (reference tests, benchmarks) and for installing from Hugging Face directly.
|
||||
#
|
||||
# The weights are not part of LightCraft: they are Meta's facebook/sam3 checkpoint on Hugging
|
||||
# Face, under the SAM License. Access is gated: open https://huggingface.co/facebook/sam3, accept
|
||||
# the license, wait for approval, then run this with a token that can read it:
|
||||
#
|
||||
# HF_TOKEN=hf_... tools/install-sam3.sh
|
||||
#
|
||||
# (or log in once with `hf auth login`; the token it saves is used). Downloads resume when
|
||||
# interrupted, and the 3.4 GB weights are checked against the official SHA-256.
|
||||
#
|
||||
# Options:
|
||||
# --dir DIR where to install (default: the folder LightCraft looks in; or LIGHTCRAFT_SAM3_DIR)
|
||||
# --repo NAME another Hugging Face repo with the same files (default: facebook/sam3)
|
||||
# --check only verify an existing installation
|
||||
set -eu
|
||||
|
||||
REPO="facebook/sam3"
|
||||
CHECK_ONLY=0
|
||||
SHA256="6d06f0a5f84e435071fe6603e61d0b4cc7b40e0d39d487cfd4d67d8cc11cc14a"
|
||||
SIZE=3439938512
|
||||
FILES="config.json vocab.json merges.txt tokenizer.json tokenizer_config.json special_tokens_map.json processor_config.json"
|
||||
|
||||
case "$(uname -s)" in
|
||||
Darwin) DEFAULT_DIR="$HOME/Library/Application Support/LightCraft/models/sam3" ;;
|
||||
*) DEFAULT_DIR="${XDG_CONFIG_HOME:-$HOME/.config}/lightcraft/models/sam3" ;;
|
||||
esac
|
||||
DIR="${LIGHTCRAFT_SAM3_DIR:-$DEFAULT_DIR}"
|
||||
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--dir) DIR="$2"; shift 2 ;;
|
||||
--repo) REPO="$2"; shift 2 ;;
|
||||
--check) CHECK_ONLY=1; shift ;;
|
||||
-h|--help) sed -n '2,20p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
|
||||
*) echo "unknown option: $1 (see --help)" >&2; exit 2 ;;
|
||||
esac
|
||||
done
|
||||
|
||||
sha256() {
|
||||
if command -v shasum >/dev/null 2>&1; then shasum -a 256 "$1" | cut -d' ' -f1
|
||||
else sha256sum "$1" | cut -d' ' -f1; fi
|
||||
}
|
||||
|
||||
verify() {
|
||||
for f in model.safetensors vocab.json merges.txt; do
|
||||
[ -s "$DIR/$f" ] || { echo "missing: $DIR/$f" >&2; return 1; }
|
||||
done
|
||||
echo "checking model.safetensors (SHA-256 of 3.4 GB, takes a few seconds)…"
|
||||
got=$(sha256 "$DIR/model.safetensors")
|
||||
if [ "$got" != "$SHA256" ]; then
|
||||
echo "model.safetensors does not match the official checkpoint (got $got)" >&2
|
||||
return 1
|
||||
fi
|
||||
echo "SAM 3 is installed in $DIR"
|
||||
}
|
||||
|
||||
if [ "$CHECK_ONLY" = 1 ]; then verify; exit $?; fi
|
||||
|
||||
command -v curl >/dev/null 2>&1 || { echo "curl is required" >&2; exit 1; }
|
||||
|
||||
TOKEN="${HF_TOKEN:-${HUGGING_FACE_HUB_TOKEN:-}}"
|
||||
if [ -z "$TOKEN" ]; then
|
||||
for t in "${HF_HOME:-$HOME/.cache/huggingface}/token" "$HOME/.huggingface/token"; do
|
||||
[ -s "$t" ] && TOKEN=$(cat "$t") && break
|
||||
done
|
||||
fi
|
||||
if [ -z "$TOKEN" ] && [ "$REPO" = "facebook/sam3" ]; then
|
||||
echo "No Hugging Face token found. facebook/sam3 is gated: accept the license at" >&2
|
||||
echo "https://huggingface.co/facebook/sam3, then set HF_TOKEN or run \`hf auth login\`." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
mkdir -p "$DIR"
|
||||
BASE="https://huggingface.co/$REPO/resolve/main"
|
||||
|
||||
# fetch FILE DEST [resume]: HTTP errors never reach DEST (--fail writes no error page), and a
|
||||
# resume the server refuses (416) starts over instead of counting as done.
|
||||
fetch() {
|
||||
set -- "$1" "$2" "${3:-}"
|
||||
if [ -n "$TOKEN" ]; then auth="Authorization: Bearer $TOKEN"; else auth="X-No-Auth: 1"; fi
|
||||
code=$(curl -fsSL ${3:+-C -} -H "$auth" -w '%{http_code}' -o "$2" "$BASE/$1") || true
|
||||
case "$code" in
|
||||
200|206) return 0 ;;
|
||||
416)
|
||||
echo "the partial download of $1 doesn't match the server's; starting over" >&2
|
||||
rm -f "$2"
|
||||
code=$(curl -fsSL -H "$auth" -w '%{http_code}' -o "$2" "$BASE/$1") || true
|
||||
[ "$code" = 200 ] && return 0
|
||||
echo "download of $REPO/$1 failed ($code)" >&2; return 1 ;;
|
||||
401|403) echo "access denied to $REPO/$1 (has your request for facebook/sam3 been approved?)" >&2; return 1 ;;
|
||||
*) echo "download of $REPO/$1 failed (${code:-no response})" >&2; return 1 ;;
|
||||
esac
|
||||
}
|
||||
|
||||
for f in $FILES; do
|
||||
echo "fetching $f"
|
||||
# small files: whole, into place only when complete
|
||||
fetch "$f" "$DIR/$f.part" && mv -f "$DIR/$f.part" "$DIR/$f" || { rm -f "$DIR/$f.part"; exit 1; }
|
||||
done
|
||||
|
||||
if [ -s "$DIR/model.safetensors" ] && [ "$(sha256 "$DIR/model.safetensors")" = "$SHA256" ]; then
|
||||
echo "model.safetensors already installed"
|
||||
else
|
||||
echo "fetching model.safetensors (3.4 GB; resumes if interrupted)"
|
||||
part="$DIR/model.safetensors.part"
|
||||
# a wrong model.safetensors is never resumed from: it goes
|
||||
rm -f "$DIR/model.safetensors"
|
||||
if [ -f "$part" ] && [ "$(wc -c < "$part" | tr -d ' ')" -ge "$SIZE" ]; then rm -f "$part"; fi
|
||||
fetch model.safetensors "$part" resume || exit 1
|
||||
have=$(wc -c < "$part" | tr -d ' ')
|
||||
if [ "$have" -lt "$SIZE" ]; then
|
||||
echo "incomplete download ($have of $SIZE bytes); run the script again to resume" >&2
|
||||
exit 1
|
||||
fi
|
||||
echo "checking model.safetensors (SHA-256 of 3.4 GB, takes a few seconds)…"
|
||||
if [ "$have" != "$SIZE" ] || [ "$(sha256 "$part")" != "$SHA256" ]; then
|
||||
rm -f "$part"
|
||||
echo "the download is damaged or not the official checkpoint; deleted it, run the script again" >&2
|
||||
exit 1
|
||||
fi
|
||||
mv -f "$part" "$DIR/model.safetensors"
|
||||
fi
|
||||
|
||||
verify
|
||||
echo "Restart LightCraft; Object and Describe in the Masking panel now use it."
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
# Python packages for tools/sam3_reference.py (reference tensors for crates/segment/tests/reference.rs).
|
||||
# python3 -m venv .venv-sam3 && .venv-sam3/bin/pip install -r tools/requirements-sam3-reference.txt
|
||||
torch>=2.6
|
||||
torchvision>=0.21
|
||||
transformers>=4.57
|
||||
pillow
|
||||
safetensors
|
||||
numpy
|
||||
|
|
@ -0,0 +1,71 @@
|
|||
"""Reference tensors for the Rust SAM 3 port (crates/segment/tests/reference.rs).
|
||||
|
||||
Runs the Hugging Face transformers implementation (fp32, CPU) on one photo and writes the
|
||||
shared pixel input plus intermediate and final outputs of the click (tracker) and text
|
||||
(detector) paths. Usage (needs torch, torchvision, transformers, pillow, safetensors):
|
||||
LIGHTCRAFT_SAM3_DIR=<model dir> python tools/sam3_reference.py photo.jpg ref.safetensors
|
||||
The text prompts below suit a photo of the moon in a blue sky; change them for other photos.
|
||||
"""
|
||||
import os, sys, time
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
from safetensors.torch import save_file
|
||||
from transformers import Sam3TrackerModel, Sam3Model, CLIPTokenizer
|
||||
|
||||
MODEL = os.environ.get("LIGHTCRAFT_SAM3_DIR", os.path.expanduser("~/Library/Application Support/LightCraft/models/sam3"))
|
||||
img_path, out_path = sys.argv[1], sys.argv[2]
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
img = Image.open(img_path).convert("RGB")
|
||||
W, H = img.size
|
||||
# Same preprocessing as the Rust side: bilinear resize to 1008², /255, (x-0.5)/0.5.
|
||||
r = img.resize((1008, 1008), Image.BILINEAR)
|
||||
px = torch.from_numpy(np.asarray(r, dtype=np.float32) / 255.0).permute(2, 0, 1)[None]
|
||||
px = (px - 0.5) / 0.5
|
||||
out = {"pixel_values": px.contiguous()}
|
||||
|
||||
t = time.time()
|
||||
trk = Sam3TrackerModel.from_pretrained(MODEL).eval()
|
||||
vis = trk.vision_encoder(px)
|
||||
out["backbone"] = vis.last_hidden_state.contiguous()
|
||||
for i, f in enumerate(vis.fpn_hidden_states):
|
||||
out[f"trk_fpn{i}"] = f.contiguous()
|
||||
# points in 1008 space: one positive, one negative
|
||||
pts = torch.tensor([[[[0.5 * 1008, 0.55 * 1008], [0.2 * 1008, 0.2 * 1008]]]], dtype=torch.float32)
|
||||
lbl = torch.tensor([[[1, 0]]], dtype=torch.int32)
|
||||
o = trk(pixel_values=px, input_points=pts, input_labels=lbl, multimask_output=False)
|
||||
out["trk_points"] = pts.contiguous()
|
||||
out["trk_masks"] = o.pred_masks.contiguous()
|
||||
out["trk_iou"] = o.iou_scores.contiguous()
|
||||
o3 = trk(pixel_values=px, input_points=pts[:, :, :1], input_labels=lbl[:, :, :1], multimask_output=True)
|
||||
out["trk_multi_masks"] = o3.pred_masks.contiguous()
|
||||
out["trk_multi_iou"] = o3.iou_scores.contiguous()
|
||||
print("tracker", time.time() - t, o.pred_masks.shape, o.iou_scores, file=sys.stderr)
|
||||
del trk
|
||||
|
||||
t = time.time()
|
||||
det = Sam3Model.from_pretrained(MODEL).eval()
|
||||
tok = CLIPTokenizer.from_pretrained(MODEL)
|
||||
for i, text in enumerate(["moon", "sky", "bird"]):
|
||||
enc = tok(text, padding="max_length", max_length=32, return_tensors="pt")
|
||||
out[f"det{i}_ids"] = enc.input_ids.to(torch.int64).contiguous()
|
||||
out[f"det{i}_mask"] = enc.attention_mask.to(torch.int64).contiguous()
|
||||
tf = det.get_text_features(input_ids=enc.input_ids, attention_mask=enc.attention_mask)
|
||||
out[f"det{i}_text"] = tf.pooler_output.contiguous()
|
||||
o = det(pixel_values=px, input_ids=enc.input_ids, attention_mask=enc.attention_mask)
|
||||
out[f"det{i}_logits"] = o.pred_logits.contiguous()
|
||||
out[f"det{i}_presence"] = o.presence_logits.contiguous()
|
||||
out[f"det{i}_boxes"] = o.pred_boxes.contiguous()
|
||||
out[f"det{i}_masks"] = o.pred_masks.contiguous()
|
||||
sc = (o.pred_logits.sigmoid() * o.presence_logits.sigmoid())[0]
|
||||
print(text, "presence", o.presence_logits.sigmoid().item(), "top", sc.topk(3).values.tolist(), file=sys.stderr)
|
||||
for i, f in enumerate(det.vision_encoder(px).fpn_hidden_states):
|
||||
out[f"det_fpn{i}"] = f.contiguous()
|
||||
print("detector", time.time() - t, file=sys.stderr)
|
||||
# tokenizer cases for the Rust BPE
|
||||
cases = ["person", "Red Car", "a dog's toy", "sky, clouds!", "café 2 people", "tree-line"]
|
||||
for i, c in enumerate(cases):
|
||||
out[f"tok{i}"] = tok(c, padding="max_length", max_length=32, return_tensors="pt").input_ids.to(torch.int64).contiguous()
|
||||
save_file(out, out_path, metadata={"image": img_path, "W": str(W), "H": str(H), "cases": "|".join(cases)})
|
||||
print("saved", out_path, file=sys.stderr)
|
||||
|
|
@ -49,6 +49,7 @@ pub const TABLE: &[(&str, Class)] = &[
|
|||
("preview", Class::Layer(3)),
|
||||
("export", Class::Layer(3)),
|
||||
("merge", Class::Layer(3)),
|
||||
("segment", Class::Layer(3)),
|
||||
("engine", Class::Layer(4)),
|
||||
("ui-egui", Class::Layer(5)),
|
||||
("mcp", Class::Layer(5)),
|
||||
|
|
|
|||
Loading…
Reference in New Issue