Update input handling logic in modeld.py

Refined how input keys are processed to ensure proper assignment, excluding 'desire'. This improves flexibility and correctness when managing input buffers in the model pipeline.

sunnypilot/modeld/modeld.py

@@ -21,8 +21,10 @@ from openpilot.sunnypilot.modeld.parse_model_outputs import Parser
 from openpilot.sunnypilot.modeld.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState
 from openpilot.sunnypilot.modeld.constants import ModelConstants
 from openpilot.sunnypilot.modeld.models.commonmodel_pyx import DrivingModelFrame, CLContext
+from openpilot.common.realtime import DT_MDL
+from openpilot.common.numpy_fast import interp
 
-from openpilot.sunnypilot.modeld.runners.run_helpers import load_model, load_metadata, prepare_inputs, parse_runner_inputs
+from openpilot.sunnypilot.modeld.runners.run_helpers import load_model, load_metadata, prepare_inputs
 
 PROCESS_NAME = "sunnypilot.modeld.modeld"
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
@@ -52,14 +54,12 @@ class ModelState:
     self.full_features_20Hz = np.zeros((ModelConstants.FULL_HISTORY_BUFFER_LEN, ModelConstants.FEATURE_LEN), dtype=np.float32)
     self.desire_20Hz =  np.zeros((ModelConstants.FULL_HISTORY_BUFFER_LEN + 1, ModelConstants.DESIRE_LEN), dtype=np.float32)
 
-    # img buffers are managed in openCL transform code
-    self.inputs = prepare_inputs()
-
     model_paths = load_model()
-    model_metadata = load_metadata()
+    self.model_metadata = load_metadata()
+    self.inputs = prepare_inputs(self.model_metadata)
 
-    self.output_slices = model_metadata['output_slices']
-    net_output_size = model_metadata['output_shapes']['outputs'][1]
+    self.output_slices = self.model_metadata['output_slices']
+    net_output_size = self.model_metadata['output_shapes']['outputs'][1]
     self.output = np.zeros(net_output_size, dtype=np.float32)
     self.parser = Parser()
 
@@ -69,6 +69,14 @@ class ModelState:
     for k,v in self.inputs.items():
       self.model.addInput(k, v)
 
+    num_elements = self.model_metadata['input_shapes']['features_buffer'][1]
+    step_size = int(-100 / num_elements)
+    self.feature_buffer_idxs = np.arange(step_size, step_size * (num_elements + 1), step_size)[::-1]
+
+    desired_shape = self.model_metadata["input_shapes"]["desire"][1]
+    middle_dim = int(self.desire_20Hz.shape[0] / desired_shape)
+    self.desire_reshape_dims = (desired_shape, middle_dim, -1)
+
   def slice_outputs(self, model_outputs: np.ndarray) -> dict[str, np.ndarray]:
     parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in self.output_slices.items()}
     if SEND_RAW_PRED:
@@ -84,7 +92,11 @@ class ModelState:
 
     self.desire_20Hz[:-1] = self.desire_20Hz[1:]
     self.desire_20Hz[-1] = new_desire
-    self.inputs['desire'][:] = self.desire_20Hz.reshape((25,4,-1)).max(axis=1).flatten()
+    self.inputs['desire'][:] = self.desire_20Hz.reshape(self.desire_reshape_dims).max(axis=1).flatten()
+
+    for key in self.inputs:
+      if key in inputs and key not in ['desire']:
+        self.inputs[key][:] = inputs[key]
 
     self.inputs['traffic_convention'][:] = inputs['traffic_convention']
 
@@ -95,13 +107,29 @@ class ModelState:
       return None
 
     self.model.execute()
-    outputs = self.parser.parse_outputs(self.slice_outputs(self.output))
+    outputs = self.parser.parse_outputs(self.slice_outputs(self.output), self.inputs.keys())
 
     self.full_features_20Hz[:-1] = self.full_features_20Hz[1:]
     self.full_features_20Hz[-1] = outputs['hidden_state'][0, :]
 
