diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index bad9fd5ff..a584aebbb 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -63,7 +63,7 @@ impl Models { self.face_model.lock().ok().and_then(|m| m.clone()) } - fn set_active(&self, task: Task, model: Option) { + pub(crate) fn set_active(&self, task: Task, model: Option) { activate(self.lane(task), &self.face_model, task, model); } } @@ -128,6 +128,9 @@ fn kind(task: Task) -> &'static Kind { } } +/// How long a Roto Brush model may sit unused before its memory is freed. +const ROTO_MODEL_IDLE: Duration = Duration::from_secs(600); + /// How long a job that needs the Roto Brush model waits for it to load (or download) before it /// gives up. A worker that never reports back must not hang the command. const ROTO_MODEL_WAIT: Duration = Duration::from_secs(600); @@ -290,8 +293,35 @@ impl Session { Err(format!("{} is still loading: try again when it is ready, or pass wait: true", info.name)) } + /// Free a Roto Brush model that has been idle for [`ROTO_MODEL_IDLE`] (and is not under a job). The + /// choice stays: it loads again when a job or a stroke needs it. + fn unload_idle_roto_model(&mut self) { + if self.is_roto_running() || self.models.status(Task::Mask).busy.is_some() { + return; + } + let Some(m) = crate::effects::roto::installed_model() else { return }; + if m.idle_secs().is_some_and(|s| s >= ROTO_MODEL_IDLE.as_secs_f64()) { + let name = m.info().name; + drop(m); + self.models.set_active(Task::Mask, None); + self.events.push(crate::Event::Toast { + message: format!("{name} unloaded after {} minutes idle; it loads again when it is used", ROTO_MODEL_IDLE.as_secs() / 60), + error: false, + }); + } + } + + /// Load the chosen Roto Brush model again if it was unloaded for being idle (before a stroke uses it). + pub fn reload_roto_model(&mut self) { + let chosen = self.chosen_model(Task::Mask).to_string(); + if chosen != seg::CLASSICAL && self.models.status(Task::Mask).active.is_none() && self.models.status(Task::Mask).busy.is_none() { + self.apply_model(Task::Mask, true); + } + } + /// Finish what the model workers did: a toast, and loading a model that just arrived. pub fn poll_models(&mut self) { + self.unload_idle_roto_model(); for task in [Task::Mask, Task::Face] { let done = self.models.lane(task).lock().ok().and_then(|mut s| s.done.take()); if let Some((message, error)) = done { diff --git a/crates/engine/src/cmd/roto_cmds.rs b/crates/engine/src/cmd/roto_cmds.rs index 0931236a2..481a23b5d 100644 --- a/crates/engine/src/cmd/roto_cmds.rs +++ b/crates/engine/src/cmd/roto_cmds.rs @@ -196,6 +196,8 @@ fn points_p(p: &Value, cmd: &str) -> Result> { fn stroke(s: &mut Session, p: &Value) -> Result { let cmd = "roto.stroke"; + // A model unloaded for being idle comes back before the stroke is segmented with it. + s.reload_roto_model(); let (cid, lid) = match roto_p(s, p, cmd) { Ok((c, l, _)) => (c, l), Err(_) => layer_p(s, p, cmd)?, diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index 0b714d717..fab025a3a 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -105,3 +105,40 @@ fn models_are_listed_chosen_and_verified() { } let _ = std::fs::remove_dir_all(d); } + +/// A Roto Brush model that has been idle for ten minutes is unloaded (the choice stays), one in use is not, +/// and a model that does not track use is never unloaded. +#[test] +fn an_idle_roto_model_is_unloaded() { + use effectcraft_segment::{Loaded, MaskModel, ModelInfo, Prompt}; + struct Fake(Option); + impl MaskModel for Fake { + fn info(&self) -> &'static ModelInfo { + &effectcraft_segment::SAM2_TINY + } + fn segment(&self, _: &[[f32; 3]], _: usize, _: usize, _: &Prompt) -> effectcraft_segment::Result> { + Err("no weights".into()) + } + fn idle_secs(&self) -> Option { + self.0 + } + } + crate::tests_roto::hold_roto_cache_test_lock(); + let (mut s, d) = session("idle"); + for (idle, unloaded) in [(Some(5.0), false), (None, false), (Some(601.0), true)] { + let m: Arc = Arc::new(Fake(idle)); + s.models.set_active(Task::Mask, Some(Loaded::Mask(m))); + assert_eq!(s.models.status(Task::Mask).active.as_deref(), Some("sam2.1-hiera-tiny")); + s.poll_models(); + assert_eq!(s.models.status(Task::Mask).active.is_none(), unloaded, "idle {idle:?}"); + assert_eq!(crate::effects::roto::installed_model().is_none(), unloaded, "the effect's model follows"); + if unloaded { + assert!(s.events.iter().any(|e| matches!(e, crate::Event::Toast { message, .. } if message.contains("unloaded after 10 minutes")))); + } + } + // the choice survives, and with the classic engine chosen reloading does nothing + s.reload_roto_model(); + assert!(s.models.status(Task::Mask).active.is_none()); + crate::effects::roto::set_model(None); + let _ = std::fs::remove_dir_all(d); +} diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index a0eac9cd6..f917e29ef 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -54,11 +54,17 @@ fn fingerprint(rgb: &[[f32; 3]], w: usize, h: usize) -> u64 { pub struct Sam2Mask { model: Arc, last: Mutex)>>, + used: Mutex, } impl Sam2Mask { pub fn new(model: Arc) -> Sam2Mask { - Sam2Mask { model, last: Mutex::new(None) } + Sam2Mask { model, last: Mutex::new(None), used: Mutex::new(std::time::Instant::now()) } + } + + /// Note that the model is being used (for [`MaskModel::idle_secs`]). + pub(crate) fn touch(&self) { + *self.used.lock().unwrap_or_else(PoisonError::into_inner) = std::time::Instant::now(); } /// The video model, for trackers built on the same weights. @@ -103,12 +109,17 @@ impl MaskModel for Sam2Mask { } fn segment(&self, rgb: &[[f32; 3]], w: usize, h: usize, prompt: &effectcraft_segment::Prompt) -> effectcraft_segment::Result> { + self.touch(); self.probability(rgb, w, h, prompt).map_err(|e| e.to_string()) } fn tracker(&self) -> Option<&dyn effectcraft_segment::ObjectTracker> { Some(self) } + + fn idle_secs(&self) -> Option { + Some(self.used.lock().unwrap_or_else(PoisonError::into_inner).elapsed().as_secs_f64()) + } } /// Verify the installed file against the registry and build the model. diff --git a/crates/sam2/src/track.rs b/crates/sam2/src/track.rs index 2b06d08cd..2476ddb74 100644 --- a/crates/sam2/src/track.rs +++ b/crates/sam2/src/track.rs @@ -151,6 +151,9 @@ impl ObjectTracker for Sam2Mask { cancel: &dyn Fn() -> bool, out: &mut dyn FnMut(usize, i64, TrackedFrame), ) -> effectcraft_segment::Result<()> { - self.run_track(req, frames, cancel, out).map_err(|e| e.to_string()) + self.touch(); + let r = self.run_track(req, frames, cancel, out).map_err(|e| e.to_string()); + self.touch(); + r } } diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 0a9a0494a..bd0b1e672 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -49,6 +49,11 @@ pub trait MaskModel: Send + Sync { fn tracker(&self) -> Option<&dyn ObjectTracker> { None } + /// Seconds since the model was last used (`None`: it does not track use, so it is never unloaded + /// for being idle). Large models free their memory after a quiet spell and load again on use. + fn idle_secs(&self) -> Option { + None + } } /// One decoded frame handed to a tracker (straight RGB 0–1, row-major `w×h`). diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index a4e58a36d..e61f49be7 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -38,6 +38,13 @@ frames before the corrected frame are not rebuilt, so frames just after it can d Tracked frames of a session are in memory; a project reopened in a new process re-tracks the whole span after a stroke change (its saved tracking is stale as a whole). +## Memory and limits +The weights take about 156 MB. A model idle for 10 minutes is unloaded (a toast says so; the choice stays) and loads +again when a Propagate or a stroke needs it. The tracker holds the memory banks of the objects it follows (about 2 MB per +stored frame, pruned to the frames later ones can still look at) and refuses to grow past 2 GiB, or past a time limit when +one is set, with a clear error instead of running out of memory. Two runs of the same prompts give bit-identical masks on +the CPU. + ## How it works `crates/sam2` is a Rust port (on [candle](https://github.com/huggingface/candle), native only) of the Hugging Face `transformers` SAM 2 implementation, Apache-2.0 (see its `NOTICE`). Every stage is checked against PyTorch by