load model and metadata dynamically

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
@@ -24,15 +22,11 @@ from openpilot.sunnypilot.modeld.fill_model_msg import fill_model_msg, fill_pose
 from openpilot.sunnypilot.modeld.constants import ModelConstants
 from openpilot.sunnypilot.modeld.models.commonmodel_pyx import DrivingModelFrame, CLContext
 
-from openpilot.sunnypilot.modeld.runners.run_helpers import METADATA_PATH, prepare_inputs, parse_runner_inputs
+from openpilot.sunnypilot.modeld.runners.run_helpers import load_model, load_metadata, prepare_inputs, parse_runner_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'}
-
 
 class FrameMeta:
   frame_id: int = 0
@@ -61,15 +55,15 @@ class ModelState:
     # img buffers are managed in openCL transform code
     self.inputs = prepare_inputs()
 
-    with open(METADATA_PATH, 'rb') as f:
-      model_metadata = pickle.load(f)
+    model_paths = load_model()
+    model_metadata = load_metadata()
 
     self.output_slices = model_metadata['output_slices']
     net_output_size = 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():

sunnypilot/modeld/runners/run_helpers.py

@@ -3,15 +3,44 @@
 # 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 = int(os.getenv('USE_ONNX', str(int(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():
-  with open(METADATA_PATH, 'rb') as f:
+  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}"
+  else:
+    metadata_path = METADATA_PATH
+
+  with open(metadata_path, 'rb') as f:
     metadata = pickle.load(f)
 
   return metadata