-    idxs = np.arange(-4,-100,-4)[::-1]
-    self.inputs['features_buffer'][:] = self.full_features_20Hz[idxs].flatten()
+    self.inputs['features_buffer'][:] = self.full_features_20Hz[self.feature_buffer_idxs].flatten()
+    if "desired_curvature" in outputs:
+      input_name_prev = None
+
+      if "prev_desired_curvs" in self.inputs.keys():
+        input_name_prev = 'prev_desired_curvs'
+      elif "prev_desired_curv" in self.inputs.keys():
+        input_name_prev = 'prev_desired_curv'
+
+      if input_name_prev is not None:
+        len = outputs['desired_curvature'][0].size
+        self.inputs[input_name_prev][:-len] = self.inputs[input_name_prev][len:]
+        self.inputs[input_name_prev][-len:] = outputs['desired_curvature'][0, :]
+
+    if "lat_planner_solution" in outputs:
+      if "lat_planner_state" in self.inputs.keys():
+        self.inputs['lat_planner_state'][2] = interp(DT_MDL, ModelConstants.T_IDXS, outputs['lat_planner_solution'][0, :, 2])
+        self.inputs['lat_planner_state'][3] = interp(DT_MDL, ModelConstants.T_IDXS, outputs['lat_planner_solution'][0, :, 3])
     return outputs
 
 
@@ -239,7 +267,24 @@ def main(demo=False):
     if prepare_only:
       cloudlog.error(f"skipping model eval. Dropped {vipc_dropped_frames} frames")
 
-    inputs = parse_runner_inputs(vec_desire, traffic_convention)
+    inputs: dict[str, np.ndarray] = {
+      'desire': vec_desire,
+      'traffic_convention': traffic_convention,
+    }
+
+    if "lateral_control_params" in model.inputs.keys():
+      inputs['lateral_control_params'] = np.array([sm["carState"].vEgo, steer_delay], dtype=np.float32)
+
+    # TODO-SP: Below should be good, but I have not tested a model with it so I can't be sure until we test it
+    # if "driving_style" in model.inputs.keys():
+    #  inputs['driving_style'] = np.array([1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0], dtype=np.float32)
+    # 
+    # if "nav_features" in model.inputs.keys():
+    #   inputs['nav_features'] = np.zeros(ModelConstants.NAV_FEATURE_LEN, dtype=np.float32)  # Get size from shape
+    # 
+    # if "nav_instructions" in model.inputs.keys():
+    #   inputs['nav_instructions'] = np.zeros(ModelConstants.NAV_INSTRUCTION_LEN, dtype=np.float32)  # Get size from shape
+
 
     mt1 = time.perf_counter()
     model_output = model.run(buf_main, buf_extra, model_transform_main, model_transform_extra, inputs, prepare_only)

sunnypilot/modeld/parse_model_outputs.py

@@ -84,7 +84,8 @@ class Parser:
     outs[name] = pred_mu_final.reshape(final_shape)
     outs[name + '_stds'] = pred_std_final.reshape(final_shape)
 
-  def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
+  def parse_outputs(self, outs: dict[str, np.ndarray], input_keys: [str]) -> dict[str, np.ndarray]:
+    """ Parse the model outputs into a dictionary of numpy arrays. The input_keys are used to determine how the output should be parsed. """
     self.parse_mdn('plan', outs, in_N=ModelConstants.PLAN_MHP_N, out_N=ModelConstants.PLAN_MHP_SELECTION,
                    out_shape=(ModelConstants.IDX_N,ModelConstants.PLAN_WIDTH))
     self.parse_mdn('lane_lines', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_LANE_LINES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH))
@@ -96,6 +97,8 @@ class Parser:
                    out_shape=(ModelConstants.LEAD_TRAJ_LEN,ModelConstants.LEAD_WIDTH))
     if 'lat_planner_solution' in outs:
       self.parse_mdn('lat_planner_solution', outs, in_N=0, out_N=0, out_shape=(ModelConstants.IDX_N,ModelConstants.LAT_PLANNER_SOLUTION_WIDTH))
+    if 'desired_curvature' in outs and "prev_desired_curv" in input_keys:
+      self.parse_mdn('desired_curvature', outs, in_N=0, out_N=0, out_shape=(ModelConstants.DESIRED_CURV_WIDTH,))
     for k in ['lead_prob', 'lane_lines_prob', 'meta']:
       self.parse_binary_crossentropy(k, outs)
     self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,))

sunnypilot/modeld/runners/onnxmodel.py

@@ -40,6 +40,7 @@ class ONNXModel(RunModel):
   def __init__(self, path, output, runtime, use_tf8, cl_context):
     self.inputs = {}
     self.output = output
