f9fcc7adab
date: 2026-06-28T09:48:35 master commit: da6313dbe95b3f24bb5d8018b0e5f950f5823ca7
100 lines
4.4 KiB
Python
100 lines
4.4 KiB
Python
import os
|
|
import glob
|
|
|
|
from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
|
|
from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE
|
|
from openpilot.system.hardware import HARDWARE, PC
|
|
|
|
Import('env', 'arch', 'release')
|
|
lenv = env.Clone()
|
|
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]
|
|
|
|
|
|
def get_camera_configs():
|
|
DEVICE_RESOLUTIONS = {
|
|
"tici": (_ar_ox_fisheye.width, _ar_ox_fisheye.height),
|
|
"tizi": (_ar_ox_fisheye.width, _ar_ox_fisheye.height),
|
|
"mici": (_os_fisheye.width, _os_fisheye.height),
|
|
}
|
|
if release or PC or 'CI' in os.environ:
|
|
return set(DEVICE_RESOLUTIONS.values())
|
|
return [DEVICE_RESOLUTIONS[HARDWARE.get_device_type()]]
|
|
|
|
CAMERA_CONFIGS = get_camera_configs()
|
|
|
|
tg_flags = {
|
|
'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0',
|
|
'Darwin': f'DEV=CPU HOME={os.path.expanduser("~")}',
|
|
}.get(arch, 'DEV=CPU:LLVM')
|
|
|
|
image_flag = {
|
|
'larch64': 'IMAGE=2',
|
|
}.get(arch, 'IMAGE=0')
|
|
|
|
model_w, model_h = MEDMODEL_INPUT_SIZE
|
|
from openpilot.selfdrive.modeld.constants import ModelConstants
|
|
frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
|
|
camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in CAMERA_CONFIGS)
|
|
|
|
pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + ':' + env.Dir("#").abspath + '"'
|
|
compile_modeld_script = File("compile_modeld.py").abspath
|
|
upstream_compile_script = File(Dir("#selfdrive/modeld").File("compile_modeld.py").abspath)
|
|
script_deps = [File("compile_modeld.py"), upstream_compile_script]
|
|
|
|
def compile_combined(model_type, onnx_args, output_name):
|
|
output_pkl = File(f"models/{output_name}").abspath
|
|
cmd = (f'{pythonpath_string} {tg_flags} {image_flag} python3 {compile_modeld_script} '
|
|
f'--model-type {model_type} '
|
|
f'--model-size {model_w}x{model_h} '
|
|
f'--camera-resolutions {camera_res_args} '
|
|
f'{onnx_args} '
|
|
f'--frame-skip {frame_skip} '
|
|
f'--output {output_pkl}')
|
|
onnx_files = [f for f in onnx_args.split() if f.endswith('.onnx')]
|
|
return lenv.Command(output_pkl, tinygrad_files + script_deps + [File(f) for f in onnx_files if os.path.isfile(f)], cmd)
|
|
|
|
# Vision + Policy (stock default model)
|
|
vision_onnx = File("models/driving_vision.onnx").abspath
|
|
policy_onnx = File("models/driving_policy.onnx").abspath
|
|
if os.path.isfile(vision_onnx) and os.path.isfile(policy_onnx):
|
|
compile_combined('vision_policy',
|
|
f'--vision-onnx {vision_onnx} --policy-onnx {policy_onnx}',
|
|
'driving_combined_tinygrad.pkl')
|
|
|
|
# Vision + Off-Policy
|
|
off_policy_onnx = File("models/driving_off_policy.onnx").abspath
|
|
if os.path.isfile(vision_onnx) and os.path.isfile(off_policy_onnx):
|
|
policy_arg = f'--policy-onnx {policy_onnx}' if os.path.isfile(policy_onnx) else ''
|
|
compile_combined('vision_multi_policy',
|
|
f'--vision-onnx {vision_onnx} {policy_arg} --off-policy-onnx {off_policy_onnx}',
|
|
'driving_combined_multi_tinygrad.pkl')
|
|
|
|
# Vision + On-Policy + Off-Policy
|
|
on_policy_onnx = File("models/driving_on_policy.onnx").abspath
|
|
if os.path.isfile(vision_onnx) and os.path.isfile(on_policy_onnx) and os.path.isfile(off_policy_onnx):
|
|
compile_combined('vision_multi_policy',
|
|
f'--vision-onnx {vision_onnx} --off-policy-onnx {off_policy_onnx} --on-policy-onnx {on_policy_onnx}',
|
|
'driving_combined_tri_tinygrad.pkl')
|
|
|
|
# Supercombo
|
|
supercombo_onnx = File("models/supercombo.onnx").abspath
|
|
if os.path.isfile(supercombo_onnx):
|
|
compile_combined('supercombo',
|
|
f'--supercombo-onnx {supercombo_onnx}',
|
|
'driving_combined_supercombo_tinygrad.pkl')
|
|
|
|
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_off_policy', 'driving_on_policy', '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)
|
|
|