diff --git a/crates/engine/src/cmd/masks.rs b/crates/engine/src/cmd/masks.rs index 1615fb06..dd1fa77c 100644 --- a/crates/engine/src/cmd/masks.rs +++ b/crates/engine/src/cmd/masks.rs @@ -432,6 +432,7 @@ pub fn specs() -> Vec { "available": crate::segment::Segmenter::AVAILABLE, "installed": g.installed(), "dir": g.dir.as_ref().map(|d| d.display().to_string()), + "remote": g.remote_endpoint(), "loaded": g.loaded(), "busy": g.busy(), "analyzing": g.analyzing(), diff --git a/crates/engine/src/segment/mod.rs b/crates/engine/src/segment/mod.rs index 3cf14712..ce927878 100644 --- a/crates/engine/src/segment/mod.rs +++ b/crates/engine/src/segment/mod.rs @@ -157,7 +157,14 @@ 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 remote_endpoint(&self) -> Option { + #[cfg(feature = "sam")] + return lightcraft_segment::remote::configured(); + #[cfg(not(feature = "sam"))] + None + } + + /// Whether the local model files are in place. pub fn installed(&self) -> bool { #[cfg(feature = "sam")] if let Some(dir) = &self.dir { @@ -230,6 +237,10 @@ impl Segmenter { if !Self::AVAILABLE { return Err("AI masks are not available in this build".into()); } + #[cfg(feature = "sam")] + if lightcraft_segment::remote::configured().is_some() { + return Ok(PathBuf::new()); + } let dir = self.dir.clone().ok_or("no folder is set for the SAM 3 model")?; if self.installed() { return Ok(dir); diff --git a/crates/engine/src/segment/worker.rs b/crates/engine/src/segment/worker.rs index 9ded0fbe..cb16ab6b 100644 --- a/crates/engine/src/segment/worker.rs +++ b/crates/engine/src/segment/worker.rs @@ -17,7 +17,10 @@ use std::sync::mpsc::{self, RecvTimeoutError}; use std::time::Duration; use lightcraft_develop::SegMask; -use lightcraft_segment::{Encoded, Sam3}; +use lightcraft_segment::{ + Encoded, Sam3, + remote::{self, Operation, Request, text_logits}, +}; use super::{Input, Prompt, Tag}; use crate::media::RenderJob; @@ -175,6 +178,7 @@ impl Worker { struct State { model: Option<(PathBuf, Sam3)>, cache: Option<(u64, Encoded)>, + remote: Option, } /// Whether `later` makes `job` pointless (it asks for the same selection again). @@ -195,7 +199,7 @@ fn run(rx: &mpsc::Receiver, shared: &Shared, idle: Duration) { match rx.recv_timeout(idle) { Ok(j) => queue.push_back(j), Err(RecvTimeoutError::Timeout) => { - if state.model.is_some() { + if state.model.is_some() || state.remote.is_some() { state = State::default(); shared.loaded.store(false, Ordering::SeqCst); log::info!("SAM 3 unloaded after {} minutes without use", idle.as_secs() / 60); @@ -228,7 +232,7 @@ fn run(rx: &mpsc::Receiver, shared: &Shared, idle: Duration) { if result.is_err() && !is_detail { state.cache = None; } - shared.loaded.store(state.model.is_some(), Ordering::SeqCst); + shared.loaded.store(state.model.is_some() || state.remote.is_some(), Ordering::SeqCst); let _ = job.reply.send(Outcome { tag: job.tag, result, superseded: false }); } } @@ -280,29 +284,68 @@ fn encoded<'a>(state: &'a mut State, shared: &Shared, job: &mut Job) -> Result<( log::info!("SAM 3 encoded {}×{} in {:?}", input.w, input.h, t.elapsed()); state.cache = Some((job.key, enc)); } - let State { model, cache } = state; + 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>, String> { - let mut merged: Option> = 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(), - }); - } +fn remote_job(state: &mut State, shared: &Shared, job: &mut Job, endpoint: &str) -> Result, String> { + let _busy = Flag::raise(&shared.encoding); + if state.remote.is_none() { + state.remote = Some(remote::Client::connect(endpoint)?); } - Ok(merged) + let client = state.remote.as_mut().ok_or("remote SAM unavailable")?; + let mut region = None; + let operation = match std::mem::replace(&mut job.kind, Kind::Prepare) { + Kind::Prepare => Operation::Prepare, + Kind::Text(text) => Operation::Text(text), + Kind::Clicks(clicks) => Operation::Clicks(clicks.iter().map(|c| [c.x, c.y, if c.positive { 1.0 } else { 0.0 }]).collect()), + Kind::Detail { prompt, region: r } => { + region = Some(r); + let (rw, rh) = (r[2] - r[0], r[3] - r[1]); + if !(rw > 0.0 && rh > 0.0) { + return Ok(None); + } + match prompt { + Prompt::Text(text) => Operation::Text(text), + Prompt::Clicks(clicks) => { + let inside: Vec<_> = clicks + .iter() + .map(|c| [((c.at.x - r[0]) / rw) as f32, ((c.at.y - r[1]) / rh) as f32, if c.include { 1.0 } else { 0.0 }]) + .filter(|p| (0.0..=1.0).contains(&p[0]) && (0.0..=1.0).contains(&p[1])) + .collect(); + if !inside.iter().any(|p| p[2] == 1.0) { + return Ok(None); + } + Operation::Clicks(inside) + } + } + } + #[cfg(test)] + Kind::Panic => return Err("test panic request".into()), + }; + let input = if region.is_some() || client.key != Some(job.key) { render(job)? } else { Input { rgb: Vec::new(), w: 0, h: 0 } }; + let request = Request { width: input.w, height: input.h, temporary: region.is_some(), operation }; + let logits = client.request(&request, &input.rgb)?; + if region.is_none() { + client.key = Some(job.key); + } + Ok(logits.map(|l| match region { + Some(r) => SegMask::from_logits_in(lightcraft_segment::MASK_SIDE, &l, r), + None => SegMask::from_logits(lightcraft_segment::MASK_SIDE, &l), + })) } fn run_job(state: &mut State, shared: &Shared, job: &mut Job) -> Result, String> { + if let Some(endpoint) = remote::configured() { + let result = remote_job(state, shared, job, &endpoint); + if result.is_err() { + state.remote = None; + } + return result.map_err(|e| format!("Remote SAM 3 failed: {e}. Check the worker and SSH tunnel, then retry.")); + } let side = lightcraft_segment::MASK_SIDE; match std::mem::replace(&mut job.kind, Kind::Prepare) { Kind::Prepare => { diff --git a/crates/segment/src/bin/lightcraft-sam3-worker.rs b/crates/segment/src/bin/lightcraft-sam3-worker.rs new file mode 100644 index 00000000..9929d2eb --- /dev/null +++ b/crates/segment/src/bin/lightcraft-sam3-worker.rs @@ -0,0 +1,8 @@ +fn main() -> Result<(), Box> { + let mut args = std::env::args().skip(1); + let dir = args.next().ok_or("usage: lightcraft-sam3-worker MODEL_DIR [127.0.0.1:8793]")?; + let address = args.next().unwrap_or_else(|| "127.0.0.1:8793".into()).parse()?; + let device = candle_core::Device::new_metal(0)?; + lightcraft_segment::remote::serve(address, std::path::Path::new(&dir), device)?; + Ok(()) +} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 42aed538..a4c3acad 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -32,6 +32,8 @@ mod clip; mod detector; #[cfg(not(target_arch = "wasm32"))] pub mod fetch; +#[cfg(not(target_arch = "wasm32"))] +pub mod remote; pub mod mask; #[cfg(not(target_arch = "wasm32"))] mod neck; diff --git a/crates/segment/src/remote.rs b/crates/segment/src/remote.rs new file mode 100644 index 00000000..cdee096f --- /dev/null +++ b/crates/segment/src/remote.rs @@ -0,0 +1,207 @@ +//! Private SAM 3 transport. Bind the worker to loopback and forward it over SSH. +use std::io::{self, Read, Write}; +use std::net::{SocketAddr, TcpListener, TcpStream}; +use std::path::Path; +use std::time::Duration; + +use crate::{Click, Encoded, INPUT, MASK_SIDE, Sam3}; +use serde::{Deserialize, Serialize}; + +const CONTROL_LIMIT: usize = 16 * 1024; +const RESPONSE_LIMIT: usize = 4 * 1024 * 1024; +const TIMEOUT: Duration = Duration::from_secs(9 * 60); +pub const ENV: &str = "LIGHTCRAFT_SAM3_REMOTE"; + +#[derive(Serialize, Deserialize)] +pub enum Operation { + Prepare, + Clicks(Vec<[f32; 3]>), + Text(String), +} + +#[derive(Serialize, Deserialize)] +pub struct Request { + pub width: usize, + pub height: usize, + pub temporary: bool, + pub operation: Operation, +} + +impl Request { + pub fn validate(&self) -> Result<(), String> { + if self.width > INPUT || self.height > INPUT || (self.width == 0) != (self.height == 0) { + return Err("invalid SAM image dimensions".into()); + } + if self.temporary && self.width == 0 { + return Err("detail pass requires an image".into()); + } + match &self.operation { + Operation::Prepare => {} + Operation::Text(text) if text.len() <= 4096 && text.split([',', ';']).any(|p| !p.trim().is_empty()) => {} + Operation::Clicks(clicks) + if !clicks.is_empty() + && clicks.len() <= 64 + && clicks.iter().all(|c| { + c.iter().all(|v| v.is_finite()) && (0.0..=1.0).contains(&c[0]) && (0.0..=1.0).contains(&c[1]) && (c[2] == 0.0 || c[2] == 1.0) + }) + && clicks.iter().any(|c| c[2] == 1.0) => {} + _ => return Err("invalid SAM prompt".into()), + } + Ok(()) + } +} + +pub fn configured() -> Option { + std::env::var(ENV).ok().filter(|v| !v.trim().is_empty()) +} + +pub fn read_frame(reader: &mut impl Read, limit: usize) -> io::Result> { + let mut size = [0; 4]; + reader.read_exact(&mut size)?; + let size = u32::from_be_bytes(size) as usize; + if size > limit { + return Err(io::Error::new(io::ErrorKind::InvalidData, "SAM frame exceeds limit")); + } + let mut data = vec![0; size]; + reader.read_exact(&mut data)?; + Ok(data) +} + +fn write_frame(writer: &mut impl Write, data: &[u8]) -> io::Result<()> { + let size = u32::try_from(data.len()).map_err(io::Error::other)?; + writer.write_all(&size.to_be_bytes())?; + writer.write_all(data) +} + +pub struct Client { + stream: TcpStream, + pub key: Option, +} + +impl Client { + pub fn connect(endpoint: &str) -> Result { + let address: SocketAddr = endpoint.parse().map_err(|e| format!("SAM endpoint must be an IP:port: {e}"))?; + if !address.ip().is_loopback() { + return Err("SAM endpoint must be loopback; use an SSH tunnel for remote access".into()); + } + let stream = TcpStream::connect_timeout(&address, Duration::from_secs(5)).map_err(|e| e.to_string())?; + stream.set_read_timeout(Some(TIMEOUT)).map_err(|e| e.to_string())?; + stream.set_write_timeout(Some(Duration::from_secs(30))).map_err(|e| e.to_string())?; + Ok(Self { stream, key: None }) + } + + pub fn request(&mut self, request: &Request, rgb: &[u8]) -> Result>, String> { + request.validate()?; + if rgb.len() != request.width * request.height * 3 { + return Err("SAM RGB buffer size does not match dimensions".into()); + } + let bytes = serde_json::to_vec(request).map_err(|e| e.to_string())?; + write_frame(&mut self.stream, &bytes).map_err(|e| e.to_string())?; + write_frame(&mut self.stream, rgb).map_err(|e| e.to_string())?; + let response = read_frame(&mut self.stream, RESPONSE_LIMIT).map_err(|e| e.to_string())?; + let result: Result>, String> = serde_json::from_slice(&response).map_err(|e| e.to_string())?; + let logits = result?; + if logits.as_ref().is_some_and(|v| v.len() != MASK_SIDE * MASK_SIDE || v.iter().any(|n| !n.is_finite())) { + return Err("invalid SAM mask response".into()); + } + Ok(logits) + } +} + +pub fn text_logits(model: &mut Sam3, enc: &mut Encoded, text: &str) -> Result>, String> { + let mut merged: Option> = 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())? { + if let Some(m) = &mut merged { + for (a, b) in m.iter_mut().zip(p.logits) { + *a = a.max(b); + } + } else { + merged = Some(p.logits); + } + } + } + Ok(merged) +} + +fn predict(model: &mut Sam3, enc: &mut Encoded, operation: Operation) -> Result>, String> { + match operation { + Operation::Prepare => { + model.prepare_clicks(enc).map_err(|e| e.to_string())?; + Ok(None) + } + Operation::Text(text) => text_logits(model, enc, &text), + Operation::Clicks(points) => { + let clicks: Vec<_> = points.into_iter().map(|p| Click { x: p[0], y: p[1], positive: p[2] == 1.0 }).collect(); + Ok(Some(model.segment_clicks(enc, &clicks).map_err(|e| e.to_string())?.logits)) + } + } +} + +fn serve_client(stream: &mut TcpStream, dir: &Path, device: candle_core::Device) -> Result<(), String> { + stream.set_read_timeout(Some(Duration::from_secs(10 * 60))).map_err(|e| e.to_string())?; + stream.set_write_timeout(Some(Duration::from_secs(30))).map_err(|e| e.to_string())?; + let mut model = None; + let mut cache = None; + loop { + let frame = read_frame(stream, CONTROL_LIMIT).map_err(|e| e.to_string())?; + let request: Request = serde_json::from_slice(&frame).map_err(|e| e.to_string())?; + request.validate()?; + let rgb = read_frame(stream, request.width * request.height * 3).map_err(|e| e.to_string())?; + if rgb.len() != request.width * request.height * 3 { + return Err("invalid RGB payload".into()); + } + let started = std::time::Instant::now(); + let result = (|| { + if model.is_none() { + model = Some(Sam3::load_on(dir, device.clone()).map_err(|e| e.to_string())?); + } + let m = model.as_mut().ok_or("SAM model unavailable")?; + let mut temporary = None; + if request.width > 0 { + let encoded = m.encode(&rgb, request.width, request.height).map_err(|e| e.to_string())?; + if request.temporary { + temporary = Some(encoded); + } else { + cache = Some(encoded); + } + } + let enc = if request.temporary { temporary.as_mut() } else { cache.as_mut() }.ok_or("SAM image cache empty")?; + predict(m, enc, request.operation) + })(); + eprintln!( + "SAM3 Metal request elapsed={:.3}s image={}x{} detail={} success={}", + started.elapsed().as_secs_f64(), + request.width, + request.height, + request.temporary, + result.is_ok() + ); + let response = serde_json::to_vec(&result).map_err(|e| e.to_string())?; + write_frame(stream, &response).map_err(|e| e.to_string())?; + if result.is_err() { + return Err("SAM inference failed; session released".into()); + } + } +} + +pub fn serve(address: SocketAddr, dir: &Path, device: candle_core::Device) -> Result<(), String> { + if !address.ip().is_loopback() || !device.is_metal() { + return Err("SAM worker requires a loopback address and a Metal device".into()); + } + if !crate::is_model_dir(dir) { + return Err("SAM model files are missing".into()); + } + let listener = TcpListener::bind(address).map_err(|e| e.to_string())?; + eprintln!("SAM3 worker listening on {address}; device=Metal; idle unload=600s"); + for stream in listener.incoming() { + let mut stream = stream.map_err(|e| e.to_string())?; + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| serve_client(&mut stream, dir, device.clone()))); + match result { + Ok(Err(error)) => eprintln!("SAM3 session closed: {error}"), + Err(_) => eprintln!("SAM3 session panicked; model released"), + Ok(Ok(())) => {} + } + } + Ok(()) +} diff --git a/crates/segment/tests/remote.rs b/crates/segment/tests/remote.rs new file mode 100644 index 00000000..ee20ea20 --- /dev/null +++ b/crates/segment/tests/remote.rs @@ -0,0 +1,21 @@ +use lightcraft_segment::remote::{Operation, Request, read_frame}; +use std::io::Cursor; + +#[test] +fn rejects_unbounded_and_truncated_frames() { + assert!(read_frame(&mut Cursor::new(u32::MAX.to_be_bytes()), 1024).is_err()); + assert!(read_frame(&mut Cursor::new([0, 0, 0, 3, 1]), 1024).is_err()); +} + +#[test] +fn validates_images_and_prompt_boundaries() { + let request = |width, height, operation| Request { width, height, temporary: false, operation }; + assert!(request(1009, 1, Operation::Prepare).validate().is_err()); + assert!(request(0, 1, Operation::Prepare).validate().is_err()); + assert!(request(0, 0, Operation::Text("dog".into())).validate().is_ok()); + assert!(request(4, 3, Operation::Clicks(vec![[0.5, 0.5, 1.0]])).validate().is_ok()); + assert!(request(4, 3, Operation::Clicks(vec![[1.5, 0.5, 1.0]])).validate().is_err()); + assert!(request(4, 3, Operation::Clicks(vec![[0.5, f32::NAN, 1.0]])).validate().is_err()); + assert!(request(4, 3, Operation::Clicks(vec![[0.5, 0.5, 0.0]])).validate().is_err()); + assert!(request(4, 3, Operation::Text(" ".into())).validate().is_err()); +} diff --git a/crates/ui-egui/src/panels/masking.rs b/crates/ui-egui/src/panels/masking.rs index 2e150a60..12ddb031 100644 --- a/crates/ui-egui/src/panels/masking.rs +++ b/crates/ui-egui/src/panels/masking.rs @@ -639,7 +639,7 @@ fn component_row_menu(app: &mut LightcraftApp, ui: &mut egui::Ui, mask: u32, k: /// 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(); + let missing = lightcraft_engine::segment::Segmenter::AVAILABLE && seg.remote_endpoint().is_none() && seg.dir.is_some() && !seg.installed(); if missing { app.offer_sam_download(Some((kind, op))); } diff --git a/docs/ai-masks.md b/docs/ai-masks.md index 5a32d690..00de5016 100644 --- a/docs/ai-masks.md +++ b/docs/ai-masks.md @@ -42,6 +42,39 @@ The result is stored with the mask: the model's 288 × 288 mask logits over the 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). +## Remote Metal inference + +Native builds with `sam` can keep the editor, RAW files, library, and exports local while +a Mac runs SAM 3. Build both ends from the same revision. On the Mac, install the model +as described below, then run: + +```sh +cargo build --release --locked -p lightcraft-segment --bin lightcraft-sam3-worker +target/release/lightcraft-sam3-worker "$HOME/Library/Application Support/LightCraft/models/sam3" 127.0.0.1:8793 +``` + +On the editing computer, start an SSH tunnel in one terminal and LightCraft in another: + +```sh +ssh -N -T -o ExitOnForwardFailure=yes -L 127.0.0.1:8793:127.0.0.1:8793 user@your-mac +LIGHTCRAFT_SAM3_REMOTE=127.0.0.1:8793 lightcraft +``` + +Object clicks, Describe prompts, and detail refinement use the remote worker. Only the +rendered model input (long edge at most 1008 pixels) and prompts cross the tunnel. The +worker caches the current photo; detail crops do not replace that cache. Returned masks +stay in the local library, so later adjustments and exports work without the connection. + +The worker requires Metal and accepts one editing connection at a time. It loads weights +on demand and releases them when the client disconnects or remains idle for ten minutes. +Both endpoints require loopback addresses: SSH provides authentication and encryption. +The listener has no independent authentication; do not expose it through a public proxy. + +No local weights are required in remote mode. An unavailable worker produces an explicit +error, not a silent CPU fallback. Restore the tunnel and retry. Unset +`LIGHTCRAFT_SAM3_REMOTE` to use the local model again. This offloads segmentation only, +not RAW development, denoising, or upscaling. The model licence below still applies. + ## Getting the model ### In the app (recommended) @@ -148,10 +181,12 @@ Therefore: | `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.status` | `{}` | `{available, installed, remote, 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}` | +`remote` reports the configured endpoint, not its health. `installed` describes local weights only. + 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 diff --git a/docs/whats-new.md b/docs/whats-new.md index ea0243f4..658ff4e0 100644 --- a/docs/whats-new.md +++ b/docs/whats-new.md @@ -135,6 +135,9 @@ lost everything at quit; a temporary session shows a banner the whole time and never writes to your library. ### Editing +- Optional remote SAM 3: keep the native editor local and run Object, Describe, and detail + inference on a Mac through SSH. Saved masks still render and export offline. + See [remote Metal inference](ai-masks.md#remote-metal-inference). - 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