+    self.use_tf8 = use_tf8
 
     self.session = create_ort_session(path, fp16_to_fp32=True)
     self.input_names = [x.name for x in self.session.get_inputs()]
@@ -63,7 +64,11 @@ class ONNXModel(RunModel):
     return None
 
   def execute(self):
-    inputs = {k: v.view(self.input_dtypes[k]) for k,v in self.inputs.items()}
+    # TODO-SP: The input below causes issues because its converting the input data when in reality it doesn't need conversion as it was already the target type.
+    #    I am leaving this comment and the input down because this needs to be looked before merging. I had similar issues when trying the tinygrad runner...
+    #    Also I checked to see if I found a similar change like this on thneed but I didn't find any, so probably thneed is still working fine.
+    # inputs = {k: v.view(self.input_dtypes[k]) for k,v in self.inputs.items()}
+    inputs = {k: (v.view(np.uint8) / 255. if self.use_tf8 and k == 'input_img' else v) for k,v in self.inputs.items()}
     inputs = {k: v.reshape(self.input_shapes[k]).astype(self.input_dtypes[k]) for k,v in inputs.items()}
     outputs = self.session.run(None, inputs)
     assert len(outputs) == 1, "Only single model outputs are supported"

sunnypilot/modeld/runners/run_helpers.py

@@ -13,30 +13,30 @@ from openpilot.sunnypilot.models.helpers import get_active_bundle
 from openpilot.system.hardware import PC
 from openpilot.system.hardware.hw import Paths
 
-USE_ONNX = int(os.getenv('USE_ONNX', str(int(PC))))
+USE_ONNX = os.getenv('USE_ONNX', PC)
 
 CUSTOM_MODEL_PATH = Paths.model_root()
-METADATA_PATH = Path(__file__).parent / 'models/supercombo_metadata.pkl'
+METADATA_PATH = Path(__file__).parent / '../models/supercombo_metadata.pkl'
 
 ModelManager = custom.ModelManagerSP
 
 
 def load_model():
   if USE_ONNX:
-    model_paths = {ModelRunner.ONNX: Path(__file__).parent / 'models/supercombo.onnx'}
+    model_paths = {ModelRunner.ONNX: Path(__file__).parent / '../models/supercombo.onnx'}
   elif bundle := get_active_bundle():
     drive_model = next(model for model in bundle.models if model.type == ModelManager.Type.drive)
     model_paths = {ModelRunner.THNEED: f"{CUSTOM_MODEL_PATH}/{drive_model.fileName}"}
   else:
-    model_paths = {ModelRunner.THNEED: Path(__file__).parent / 'models/supercombo.thneed'}
+    model_paths = {ModelRunner.THNEED: Path(__file__).parent / '../models/supercombo.thneed'}
 
   return model_paths
 
 
 def load_metadata():
   if bundle := get_active_bundle():
-    drive_model = next(model for model in bundle.models if model.type == ModelManager.Type.metadata)
-    metadata_path = f"{CUSTOM_MODEL_PATH}/{drive_model.fileName}"
+    metadata_model = next(model for model in bundle.models if model.type == ModelManager.Type.metadata)
+    metadata_path = f"{CUSTOM_MODEL_PATH}/{metadata_model.fileName}"
   else:
     metadata_path = METADATA_PATH
 
@@ -46,22 +46,12 @@ def load_metadata():
   return metadata
 
 
-def prepare_inputs() -> dict[str, np.ndarray]:
-  model_metadata = load_metadata()
-
+def prepare_inputs(model_metadata) -> dict[str, np.ndarray]:
+  # img buffers are managed in openCL transform code so we don't pass them as inputs
   inputs: dict[str, np.ndarray] = {
-    key: np.zeros(shape, dtype=np.float32)
+    key: np.zeros(shape, dtype=np.float32).flatten()  # Inputs were defined flattened back then
     for key, shape in model_metadata['input_shapes'].items()
     if key not in ['input_imgs', 'big_input_imgs']
   }
 
   return inputs
-
-
-def parse_runner_inputs(vec_desire, traffic_convention) -> dict[str, np.ndarray]:
-  inputs: dict[str, np.ndarray] = {
-    'desire': vec_desire,
-    'traffic_convention': traffic_convention,
-  }
-
-  return inputs