Use upstream tinygrad disk tensors for model loading (#38956)

openpilot/selfdrive/modeld/dmonitoringmodeld.py

@@ -27,8 +27,7 @@ class ModelState:
   output: np.ndarray
 
   def __init__(self, cam_w: int, cam_h: int):
-    with open(MODEL_PKL_PATH, 'rb') as f:
-      jits = load_oob(f)
+    jits = load_oob(MODEL_PKL_PATH)
     self.DEV = jits['input_specs']['input_img'][2]
     self.input_shapes = jits['metadata']['input_shapes']
     self.output_slices = pickle.loads(base64.b64decode(jits['metadata']['metadata']['output_slices']))

openpilot/selfdrive/modeld/helpers.py

@@ -1,8 +1,7 @@
-import mmap
-import pickle
-import struct
+import sys
 from pathlib import Path
 
+from openpilot.common.hardware import AGNOS
 from openpilot.common.hardware.usb import CHESTNUT_USB_PRODUCT, USB_DEVICES_PATH, is_chestnut_usb_id
 
 MODELS_DIR = Path(__file__).resolve().parent / 'models'
@@ -12,19 +11,12 @@ def modeld_pkl_path(chestnut: bool):
   prefix = 'big_' if chestnut else ''
   return MODELS_DIR / f'{prefix}driving_tinygrad.pkl'
 
-def load_oob(f):
-  data = memoryview(mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_COPY))
-  opcode_size = struct.unpack_from('<q', data)[0]
-  def buffers():
-    offset = 8 + opcode_size
-    while offset < len(data):
-      size = struct.unpack_from('<q', data, offset)[0]
-      offset += 8
-      if offset + size > len(data):
-        raise EOFError("incomplete model buffer")
-      yield pickle.PickleBuffer(data[offset:offset + size])
-      offset += size
-  return pickle.loads(data[8:8 + opcode_size], buffers=buffers())
+def load_oob(path, chestnut=False):
+  from tinygrad import Context
+  device = 'USB+AMD:LLVM' if chestnut else 'QCOM' if AGNOS else 'METAL' if sys.platform == 'darwin' else 'CPU:LLVM'
+  with Context(DEV=device):
+    from tinygrad_repo.examples.openpilot.helpers import load_pickle
+    return load_pickle(path, out_of_band=True)
 
 def chestnut_present() -> bool:
   for d in USB_DEVICES_PATH.glob("*"):

openpilot/selfdrive/modeld/modeld.py

@@ -139,8 +139,7 @@ class ModelState:
   prev_desire: np.ndarray  # for tracking the rising edge of the pulse
 
   def __init__(self, cam_w: int, cam_h: int, chestnut: bool):
-    with open(modeld_pkl_path(chestnut), 'rb') as f:
-      jits = load_oob(f)
+    jits = load_oob(modeld_pkl_path(chestnut), chestnut)
     self.model_device = jits['input_specs']['new_img'][2]
     self.input_shapes = {name: (shape, np.dtype(dtype)) for name, (shape, dtype, _) in jits['input_specs'].items()}
     self.state_pairs = {name: f'next_{name}' for name in self.input_shapes if f'next_{name}' in jits['metadata']['output_shapes']}

openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e00a53cb9c64fe70a58224bc7879e5ea1b7e18baa7710f0929bf4f2e4537a468
-size 776612574
+oid sha256:a5fc6813f45be5673382d6571a5fc598a7cd8e47c534101fbc3a45767c169c53
+size 773082315