Roto model loading: drop a load for an old choice, key viewer frames by the active model

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RDvxCS8xcR28VArjzn1x7U

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<Loaded>) {
         activate(self.lane(task), &self.face_model, task, model);
     }
@@ -92,6 +100,30 @@ fn activate(lane: &Arc<Mutex<Status>>, 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<Mutex<Status>>, 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<Mutex<Status>>, 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);
             }
         }

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<Vec<f32>> {
+            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<dyn MaskModel> = 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]

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 {