Initial pass at onnx wiring into SpotObserver

include/spot-observer.h

@@ -94,11 +94,34 @@ UNITY_INTERFACE_EXPORT
 bool UNITY_INTERFACE_API SOb_PushNextVisionPipelineImageSetToUnityBuffers(int32_t robot_id, int32_t cam_stream_id);
 
 // Model stuff
+
+// Model families. The family cannot be inferred from the file: both the
+// single-shot and the streaming (KV-cache) depth models are .onnx.
+//   SOb_MODEL_SINGLE_SHOT - stateless model; one instance may be shared by
+//                           several pipelines (path-cached internally).
+//   SOb_MODEL_STREAMING   - autoregressive model with per-sequence state; every
+//                           load returns a fresh instance and each instance can
+//                           drive exactly one pipeline at a time.
+#define SOb_MODEL_SINGLE_SHOT 0
+#define SOb_MODEL_STREAMING   1
+
 UNITY_INTERFACE_EXPORT
 SObModel UNITY_INTERFACE_API SOb_LoadModel(const char* modelPath, const char* backend);
+// Like SOb_LoadModel, with the model family stated explicitly (SOb_MODEL_*).
+// SOb_LoadModel is equivalent to kind = SOb_MODEL_SINGLE_SHOT.
+UNITY_INTERFACE_EXPORT
+SObModel UNITY_INTERFACE_API SOb_LoadModelEx(const char* modelPath, const char* backend, int32_t kind);
 UNITY_INTERFACE_EXPORT
 void UNITY_INTERFACE_API SOb_UnloadModel(SObModel model);
 
+// Live model switch: stops the vision pipeline on the given camera stream (if
+// one is running) and relaunches it with `model`. The camera stream keeps
+// running throughout; only the inference side is swapped. Intended pattern is to
+// load every selectable model once at startup and flip between the handles --
+// no load stall and no unload hazard at switch time.
+UNITY_INTERFACE_EXPORT
+bool UNITY_INTERFACE_API SOb_SwitchVisionPipelineModel(int32_t robot_id, int32_t cam_stream_id, SObModel model);
+
 // Config calls
 UNITY_INTERFACE_EXPORT
 bool UNITY_INTERFACE_API SOb_ToggleDepthCompletion(bool enable);

onnx/set_window.py

