diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index ae15fe1d3..b66680507 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -36,6 +36,8 @@ pub struct Status { /// A finished job's message and whether it failed, for a toast (taken by /// [`Session::poll_models`]). pub done: Option<(String, bool)>, + /// Counts selection changes: a load started under an older count is stale and its model is discarded. + pub generation: u64, } /// The session's trained models: per task, the worker's status; and the face model in use. @@ -63,6 +65,12 @@ impl Models { self.face_model.lock().ok().and_then(|m| m.clone()) } + /// Activate a model that finished loading for the choice counted as `generation`; `false` (and the model is + /// dropped) when the choice changed meanwhile. + pub(crate) fn set_active_if_current(&self, task: Task, model: Loaded, generation: u64) -> bool { + activate_if_current(self.lane(task), &self.face_model, task, model, generation) + } + pub(crate) fn set_active(&self, task: Task, model: Option) { activate(self.lane(task), &self.face_model, task, model); } @@ -92,6 +100,30 @@ fn activate(lane: &Arc>, face: &FaceSlot, task: Task, model: Optio set(lane, |s| s.active = id); } +/// [`activate`] a freshly loaded model unless the choice changed since the load started (`generation`); whether it +/// was activated. The check and the activation happen under the lane's lock, so a selection can't slip between them. +fn activate_if_current(lane: &Arc>, face: &FaceSlot, task: Task, model: Loaded, generation: u64) -> bool { + let Ok(mut s) = lane.lock() else { return false }; + if s.generation != generation { + return false; + } + let id = match &model { + Loaded::Mask(m) => m.info().id.to_string(), + Loaded::Face(m) => m.info().id.to_string(), + }; + match (task, model) { + (Task::Mask, Loaded::Mask(m)) => crate::effects::roto::set_model(Some(m)), + (Task::Face, Loaded::Face(f)) => { + if let Ok(mut slot) = face.lock() { + *slot = Some(f); + } + } + _ => return false, + } + s.active = Some(id); + true +} + fn set(status: &Arc>, f: impl FnOnce(&mut Status)) { if let Ok(mut s) = status.lock() { f(&mut s); @@ -160,7 +192,11 @@ impl Session { Task::Mask => &mut self.prefs.roto.model, Task::Face => &mut self.prefs.face.model, }; - *slot = id.to_string(); + if slot != id { + *slot = id.to_string(); + // a load still running for the previous choice must not activate when it finishes + set(self.models.lane(task), |s| s.generation = s.generation.wrapping_add(1)); + } } /// Load (or drop) the models Settings chooses; called whenever settings change. @@ -190,9 +226,13 @@ impl Session { return; }; let (face_slot, label) = (self.models.face_model.clone(), kind(task).label); + let generation = st.generation; spawn(&lane, label, format!("Loading {}…", info.name), wait, move |status| { let m = load_file(&path, info.id)?; - activate(status, &face_slot, task, Some(m)); + // the choice changed while this loaded: drop the model, `poll_models` loads the new choice + if !activate_if_current(status, &face_slot, task, m, generation) { + return Ok(String::new()); + } Ok(format!("{label}: {} loaded", info.name)) }); } @@ -329,7 +369,10 @@ impl Session { 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 { - self.events.push(crate::Event::Toast { message, error }); + // a discarded (stale) load has no message + if !message.is_empty() { + self.events.push(crate::Event::Toast { message, error }); + } self.apply_model(task, false); } } diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index 84bbd2381..567b51666 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -161,6 +161,36 @@ fn an_idle_roto_model_is_unloaded() { let _ = std::fs::remove_dir_all(d); } +/// A model that finishes loading after the choice changed is dropped, not activated. +#[test] +fn a_load_for_an_old_choice_does_not_activate() { + use effectcraft_segment::{Loaded, MaskModel, ModelInfo, Prompt}; + struct Fake; + 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()) + } + } + super::roto::hold_roto_cache_test_lock(); + let (mut s, d) = session("stale-load"); + let started = s.models.status(Task::Mask).generation; + // the user picks another engine while the load runs + s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); + s.execute_checked("roto.model.select", json!({"id": "sam2.1-hiera-tiny"})).unwrap(); + assert_ne!(s.models.status(Task::Mask).generation, started, "a changed choice is counted"); + let m: Arc = Arc::new(Fake); + assert!(!s.models.set_active_if_current(Task::Mask, Loaded::Mask(m.clone()), started), "the stale load is dropped"); + assert!(s.models.status(Task::Mask).active.is_none() && crate::effects::roto::installed_model().is_none()); + let now = s.models.status(Task::Mask).generation; + assert!(s.models.set_active_if_current(Task::Mask, Loaded::Mask(m), now), "a load for the current choice activates"); + assert_eq!(s.models.status(Task::Mask).active.as_deref(), Some("sam2.1-hiera-tiny")); + crate::effects::roto::set_model(None); + let _ = std::fs::remove_dir_all(d); +} + /// A stroke reloading an unloaded model never waits for the load and never retries one that already failed. #[cfg(not(target_arch = "wasm32"))] #[test] diff --git a/crates/ui-egui/src/app/viewer_frames.rs b/crates/ui-egui/src/app/viewer_frames.rs index c878fafa4..f87db615c 100644 --- a/crates/ui-egui/src/app/viewer_frames.rs +++ b/crates/ui-egui/src/app/viewer_frames.rs @@ -84,10 +84,20 @@ impl EffectcraftApp { frame: 0, scale: (scale * 1000.0).round() as u32, view: self.view_hash(comp), - opts: frames::opts_hash(&self.frame_opts(comp, scale)), + opts: frames::opts_hash(&self.frame_opts(comp, scale)) ^ self.roto_model_stamp(), } } + /// Changes when the Roto Brush model in use does: a frame drawn with the classic fallback while the model loaded + /// is not the frame the model gives, so the two must not share a cache entry. + fn roto_model_stamp(&self) -> u64 { + use std::hash::{Hash, Hasher}; + let Some(id) = self.session.models_roto_active() else { return 0 }; + let mut h = std::collections::hash_map::DefaultHasher::new(); + id.hash(&mut h); + h.finish() | 1 + } + /// The frame series the viewer of `comp` shows now (its scale, view and options, for the /// comp's current content): what the cache bars count. pub fn shown_series(&self, comp: ItemId) -> FrameKey {