IQ.Pilot Release Commit @ f2a861c

iqpilot/selfdrive/iqmodeld/models/runners/tinygrad/fused_runner.py

@@ -9,12 +9,12 @@ from typing import Any
 
 import numpy as np
 
-from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import (
+from iqpilot.selfdrive.iqmodeld.models.runners.model_runner import (
     CUSTOM_MODEL_PATH, NumpyDict, ShapeDict, SliceDict,
 )
-from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelRunner
-from openpilot.iqpilot.selfdrive.iqmodeld.models.split_model_constants import SplitModelConstants
-from openpilot.iqpilot.selfdrive.iqmodeld.parser import PhaseParser
+from iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelRunner
+from iqpilot.selfdrive.iqmodeld.models.split_model_constants import SplitModelConstants
+from iqpilot.selfdrive.iqmodeld.parser import PhaseParser
 
 
 def _tinygrad_imports():
@@ -27,8 +27,6 @@ WARP_DEV = os.getenv('WARP_DEV')
 
 
 class TinygradFusedRunner(ModelRunner):
-    """Runs a fused warp+vision+policy pkl. Bundle ships one `driving_fused_*` artifact."""
-
     uses_opencl_warp: bool = False
 
     def __init__(self):
@@ -110,19 +108,30 @@ class TinygradFusedRunner(ModelRunner):
             'feat_q':    zeros_f32((self._frame_skip * (fb[1] - 1) + 1, fb[0], fb[2])),
             'desire_q':  zeros_f32((self._frame_skip * dp[1], dp[0], dp[2])),
         }
-        # shapes must match the captured run_policy JIT inputs
         on_shapes = self._on_meta['input_shapes']
+        captured = self._run_policy.captured
+        jit_shapes = {
+            name: tuple(int(s) for s in view.shape)
+            for name, (view, _vars, _dtype, _device) in zip(captured.expected_names, captured.expected_input_info)
+        }
+
+        def policy_input_shape(name):
+            shape = on_shapes.get(name, jit_shapes.get(name))
+            if shape is None:
+                raise ValueError(f"fused pkl declares no shape for policy input {name}")
+            return shape
+
         self._npy_buffers = {
             'desire':             np.zeros(dp[2], dtype=np.float32),
-            'traffic_convention': np.zeros(on_shapes['traffic_convention'], dtype=np.float32),
-            'action_t':           np.zeros(on_shapes['action_t'], dtype=np.float32),
+            'traffic_convention': np.zeros(policy_input_shape('traffic_convention'), dtype=np.float32),
             'tfm':                np.zeros((3, 3), dtype=np.float32),
             'big_tfm':            np.zeros((3, 3), dtype=np.float32),
         }
+        if 'action_t' in jit_shapes:
+            self._npy_buffers['action_t'] = np.zeros(policy_input_shape('action_t'), dtype=np.float32)
         self._cam_resolution = (cam_w, cam_h)
 
     def run_fused(self, bufs: dict, transforms: dict[str, np.ndarray], numpy_inputs: NumpyDict) -> NumpyDict:
-        """warp + vision + policy in one pass from raw NV12 bufs + transform matrices."""
         Tensor, Device = _tinygrad_imports()
 
         main_buf = bufs['img']
@@ -134,14 +143,13 @@ class TinygradFusedRunner(ModelRunner):
             self._npy_buffers['desire'][:] = numpy_inputs[desire_key]
         if 'traffic_convention' in numpy_inputs:
             self._npy_buffers['traffic_convention'][:] = numpy_inputs['traffic_convention']
-        if 'action_t' in numpy_inputs:
+        if 'action_t' in numpy_inputs and 'action_t' in self._npy_buffers:
             self._npy_buffers['action_t'][:] = numpy_inputs['action_t']
         self._npy_buffers['tfm'][:] = transforms['img']
         self._npy_buffers['big_tfm'][:] = transforms['big_img']
 
         npy = lambda key: Tensor(self._npy_buffers[key], device='NPY')
 
-        # frames go on the compute device to match the captured warp JIT
         frame = self._frame_tensor('img', bufs['img'])
         big_frame = self._frame_tensor('big_img', bufs['big_img'])
 
@@ -149,12 +157,13 @@ class TinygradFusedRunner(ModelRunner):
         img, big_img = warp_jit(img_q=self._queues['img_q'], big_img_q=self._queues['big_img_q'],
                                 tfm=npy('tfm'), big_tfm=npy('big_tfm'), frame=frame, big_frame=big_frame)
 
-        vision_out_t, on_out_t, off_out_t = self._run_policy(
+        policy_inputs = dict(
             img=img, big_img=big_img, feat_q=self._queues['feat_q'], desire_q=self._queues['desire_q'],
-            desire=npy('desire'), traffic_convention=npy('traffic_convention'), action_t=npy('action_t'))
+            desire=npy('desire'), traffic_convention=npy('traffic_convention'))
+        if 'action_t' in self._npy_buffers:
+            policy_inputs['action_t'] = npy('action_t')
+        vision_out_t, on_out_t, off_out_t = self._run_policy(**policy_inputs)
 
-        # parse each model's output on its own sliced dict; parsing a merged dict
-        # would run parse_dynamic_outputs twice and double-parse plan/lead
         def _slice(tensor_out, meta) -> NumpyDict:
             flat = tensor_out.numpy().flatten()
             return {k: flat[np.newaxis, sl] for k, sl in meta['output_slices'].items() if k != 'pad'}