@@ -0,0 +1,129 @@
+"""Retarget the streaming model's KV-cache retention window.
+
+The window is baked into the graph as 48 Constant nodes feeding the `starts`
+input of the Slice that trims each new_* cache output. Rewriting them is a
+graph-only edit: the 7.5 GB .onnx.data blob is referenced by filename and is
+neither read nor rewritten, so each variant costs ~8 MB and shares the weights.
+
+VRAM saved is linear in the window, and doubled at peak because the graph reads
+`past_*` while writing `new_*`:
+
+    per retained frame = 48 * 16 * 1041 * 64 * 4 B = 195.2 MiB   (fp32)
+    peak cache         = 2 * window * per_frame
+
+Usage:
+    python set_window.py 8
+    python set_window.py 8 --in mae_model_step_consolidated.onnx
+"""
+import argparse, pathlib, sys
+import onnx
+from onnx import numpy_helper
+import numpy as np
+
+HEADS, TOKENS, HEAD_DIM, N_CACHE = 16, 1041, 64, 48
+BYTES_PER_FRAME = N_CACHE * HEADS * TOKENS * HEAD_DIM * 4
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("window", type=int, help="frames of KV cache to retain")
+    ap.add_argument("--in", dest="src", default="mae_model_step_consolidated.onnx")
+    ap.add_argument("--out", dest="dst", default=None)
+    args = ap.parse_args()
+
+    if args.window < 1:
+        print("window must be >= 1", file=sys.stderr)
+        return 2
+
+    here = pathlib.Path(__file__).parent
+    src = here / args.src
+    dst = here / (args.dst or f"{src.stem}_w{args.window}.onnx")
+
+    # load_external_data=False keeps the 7.5 GB blob out of memory entirely.
+    model = onnx.load(src, load_external_data=False)
+    graph = model.graph
+
+    # Preflight: every external-data file the graph references must sit next to
+    # the graph, under the exact name recorded inside it. Fail with the list
+    # rather than letting the checker produce a cryptic CWD-relative error.
+    referenced = {e.value for i in graph.initializer
+                  for e in i.external_data if e.key == "location"}
+    missing = sorted(loc for loc in referenced if not (src.parent / loc).exists())
+    if missing:
+        print("graph references external data files that are not next to it:", file=sys.stderr)
+        for loc in missing:
+            print(f"  expected: {src.parent / loc}", file=sys.stderr)
+        print("copy the .data file(s) into that directory with exactly those names", file=sys.stderr)
+        return 1
+
+    cache_outputs = {v.name for v in graph.output if v.name.startswith("new_")}
+    producer = {o: n for n in graph.node for o in n.output}
+
+    # The Slice may not feed the graph output directly: fp16 conversion and
+    # re-exports append pass-through ops (Cast, Identity) after it. Walk back
+    # from each new_* output through those to the producing Slice.
+    PASSTHROUGH = {"Identity", "Cast", "Squeeze", "Unsqueeze"}
+    def find_slice(output_name):
+        node, hops = producer.get(output_name), 0
+        while node is not None and node.op_type in PASSTHROUGH and hops < 8:
+            node = producer.get(node.input[0])
+            hops += 1
+        return node if node is not None and node.op_type == "Slice" else None
+
+    slice_nodes = [s for s in (find_slice(name) for name in sorted(cache_outputs)) if s is not None]
+    starts_inputs = {n.input[1] for n in slice_nodes}
+
+    if len(starts_inputs) != N_CACHE:
+        # Fallback: find Slices by pattern -- 1-element negative `starts` -- and
+        # report what actually feeds the outputs so a failure is diagnosable.
+        feeders = sorted({(producer[n].op_type if n in producer else "<graph input>")
+                          for n in cache_outputs})
+        print(f"expected {N_CACHE} cache Slice nodes, found {len(starts_inputs)}", file=sys.stderr)
+        print(f"ops feeding new_* outputs: {feeders}", file=sys.stderr)
+        print("(if these are not Slice/Cast/Identity, the export's window structure "
+              "changed -- send this output back)", file=sys.stderr)
+        return 1
+
+    patched, old_values = 0, set()
+    new_val = np.array([-args.window], dtype=np.int64)
+
+    for node in graph.node:
+        if node.op_type == "Constant" and node.output and node.output[0] in starts_inputs:
+            for attr in node.attribute:
+                if attr.name == "value":
+                    old_values.add(int(numpy_helper.to_array(attr.t)[0]))
+                    attr.t.CopyFrom(numpy_helper.from_array(new_val, attr.t.name))
+                    patched += 1
+    for init in graph.initializer:
+        if init.name in starts_inputs:
+            old_values.add(int(numpy_helper.to_array(init)[0]))
+            init.CopyFrom(numpy_helper.from_array(new_val, init.name))
+            patched += 1
+
+    if patched != N_CACHE:
+        print(f"patched {patched} of {N_CACHE} window constants -- aborting", file=sys.stderr)
+        return 1
+
+    # The checker resolves external-data locations relative to the process CWD,
+    # not the graph -- pin it to the graph's directory so running the script
+    # from anywhere behaves the same.
+    import os
+    prev_cwd = os.getcwd()
+    os.chdir(src.parent)
+    try:
+        onnx.checker.check_model(model)
+    finally:
+        os.chdir(prev_cwd)
+    onnx.save(model, dst)
+
+    old = ", ".join(str(-v) for v in sorted(old_values))
+    peak = 2 * args.window * BYTES_PER_FRAME / (1024 ** 3)
+    print(f"patched {patched} window constants: {old} -> {args.window}")
+    print(f"wrote {dst.name} ({dst.stat().st_size / 1024**2:.1f} MB; shares {src.stem}.onnx.data)")
+    print(f"fp32 peak KV cache at window {args.window}: {peak:.2f} GiB "
+          f"(steady {peak/2:.2f} GiB)")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/cuda_kernels.cu

@@ -1421,4 +1421,116 @@ cudaError_t postprocess_depth_image(
     return cudaGetLastError();
 }
 
