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

@@ -1,11 +1,9 @@
 #!/usr/bin/env python3
 import os
 import time
-import pickle
 import numpy as np
 import cereal.messaging as messaging
 from cereal import car, log
-from pathlib import Path
 from setproctitle import setproctitle
 from cereal.messaging import PubMaster, SubMaster
 from msgq.visionipc import VisionIpcClient, VisionStreamType, VisionBuf
@@ -23,16 +21,14 @@ 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
 
 PROCESS_NAME = "sunnypilot.modeld.modeld"
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
 
-MODEL_PATHS = {
-  ModelRunner.THNEED: Path(__file__).parent / 'models/supercombo.thneed',
-  ModelRunner.ONNX: Path(__file__).parent / 'models/supercombo.onnx'}
-
-METADATA_PATH = Path(__file__).parent / 'models/supercombo_metadata.pkl'
-
 
 class FrameMeta:
   frame_id: int = 0
@@ -58,27 +54,29 @@ 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 = {
-      'desire': np.zeros(ModelConstants.DESIRE_LEN * (ModelConstants.HISTORY_BUFFER_LEN+1), dtype=np.float32),
-      'traffic_convention': np.zeros(ModelConstants.TRAFFIC_CONVENTION_LEN, dtype=np.float32),
-      'features_buffer': np.zeros(ModelConstants.HISTORY_BUFFER_LEN * ModelConstants.FEATURE_LEN, dtype=np.float32),
-    }
+    model_paths = load_model()
+    self.model_metadata = load_metadata()
+    self.inputs = prepare_inputs(self.model_metadata)
 
-    with open(METADATA_PATH, 'rb') as f:
-      model_metadata = pickle.load(f)
-
-    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()
 
-    self.model = ModelRunner(MODEL_PATHS, self.output, Runtime.GPU, False, context)
+    self.model = ModelRunner(model_paths, self.output, Runtime.GPU, False, context)
     self.model.addInput("input_imgs", None)
     self.model.addInput("big_input_imgs", None)
     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:
@@ -94,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']
 
@@ -105,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
 
 
@@ -249,10 +267,24 @@ def main(demo=False):
     if prepare_only:
       cloudlog.error(f"skipping model eval. Dropped {vipc_dropped_frames} frames")
 
-    inputs:dict[str, np.ndarray] = {
+    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

@@ -0,0 +1,57 @@
+# Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
+#
+# This file is part of sunnypilot and is licensed under the MIT License.
+# See the LICENSE.md file in the root directory for more details.
+
+import os
+import pickle
+import numpy as np
+from pathlib import Path
+from cereal import custom
+from openpilot.sunnypilot.modeld.runners import ModelRunner
+from openpilot.sunnypilot.models.helpers import get_active_bundle
+from openpilot.system.hardware import PC
+from openpilot.system.hardware.hw import Paths
+
+USE_ONNX = os.getenv('USE_ONNX', PC)
+
+CUSTOM_MODEL_PATH = Paths.model_root()
+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'}
+  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'}
+
+  return model_paths
+
+
+def load_metadata():
+  if bundle := get_active_bundle():
+    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
+
+  with open(metadata_path, 'rb') as f:
+    metadata = pickle.load(f)
+
+  return 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).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

sunnypilot/models/helpers.py

@@ -22,7 +22,7 @@ async def verify_file(file_path: str, expected_hash: str) -> bool:
   return sha256_hash.hexdigest().lower() == expected_hash.lower()
 
 
-def get_active_bundle(params: Params) -> custom.ModelManagerSP.ModelBundle:
+def get_active_bundle(params: Params = None) -> custom.ModelManagerSP.ModelBundle:
   """Gets the active model bundle from cache"""
   if params is None:
     params = Params()