modeld_v2: conditional model compilation for metadrive testing (#1623)

* modeld_v2: conditional model compilation for PC

* full send

* shebang

---------

Co-authored-by: Jason Wen <haibin.wen3@gmail.com>

sunnypilot/modeld_v2/SConscript

@@ -1,3 +1,4 @@
+import os
 import glob
 
 Import('env', 'envCython', 'arch', 'cereal', 'messaging', 'common', 'visionipc', 'transformations')
@@ -28,3 +29,38 @@ for pathdef, fn in {'TRANSFORM': 'transforms/transform.cl', 'LOADYUV': 'transfor
 cython_libs = envCython["LIBS"] + libs
 commonmodel_lib = lenv.Library('commonmodel', common_src)
 lenvCython.Program('models/commonmodel_pyx.so', 'models/commonmodel_pyx.pyx', LIBS=[commonmodel_lib, *cython_libs], FRAMEWORKS=frameworks)
+tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=env.Dir("#").abspath) if 'pycache' not in x]
+
+# Get model metadata
+PC = not os.path.isfile('/TICI')
+if PC:
+  inputs = tinygrad_files + [File(Dir("#sunnypilot/modeld_v2").File("install_models_pc.py").abspath)]
+  outputs = []
+  model_dir = Dir("models").abspath
+  cmd = f'python3 {Dir("#sunnypilot/modeld_v2").abspath}/install_models_pc.py {model_dir}'
+
+  for model_name in ['supercombo', 'driving_vision', 'driving_policy']:
+    if File(f"models/{model_name}.onnx").exists():
+        inputs.append(File(f"models/{model_name}.onnx"))
+        inputs.append(File(f"models/{model_name}_tinygrad.pkl"))
+        outputs.append(File(f"models/{model_name}_metadata.pkl"))
+  if outputs:
+    lenv.Command(outputs, inputs, cmd)
+
+def tg_compile(flags, model_name):
+  pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
+  fn = File(f"models/{model_name}").abspath
+  return lenv.Command(
+    fn + "_tinygrad.pkl",
+    [fn + ".onnx"] + tinygrad_files,
+    f'{pythonpath_string} {flags} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {fn}_tinygrad.pkl'
+  )
+
+# Compile small models
+for model_name in ['supercombo', 'driving_vision', 'driving_policy']:
+  if File(f"models/{model_name}.onnx").exists():
+    flags = {
+      'larch64': 'DEV=QCOM',
+      'Darwin': f'DEV=CPU HOME={os.path.expanduser("~")} IMAGE=0', # tinygrad calls brew which needs a $HOME in the env
+    }.get(arch, 'DEV=CPU CPU_LLVM=1 IMAGE=0')
+    tg_compile(flags, model_name)

sunnypilot/modeld_v2/install_models_pc.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+import sys
+import shutil
+import pickle
+import codecs
+import onnx
+from pathlib import Path
+
+from openpilot.system.hardware.hw import Paths
+
+
+def get_name_and_shape(value_info):
+  shape = tuple([int(dim.dim_value) for dim in value_info.type.tensor_type.shape.dim])
+  return value_info.name, shape
+
+
+def get_metadata_value_by_name(model, name):
+  for prop in model.metadata_props:
+    if prop.key == name:
+      return prop.value
+  return None
+
+
+def generate_metadata_pkl(model_path, output_path):
+  try:
+    model = onnx.load(str(model_path))
+    output_slices = get_metadata_value_by_name(model, 'output_slices')
+
+    if output_slices:
+      metadata = {
+        'model_checkpoint': get_metadata_value_by_name(model, 'model_checkpoint'),
+        'output_slices': pickle.loads(codecs.decode(output_slices.encode(), "base64")),
+        'input_shapes': dict([get_name_and_shape(x) for x in model.graph.input]),
+        'output_shapes': dict([get_name_and_shape(x) for x in model.graph.output])
+      }
+      with open(output_path, 'wb') as f:
+        pickle.dump(metadata, f)
+      return True
+    else:
+      return False
+  except Exception:
+    return False
+
+
+def install_models(model_dir):
+  model_dir = Path(model_dir)
+  models = ["driving_policy", "driving_vision"]
+  found_models = []
+
+  for model in models:
+    if (model_dir / f"{model}.onnx").exists():
+      found_models.append(model)
+
+  if not found_models:
+    return
+
+  try:
+    custom_name = input(f"Found models ({', '.join(found_models)}). Enter model short name (e.g. wmiv4): ").strip()
+  except EOFError:
+    return
+
+  if not custom_name:
+    print("No name provided, skipping installation.")
+    return
+
+  dest_dir = Path(Paths.model_root())
+  dest_dir.mkdir(parents=True, exist_ok=True)
+
+  for model in found_models:
+    onnx_path = model_dir / f"{model}.onnx"
+    tinygrad_pkl = model_dir / f"{model}_tinygrad.pkl"
+    metadata_pkl = model_dir / f"{model}_metadata.pkl"
+
+    if not metadata_pkl.exists():
+      generate_metadata_pkl(onnx_path, metadata_pkl)
+
+    dest_tinygrad = dest_dir / f"{model}_{custom_name}_tinygrad.pkl"
+    dest_metadata = dest_dir / f"{model}_{custom_name}_metadata.pkl"
+
+    if tinygrad_pkl.exists():
+      shutil.move(str(tinygrad_pkl), str(dest_tinygrad))
+    if metadata_pkl.exists():
+      shutil.move(str(metadata_pkl), str(dest_metadata))
+
+if __name__ == "__main__":
+  if len(sys.argv) < 2:
+    print("Usage: install_models_pc.py <model_dir>")
+    sys.exit(1)
+  install_models(sys.argv[1])

sunnypilot/models/runners/helpers.py

@@ -2,25 +2,17 @@ from openpilot.sunnypilot.models.helpers import get_active_bundle
 from openpilot.sunnypilot.models.runners.model_runner import ModelRunner
 from openpilot.sunnypilot.models.runners.tinygrad.tinygrad_runner import TinygradRunner, TinygradSplitRunner
 from openpilot.sunnypilot.models.runners.constants import ModelType
-from openpilot.system.hardware import TICI
 
-if not TICI:
-  from openpilot.sunnypilot.models.runners.onnx.onnx_runner import ONNXRunner
 
 def get_model_runner() -> ModelRunner:
   """
   Factory function to create and return the appropriate ModelRunner instance.
 
-  Selects between ONNXRunner (for non-TICI platforms) and TinygradRunner
-  (for TICI platforms), choosing TinygradSplitRunner if separate vision/policy
+  Selects TinygradRunner, choosing TinygradSplitRunner if separate vision/policy
   models are detected in the active bundle.
 
   :return: An instance of a ModelRunner subclass (ONNXRunner, TinygradRunner, or TinygradSplitRunner).
   """
-  if not TICI:
-    return ONNXRunner()
-
-  # On TICI platforms, use Tinygrad runners
   bundle = get_active_bundle()
   if bundle and bundle.models:
     model_types = {m.type.raw for m in bundle.models}