diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index 9ff8a5ca7..8d45b5473 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -449,9 +449,17 @@ fn install_bytes(dir: &Path, m: &ModelInfo, bytes: &[u8]) -> std::result::Result Ok(path) } +fn models_for_platform(task: Task, wasm: bool) -> impl Iterator { + seg::models(task).filter(move |m| !wasm || !m.id.starts_with("sam2.1-hiera-")) +} + +fn target_models(task: Task) -> impl Iterator { + models_for_platform(task, cfg!(target_arch = "wasm32")) +} + fn model_p(p: &Value, cmd: &str, task: Task) -> Result<&'static ModelInfo> { let id = str_p(p, "id").ok_or_else(|| bad(cmd, "missing `id`"))?; - seg::info(id).filter(|m| m.task == task).ok_or_else(|| bad(cmd, format!("unknown model `{id}` (see {}.models)", kind(task).prefix))) + target_models(task).find(|m| m.id == id).ok_or_else(|| bad(cmd, format!("unknown model `{id}` (see {}.models)", kind(task).prefix))) } fn list(s: &mut Session, task: Task) -> Result { @@ -467,7 +475,7 @@ fn list(s: &mut Session, task: Task) -> Result { "selected": chosen == seg::CLASSICAL, "active": st.active.is_none(), })]; - for m in seg::models(task) { + for m in target_models(task) { #[cfg(not(target_arch = "wasm32"))] let (recommended, hint) = (effectcraft_sam2::mask::recommended().id == m.id, effectcraft_sam2::mask::hint(m.id)); #[cfg(target_arch = "wasm32")] @@ -528,7 +536,9 @@ fn install(s: &mut Session, p: &Value, task: Task) -> Result { Some(_) => model_p(p, &cmd, task)?, None => { let h = seg::sha256::hex(&bytes); - seg::models(task).find(|m| m.sha256 == h).ok_or_else(|| bad(&cmd, "this file is not a published model in the registry (SHA-256 differs)"))? + target_models(task) + .find(|m| m.sha256 == h) + .ok_or_else(|| bad(&cmd, "this file is not a supported published model in the registry (SHA-256 differs)"))? } }; let dest = install_bytes(&dir, m, &bytes).map_err(|e| bad(&cmd, e))?; @@ -652,3 +662,19 @@ pub fn specs() -> Vec { cmd!("face.model.remove", "Remove Face Tracking Model", [], None, "{id}", always, |s, p| remove(s, p, Task::Face)), ] } + +#[cfg(test)] +mod availability_tests { + use super::*; + + #[test] + fn browser_models_exclude_native_trackers_and_keep_portable_models() { + let browser: Vec<_> = models_for_platform(Task::Mask, true).map(|m| m.id).collect(); + assert_eq!(browser, ["mobilesam"]); + let desktop: Vec<_> = models_for_platform(Task::Mask, false).map(|m| m.id).collect(); + assert_eq!(desktop, ["mobilesam", "sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"]); + for wasm in [false, true] { + assert_eq!(models_for_platform(Task::Face, wasm).map(|m| m.id).collect::>(), ["mediapipe-face"]); + } + } +} diff --git a/crates/engine/src/tests/models.rs b/crates/engine/src/tests/models.rs index 567b51666..28ddebf55 100644 --- a/crates/engine/src/tests/models.rs +++ b/crates/engine/src/tests/models.rs @@ -15,6 +15,32 @@ fn session(tag: &str) -> (Session, std::path::PathBuf) { (Session { config: Some(store), ..Default::default() }, d) } +#[test] +fn model_list_and_direct_selection_agree_on_target_availability() { + super::roto::hold_roto_cache_test_lock(); + let mut s = Session::default(); + let list = s.execute_checked("roto.models", json!({})).unwrap(); + let models = list["models"].as_array().unwrap(); + for id in ["sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"] { + let chosen = s.chosen_model(Task::Mask).to_string(); + let generation = s.models.status(Task::Mask).generation; + let selected = s.execute_checked("roto.model.select", json!({"id": id})); + if cfg!(target_arch = "wasm32") { + assert!(!models.iter().any(|m| m["id"] == id)); + assert!(selected.is_err()); + assert_eq!(s.chosen_model(Task::Mask), chosen); + assert_eq!(s.models.status(Task::Mask).generation, generation, "rejection must not alter model state"); + } else { + assert!(models.iter().any(|m| m["id"] == id)); + selected.unwrap(); + assert_eq!(s.chosen_model(Task::Mask), id); + } + } + s.execute_checked("roto.model.select", json!({"id": "classical"})).unwrap(); + assert_eq!(s.chosen_model(Task::Mask), "classical"); + assert!(models.iter().any(|m| m["id"] == "mobilesam")); +} + /// Poll until `task`'s model is in use (or a few seconds pass). fn wait_loaded(s: &mut Session, task: Task) { for _ in 0..400 { @@ -37,7 +63,7 @@ fn models_are_listed_chosen_and_verified() { let list = s.execute_checked(&format!("{prefix}.models"), json!({})).unwrap(); let ids: Vec<&str> = list["models"].as_array().unwrap().iter().map(|m| m["id"].as_str().unwrap()).collect(); // Roto Brush also offers SAM 2.1 in four sizes (video trackers; see crates/sam2). - let expect: Vec<&str> = if task == Task::Mask { + let expect: Vec<&str> = if task == Task::Mask && !cfg!(target_arch = "wasm32") { vec!["classical", model, "sam2.1-hiera-tiny", "sam2.1-hiera-small", "sam2.1-hiera-base-plus", "sam2.1-hiera-large"] } else { vec!["classical", model]