mirror of
https://github.com/sunnypilot/sunnypilot.git
synced 2026-10-19 16:13:45 +08:00
Adeeb Shihadeh
72 lines
3.4 KiB
Python
72 lines
3.4 KiB
Python
import ctypes.util
|
|
import glob
|
|
import os
|
|
from SCons.Script import Action, Value
|
|
from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
|
|
from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE
|
|
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
|
|
|
|
|
Import('env', 'arch')
|
|
|
|
# TODO: add GCC support to tinygrad?
|
|
# tinygrad's CPU:LLVM backend requires libLLVM, skip the build on PC if it's not available
|
|
if arch not in ('comma_arm64', 'Darwin') and ctypes.util.find_library('LLVM') is None:
|
|
print("\033[93mmodeld: LLVM not found, skipping tinygrad model compilation. Install LLVM to compile models.\033[0m")
|
|
Return()
|
|
|
|
lenv = env.Clone()
|
|
lenv.PrependENVPath('PYTHONPATH', Dir('#tinygrad_repo').abspath)
|
|
|
|
tinygrad_root = env.Dir("#").abspath
|
|
tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=tinygrad_root)
|
|
if 'pycache' not in x and os.path.isfile(os.path.join(tinygrad_root, x))]
|
|
|
|
camera_configs = [(c.width, c.height) for c in (_ar_ox_fisheye, _os_fisheye)]
|
|
|
|
if arch == 'comma_arm64':
|
|
tg_flags = 'DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1'
|
|
else:
|
|
# JIT=2 disables graph batching, which produces incorrect outputs after buffers change.
|
|
tg_flags = 'DEV=METAL JIT=2' if arch == 'Darwin' else 'DEV=CPU:LLVM'
|
|
|
|
# tinygrad calls brew which needs a $HOME in the env
|
|
mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else ''
|
|
|
|
warp_deps = [File("#openpilot/system/camerad/cameras/nv12_info.py")]
|
|
compiler = Dir('#tinygrad_repo/examples/openpilot').abspath
|
|
# build on all cores
|
|
taskset = 'taskset -c 0-7 ' if arch == 'comma_arm64' else ''
|
|
|
|
def compile_model(onnx_path, pkl_path):
|
|
onnx_path, target_pkl_path = File(onnx_path).abspath, File(pkl_path).abspath
|
|
cmd = (f'{tg_flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_onnx.py" '
|
|
f'"{onnx_path}" "{target_pkl_path}" --out-of-band --benchmark-runs 1')
|
|
lenv.Command(
|
|
target_pkl_path,
|
|
tinygrad_files + [onnx_path, Value(cmd)],
|
|
Action(cmd, " [ONNX] $TARGET"),
|
|
)
|
|
|
|
compile_model('models/dmonitoring_model.onnx', 'models/dmonitoring_model_tinygrad.pkl')
|
|
compile_model('models/driving_supercombo.onnx', 'models/driving_tinygrad.pkl')
|
|
|
|
model_w, model_h = MEDMODEL_INPUT_SIZE
|
|
for cam_w, cam_h in camera_configs:
|
|
warp_pkl_path = File(f"models/driving_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
|
|
stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h)
|
|
cmd = (f'{tg_flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_warp.py" '
|
|
f'--frame {cam_w},{cam_h},{stride},{y_height},{uv_height},{stride * (y_height + uv_height)} '
|
|
f'--warp-to {model_w}x{model_h} --layout yuv420 --frames 2 '
|
|
f'--output {warp_pkl_path}')
|
|
lenv.Command(warp_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], cmd)
|
|
|
|
dm_w, dm_h = DM_INPUT_SIZE
|
|
for cam_w, cam_h in camera_configs:
|
|
dm_pkl_path = File(f"models/dm_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
|
|
stride, y_height, uv_height, frame_size = get_nv12_info(cam_w, cam_h)
|
|
cmd = (f'{tg_flags} {mac_brew_string} python3 "{compiler}/compile_warp.py" '
|
|
f'--frame {cam_w},{cam_h},{stride},{y_height},{uv_height},{frame_size} --warp-to {dm_w}x{dm_h} '
|
|
f'--layout luma --border-fill 16 --transform-device NPY --output {dm_pkl_path}')
|
|
lenv.Command(dm_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], cmd)
|