Run misc in modeld

selfdrive/modeld/modeld.py

@@ -38,8 +38,10 @@ SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
 
 VISION_PKL_PATH = Path(__file__).parent / 'models/driving_vision_tinygrad.pkl'
 POLICY_PKL_PATH = Path(__file__).parent / 'models/driving_policy_tinygrad.pkl'
+MISC_PKL_PATH = Path(__file__).parent / 'models/driving_misc_tinygrad.pkl'
 VISION_METADATA_PATH = Path(__file__).parent / 'models/driving_vision_metadata.pkl'
 POLICY_METADATA_PATH = Path(__file__).parent / 'models/driving_policy_metadata.pkl'
+MISC_METADATA_PATH = Path(__file__).parent / 'models/driving_misc_metadata.pkl'
 
 LAT_SMOOTH_SECONDS = 0.3
 LONG_SMOOTH_SECONDS = 0.3
@@ -113,11 +115,19 @@ class ModelState:
       self.policy_output_slices = policy_metadata['output_slices']
       policy_output_size = policy_metadata['output_shapes']['outputs'][1]
 
+    with open(MISC_METADATA_PATH, 'rb') as f:
+      misc_metadata = pickle.load(f)
+      self.misc_input_shapes =  misc_metadata['input_shapes']
+      self.misc_output_slices = misc_metadata['output_slices']
+      misc_output_size = misc_metadata['output_shapes']['outputs'][1]
+
     # img buffers are managed in openCL transform code
     self.vision_inputs: dict[str, Tensor] = {}
     self.vision_output = np.zeros(vision_output_size, dtype=np.float32)
     self.policy_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()}
     self.policy_output = np.zeros(policy_output_size, dtype=np.float32)
+    self.misc_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items() if k not in {"prev_desired_curv", "lateral_control_params"}}
+    self.misc_output = np.zeros(misc_output_size, dtype=np.float32)
     self.parser = Parser()
 
     with open(VISION_PKL_PATH, "rb") as f:
@@ -126,6 +136,9 @@ class ModelState:
     with open(POLICY_PKL_PATH, "rb") as f:
       self.policy_run = pickle.load(f)
 
+    with open(MISC_PKL_PATH, "rb") as f:
+      self.misc_run = pickle.load(f)
+
   def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]:
     parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()}
     return parsed_model_outputs
@@ -169,12 +182,15 @@ class ModelState:
     self.policy_output = self.policy_run(**self.policy_inputs).numpy().flatten()
     policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(self.policy_output, self.policy_output_slices))
 
+    self.misc_output = self.misc_run(**self.misc_inputs).numpy().flatten()
+    misc_outputs_dict = self.parser.parse_misc_outputs(self.slice_outputs(self.misc_output, self.misc_output_slices))
+
     # TODO model only uses last value now
     self.full_prev_desired_curv[0,:-1] = self.full_prev_desired_curv[0,1:]
     self.full_prev_desired_curv[0,-1,:] = policy_outputs_dict['desired_curvature'][0, :]
     self.numpy_inputs['prev_desired_curv'][:] = 0*self.full_prev_desired_curv[0, self.temporal_idxs]
 
-    combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict}
+    combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict, **misc_outputs_dict}
     if SEND_RAW_PRED:
       combined_outputs_dict['raw_pred'] = np.concatenate([self.vision_output.copy(), self.policy_output.copy()])