From 339d9d70dc0944b2bbdf2c6e1f11192a711a272b Mon Sep 17 00:00:00 2001 From: lockewerks <59770696+lockewerks@users.noreply.github.com> Date: Mon, 25 May 2026 16:50:56 -0600 Subject: [PATCH] feat(cog-pose): wire pose_v1.safetensors into inference path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace silent stub fallback in InferenceEngine with a tracing-aware loader: - emit `tracing::info!` when pose_v1.safetensors loads into the candle backend, including the resolved on-disk path - emit `tracing::warn!` (instead of silently falling through) when the configured weights file is missing or no path is set at all, so the appliance log surfaces "no model — running stub" instead of just emitting confidence=0 frames forever - inline unit tests covering both branches: stub fallback shape + confidence, and a real-weights forward pass that asserts the output decomposes into 17 (x, y) pairs all finite and in [0, 1] The forward pass itself already matched the training script (Conv1d 56->64->128->128 with dilations [1, 2, 4], GlobalMeanPool, Linear 128->256->34, sigmoid) and the tensor names in the safetensors file (`enc.c1/c2/c3.weight|bias`, `head.fc1/fc2.weight|bias`) line up with the VarBuilder prefixes — no architecture changes required, only loader hygiene. Confidence is the published validation PCK@50 (0.185); v0.0.1 doesn't emit per-frame confidence and we're not fudging that here. --- .../cog-pose-estimation/src/inference.rs | 120 +++++++++++++++++- 1 file changed, 117 insertions(+), 3 deletions(-) diff --git a/v2/crates/cog-pose-estimation/src/inference.rs b/v2/crates/cog-pose-estimation/src/inference.rs index 2e1623ed..e476acf1 100644 --- a/v2/crates/cog-pose-estimation/src/inference.rs +++ b/v2/crates/cog-pose-estimation/src/inference.rs @@ -145,8 +145,10 @@ impl InferenceEngine { } /// Create an engine with a specific weights path (used by `--config` - /// in `cog-pose-estimation run`). If `weights_path` is `None`, the - /// stub fallback is used. + /// in `cog-pose-estimation run`). If `weights_path` is `None` or the + /// file does not exist on disk, the engine falls back to the + /// centred-skeleton stub and emits a `tracing::warn!` so the + /// appliance log shows why no real keypoints are coming through. pub fn with_weights(weights_path: Option<&Path>) -> Result> { let device = pick_device(); let inner = match weights_path { @@ -159,9 +161,26 @@ impl InferenceEngine { VarBuilder::from_mmaped_safetensors(&[p.to_path_buf()], DType::F32, &device)? }; let net = PoseNet::new(vb)?; + tracing::info!( + weights = %p.display(), + "loaded pose_v1.safetensors into candle backend" + ); Some(Arc::new(LoadedModel { net })) } - _ => None, + Some(p) => { + tracing::warn!( + weights = %p.display(), + "pose weights file not found; falling back to centred-skeleton stub (confidence=0)" + ); + None + } + None => { + tracing::warn!( + "no pose weights path configured and no default weights found on disk; \ + falling back to centred-skeleton stub (confidence=0)" + ); + None + } }; Ok(Self { inner, device }) } @@ -255,3 +274,98 @@ fn default_weights_path() -> Option { ]; candidates.into_iter().find(|p| p.exists()) } + +// --------------------------------------------------------------------------- +// Unit tests — exercise the safetensors → forward-pass path. Integration-level +// assertions (CLI surface, manifest round-trip, etc.) live in `tests/smoke.rs`. +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + /// Locate `pose_v1.safetensors` from any of the cwds a `cargo test` + /// invocation might land in (workspace root, `v2/`, or the crate dir). + fn locate_weights() -> Option { + let candidates = [ + std::path::PathBuf::from("cog/artifacts/pose_v1.safetensors"), + std::path::PathBuf::from("crates/cog-pose-estimation/cog/artifacts/pose_v1.safetensors"), + std::path::PathBuf::from( + "v2/crates/cog-pose-estimation/cog/artifacts/pose_v1.safetensors", + ), + ]; + candidates.into_iter().find(|p| p.exists()) + } + + #[test] + fn stub_fallback_when_weights_missing() { + // `with_weights(None)` must never panic and must produce a + // finite, well-shaped output so the runtime loop keeps making + // progress while the operator notices the warn log. + let engine = InferenceEngine::with_weights(None).expect("engine init"); + assert_eq!(engine.backend(), "stub"); + let out = engine.infer(&SyntheticInput.as_window()).expect("infer"); + assert!(out.is_finite()); + assert_eq!(out.keypoints.len(), OUTPUT_KEYPOINTS * 2); + assert_eq!(out.confidence, 0.0); + } + + #[test] + fn weights_load_and_forward_produces_seventeen_keypoint_pairs() { + let Some(weights) = locate_weights() else { + eprintln!( + "(skipping — pose_v1.safetensors not on disk; run from the cog crate or repo root)" + ); + return; + }; + let engine = InferenceEngine::with_weights(Some(&weights)).expect("load real weights"); + assert!( + engine.backend().starts_with("candle-"), + "expected candle backend, got {}", + engine.backend() + ); + + // Synthetic [56, 20] zero-input window — the documented "no-op" + // test signal. Anything finite and well-shaped proves the + // safetensors weights flowed through the forward pass. + let out = engine.infer(&SyntheticInput.as_window()).expect("infer"); + + // Shape: 17 (x, y) pairs = 34 scalars, no NaN, no Inf, all in + // sigmoid's [0, 1] range. + assert_eq!( + out.keypoints.len(), + OUTPUT_KEYPOINTS * 2, + "expected {} scalars for 17 keypoint pairs", + OUTPUT_KEYPOINTS * 2 + ); + let pairs: Vec<[f32; 2]> = out + .keypoints + .chunks_exact(2) + .map(|c| [c[0], c[1]]) + .collect(); + assert_eq!(pairs.len(), OUTPUT_KEYPOINTS, "expected 17 (x, y) pairs"); + for (i, [x, y]) in pairs.iter().enumerate() { + assert!( + x.is_finite() && y.is_finite(), + "keypoint {i} not finite: ({x}, {y})" + ); + assert!(!x.is_nan() && !y.is_nan(), "keypoint {i} is NaN"); + assert!( + (0.0..=1.0).contains(x) && (0.0..=1.0).contains(y), + "keypoint {i} out of [0,1]: ({x}, {y})" + ); + } + + // Confidence is the published PCK@50 (constant for v0.0.1), so + // anything > 0 proves we didn't silently fall through to the stub. + assert!(out.confidence > 0.0); + assert!(out.confidence.is_finite()); + } + + #[test] + fn rejects_wrong_shape_input_before_any_forward_pass() { + let engine = InferenceEngine::with_weights(None).expect("engine init"); + let bad = CsiWindow { data: vec![0.0; 7] }; + assert!(engine.infer(&bad).is_err()); + } +}