PR 26 of #508 — device-agnostic layer buffers + Linux spawn-context for PAF pool

[gitttt-1234]

Stacked on #562 (PR 25). Closes two bug families surfaced by the CUDA bench against PR 25's tip on an NVIDIA A40 box.

What the bench surfaced

Bug	Where	Reproduces on
Top-down predict_streaming device mismatch	cuda:0 vs cpu and mps:0 vs cpu	CUDA + MPS (not CPU)
paf_workers > 0 hangs forever	first worker call	Linux + CUDA (fork start-method default)

The first one I'd assumed was CUDA-specific. Testing on Mac MPS reproduced it with the exact same shape of error — confirming the root cause is device-agnostic anti-pattern in the new flow, not a CUDA quirk.

Root cause 1 — CPU-allocated output buffers across the layers

Several InferenceLayer subclasses allocated output buffers with bare torch.full((...), float("nan")) / torch.ones(B) calls. No device= kwarg. On CPU this is silent (everything is CPU); on any non-CPU device the scatter from the layer's device-resident tensors into a CPU buffer raises RuntimeError: Expected all tensors to be on the same device.

Files touched (the same anti-pattern fixed everywhere):

• sleap_nn/inference/layers/topdown.py — 3 torch.full allocations now device-aware (stage2_kpts_img.device).
• sleap_nn/inference/layers/centroid.py — padded_peaks, padded_vals, centroid_vals, both eff_scale allocations now device-aware in both the GT branch and the postprocess. This was the actual upstream source of the topdown failure — the centroid layer was returning CPU-resident pred_centroids even when running on MPS/CUDA, which then poisoned the scatter in topdown stage 2.
• sleap_nn/inference/layers/centered_instance.py — b_idx, matched_vals, pred_centroid_values, eff_scale device-aware.
• sleap_nn/inference/layers/single_instance.py — eff_scale device-aware (uses x.device).
• sleap_nn/inference/layers/{bottomup, bottomup_multiclass, topdown_multiclass}.py — eff_scale device-aware via scaled.device.

Root cause 2 — PafGroupingPool fork-on-Linux deadlocks against CUDA

sleap_nn/inference/streaming.py:328 constructed its ProcessPoolExecutor without an explicit mp_context. On Linux this defaults to fork, which inherits the parent's already-initialized CUDA context and deadlocks the first worker call. Pinned mp_context=multiprocessing.get_context(\"spawn\") — matches the existing default on macOS / Windows.

Why the existing test suite missed all of this

Every pre-existing inference test either (a) used _StubLayer instead of a real backend, (b) used NumpyProvider with synthetic frames, or (c) mocked the factory. None exercised the actual video → preprocess → backend.forward → postprocess → Outputs chain on a real fixture. The new tests/inference/test_e2e_video.py plugs that gap.

10 parametrized tests = 5 CPU + 5 MPS-gated, one per supported model type (single_instance, centroid_only, topdown, bottomup, multiclass_bottomup). Pre-fix the MPS topdown case raised the device-mismatch error; post-fix all 10 pass.

Out of scope for this PR

The CUDA bench also showed channel-mismatch failures for centroid-only + bottom-up predict_streaming on a real video (weight=[36, 72, 3, 3], expected 72 channels, got 36). These reproduce only on CUDA — clean on CPU and MPS with the same code + checkpoint + video. Probably cuDNN strictness or a torch 2.9.1 quirk with UNet skip-connection alignment on non-square inputs. Need a CUDA box to fix; filing as a separate follow-up issue.

The standalone centered_instance factory dispatch (raised Unsupported model_paths combination in the bench) is also intentionally out of scope — it's a missing feature symmetric to PR 25's centroid-only support, with a similar footprint. Will be filed as PR 27.

Test plan

• pytest tests/inference/test_e2e_video.py — 10 passed (5 CPU + 5 MPS) on this Mac.
• pytest tests/inference/test_paf_worker_pool.py — 8 passed (spawn fix intact, no regressions).
• pytest tests/inference/ tests/cli/ tests/data/test_instance_centroids.py — 414 passed, 23 skipped (CUDA-gated).
• black --check sleap_nn tests — clean.
• ruff check sleap_nn/ — clean.
• Re-run the CUDA bench on the A40 box (section D + topdown C should both now pass; centroid + bottomup channel mismatches will persist — that's the separate issue).

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 68.41398% with 235 lines in your changes missing coverage. Please review.
✅ Project coverage is 64.46%. Comparing base (60575f1) to head (bff2525).

Files with missing lines	Patch %	Lines
sleap_nn/inference/predictor.py	48.37%	143 Missing ⚠️
sleap_nn/inference/run.py	14.28%	30 Missing ⚠️
sleap_nn/inference/loaders.py	88.26%	25 Missing ⚠️
sleap_nn/inference/layers/base.py	68.18%	21 Missing ⚠️
sleap_nn/cli.py	89.65%	6 Missing ⚠️
sleap_nn/inference/layers/centroid.py	70.58%	5 Missing ⚠️
sleap_nn/inference/layers/topdown.py	88.88%	5 Missing ⚠️

Additional details and impacted files

@@                                Coverage Diff                                @@
##           divya/inf-refactor-25-centroid-only-inference     #563      +/-   ##
=================================================================================
+ Coverage                                          64.20%   64.46%   +0.26%     
=================================================================================
  Files                                                124      125       +1     
  Lines                                              19066    19431     +365     
=================================================================================
+ Hits                                               12241    12527     +286     
- Misses                                              6825     6904      +79     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.