diff --git a/sunnypilot/modeld_v2/SConscript b/sunnypilot/modeld_v2/SConscript index 4c04b7382..750a242ad 100644 --- a/sunnypilot/modeld_v2/SConscript +++ b/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) diff --git a/sunnypilot/modeld_v2/install_models_pc.py b/sunnypilot/modeld_v2/install_models_pc.py new file mode 100755 index 000000000..3f964dc28 --- /dev/null +++ b/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 ") + sys.exit(1) + install_models(sys.argv[1]) diff --git a/sunnypilot/models/runners/helpers.py b/sunnypilot/models/runners/helpers.py index 6a128b340..8f9d8fc2f 100644 --- a/sunnypilot/models/runners/helpers.py +++ b/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}