+///////////////////////////////////////////////////////////////////////////////////////////////////
+// Resampling between camera resolution and a model's native input size.
+
+__global__ void resize_bilinear_chw_kernel(
+    const float* __restrict__ src,
+    float* __restrict__ dst,
+    int in_h, int in_w,
+    int out_h, int out_w,
+    int channels
+) {
+    const int x = blockIdx.x * blockDim.x + threadIdx.x;
+    const int y = blockIdx.y * blockDim.y + threadIdx.y;
+    const int c = blockIdx.z;
+    if (x >= out_w || y >= out_h || c >= channels) return;
+
+    // Half-pixel centres (equivalent to align_corners=False), so the resample
+    // stays centred and does not shift the image by half a source pixel.
+    const float fy = (y + 0.5f) * in_h / out_h - 0.5f;
+    const float fx = (x + 0.5f) * in_w / out_w - 0.5f;
+
+    int y0 = static_cast<int>(floorf(fy));
+    int x0 = static_cast<int>(floorf(fx));
+    const float wy = fy - y0;
+    const float wx = fx - x0;
+    const int y1 = min(max(y0 + 1, 0), in_h - 1);
+    const int x1 = min(max(x0 + 1, 0), in_w - 1);
+    y0 = min(max(y0, 0), in_h - 1);
+    x0 = min(max(x0, 0), in_w - 1);
+
+    const float* plane = src + static_cast<size_t>(c) * in_h * in_w;
+    const float v00 = plane[static_cast<size_t>(y0) * in_w + x0];
+    const float v01 = plane[static_cast<size_t>(y0) * in_w + x1];
+    const float v10 = plane[static_cast<size_t>(y1) * in_w + x0];
+    const float v11 = plane[static_cast<size_t>(y1) * in_w + x1];
+
+    dst[static_cast<size_t>(c) * out_h * out_w + static_cast<size_t>(y) * out_w + x] =
+        (1.f - wy) * ((1.f - wx) * v00 + wx * v01) +
+                wy * ((1.f - wx) * v10 + wx * v11);
+}
+
+__global__ void resize_sparse_depth_kernel(
+    const float* __restrict__ src,
+    float* __restrict__ dst,
+    int in_h, int in_w,
+    int out_h, int out_w
+) {
+    const int x = blockIdx.x * blockDim.x + threadIdx.x;
+    const int y = blockIdx.y * blockDim.y + threadIdx.y;
+    if (x >= out_w || y >= out_h) return;
+
+    // Source footprint of this output pixel, plus its centre for tie-breaking.
+    const float cy = (y + 0.5f) * in_h / out_h;
+    const float cx = (x + 0.5f) * in_w / out_w;
+    int y0 = max(static_cast<int>(floorf(static_cast<float>(y)     * in_h / out_h)), 0);
+    int y1 = min(static_cast<int>(ceilf (static_cast<float>(y + 1) * in_h / out_h)), in_h);
+    int x0 = max(static_cast<int>(floorf(static_cast<float>(x)     * in_w / out_w)), 0);
+    int x1 = min(static_cast<int>(ceilf (static_cast<float>(x + 1) * in_w / out_w)), in_w);
+    // Upscaling can leave an empty footprint; always sample at least one pixel.
+    if (y1 <= y0) y1 = min(y0 + 1, in_h);
+    if (x1 <= x0) x1 = min(x0 + 1, in_w);
+
+    float best = 0.f;
+    float best_dist2 = CUDART_INF_F;
+    for (int yy = y0; yy < y1; ++yy) {
+        for (int xx = x0; xx < x1; ++xx) {
+            const float v = src[static_cast<size_t>(yy) * in_w + xx];
+            // 0 means "no sample". Negated compare so NaN is rejected too.
+            // Range validation is the model's, not the resample's.
+            if (!(v > 0.f)) continue;
+            const float dy = (yy + 0.5f) - cy;
+            const float dx = (xx + 0.5f) - cx;
+            const float d2 = dy * dy + dx * dx;
+            if (d2 < best_dist2) { best_dist2 = d2; best = v; }
+        }
+    }
+    dst[static_cast<size_t>(y) * out_w + x] = best;
+}
+
+cudaError_t resize_bilinear_chw(
+    const float* d_in,
+    float* d_out,
+    int in_h, int in_w,
+    int out_h, int out_w,
+    int channels,
+    cudaStream_t stream
+) {
+    dim3 block(32, 8);
+    dim3 grid((out_w + block.x - 1) / block.x,
+              (out_h + block.y - 1) / block.y,
+              channels);
+    resize_bilinear_chw_kernel<<<grid, block, 0, stream>>>(
+        d_in, d_out, in_h, in_w, out_h, out_w, channels
+    );
+    return cudaGetLastError();
+}
+
+cudaError_t resize_sparse_depth(
+    const float* d_in,
+    float* d_out,
+    int in_h, int in_w,
+    int out_h, int out_w,
+    cudaStream_t stream
+) {
+    dim3 block(32, 8);
+    dim3 grid((out_w + block.x - 1) / block.x,
+              (out_h + block.y - 1) / block.y);
+    resize_sparse_depth_kernel<<<grid, block, 0, stream>>>(
+        d_in, d_out, in_h, in_w, out_h, out_w
+    );
+    return cudaGetLastError();
+}
+
 } // namespace SOb

src/include/cuda_kernels.cuh

@@ -67,6 +67,34 @@ void convert_uint8_img_to_float_img(
 
 void loadImageToCudaFloatRGB(const std::string& path, int& outW, int& outH, float* d_image);
 
+// Resampling between the camera resolution and a model's native input size.
+// Neither axis is an integer ratio (e.g. 480->392, 640->518), so the integer
+// downscale in preprocess_depth_image2 does not cover this.
+
+// Planar CHW bilinear resample.
+cudaError_t resize_bilinear_chw(
+    const float* d_in,
+    float* d_out,
+    int in_h, int in_w,
+    int out_h, int out_w,
+    int channels,
+    cudaStream_t stream = 0
+);
+
+// Resample sparse metric depth. Bilinear is wrong here: 0 means "no sample", and
+// blending it with valid neighbours fabricates depth. Each output pixel takes the
+// nearest sampled pixel in its source footprint, which preserves sparse points
+// that plain nearest-neighbour would drop when downscaling.
+// Only 0 (and NaN) are treated as "no sample"; range validation belongs to the
+// consuming model, not to the resample.
+cudaError_t resize_sparse_depth(
+    const float* d_in,
+    float* d_out,
+    int in_h, int in_w,
+    int out_h, int out_w,
+    cudaStream_t stream = 0
+);
+
 // Running average depth maintenance
 cudaError_t prefill_invalid_depth(
     float* d_depth_data,