sunnypilot v2026.002.000

version: sunnypilot v2026.002.000 (staging)
date: 2026-07-21T07:44:02
master commit: 2c334ede443d7391d27575af8c854a095ba702a8

openpilot/sunnypilot/modeld_v2/.gitignore

@@ -0,0 +1 @@
+*_pyx.cpp

openpilot/sunnypilot/modeld_v2/__init__.py

@@ -0,0 +1,5 @@
+from pathlib import Path
+
+MODEL_PATH = Path(__file__).parent / 'models/supercombo.onnx'
+MODEL_PKL_PATH = Path(__file__).parent / 'models/supercombo_tinygrad.pkl'
+METADATA_PATH = Path(__file__).parent / 'models/supercombo_metadata.pkl'

openpilot/sunnypilot/modeld_v2/camera_offset_helper.py

@@ -0,0 +1,39 @@
+"""
+Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
+
+This file is part of sunnypilot and is licensed under the MIT License.
+See the LICENSE.md file in the root directory for more details.
+"""
+import numpy as np
+
+from openpilot.common.transformations.camera import DEVICE_CAMERAS
+
+
+class CameraOffsetHelper:
+  def __init__(self):
+    self.camera_offset = 0.0
+    self.actual_camera_offset = 0.0
+
+  @staticmethod
+  def apply_camera_offset(model_transform, intrinsics, height, offset_param):
+    cy = intrinsics[1, 2]
+    shear = np.eye(3, dtype=np.float32)
+    shear[0, 1] = offset_param / height
+    shear[0, 2] = -offset_param / height * cy
+    model_transform = (shear @ model_transform).astype(np.float32)
+    return model_transform
+
+  def set_offset(self, offset):
+    self.camera_offset = offset
+
+  def update(self, model_transform_main, model_transform_extra, sm, main_wide_camera):
+    self.actual_camera_offset = (0.9 * self.actual_camera_offset) + (0.1 * self.camera_offset)
+    dc = DEVICE_CAMERAS[(str(sm['deviceState'].deviceType), str(sm['roadCameraState'].sensor))]
+    height = sm["liveCalibration"].height[0] if sm['liveCalibration'].height else 1.22
+
+    intrinsics_main = dc.ecam.intrinsics if main_wide_camera else dc.fcam.intrinsics
+    model_transform_main = self.apply_camera_offset(model_transform_main, intrinsics_main, height, self.actual_camera_offset)
+
+    intrinsics_extra = dc.ecam.intrinsics
+    model_transform_extra = self.apply_camera_offset(model_transform_extra, intrinsics_extra, height, self.actual_camera_offset)
+    return model_transform_main, model_transform_extra

openpilot/sunnypilot/modeld_v2/compile_modeld.py

@@ -0,0 +1,480 @@
+#!/usr/bin/env python3
+"""
+Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
+
+This file is part of sunnypilot and is licensed under the MIT License.
+See the LICENSE.md file in the root directory for more details.
+"""
+
+import argparse
+import os
+import pickle
+import time
+from functools import partial
+from collections import defaultdict
+
+import numpy as np
+from tinygrad.tensor import Tensor
+from tinygrad.device import Device
+from tinygrad.engine.jit import TinyJit
+
+from openpilot.selfdrive.modeld.compile_modeld import (
+  NV12Frame, make_frame_prepare,
+  shift_and_sample, sample_skip, sample_desire,
+)
+
+MODEL_TYPES = ('vision_policy', 'supercombo', 'vision_multi_policy')
+
+
+def _detect_desire_key(policy_input_shapes):
+  for k in policy_input_shapes:
+    if k.startswith('desire'):
+      return k
+  return None
+
+
+def _detect_vision_keys(vision_input_shapes):
+  img_keys = sorted([k for k in vision_input_shapes if 'img' in k])
+  road_key = next((k for k in img_keys if 'big' not in k), None)
+  wide_key = next((k for k in img_keys if 'big' in k), None)
+  if road_key is None or wide_key is None:
+    raise ValueError(f"Cannot determine road/wide image keys from {list(vision_input_shapes.keys())}")
+  return road_key, wide_key
+
+
+def make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device):
+  road_key, _ = _detect_vision_keys(vision_input_shapes)
+  img = vision_input_shapes[road_key]
+  n_frames = img[1] // 6
+  img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3])
+
+  fb = policy_input_shapes['features_buffer']
+  desire_key = _detect_desire_key(policy_input_shapes)
+  dp = policy_input_shapes[desire_key]
+  tc = policy_input_shapes.get('traffic_convention', (1, 2))
+
+  npy = {
+    'desire': np.zeros(dp[2], dtype=np.float32),
+    'traffic_convention': np.zeros(tc, dtype=np.float32),
+    'tfm': np.zeros((3, 3), dtype=np.float32),
+    'big_tfm': np.zeros((3, 3), dtype=np.float32),
+  }
+
+  handled = {'features_buffer', desire_key, 'traffic_convention'}
+  for key, shape in policy_input_shapes.items():
+    if key in handled:
+      continue
+    npy[key] = np.zeros(shape, dtype=np.float32)
+
+  input_queues = {
+    'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
+    'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
+    'feat_q': Tensor(np.zeros((frame_skip * (fb[1] - 1) + 1, fb[0], fb[2]), dtype=np.float32), device=device).contiguous().realize(),
+    'desire_q': Tensor(np.zeros((frame_skip * dp[1], dp[0], dp[2]), dtype=np.float32), device=device).contiguous().realize(),
+    **{k: Tensor(v, device='NPY').realize() for k, v in npy.items()},
+  }
+  return input_queues, npy
+
+
+def make_run_split_policy(vision_runner, policy_runner, nv12: NV12Frame, model_w, model_h,
+                          vision_features_slice, frame_skip, desire_key, extra_policy_keys,
+                          vision_road_key, vision_wide_key, prepare_only=False):
+  frame_prepare = make_frame_prepare(nv12, model_w, model_h)
+  sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
+  sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
+
+  def run_policy(img_q, big_img_q, feat_q, desire_q, desire, traffic_convention, tfm, big_tfm, frame, big_frame, **extra):
+    npy_tensors = [tfm.to(Device.DEFAULT), big_tfm.to(Device.DEFAULT),
+                   desire.to(Device.DEFAULT), traffic_convention.to(Device.DEFAULT)]
+    extra_device = {k: extra[k].to(Device.DEFAULT) for k in extra_policy_keys}
+    Tensor.realize(*npy_tensors, *extra_device.values())
+    tfm, big_tfm, desire, traffic_convention = npy_tensors
+
+    img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn)
+    big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn)
+
+    if prepare_only:
+      return img, big_img
+
+    vision_out = next(iter(vision_runner({vision_road_key: img, vision_wide_key: big_img}).values())).cast('float32')
+
+    new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0)
+    feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn)
+    desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn)
+
+    inputs = {'features_buffer': feat_buf, desire_key: desire_buf, 'traffic_convention': traffic_convention, **extra_device}
+    policy_out = next(iter(policy_runner(inputs).values())).cast('float32')
+
+    return vision_out, policy_out
+  return run_policy
+
+
+def compile_split_policy(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip,
+                         vision_runner, policy_runner, vision_metadata, policy_metadata):
+  print(f"Compiling combined policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...")
+
+  vision_features_slice = vision_metadata['output_slices']['hidden_state']
+  vision_input_shapes = vision_metadata['input_shapes']
+  policy_input_shapes = policy_metadata['input_shapes']
+  desire_key = _detect_desire_key(policy_input_shapes)
+  extra_policy_keys = [k for k in policy_input_shapes if k not in ('features_buffer', desire_key, 'traffic_convention')]
+  vision_road_key, vision_wide_key = _detect_vision_keys(vision_input_shapes)
+
+  _run = make_run_split_policy(vision_runner, policy_runner, nv12, model_w, model_h,
+                               vision_features_slice, frame_skip, desire_key, extra_policy_keys,
+                               vision_road_key, vision_wide_key, prepare_only)
+  run_policy_jit = TinyJit(_run, prune=True)
+
+  SEED = 42
+
+  def random_inputs_run_fn(fn, seed, test_val=None, test_buffers=None, expect_match=True):
+    input_queues, npy = make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, Device.DEFAULT)
+    np.random.seed(seed)
+    Tensor.manual_seed(seed)
+
+    testing = test_val is not None or test_buffers is not None
+    n_runs = 1 if testing else 3
+
+    for i in range(n_runs):
+      frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize()
+      big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize()
+      for v in npy.values():
+        v[:] = np.random.randn(*v.shape).astype(v.dtype)
+      Device.default.synchronize()
+      st = time.perf_counter()
+      outs = fn(**input_queues, frame=frame, big_frame=big_frame)
+      mt = time.perf_counter()
+      Device.default.synchronize()
+      et = time.perf_counter()
+      print(f"  [{i+1}/{n_runs}] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms")
+
+      if i == 0:
+        val = [np.copy(v.numpy()) for v in outs]
+        buffers = [np.copy(v.numpy().copy()) for v in input_queues.values()]
+
+    if test_val is not None:
+      match = all(np.array_equal(a, b) for a, b in zip(val, test_val, strict=True))
+      assert match == expect_match, f"outputs {'differ from' if expect_match else 'match'} baseline (seed={seed})"
+    if test_buffers is not None:
+      match = all(np.array_equal(a, b) for a, b in zip(buffers, test_buffers, strict=True))
+      assert match == expect_match, f"buffers {'differ from' if expect_match else 'match'} baseline (seed={seed})"
+    return fn, val, buffers
+
+  print('capture + replay')
+  run_policy_jit, test_val, test_buffers = random_inputs_run_fn(run_policy_jit, SEED)
+
+  print('pickle round trip')
+  run_policy_jit = pickle.loads(pickle.dumps(run_policy_jit))
+  random_inputs_run_fn(run_policy_jit, SEED, test_val, test_buffers, expect_match=True)
+  random_inputs_run_fn(run_policy_jit, SEED+1, test_val, test_buffers, expect_match=False)
+  return run_policy_jit
+
+
+def derive_frame_skip(vision_input_shapes, policy_input_shapes):
+  fb = policy_input_shapes.get('features_buffer')
+  if fb is None:
+    return 1
+  fb_history = fb[1]
+  if fb_history >= 99:
+    return 1
+  return 4
+
+
+def make_supercombo_input_queues(input_shapes, frame_skip, device):
+  img_shape = input_shapes.get('img', input_shapes.get('input_imgs'))
+  if img_shape is None:
+    raise ValueError("No img input found in model shapes")
+
+  n_frames = img_shape[1] // 6
+  img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img_shape[2], img_shape[3])
+
+  numpy_keys = {}
+  queue_keys = {}
+
+  for key, shape in input_shapes.items():
+    if 'img' in key:
+      continue
+    if len(shape) == 3 and shape[1] > 1:
+      if key.startswith('desire'):
+        numpy_keys[key] = np.zeros(shape[2], dtype=np.float32)
+        queue_keys[f'{key}_q'] = Tensor(
+          np.zeros((frame_skip * shape[1], shape[0], shape[2]), dtype=np.float32),
+          device=device).contiguous().realize()
+      elif key == 'features_buffer':
+        queue_keys['feat_q'] = Tensor(
+          np.zeros((frame_skip * (shape[1] - 1) + 1, shape[0], shape[2]), dtype=np.float32),
+          device=device).contiguous().realize()
+      else:
+        numpy_keys[key] = np.zeros(shape, dtype=np.float32)
+    elif len(shape) == 2:
+      numpy_keys[key] = np.zeros(shape, dtype=np.float32)
+
+  if 'traffic_convention' not in numpy_keys:
+    tc_shape = input_shapes.get('traffic_convention', (1, 2))
+    numpy_keys['traffic_convention'] = np.zeros(tc_shape, dtype=np.float32)
+
+  numpy_keys['tfm'] = np.zeros((3, 3), dtype=np.float32)
+  numpy_keys['big_tfm'] = np.zeros((3, 3), dtype=np.float32)
+
+  input_queues = {
+    'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
+    'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
+    **queue_keys,
+    **{k: Tensor(v, device='NPY').realize() for k, v in numpy_keys.items()},
+  }
+  return input_queues, numpy_keys
+
+
+def make_run_supercombo(model_runner, nv12: NV12Frame, model_w, model_h,
+                        features_slice, frame_skip, input_shapes, prepare_only=False):
+  frame_prepare = make_frame_prepare(nv12, model_w, model_h)
+  sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
+  sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
+
+  desire_key = _detect_desire_key(input_shapes)
+  if desire_key is None:
+    raise ValueError(f"No desire* key found in input_shapes: {list(input_shapes.keys())}")
+  road_img_key, wide_img_key = _detect_vision_keys(input_shapes)
+  extra_policy_keys = [k for k in input_shapes
+                       if k not in (desire_key, 'features_buffer', 'traffic_convention')
+                       and 'img' not in k]
+
+  def run_supercombo(img_q, big_img_q, feat_q, desire_q,
+                     frame, big_frame, **kwargs):
+    desire = kwargs.get(desire_key)
+    traffic_convention = kwargs.get('traffic_convention')
+    tfm = kwargs['tfm']
+    big_tfm = kwargs['big_tfm']
+
+    tfm = tfm.to(Device.DEFAULT)
+    big_tfm = big_tfm.to(Device.DEFAULT)
+    desire = desire.to(Device.DEFAULT)
+    traffic_convention = traffic_convention.to(Device.DEFAULT)
+    Tensor.realize(tfm, big_tfm, desire, traffic_convention)
+
+    img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn)
+    big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn)
+
+    if prepare_only:
+      return img, big_img
+
+    desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn)
+    feat_buf = sample_skip_fn(feat_q)
+
+    inputs = {road_img_key: img, wide_img_key: big_img,
+              desire_key: desire_buf, 'features_buffer': feat_buf,
+              'traffic_convention': traffic_convention}
+    for k in extra_policy_keys:
+      if k in kwargs:
+        inputs[k] = kwargs[k].to(Device.DEFAULT)
+
+    model_out = next(iter(model_runner(inputs).values())).cast('float32')
+
+    new_feat = model_out[:, features_slice].reshape(1, -1).unsqueeze(0)
+    shift_and_sample(feat_q, new_feat, sample_skip_fn)
+
+    return model_out
+
+  return run_supercombo
+
+
+def make_run_vision_multi_policy(vision_runner, policy_runners, nv12: NV12Frame, model_w, model_h,
+                                 vision_features_slice, frame_skip, desire_key, extra_policy_keys,
+                                 vision_road_key, vision_wide_key, prepare_only=False):
+  frame_prepare = make_frame_prepare(nv12, model_w, model_h)
+  sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
+  sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
+
+  def run_multi_policy(img_q, big_img_q, feat_q, desire_q, desire,
+                       traffic_convention, tfm, big_tfm, frame, big_frame, **extra):
+    npy_tensors = [tfm.to(Device.DEFAULT), big_tfm.to(Device.DEFAULT),
+                   desire.to(Device.DEFAULT), traffic_convention.to(Device.DEFAULT)]
+    extra_device = {k: extra[k].to(Device.DEFAULT) for k in extra_policy_keys}
+    Tensor.realize(*npy_tensors, *extra_device.values())
+    tfm, big_tfm, desire, traffic_convention = npy_tensors
+
+    img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn)
+    big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn)
+
+    if prepare_only:
+      return img, big_img
+
+    vision_out = next(iter(vision_runner({vision_road_key: img, vision_wide_key: big_img}).values())).cast('float32')
+
+    new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0)
+    feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn)
+    desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn)
+
+    inputs = {'features_buffer': feat_buf, desire_key: desire_buf, 'traffic_convention': traffic_convention, **extra_device}
+
+    policy_outputs = []
+    for runner in policy_runners:
+      policy_out = next(iter(runner(inputs).values())).cast('float32')
+      policy_outputs.append(policy_out)
+
+    return (vision_out, *policy_outputs)
+
+  return run_multi_policy
+
+
+def _warmup_and_serialize(run_jit, input_queues, npy, nv12):
+  for i in range(3):
+    np.random.seed(42 + i)
+    frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize()
+    big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize()
+    for v in npy.values():
+      v[:] = np.random.randn(*v.shape).astype(v.dtype)
+    Device.default.synchronize()
+    st = time.perf_counter()
+    run_jit(**input_queues, frame=frame, big_frame=big_frame)
+    mt = time.perf_counter()
+    Device.default.synchronize()
+    et = time.perf_counter()
+    print(f"  [{i + 1}/3] enqueue {(mt - st) * 1e3:6.2f} ms -- total {(et - st) * 1e3:6.2f} ms")
+  return pickle.loads(pickle.dumps(run_jit))
+
+
+def compile_supercombo(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip,
+                       model_runner, metadata):
+  print(f"Compiling combined supercombo JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...")
+
+  features_slice = metadata['output_slices']['hidden_state']
+  input_shapes = metadata['input_shapes']
+
+  _run = make_run_supercombo(model_runner, nv12, model_w, model_h,
+                             features_slice, frame_skip, input_shapes, prepare_only)
+  run_jit = TinyJit(_run, prune=True)
+
+  input_queues, npy = make_supercombo_input_queues(input_shapes, frame_skip, Device.DEFAULT)
+
+  run_jit = _warmup_and_serialize(run_jit, input_queues, npy, nv12)
+  return run_jit
+
+
+def compile_multi_policy(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip,
+                         vision_runner, policy_runners, vision_metadata, policy_metadata):
+  print(f"Compiling combined multi-policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...")
+
+  vision_features_slice = vision_metadata['output_slices']['hidden_state']
+  vision_input_shapes = vision_metadata['input_shapes']
+  policy_input_shapes = policy_metadata['input_shapes']
+  desire_key = _detect_desire_key(policy_input_shapes)
+  extra_policy_keys = [k for k in policy_input_shapes if k not in ('features_buffer', desire_key, 'traffic_convention')]
+  vision_road_key, vision_wide_key = _detect_vision_keys(vision_input_shapes)
+
+  _run = make_run_vision_multi_policy(vision_runner, policy_runners, nv12, model_w, model_h,
+                                      vision_features_slice, frame_skip, desire_key, extra_policy_keys,
+                                      vision_road_key, vision_wide_key, prepare_only)
+  run_jit = TinyJit(_run, prune=True)
+
+  input_queues, npy = make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, Device.DEFAULT)
+
+  run_jit = _warmup_and_serialize(run_jit, input_queues, npy, nv12)
+  return run_jit
+
+
+def _parse_size(s):
+  w, h = s.lower().split('x')
+  return int(w), int(h)
+
+
+if __name__ == "__main__":
+  from tinygrad.nn.onnx import OnnxRunner
+  from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
+  from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict
+
+  p = argparse.ArgumentParser(description="Compile combined JIT pkl for sunnypilot modeld_v2")
+  p.add_argument('--model-type', choices=MODEL_TYPES, required=True)
+  p.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH')
+  p.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True)
+  p.add_argument('--frame-skip', type=int, default=None, help='frame skip value (auto-derived if not provided)')
+  p.add_argument('--output', required=True)
+
+  p.add_argument('--vision-onnx', help='vision ONNX (for split models)')
+  p.add_argument('--policy-onnx', help='policy ONNX (for vision_policy)')
+  p.add_argument('--off-policy-onnx', help='off-policy ONNX (for vision_multi_policy)')
+  p.add_argument('--on-policy-onnx', help='on-policy ONNX (for vision_multi_policy)')
+  p.add_argument('--supercombo-onnx', help='supercombo ONNX (for supercombo)')
+
+  args = p.parse_args()
+  out = defaultdict(dict)
+
+  if args.model_type == 'vision_policy':
+    assert args.vision_onnx and args.policy_onnx
+    vision_runner = OnnxRunner(args.vision_onnx)
+    policy_runner = OnnxRunner(args.policy_onnx)
+    out['metadata']['vision'] = make_metadata_dict(args.vision_onnx)
+    out['metadata']['policy'] = make_metadata_dict(args.policy_onnx)
+
+    frame_skip = args.frame_skip if args.frame_skip is not None else derive_frame_skip(out['metadata']['vision']['input_shapes'],
+                                                                                       out['metadata']['policy']['input_shapes'])
+
+    for cam_w, cam_h in args.camera_resolutions:
+      nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
+      model_w, model_h = args.model_size
+      out[(cam_w, cam_h)] = {
+        name: compile_split_policy(nv12, model_w, model_h, prepare_only, frame_skip,
+                                   vision_runner, policy_runner,
+                                   out['metadata']['vision'], out['metadata']['policy'])
+        for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)]
+      }
+
+  elif args.model_type == 'supercombo':
+    assert args.supercombo_onnx
+    model_runner = OnnxRunner(args.supercombo_onnx)
+    out['metadata']['model'] = make_metadata_dict(args.supercombo_onnx)
+
+    frame_skip = args.frame_skip if args.frame_skip is not None else derive_frame_skip({}, out['metadata']['model']['input_shapes'])
+
+    for cam_w, cam_h in args.camera_resolutions:
+      nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
+      model_w, model_h = args.model_size
+      out[(cam_w, cam_h)] = {
+        name: compile_supercombo(nv12, model_w, model_h, prepare_only, frame_skip,
+                                 model_runner, out['metadata']['model'])
+        for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)]
+      }
+
+  elif args.model_type == 'vision_multi_policy':
+    assert args.vision_onnx
+    vision_runner = OnnxRunner(args.vision_onnx)
+    out['metadata']['vision'] = make_metadata_dict(args.vision_onnx)
+
+    policy_runners = []
+    policy_onnxes = []
+    if args.policy_onnx:
+      policy_onnxes.append(('policy', args.policy_onnx))
+    if args.off_policy_onnx:
+      policy_onnxes.append(('off_policy', args.off_policy_onnx))
+    if args.on_policy_onnx:
+      policy_onnxes.append(('on_policy', args.on_policy_onnx))
+
+    for name, onnx_path in policy_onnxes:
+      runner = OnnxRunner(onnx_path)
+      policy_runners.append(runner)
+      out['metadata'][name] = make_metadata_dict(onnx_path)
+
+    first_policy_key = policy_onnxes[0][0]
+    frame_skip = args.frame_skip if args.frame_skip is not None else derive_frame_skip(out['metadata']['vision']['input_shapes'],
+                                                                                       out['metadata'][first_policy_key]['input_shapes'])
+
+    for cam_w, cam_h in args.camera_resolutions:
+      nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
+      model_w, model_h = args.model_size
+      out[(cam_w, cam_h)] = {
+        name: compile_multi_policy(nv12, model_w, model_h, prepare_only, frame_skip,
+                                   vision_runner, policy_runners,
+                                   out['metadata']['vision'], out['metadata'][first_policy_key])
+        for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)]
+      }
+
+  with open(args.output, "wb") as f:
+    pickle.dump(out, f)
+  pkl_size = os.path.getsize(args.output)
+  print(f"Saved combined JIT to {args.output} ({pkl_size / 1e6:.2f} MB)")
+
+  from openpilot.common.file_chunker import chunk_file, get_chunk_targets
+  chunk_targets = get_chunk_targets(args.output, pkl_size)
+  chunk_file(args.output, chunk_targets)
+  num_chunks = len(chunk_targets) - 1
+  print(f"Chunked into {num_chunks} file(s)")

openpilot/sunnypilot/modeld_v2/constants.py

@@ -0,0 +1,84 @@
+import numpy as np
+
+def index_function(idx, max_val=192, max_idx=32):
+  return (max_val) * ((idx/max_idx)**2)
+
+class ModelConstants:
+  # time and distance indices
+  IDX_N = 33
+  T_IDXS = [index_function(idx, max_val=10.0) for idx in range(IDX_N)]
+  X_IDXS = [index_function(idx, max_val=192.0) for idx in range(IDX_N)]
+  LEAD_T_IDXS = [0., 2., 4., 6., 8., 10.]
+  LEAD_T_OFFSETS = [0., 2., 4.]
+  META_T_IDXS = [2., 4., 6., 8., 10.]
+
+  # model inputs constants
+  MODEL_FREQ = 20
+  FEATURE_LEN = 512
+  FULL_HISTORY_BUFFER_LEN = 99
+  DESIRE_LEN = 8
+  TRAFFIC_CONVENTION_LEN = 2
+  LAT_PLANNER_STATE_LEN = 4
+  LATERAL_CONTROL_PARAMS_LEN = 2
+  PREV_DESIRED_CURV_LEN = 1
+
+  # model outputs constants
+  FCW_THRESHOLDS_5MS2 = np.array([.05, .05, .15, .15, .15], dtype=np.float32)
+  FCW_THRESHOLDS_3MS2 = np.array([.7, .7], dtype=np.float32)
+  FCW_5MS2_PROBS_WIDTH = 5
+  FCW_3MS2_PROBS_WIDTH = 2
+
+  DISENGAGE_WIDTH = 5
+  POSE_WIDTH = 6
+  WIDE_FROM_DEVICE_WIDTH = 3
+  LEAD_WIDTH = 4
+  LANE_LINES_WIDTH = 2
+  ROAD_EDGES_WIDTH = 2
+  PLAN_WIDTH = 15
+  DESIRE_PRED_WIDTH = 8
+  LAT_PLANNER_SOLUTION_WIDTH = 4
+  DESIRED_CURV_WIDTH = 1
+
+  NUM_LANE_LINES = 4
+  NUM_ROAD_EDGES = 2
+
+  LEAD_TRAJ_LEN = 6
+  DESIRE_PRED_LEN = 4
+
+  PLAN_MHP_N = 5
+  LEAD_MHP_N = 2
+  PLAN_MHP_SELECTION = 1
+  LEAD_MHP_SELECTION = 3
+
+  FCW_THRESHOLD_5MS2_HIGH = 0.15
+  FCW_THRESHOLD_5MS2_LOW = 0.05
+  FCW_THRESHOLD_3MS2 = 0.7
+
+  CONFIDENCE_BUFFER_LEN = 5
+  RYG_GREEN = 0.01165
+  RYG_YELLOW = 0.06157
+
+  POLY_PATH_DEGREE = 4
+
+# model outputs slices
+class Plan:
+  POSITION = slice(0, 3)
+  VELOCITY = slice(3, 6)
+  ACCELERATION = slice(6, 9)
+  T_FROM_CURRENT_EULER = slice(9, 12)
+  ORIENTATION_RATE = slice(12, 15)
+
+class Meta:
+  ENGAGED = slice(0, 1)
+  # next 2, 4, 6, 8, 10 seconds
+  GAS_DISENGAGE = slice(1, 31, 6)
+  BRAKE_DISENGAGE = slice(2, 31, 6)
+  STEER_OVERRIDE = slice(3, 31, 6)
+  HARD_BRAKE_3 = slice(4, 31, 6)
+  HARD_BRAKE_4 = slice(5, 31, 6)
+  HARD_BRAKE_5 = slice(6, 31, 6)
+  # next 0, 2, 4, 6, 8, 10 seconds
+  GAS_PRESS = slice(31, 55, 4)
+  BRAKE_PRESS = slice(32, 55, 4)
+  LEFT_BLINKER = slice(33, 55, 4)
+  RIGHT_BLINKER = slice(34, 55, 4)

openpilot/sunnypilot/modeld_v2/fill_model_msg.py

@@ -0,0 +1,218 @@
+import os
+import capnp
+import numpy as np
+from openpilot.cereal import log
+from openpilot.sunnypilot.modeld_v2.constants import ModelConstants, Plan
+from openpilot.sunnypilot.models.helpers import plan_x_idxs_helper
+from openpilot.selfdrive.controls.lib.drive_helpers import get_curvature_from_plan
+
+SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
+
+ConfidenceClass = log.ModelDataV2.ConfidenceClass
+
+
+def get_curvature_from_output(output, plan, vego, lat_action_t, mlsim):
+  if not mlsim:
+    if desired_curv := output.get('desired_curvature'):  # If the model outputs the desired curvature, use that directly
+      return float(desired_curv[0, 0])
+
+  return float(get_curvature_from_plan(plan[:, Plan.T_FROM_CURRENT_EULER][:, 2], plan[:, Plan.ORIENTATION_RATE][:, 2],
+                                       ModelConstants.T_IDXS, vego, lat_action_t))
+
+
+class PublishState:
+  def __init__(self):
+    self.disengage_buffer = np.zeros(ModelConstants.CONFIDENCE_BUFFER_LEN*ModelConstants.DISENGAGE_WIDTH, dtype=np.float32)
+    self.prev_brake_5ms2_probs = np.zeros(ModelConstants.FCW_5MS2_PROBS_WIDTH, dtype=np.float32)
+    self.prev_brake_3ms2_probs = np.zeros(ModelConstants.FCW_3MS2_PROBS_WIDTH, dtype=np.float32)
+
+def fill_xyzt(builder, t, x, y, z, x_std=None, y_std=None, z_std=None):
+  builder.t = t
+  builder.x = x.tolist()
+  builder.y = y.tolist()
+  builder.z = z.tolist()
+  if x_std is not None:
+    builder.xStd = x_std.tolist()
+  if y_std is not None:
+    builder.yStd = y_std.tolist()
+  if z_std is not None:
+    builder.zStd = z_std.tolist()
+
+def fill_xyvat(builder, t, x, y, v, a, x_std=None, y_std=None, v_std=None, a_std=None):
+  builder.t = t
+  builder.x = x.tolist()
+  builder.y = y.tolist()
+  builder.v = v.tolist()
+  builder.a = a.tolist()
+  if x_std is not None:
+    builder.xStd = x_std.tolist()
+  if y_std is not None:
+    builder.yStd = y_std.tolist()
+  if v_std is not None:
+    builder.vStd = v_std.tolist()
+  if a_std is not None:
+    builder.aStd = a_std.tolist()
+
+def fill_xyz_poly(builder, degree, x, y, z):
+  xyz = np.stack([x, y, z], axis=1)
+  coeffs = np.polynomial.polynomial.polyfit(ModelConstants.T_IDXS, xyz, deg=degree)
+  builder.xCoefficients = coeffs[:, 0].tolist()
+  builder.yCoefficients = coeffs[:, 1].tolist()
+  builder.zCoefficients = coeffs[:, 2].tolist()
+
+def fill_lane_line_meta(builder, lane_lines, lane_line_probs):
+  builder.leftY = lane_lines[1].y[0]
+  builder.leftProb = lane_line_probs[1]
+  builder.rightY = lane_lines[2].y[0]
+  builder.rightProb = lane_line_probs[2]
+
+def fill_model_msg(base_msg: capnp._DynamicStructBuilder, extended_msg: capnp._DynamicStructBuilder,
+                   net_output_data: dict[str, np.ndarray], action: log.ModelDataV2.Action,
+                   publish_state: PublishState, vipc_frame_id: int, vipc_frame_id_extra: int,
+                   frame_id: int, frame_drop: float, timestamp_eof: int, model_execution_time: float,
+                   valid: bool, model_meta) -> None:
+  frame_age = frame_id - vipc_frame_id if frame_id > vipc_frame_id else 0
+  frame_drop_perc = frame_drop * 100
+  extended_msg.valid = valid
+  base_msg.valid = valid
+
+  driving_model_data = base_msg.drivingModelData
+
+  driving_model_data.frameId = vipc_frame_id
+  driving_model_data.frameIdExtra = vipc_frame_id_extra
+  driving_model_data.frameDropPerc = frame_drop_perc
+  driving_model_data.modelExecutionTime = model_execution_time
+  driving_model_data.action = action
+
+  modelV2 = extended_msg.modelV2
+  modelV2.frameId = vipc_frame_id
+  modelV2.frameIdExtra = vipc_frame_id_extra
+  modelV2.frameAge = frame_age
+  modelV2.frameDropPerc = frame_drop_perc
+  modelV2.timestampEof = timestamp_eof
+  modelV2.modelExecutionTime = model_execution_time
+
+  # plan
+  fill_xyzt(modelV2.position, ModelConstants.T_IDXS, *net_output_data['plan'][0,:,Plan.POSITION].T, *net_output_data['plan_stds'][0,:,Plan.POSITION].T)
+  fill_xyzt(modelV2.velocity, ModelConstants.T_IDXS, *net_output_data['plan'][0,:,Plan.VELOCITY].T)
+  fill_xyzt(modelV2.acceleration, ModelConstants.T_IDXS, *net_output_data['plan'][0,:,Plan.ACCELERATION].T)
+  fill_xyzt(modelV2.orientation, ModelConstants.T_IDXS, *net_output_data['plan'][0,:,Plan.T_FROM_CURRENT_EULER].T)
+  fill_xyzt(modelV2.orientationRate, ModelConstants.T_IDXS, *net_output_data['plan'][0,:,Plan.ORIENTATION_RATE].T)
+
+  # temporal pose
+  temporal_pose = modelV2.deprecated.temporalPose
+  if 'sim_pose' in net_output_data:
+    temporal_pose.trans = net_output_data['sim_pose'][0,:ModelConstants.POSE_WIDTH//2].tolist()
+    temporal_pose.transStd = net_output_data['sim_pose_stds'][0,:ModelConstants.POSE_WIDTH//2].tolist()
+    temporal_pose.rot = net_output_data['sim_pose'][0,ModelConstants.POSE_WIDTH//2:].tolist()
+    temporal_pose.rotStd = net_output_data['sim_pose_stds'][0,ModelConstants.POSE_WIDTH//2:].tolist()
+  else:
+    temporal_pose.trans = net_output_data['plan'][0,0,Plan.VELOCITY].tolist()
+    temporal_pose.transStd = net_output_data['plan_stds'][0,0,Plan.VELOCITY].tolist()
+    temporal_pose.rot = net_output_data['plan'][0,0,Plan.ORIENTATION_RATE].tolist()
+    temporal_pose.rotStd = net_output_data['plan_stds'][0,0,Plan.ORIENTATION_RATE].tolist()
+
+  # poly path
+  fill_xyz_poly(driving_model_data.path, ModelConstants.POLY_PATH_DEGREE, *net_output_data['plan'][0,:,Plan.POSITION].T)
+
+  # action (includes lateral planning now)
+  modelV2.action = action
+
+  # times at X_IDXS of edges and lines
+  LINE_T_IDXS: list[float] = plan_x_idxs_helper(ModelConstants, Plan, net_output_data)
+
+  # lane lines
+  modelV2.init('laneLines', 4)
+  for i in range(4):
+    lane_line = modelV2.laneLines[i]
+    fill_xyzt(lane_line, LINE_T_IDXS, np.array(ModelConstants.X_IDXS), net_output_data['lane_lines'][0,i,:,0], net_output_data['lane_lines'][0,i,:,1])
+  modelV2.laneLineStds = net_output_data['lane_lines_stds'][0,:,0,0].tolist()
+  modelV2.laneLineProbs = net_output_data['lane_lines_prob'][0,1::2].tolist()
+
+  fill_lane_line_meta(driving_model_data.laneLineMeta, modelV2.laneLines, modelV2.laneLineProbs)
+
+  # road edges
+  modelV2.init('roadEdges', 2)
+  for i in range(2):
+    road_edge = modelV2.roadEdges[i]
+    fill_xyzt(road_edge, LINE_T_IDXS, np.array(ModelConstants.X_IDXS), net_output_data['road_edges'][0,i,:,0], net_output_data['road_edges'][0,i,:,1])
+  modelV2.roadEdgeStds = net_output_data['road_edges_stds'][0,:,0,0].tolist()
+
+  # leads
+  modelV2.init('leadsV3', 3)
+  for i in range(3):
+    lead = modelV2.leadsV3[i]
+    fill_xyvat(lead, ModelConstants.LEAD_T_IDXS, *net_output_data['lead'][0,i].T, *net_output_data['lead_stds'][0,i].T)
+    lead.prob = net_output_data['lead_prob'][0,i].tolist()
+    lead.probTime = ModelConstants.LEAD_T_OFFSETS[i]
+
+  # meta
+  meta = modelV2.meta
+  meta.desireState = net_output_data['desire_state'][0].reshape(-1).tolist()
+  meta.desirePrediction = net_output_data['desire_pred'][0].reshape(-1).tolist()
+  meta.engagedProb = net_output_data['meta'][0,model_meta.ENGAGED].item()
+  meta.init('disengagePredictions')
+  disengage_predictions = meta.disengagePredictions
+  disengage_predictions.t = ModelConstants.META_T_IDXS
+  disengage_predictions.brakeDisengageProbs = net_output_data['meta'][0,model_meta.BRAKE_DISENGAGE].tolist()
+  disengage_predictions.gasDisengageProbs = net_output_data['meta'][0,model_meta.GAS_DISENGAGE].tolist()
+  disengage_predictions.steerOverrideProbs = net_output_data['meta'][0,model_meta.STEER_OVERRIDE].tolist()
+  disengage_predictions.brake3MetersPerSecondSquaredProbs = net_output_data['meta'][0,model_meta.HARD_BRAKE_3].tolist()
+  disengage_predictions.brake4MetersPerSecondSquaredProbs = net_output_data['meta'][0,model_meta.HARD_BRAKE_4].tolist()
+  disengage_predictions.brake5MetersPerSecondSquaredProbs = net_output_data['meta'][0,model_meta.HARD_BRAKE_5].tolist()
+
+  if hasattr(model_meta, 'GAS_PRESS') and hasattr(model_meta, 'BRAKE_PRESS'):
+    disengage_predictions.gasPressProbs = net_output_data['meta'][0,model_meta.GAS_PRESS].tolist()
+    disengage_predictions.brakePressProbs = net_output_data['meta'][0,model_meta.BRAKE_PRESS].tolist()
+
+  publish_state.prev_brake_5ms2_probs[:-1] = publish_state.prev_brake_5ms2_probs[1:]
+  publish_state.prev_brake_5ms2_probs[-1] = net_output_data['meta'][0,model_meta.HARD_BRAKE_5][0]
+  publish_state.prev_brake_3ms2_probs[:-1] = publish_state.prev_brake_3ms2_probs[1:]
+  publish_state.prev_brake_3ms2_probs[-1] = net_output_data['meta'][0,model_meta.HARD_BRAKE_3][0]
+  hard_brake_predicted = (publish_state.prev_brake_5ms2_probs > ModelConstants.FCW_THRESHOLDS_5MS2).all() and \
+    (publish_state.prev_brake_3ms2_probs > ModelConstants.FCW_THRESHOLDS_3MS2).all()
+  meta.hardBrakePredicted = hard_brake_predicted.item()
+
+  # confidence
+  if vipc_frame_id % (2*ModelConstants.MODEL_FREQ) == 0:
+    # any disengage prob
+    brake_disengage_probs = net_output_data['meta'][0,model_meta.BRAKE_DISENGAGE]
+    gas_disengage_probs = net_output_data['meta'][0,model_meta.GAS_DISENGAGE]
+    steer_override_probs = net_output_data['meta'][0,model_meta.STEER_OVERRIDE]
+    any_disengage_probs = 1-((1-brake_disengage_probs)*(1-gas_disengage_probs)*(1-steer_override_probs))
+    # independent disengage prob for each 2s slice
+    ind_disengage_probs = np.r_[any_disengage_probs[0], np.diff(any_disengage_probs) / (1 - any_disengage_probs[:-1])]
+    # rolling buf for 2, 4, 6, 8, 10s
+    publish_state.disengage_buffer[:-ModelConstants.DISENGAGE_WIDTH] = publish_state.disengage_buffer[ModelConstants.DISENGAGE_WIDTH:]
+    publish_state.disengage_buffer[-ModelConstants.DISENGAGE_WIDTH:] = ind_disengage_probs
+
+  score = 0.
+  for i in range(ModelConstants.DISENGAGE_WIDTH):
+    score += publish_state.disengage_buffer[i*ModelConstants.DISENGAGE_WIDTH+ModelConstants.DISENGAGE_WIDTH-1-i].item() / ModelConstants.DISENGAGE_WIDTH
+  if score < ModelConstants.RYG_GREEN:
+    modelV2.confidence = ConfidenceClass.green
+  elif score < ModelConstants.RYG_YELLOW:
+    modelV2.confidence = ConfidenceClass.yellow
+  else:
+    modelV2.confidence = ConfidenceClass.red
+
+  # raw prediction if enabled
+  if SEND_RAW_PRED:
+    modelV2.rawPredictions = net_output_data['raw_pred'].tobytes()
+
+def fill_pose_msg(msg: capnp._DynamicStructBuilder, net_output_data: dict[str, np.ndarray],
+                  vipc_frame_id: int, vipc_dropped_frames: int, timestamp_eof: int, live_calib_seen: bool) -> None:
+  msg.valid = live_calib_seen & (vipc_dropped_frames < 1)
+  cameraOdometry = msg.cameraOdometry
+
+  cameraOdometry.frameId = vipc_frame_id
+  cameraOdometry.timestampEof = timestamp_eof
+
+  cameraOdometry.trans = net_output_data['pose'][0,:3].tolist()
+  cameraOdometry.rot = net_output_data['pose'][0,3:].tolist()
+  cameraOdometry.wideFromDeviceEuler = net_output_data['wide_from_device_euler'][0,:].tolist()
+  cameraOdometry.roadTransformTrans = net_output_data['road_transform'][0,:3].tolist()
+  cameraOdometry.transStd = net_output_data['pose_stds'][0,:3].tolist()
+  cameraOdometry.rotStd = net_output_data['pose_stds'][0,3:].tolist()
+  cameraOdometry.wideFromDeviceEulerStd = net_output_data['wide_from_device_euler_stds'][0,:].tolist()
+  cameraOdometry.roadTransformTransStd = net_output_data['road_transform_stds'][0,:3].tolist()

openpilot/sunnypilot/modeld_v2/get_model_metadata.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+import sys
+import pathlib
+import codecs
+import pickle
+from typing import Any
+
+from tinygrad.nn.onnx import OnnxPBParser
+
+
+class MetadataOnnxPBParser(OnnxPBParser):
+  def _parse_ModelProto(self) -> dict:
+    obj: dict[str, Any] = {"graph": {"input": [], "output": []}, "metadata_props": []}
+    for fid, wire_type in self._parse_message(self.reader.len):
+      match fid:
+        case 7:
+          obj["graph"] = self._parse_GraphProto()
+        case 14:
+          obj["metadata_props"].append(self._parse_StringStringEntryProto())
+        case _:
+          self.reader.skip_field(wire_type)
+    return obj
+
+
+def get_name_and_shape(value_info: dict[str, Any]) -> tuple[str, tuple[int, ...]]:
+  shape = tuple(int(dim) if isinstance(dim, int) else 0 for dim in value_info["parsed_type"].shape)
+  name = value_info["name"]
+  return name, shape
+
+
+def get_metadata_value_by_name(model: dict[str, Any], name: str) -> str | Any:
+  for prop in model["metadata_props"]:
+    if prop["key"] == name:
+      return prop["value"]
+  return None
+
+
+if __name__ == "__main__":
+  model_path = pathlib.Path(sys.argv[1])
+  model = MetadataOnnxPBParser(model_path).parse()
+  output_slices = get_metadata_value_by_name(model, 'output_slices')
+  assert output_slices is not None, 'output_slices not found in metadata'
+
+  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"]),
+  }
+
+  metadata_path = model_path.parent / (model_path.stem + '_metadata.pkl')
+  with open(metadata_path, 'wb') as f:
+    pickle.dump(metadata, f)
+
+  print(f'saved metadata to {metadata_path}')

openpilot/sunnypilot/modeld_v2/install_models_pc.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+import sys
+import shutil
+import pickle
+import codecs
+from pathlib import Path
+
+from openpilot.common.hardware.hw import Paths
+from openpilot.sunnypilot.modeld_v2.get_model_metadata import MetadataOnnxPBParser, get_name_and_shape, get_metadata_value_by_name
+
+
+def generate_metadata_pkl(model_path, output_path):
+  try:
+    model = MetadataOnnxPBParser(model_path).parse()
+    output_slices = get_metadata_value_by_name(model, 'output_slices')
+    if not output_slices:
+      return False
+    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
+  except Exception:
+    return False
+
+
+def install_models(model_dir):
+  model_dir = Path(model_dir)
+  models = ["driving_off_policy", "driving_on_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])

openpilot/sunnypilot/modeld_v2/meta_20hz.py

@@ -0,0 +1,17 @@
+from openpilot.sunnypilot.modeld_v2.constants import Meta
+
+
+class Meta20hz(Meta):
+  ENGAGED = slice(0, 1)
+  # next 2, 4, 6, 8, 10 seconds
+  GAS_DISENGAGE = slice(1, 31, 6)
+  BRAKE_DISENGAGE = slice(2, 31, 6)
+  STEER_OVERRIDE = slice(3, 31, 6)
+  HARD_BRAKE_3 = slice(4, 31, 6)
+  HARD_BRAKE_4 = slice(5, 31, 6)
+  HARD_BRAKE_5 = slice(6, 31, 6)
+  # next 0, 2, 4, 6, 8, 10 seconds
+  GAS_PRESS = slice(31, 55, 4)
+  BRAKE_PRESS = slice(32, 55, 4)
+  LEFT_BLINKER = slice(33, 55, 4)
+  RIGHT_BLINKER = slice(34, 55, 4)

openpilot/sunnypilot/modeld_v2/meta_helper.py

@@ -0,0 +1,26 @@
+from openpilot.sunnypilot.modeld_v2.constants import Meta
+from openpilot.cereal import custom
+from openpilot.sunnypilot.modeld_v2.meta_20hz import Meta20hz
+from openpilot.sunnypilot.models.helpers import get_active_bundle
+
+ModelBundle = custom.ModelManagerSP.ModelBundle
+
+
+def load_meta_constants():
+  """
+  Determines and loads the appropriate meta model class based on the metadata provided. The function checks
+  specific keys and conditions within the provided metadata dictionary to identify the corresponding meta
+  model class to return.
+
+  :param model_metadata: Dictionary containing metadata about the model. It includes
+      details such as input shapes, output slices, and other configurations for identifying
+      metadata-dependent meta model classes.
+  :type model_metadata: dict
+  :return: The appropriate meta model class (Meta, MetaSimPose, or MetaTombRaider)
+      based on the conditions and metadata provided.
+  :rtype: type
+  """
+  if (bundle := get_active_bundle()) and bundle.is20hz:
+    return Meta20hz
+
+  return Meta  # Default

openpilot/sunnypilot/modeld_v2/modeld

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+
+DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null && pwd)"
+exec "$DIR/modeld.py" "$@"

openpilot/sunnypilot/modeld_v2/modeld.py

@@ -0,0 +1,458 @@
+#!/usr/bin/env python3
+"""
+Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
+
+This file is part of sunnypilot and is licensed under the MIT License.
+See the LICENSE.md file in the root directory for more details.
+"""
+
+import os
+from openpilot.common.hardware import TICI
+os.environ['DEV'] = 'QCOM' if TICI else 'CPU'
+USBGPU = "USBGPU" in os.environ
+if USBGPU:
+  os.environ['DEV'] = 'AMD'
+  os.environ['AMD_IFACE'] = 'USB'
+import pickle
+import time
+import numpy as np
+import openpilot.cereal.messaging as messaging
+from openpilot.cereal import log
+from opendbc.car.structs import car
+from setproctitle import setproctitle
+from openpilot.cereal.messaging import PubMaster, SubMaster
+from msgq.visionipc import VisionIpcClient, VisionStreamType, VisionBuf
+from opendbc.car.car_helpers import get_demo_car_params
+from openpilot.common.swaglog import cloudlog
+from openpilot.common.params import Params
+from openpilot.common.filter_simple import FirstOrderFilter
+from openpilot.common.realtime import config_realtime_process, DT_MDL
+from openpilot.common.transformations.camera import DEVICE_CAMERAS
+from openpilot.common.transformations.model import get_warp_matrix
+from openpilot.system import sentry
+from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
+from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value
+
+from openpilot.sunnypilot.modeld_v2.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState, get_curvature_from_output
+from openpilot.sunnypilot.modeld_v2.constants import Plan
+from openpilot.sunnypilot.modeld_v2.meta_helper import load_meta_constants
+from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelper
+
+from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
+from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
+from openpilot.sunnypilot.models.helpers import get_active_bundle
+
+PROCESS_NAME = "openpilot.selfdrive.modeld.modeld_tinygrad"
+
+
+def _pkl_exists(path):
+  from openpilot.common.file_chunker import get_manifest_path
+  return os.path.exists(path) or os.path.exists(get_manifest_path(path))
+
+
+def _find_driving_pkl(bundle):
+  if (override := os.environ.get('COMBINED_MODEL_PKL')) and _pkl_exists(override):
+    return override
+  if bundle is None or not bundle.models:
+    return None
+  from openpilot.common.hardware.hw import Paths
+  model_root = Paths.model_root()
+
+  pkl_name = bundle.models[0].artifact.fileName
+  pkl_path = os.path.join(model_root, pkl_name)
+  if _pkl_exists(pkl_path):
+    return pkl_path
+
+
+class FrameMeta:
+  frame_id: int = 0
+  timestamp_sof: int = 0
+  timestamp_eof: int = 0
+
+  def __init__(self, vipc=None):
+    if vipc is not None:
+      self.frame_id, self.timestamp_sof, self.timestamp_eof = vipc.frame_id, vipc.timestamp_sof, vipc.timestamp_eof
+
+
+class ModelState(ModelStateBase):
+  inputs: dict[str, np.ndarray]
+  prev_desire: np.ndarray
+
+  def __init__(self, cam_w: int, cam_h: int):
+    ModelStateBase.__init__(self)
+
+    env_pkl = os.environ.get('COMBINED_MODEL_PKL')
+    if env_pkl and os.path.exists(env_pkl):
+      model_bundle = None
+    else:
+      model_bundle = get_active_bundle()
+    self.generation = model_bundle.generation if model_bundle is not None else None
+    overrides = {override.key: override.value for override in model_bundle.overrides} if model_bundle else {}
+
+    self.LAT_SMOOTH_SECONDS = float(overrides.get('lat', ".0"))
+    self.LONG_SMOOTH_SECONDS = float(overrides.get('long', ".0"))
+    self.MIN_LAT_CONTROL_SPEED = 0.3
+    self.PLANPLUS_CONTROL: float = 1.0
+
+    pkl_path = _find_driving_pkl(model_bundle)
+    assert pkl_path is not None, "No driving pkl found — all models must be compiled with compile_modeld.py"
+    self._init_combined(pkl_path, cam_w, cam_h, model_bundle)
+
+  def _init_combined(self, pkl_path, cam_w, cam_h, bundle):
+    from tinygrad.tensor import Tensor
+    from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
+    from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues
+    from tinygrad.device import Device
+
+    from openpilot.common.file_chunker import open_file_chunked
+
+    cloudlog.warning(f"loading combined pkl: {pkl_path}")
+    jits = pickle.load(open_file_chunked(pkl_path))
+
+    self.DEV = Device.DEFAULT
+
+    metadata = jits['metadata']
+    if 'model' in metadata:
+      model_metadata = metadata['model']
+      self.vision_output_slices = model_metadata['output_slices']
+      self.policy_output_slices = {}
+      self._policy_slices_list = []
+      self._combined_model_type = 'supercombo'
+      self._vision_input_names = [k for k in model_metadata['input_shapes'] if 'img' in k]
+      from openpilot.sunnypilot.modeld_v2.compile_modeld import make_supercombo_input_queues
+      frame_skip = derive_frame_skip({}, model_metadata['input_shapes'])
+      self.input_queues, self.numpy_inputs = make_supercombo_input_queues(model_metadata['input_shapes'], frame_skip, device=self.DEV)
+    else:
+      vision_metadata = metadata['vision']
+      policy_keys = [k for k in metadata if k != 'vision']
+      if policy_keys == ['policy']:
+        self._combined_model_type = 'split'
+      else:
+        self._combined_model_type = 'multi_policy'
+      self.vision_output_slices = vision_metadata['output_slices']
+      self._policy_keys = policy_keys
+      self._policy_slices_list = [metadata[k]['output_slices'] for k in policy_keys]
+      self.policy_output_slices = self._policy_slices_list[0]
+      self._has_on_policy = any('on' in k.lower() for k in policy_keys)
+      first_policy_metadata = metadata[policy_keys[0]]
+      vision_input_shapes = vision_metadata['input_shapes']
+      policy_input_shapes = first_policy_metadata['input_shapes']
+      self._vision_input_names = [k for k in vision_input_shapes if 'img' in k]
+      frame_skip = derive_frame_skip(vision_input_shapes, policy_input_shapes)
+      self.input_queues, self.numpy_inputs = make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device=self.DEV)
+
+    from openpilot.sunnypilot.modeld_v2.parse_model_outputs_split import Parser as SplitParser
+    from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser as CombinedParser
+    self.parser = SplitParser() if self._combined_model_type != 'supercombo' else CombinedParser()
+
+    is_20hz = bundle.is20hz if bundle else self._combined_model_type in ('split', 'multi_policy')
+    if is_20hz:
+      from openpilot.sunnypilot.models.split_model_constants import SplitModelConstants
+      self.constants = SplitModelConstants()
+    else:
+      from openpilot.sunnypilot.modeld_v2.constants import ModelConstants
+      self.constants = ModelConstants()
+
+    self.prev_desire = np.zeros(self.constants.DESIRE_LEN, dtype=np.float32)
+    self.full_frames: dict = {}
+    self._blob_cache: dict = {}
+    nv12_info = get_nv12_info(cam_w, cam_h)
+    self.frame_buf_params = dict.fromkeys(self._vision_input_names, nv12_info)
+
+    self._run_policy = jits[(cam_w, cam_h)]['run_policy']
+    self._warp_enqueue = jits[(cam_w, cam_h)]['warp_enqueue']
+    road_name = next(k for k in self._vision_input_names if 'big' not in k)
+    yuv_size = self.frame_buf_params[road_name][3]
+    self._warp_enqueue(
+      **self.input_queues,
+      frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.DEV).contiguous().realize(),
+      big_frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.DEV).contiguous().realize())
+
+
+  @property
+  def mlsim(self) -> bool:
+    return bool(self.generation is not None and self.generation >= 11)
+
+  @property
+  def vision_input_names(self) -> list[str]:
+    return self._vision_input_names
+
+  @property
+  def desire_key(self) -> str:
+    return next(k for k in self.numpy_inputs if k.startswith('desire'))
+
+  def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray],
+                inputs: dict[str, np.ndarray], prepare_only: bool) -> dict[str, np.ndarray] | None:
+    from tinygrad.tensor import Tensor
+
+    for key in bufs.keys():
+      ptr = np.frombuffer(bufs[key].data, dtype=np.uint8).ctypes.data
+      yuv_size = self.frame_buf_params[key][3]
+      cache_key = (key, ptr)
+      if cache_key not in self._blob_cache:
+        self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8', device=self.DEV)
+      self.full_frames[key] = self._blob_cache[cache_key]
+
+    desire_key = self.desire_key
+    inputs[desire_key][0] = 0
+    self.numpy_inputs[desire_key][:] = np.where(inputs[desire_key] - self.prev_desire > .99, inputs[desire_key], 0)
+    self.prev_desire[:] = inputs[desire_key]
+    for key in ('traffic_convention', 'lateral_control_params'):
+      if key in self.numpy_inputs and key in inputs:
+        self.numpy_inputs[key][:] = inputs[key]
+
+    road_key = next(n for n in bufs if 'big' not in n)
+    wide_key = next(n for n in bufs if 'big' in n)
+    self.numpy_inputs['tfm'][:, :] = transforms[road_key].reshape(3, 3)
+    self.numpy_inputs['big_tfm'][:, :] = transforms[wide_key].reshape(3, 3)
+
+    if prepare_only:
+      self._warp_enqueue(**self.input_queues, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key])
+      return None
+
+    raw_outputs = self._run_policy(**self.input_queues, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key])
+
+    if self._combined_model_type == 'supercombo':
+      model_output = raw_outputs.numpy().flatten()
+      sliced = {k: model_output[np.newaxis, v] for k, v in self.vision_output_slices.items()}
+      outputs = self.parser.parse_outputs(sliced)
+    else:
+      vision_output = raw_outputs[0].numpy().flatten()
+      vision_sliced = {k: vision_output[np.newaxis, v] for k, v in self.vision_output_slices.items()}
+      outputs = self.parser.parse_vision_outputs(vision_sliced)
+
+      for i, policy_slices in enumerate(self._policy_slices_list):
+        policy_output = raw_outputs[i + 1].numpy().flatten()
+        policy_sliced = {k: policy_output[np.newaxis, v] for k, v in policy_slices.items()}
+        parsed = self.parser.parse_policy_outputs(policy_sliced)
+        if 'off' in self._policy_keys[i] and self._has_on_policy:
+          parsed.pop('plan', None)
+        outputs.update(parsed)
+
+      if 'planplus' in outputs and 'plan' in outputs:
+        outputs['plan'] = outputs['plan'] + outputs['planplus']
+
+    if 'desired_curvature' in outputs and 'prev_desired_curv' in self.numpy_inputs:
+      buf = self.numpy_inputs['prev_desired_curv']
+      buf[0, :-1] = buf[0, 1:]
+      buf[0, -1, :] = outputs['desired_curvature'][0, :] if not self.mlsim else 0
+
+    return outputs
+
+  def get_action_from_model(self, model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action,
+                            lat_action_t: float, long_action_t: float, v_ego: float) -> log.ModelDataV2.Action:
+    plan = model_output['plan'][0]
+    desired_accel, should_stop = get_accel_from_plan(plan[:, Plan.VELOCITY][:, 0], plan[:, Plan.ACCELERATION][:, 0], self.constants.T_IDXS,
+                                                     action_t=long_action_t)
+    desired_accel = smooth_value(desired_accel, prev_action.desiredAcceleration, self.LONG_SMOOTH_SECONDS)
+
+    curvature_plan = plan + (self.PLANPLUS_CONTROL - 1.0) * model_output['planplus'][0] if 'planplus' in model_output and self.PLANPLUS_CONTROL != 1.0 else plan
+    desired_curvature = get_curvature_from_output(model_output, curvature_plan, v_ego, lat_action_t, self.mlsim)
+    if self.generation is not None and self.generation >= 10: # smooth curvature for post FOF models
+      if v_ego > self.MIN_LAT_CONTROL_SPEED:
+        desired_curvature = smooth_value(desired_curvature, prev_action.desiredCurvature, self.LAT_SMOOTH_SECONDS)
+      else:
+        desired_curvature = prev_action.desiredCurvature
+
+    return log.ModelDataV2.Action(desiredCurvature=float(desired_curvature),desiredAcceleration=float(desired_accel), shouldStop=bool(should_stop))
+
+
+def main(demo=False):
+  cloudlog.warning("modeld init")
+
+  sentry.set_tag("daemon", PROCESS_NAME)
+  cloudlog.bind(daemon=PROCESS_NAME)
+  setproctitle(PROCESS_NAME)
+  config_realtime_process(7, 54)
+
+  # visionipc clients
+  while True:
+    available_streams = VisionIpcClient.available_streams("camerad", block=False)
+    if available_streams:
+      use_extra_client = VisionStreamType.VISION_STREAM_WIDE_ROAD in available_streams and VisionStreamType.VISION_STREAM_ROAD in available_streams
+      main_wide_camera = VisionStreamType.VISION_STREAM_ROAD not in available_streams
+      break
+    time.sleep(.1)
+
+  vipc_client_main_stream = VisionStreamType.VISION_STREAM_WIDE_ROAD if main_wide_camera else VisionStreamType.VISION_STREAM_ROAD
+  vipc_client_main = VisionIpcClient("camerad", vipc_client_main_stream, True)
+  vipc_client_extra = VisionIpcClient("camerad", VisionStreamType.VISION_STREAM_WIDE_ROAD, False)
+  cloudlog.warning(f"vision stream set up, main_wide_camera: {main_wide_camera}, use_extra_client: {use_extra_client}")
+
+  while not vipc_client_main.connect(False):
+    time.sleep(0.1)
+  while use_extra_client and not vipc_client_extra.connect(False):
+    time.sleep(0.1)
+
+  cloudlog.warning(f"connected main cam with buffer size: {vipc_client_main.buffer_len} ({vipc_client_main.width} x {vipc_client_main.height})")
+  if use_extra_client:
+    cloudlog.warning(f"connected extra cam with buffer size: {vipc_client_extra.buffer_len} ({vipc_client_extra.width} x {vipc_client_extra.height})")
+
+  cloudlog.warning("loading model")
+  model = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height)
+  cloudlog.warning("models loaded, modeld starting")
+
+  # messaging
+  pm = PubMaster(["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"])
+  sm = SubMaster(["deviceState", "carState", "roadCameraState", "liveCalibration", "driverMonitoringState", "carControl", "liveDelay"])
+
+  publish_state = PublishState()
+  params = Params()
+
+  # setup filter to track dropped frames
+  frame_dropped_filter = FirstOrderFilter(0., 10., 1. / model.constants.MODEL_FREQ)
+  frame_id = 0
+  last_vipc_frame_id = 0
+  run_count = 0
+
+  model_transform_main = np.zeros((3, 3), dtype=np.float32)
+  model_transform_extra = np.zeros((3, 3), dtype=np.float32)
+  live_calib_seen = False
+  buf_main, buf_extra = None, None
+  meta_main = FrameMeta()
+  meta_extra = FrameMeta()
+  camera_offset_helper = CameraOffsetHelper()
+
+
+  if demo:
+    CP = get_demo_car_params()
+  else:
+    CP = messaging.log_from_bytes(params.get("CarParams", block=True), car.CarParams)
+  cloudlog.info("modeld got CarParams: %s", CP.brand)
+
+  # TODO Move smooth seconds to action function
+  long_delay = CP.longitudinalActuatorDelay + model.LONG_SMOOTH_SECONDS
+  prev_action = log.ModelDataV2.Action()
+
+  DH = DesireHelper()
+  meta_constants = load_meta_constants()
+
+  while True:
+    # Keep receiving frames until we are at least 1 frame ahead of previous extra frame
+    while meta_main.timestamp_sof < meta_extra.timestamp_sof + 25000000:
+      buf_main = vipc_client_main.recv()
+      meta_main = FrameMeta(vipc_client_main)
+      if buf_main is None:
+        break
+
+    if buf_main is None:
+      cloudlog.debug("vipc_client_main no frame")
+      continue
+
+    if use_extra_client:
+      # Keep receiving extra frames until frame id matches main camera
+      while True:
+        buf_extra = vipc_client_extra.recv()
+        meta_extra = FrameMeta(vipc_client_extra)
+        if buf_extra is None or meta_main.timestamp_sof < meta_extra.timestamp_sof + 25000000:
+          break
+
+      if buf_extra is None:
+        cloudlog.debug("vipc_client_extra no frame")
+        continue
+
+      if abs(meta_main.timestamp_sof - meta_extra.timestamp_sof) > 10000000:
+        cloudlog.error(f"frames out of sync! main: {meta_main.frame_id} ({meta_main.timestamp_sof / 1e9:.5f}),\
+                       extra: {meta_extra.frame_id} ({meta_extra.timestamp_sof / 1e9:.5f})")
+
+    else:
+      # Use single camera
+      buf_extra = buf_main
+      meta_extra = meta_main
+
+    sm.update(0)
+    desire = DH.desire
+    is_rhd = sm["driverMonitoringState"].isRHD
+    frame_id = sm["roadCameraState"].frameId
+    v_ego = max(sm["carState"].vEgo, 0.)
+    if sm.frame % 60 == 0:
+      model.lat_delay = get_lat_delay(params, sm["liveDelay"].lateralDelay)
+      model.PLANPLUS_CONTROL = params.get("PlanplusControl", return_default=True)
+      camera_offset_helper.set_offset(params.get("CameraOffset", return_default=True))
+    lat_delay = model.lat_delay + model.LAT_SMOOTH_SECONDS
+    if sm.updated["liveCalibration"] and sm.seen['roadCameraState'] and sm.seen['deviceState']:
+      device_from_calib_euler = np.array(sm["liveCalibration"].rpyCalib, dtype=np.float32)
+      dc = DEVICE_CAMERAS[(str(sm['deviceState'].deviceType), str(sm['roadCameraState'].sensor))]
+      model_transform_main = get_warp_matrix(device_from_calib_euler, dc.ecam.intrinsics if main_wide_camera else dc.fcam.intrinsics, False).astype(np.float32)
+      model_transform_extra = get_warp_matrix(device_from_calib_euler, dc.ecam.intrinsics, True).astype(np.float32)
+      model_transform_main, model_transform_extra = camera_offset_helper.update(model_transform_main, model_transform_extra, sm, main_wide_camera)
+      live_calib_seen = True
+
+    traffic_convention = np.zeros(2)
+    traffic_convention[int(is_rhd)] = 1
+
+    vec_desire = np.zeros(model.constants.DESIRE_LEN, dtype=np.float32)
+    if desire >= 0 and desire < model.constants.DESIRE_LEN:
+      vec_desire[desire] = 1
+
+    # tracked dropped frames
+    vipc_dropped_frames = max(0, meta_main.frame_id - last_vipc_frame_id - 1)
+    frames_dropped = frame_dropped_filter.update(min(vipc_dropped_frames, 10))
+    if run_count < 10: # let frame drops warm up
+      frame_dropped_filter.x = 0.
+      frames_dropped = 0.
+    run_count = run_count + 1
+
+    frame_drop_ratio = frames_dropped / (1 + frames_dropped)
+    prepare_only = vipc_dropped_frames > 0
+    if prepare_only:
+      cloudlog.error(f"skipping model eval. Dropped {vipc_dropped_frames} frames")
+
+    bufs = {name: buf_extra if 'big' in name else buf_main for name in model.vision_input_names}
+    transforms = {name: model_transform_extra if 'big' in name else model_transform_main for name in model.vision_input_names}
+    inputs:dict[str, np.ndarray] = {
+      model.desire_key: vec_desire,
+      'traffic_convention': traffic_convention,
+    }
+
+    if 'lateral_control_params' in model.numpy_inputs:
+      inputs['lateral_control_params'] = np.array([v_ego, lat_delay], dtype=np.float32)
+
+    mt1 = time.perf_counter()
+    model_output = model.run(bufs, transforms, inputs, prepare_only)
+    mt2 = time.perf_counter()
+    model_execution_time = mt2 - mt1
+
+    if model_output is not None:
+      modelv2_send = messaging.new_message('modelV2')
+      drivingdata_send = messaging.new_message('drivingModelData')
+      posenet_send = messaging.new_message('cameraOdometry')
+      mdv2sp_send = messaging.new_message('modelDataV2SP')
+
+      action = model.get_action_from_model(model_output, prev_action, lat_delay + DT_MDL, long_delay + DT_MDL, v_ego)
+      prev_action = action
+      fill_model_msg(drivingdata_send, modelv2_send, model_output, action,
+                     publish_state, meta_main.frame_id, meta_extra.frame_id, frame_id,
+                     frame_drop_ratio, meta_main.timestamp_eof, model_execution_time, live_calib_seen, meta_constants)
+
+      desire_state = modelv2_send.modelV2.meta.desireState
+      l_lane_change_prob = desire_state[log.Desire.laneChangeLeft]
+      r_lane_change_prob = desire_state[log.Desire.laneChangeRight]
+      lane_change_prob = l_lane_change_prob + r_lane_change_prob
+      DH.update(sm['carState'], sm['carControl'].latActive, lane_change_prob)
+      modelv2_send.modelV2.meta.laneChangeState = DH.lane_change_state
+      modelv2_send.modelV2.meta.laneChangeDirection = DH.lane_change_direction
+      mdv2sp_send.modelDataV2SP.laneTurnDirection = DH.lane_turn_direction
+      drivingdata_send.drivingModelData.meta.laneChangeState = DH.lane_change_state
+      drivingdata_send.drivingModelData.meta.laneChangeDirection = DH.lane_change_direction
+
+      fill_pose_msg(posenet_send, model_output, meta_main.frame_id, vipc_dropped_frames, meta_main.timestamp_eof, live_calib_seen)
+      pm.send('modelV2', modelv2_send)
+      pm.send('drivingModelData', drivingdata_send)
+      pm.send('cameraOdometry', posenet_send)
+      pm.send('modelDataV2SP', mdv2sp_send)
+    last_vipc_frame_id = meta_main.frame_id
+
+
+if __name__ == "__main__":
+  try:
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--demo', action='store_true', help='A boolean for demo mode.')
+    args = parser.parse_args()
+    main(demo=args.demo)
+  except KeyboardInterrupt:
+    cloudlog.warning(f"child {PROCESS_NAME} got SIGINT")
+  except Exception:
+    sentry.capture_exception()
+    raise

openpilot/sunnypilot/modeld_v2/modeld_base.py

@@ -0,0 +1,12 @@
+"""
+Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
+
+This file is part of sunnypilot and is licensed under the MIT License.
+See the LICENSE.md file in the root directory for more details.
+"""
+from openpilot.common.params import Params
+
+
+class ModelStateBase:
+  def __init__(self):
+    self.lat_delay = Params().get("LagdValueCache", return_default=True)

openpilot/sunnypilot/modeld_v2/parse_model_outputs.py

@@ -0,0 +1,107 @@
+import numpy as np
+from openpilot.sunnypilot.modeld_v2.constants import ModelConstants
+
+def safe_exp(x, out=None):
+  # -11 is around 10**14, more causes float16 overflow
+  return np.exp(np.clip(x, -np.inf, 11), out=out)
+
+def sigmoid(x):
+  return 1. / (1. + safe_exp(-x))
+
+def softmax(x, axis=-1):
+  x -= np.max(x, axis=axis, keepdims=True)
+  if x.dtype == np.float32 or x.dtype == np.float64:
+    safe_exp(x, out=x)
+  else:
+    x = safe_exp(x)
+  x /= np.sum(x, axis=axis, keepdims=True)
+  return x
+
+class Parser:
+  def __init__(self, ignore_missing=False):
+    self.ignore_missing = ignore_missing
+
+  def check_missing(self, outs, name):
+    if name not in outs and not self.ignore_missing:
+      raise ValueError(f"Missing output {name}")
+    return name not in outs
+
+  def parse_categorical_crossentropy(self, name, outs, out_shape=None):
+    if self.check_missing(outs, name):
+      return
+    raw = outs[name]
+    if out_shape is not None:
+      raw = raw.reshape((raw.shape[0],) + out_shape)
+    outs[name] = softmax(raw, axis=-1)
+
+  def parse_binary_crossentropy(self, name, outs):
+    if self.check_missing(outs, name):
+      return
+    raw = outs[name]
+    outs[name] = sigmoid(raw)
+
+  def parse_mdn(self, name, outs, in_N=0, out_N=1, out_shape=None):
+    if self.check_missing(outs, name):
+      return
+    raw = outs[name]
+    raw = raw.reshape((raw.shape[0], max(in_N, 1), -1))
+
+    n_values = (raw.shape[2] - out_N)//2
+    pred_mu = raw[:,:,:n_values]
+    pred_std = safe_exp(raw[:,:,n_values: 2*n_values])
+
+    if in_N > 1:
+      weights = np.zeros((raw.shape[0], in_N, out_N), dtype=raw.dtype)
+      for i in range(out_N):
+        weights[:,:,i - out_N] = softmax(raw[:,:,i - out_N], axis=-1)
+
+      if out_N == 1:
+        for fidx in range(weights.shape[0]):
+          idxs = np.argsort(weights[fidx][:,0])[::-1]
+          weights[fidx] = weights[fidx][idxs]
+          pred_mu[fidx] = pred_mu[fidx][idxs]
+          pred_std[fidx] = pred_std[fidx][idxs]
+      full_shape = tuple([raw.shape[0], in_N] + list(out_shape))
+      outs[name + '_weights'] = weights
+      outs[name + '_hypotheses'] = pred_mu.reshape(full_shape)
+      outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape)
+
+      pred_mu_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
+      pred_std_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
+      for fidx in range(weights.shape[0]):
+        for hidx in range(out_N):
+          idxs = np.argsort(weights[fidx,:,hidx])[::-1]
+          pred_mu_final[fidx, hidx] = pred_mu[fidx, idxs[0]]
+          pred_std_final[fidx, hidx] = pred_std[fidx, idxs[0]]
+    else:
+      pred_mu_final = pred_mu
+      pred_std_final = pred_std
+
+    if out_N > 1:
+      final_shape = tuple([raw.shape[0], out_N] + list(out_shape))
+    else:
+      final_shape = tuple([raw.shape[0],] + list(out_shape))
+    outs[name] = pred_mu_final.reshape(final_shape)
+    outs[name + '_stds'] = pred_std_final.reshape(final_shape)
+
+  def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
+    self.parse_mdn('plan', outs, in_N=ModelConstants.PLAN_MHP_N, out_N=ModelConstants.PLAN_MHP_SELECTION,
+                   out_shape=(ModelConstants.IDX_N,ModelConstants.PLAN_WIDTH))
+    self.parse_mdn('lane_lines', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_LANE_LINES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH))
+    self.parse_mdn('road_edges', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_ROAD_EDGES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH))
+    self.parse_mdn('pose', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,))
+    self.parse_mdn('road_transform', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,))
+    if 'sim_pose' in outs:
+      self.parse_mdn('sim_pose', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,))
+    self.parse_mdn('wide_from_device_euler', outs, in_N=0, out_N=0, out_shape=(ModelConstants.WIDE_FROM_DEVICE_WIDTH,))
+    self.parse_mdn('lead', outs, in_N=ModelConstants.LEAD_MHP_N, out_N=ModelConstants.LEAD_MHP_SELECTION,
+                   out_shape=(ModelConstants.LEAD_TRAJ_LEN,ModelConstants.LEAD_WIDTH))
+    if 'lat_planner_solution' in outs:
+      self.parse_mdn('lat_planner_solution', outs, in_N=0, out_N=0, out_shape=(ModelConstants.IDX_N,ModelConstants.LAT_PLANNER_SOLUTION_WIDTH))
+    if 'desired_curvature' in outs:
+      self.parse_mdn('desired_curvature', outs, in_N=0, out_N=0, out_shape=(ModelConstants.DESIRED_CURV_WIDTH,))
+    for k in ['lead_prob', 'lane_lines_prob', 'meta']:
+      self.parse_binary_crossentropy(k, outs)
+    self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,))
+    self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH))
+    return outs

openpilot/sunnypilot/modeld_v2/parse_model_outputs_split.py

@@ -0,0 +1,154 @@
+import numpy as np
+from openpilot.sunnypilot.models.split_model_constants import SplitModelConstants
+
+
+def safe_exp(x, out=None):
+  # -11 is around 10**14, more causes float16 overflow
+  return np.exp(np.clip(x, -np.inf, 11), out=out)
+
+
+def sigmoid(x):
+  return 1. / (1. + safe_exp(-x))
+
+
+def softmax(x, axis=-1):
+  x -= np.max(x, axis=axis, keepdims=True)
+  if x.dtype == np.float32 or x.dtype == np.float64:
+    safe_exp(x, out=x)
+  else:
+    x = safe_exp(x)
+  x /= np.sum(x, axis=axis, keepdims=True)
+  return x
+
+
+class Parser:
+  def __init__(self, ignore_missing=False):
+    self.ignore_missing = ignore_missing
+
+  def check_missing(self, outs, name):
+    if name not in outs and not self.ignore_missing:
+      raise ValueError(f"Missing output {name}")
+    return name not in outs
+
+  def parse_categorical_crossentropy(self, name, outs, out_shape=None):
+    if self.check_missing(outs, name):
+      return
+    raw = outs[name]
+    if out_shape is not None:
+      raw = raw.reshape((raw.shape[0],) + out_shape)
+    outs[name] = softmax(raw, axis=-1)
+
+  def parse_binary_crossentropy(self, name, outs):
+    if self.check_missing(outs, name):
+      return
+    raw = outs[name]
+    outs[name] = sigmoid(raw)
+
+  def parse_mdn(self, name, outs, in_N=0, out_N=1, out_shape=None):
+    if self.check_missing(outs, name):
+      return
+    raw = outs[name]
+    raw = raw.reshape((raw.shape[0], max(in_N, 1), -1))
+
+    n_values = (raw.shape[2] - out_N)//2
+    pred_mu = raw[:,:,:n_values]
+    pred_std = safe_exp(raw[:,:,n_values: 2*n_values])
+
+    if in_N > 1:
+      weights = np.zeros((raw.shape[0], in_N, out_N), dtype=raw.dtype)
+      for i in range(out_N):
+        weights[:,:,i - out_N] = softmax(raw[:,:,i - out_N], axis=-1)
+
+      if out_N == 1:
+        for fidx in range(weights.shape[0]):
+          idxs = np.argsort(weights[fidx][:,0])[::-1]
+          weights[fidx] = weights[fidx][idxs]
+          pred_mu[fidx] = pred_mu[fidx][idxs]
+          pred_std[fidx] = pred_std[fidx][idxs]
+      full_shape = tuple([raw.shape[0], in_N] + list(out_shape))
+      outs[name + '_weights'] = weights
+      outs[name + '_hypotheses'] = pred_mu.reshape(full_shape)
+      outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape)
+
+      pred_mu_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
+      pred_std_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype)
+      for fidx in range(weights.shape[0]):
+        for hidx in range(out_N):
+          idxs = np.argsort(weights[fidx,:,hidx])[::-1]
+          pred_mu_final[fidx, hidx] = pred_mu[fidx, idxs[0]]
+          pred_std_final[fidx, hidx] = pred_std[fidx, idxs[0]]
+    else:
+      pred_mu_final = pred_mu
+      pred_std_final = pred_std
+
+    if out_N > 1:
+      final_shape = tuple([raw.shape[0], out_N] + list(out_shape))
+    else:
+      final_shape = tuple([raw.shape[0],] + list(out_shape))
+    outs[name] = pred_mu_final.reshape(final_shape)
+    outs[name + '_stds'] = pred_std_final.reshape(final_shape)
+
+  def is_mhp(self, outs, name, shape):
+    if self.check_missing(outs, name):
+      return False
+    if outs[name].shape[1] == 2 * shape:
+      return False
+    return True
+
+  def parse_dynamic_outputs(self, outs: dict[str, np.ndarray]) -> None:
+    if 'lead' in outs:
+      lead_mhp = self.is_mhp(outs, 'lead',
+                             SplitModelConstants.LEAD_MHP_SELECTION * SplitModelConstants.LEAD_TRAJ_LEN * SplitModelConstants.LEAD_WIDTH)
+      lead_in_N, lead_out_N = (SplitModelConstants.LEAD_MHP_N, SplitModelConstants.LEAD_MHP_SELECTION) if lead_mhp else (0, 0)
+      lead_out_shape = (SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH) if lead_mhp else \
+        (SplitModelConstants.LEAD_MHP_SELECTION, SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH)
+      self.parse_mdn('lead', outs, in_N=lead_in_N, out_N=lead_out_N, out_shape=lead_out_shape)
+    if 'plan' in outs:
+      plan_mhp = self.is_mhp(outs, 'plan', SplitModelConstants.IDX_N * SplitModelConstants.PLAN_WIDTH)
+      plan_in_N, plan_out_N = (SplitModelConstants.PLAN_MHP_N, SplitModelConstants.PLAN_MHP_SELECTION) if plan_mhp else (0, 0)
+      self.parse_mdn('plan', outs, in_N=plan_in_N, out_N=plan_out_N,
+                     out_shape=(SplitModelConstants.IDX_N, SplitModelConstants.PLAN_WIDTH))
+    if 'planplus' in outs:
+      self.parse_mdn('planplus', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.IDX_N, SplitModelConstants.PLAN_WIDTH))
+
+  def split_outputs(self, outs: dict[str, np.ndarray]) -> None:
+    if 'desired_curvature' in outs:
+      self.parse_mdn('desired_curvature', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.DESIRED_CURV_WIDTH,))
+    if 'desire_pred' in outs:
+      self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(SplitModelConstants.DESIRE_PRED_LEN,SplitModelConstants.DESIRE_PRED_WIDTH))
+    if 'desire_state' in outs:
+      self.parse_categorical_crossentropy('desire_state', outs, out_shape=(SplitModelConstants.DESIRE_PRED_WIDTH,))
+    if 'lane_lines' in outs:
+      self.parse_mdn('lane_lines', outs, in_N=0, out_N=0,
+                    out_shape=(SplitModelConstants.NUM_LANE_LINES,SplitModelConstants.IDX_N,SplitModelConstants.LANE_LINES_WIDTH))
+    if 'lane_lines_prob' in outs:
+      self.parse_binary_crossentropy('lane_lines_prob', outs)
+    if 'lead_prob' in outs:
+      self.parse_binary_crossentropy('lead_prob', outs)
+    if 'lat_planner_solution' in outs:
+      self.parse_mdn('lat_planner_solution', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.IDX_N,SplitModelConstants.LAT_PLANNER_SOLUTION_WIDTH))
+    if 'meta' in outs:
+      self.parse_binary_crossentropy('meta', outs)
+    if 'road_edges' in outs:
+      self.parse_mdn('road_edges', outs, in_N=0, out_N=0,
+                    out_shape=(SplitModelConstants.NUM_ROAD_EDGES,SplitModelConstants.IDX_N,SplitModelConstants.LANE_LINES_WIDTH))
+    if 'sim_pose' in outs:
+      self.parse_mdn('sim_pose', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,))
+
+  def parse_vision_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
+    self.parse_mdn('pose', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,))
+    self.parse_mdn('wide_from_device_euler', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.WIDE_FROM_DEVICE_WIDTH,))
+    self.parse_mdn('road_transform', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,))
+    self.parse_dynamic_outputs(outs)
+    self.split_outputs(outs)
+    return outs
+
+  def parse_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
+    self.parse_dynamic_outputs(outs)
+    self.split_outputs(outs)
+    return outs
+
+  def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
+    outs = self.parse_vision_outputs(outs)
+    outs = self.parse_policy_outputs(outs)
+    return outs

openpilot/sunnypilot/modeld_v2/tests/conftest.py

@@ -0,0 +1,204 @@
+"""
+Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
+
+This file is part of sunnypilot and is licensed under the MIT License.
+See the LICENSE.md file in the root directory for more details.
+"""
+
+import pickle
+import pytest
+
+import openpilot.sunnypilot.models.helpers as helpers
+import openpilot.sunnypilot.modeld_v2.modeld as modeld_module
+from openpilot.sunnypilot.modeld_v2.constants import ModelConstants
+from openpilot.sunnypilot.models.split_model_constants import SplitModelConstants
+
+
+class DummyOverride:
+  def __init__(self, key, value):
+    self.key = key
+    self.value = value
+
+
+class DummyArtifact:
+  def __init__(self, file_name):
+    self.fileName = file_name
+
+
+class DummyModelType:
+  def __init__(self, raw):
+    self.raw = raw
+
+
+class DummyModel:
+  def __init__(self, type_str, artifact_file):
+    self.type = DummyModelType(type_str)
+    self.artifact = DummyArtifact(artifact_file)
+
+
+class DummyBundle:
+  def __init__(self, is_20hz=False, models=None, generation=10):
+    self.overrides = [DummyOverride('lat', '.1'), DummyOverride('long', '.3')]
+    self.generation = generation
+    self.is20hz = is_20hz
+    self.models = models or []
+
+
+class Archetype:
+  def __init__(self, name, metadata_structure, model_stubs, is_20hz,
+               expected_model_type, expected_constants_class, expected_parser_module,
+               expected_desire_key):
+    self.name = name
+    self.metadata_structure = metadata_structure
+    self.model_stubs = model_stubs
+    self.is_20hz = is_20hz
+    self.expected_model_type = expected_model_type
+    self.expected_constants_class = expected_constants_class
+    self.expected_parser_module = expected_parser_module
+    self.expected_desire_key = expected_desire_key
+
+
+def _noop_jit(**kwargs):
+  pass
+
+
+def _make_vision_policy_metadata(vision_input_shapes, policy_input_shapes,
+                                 vision_output_slices, policy_output_slices):
+  return {
+    'vision': {'input_shapes': vision_input_shapes, 'output_slices': vision_output_slices},
+    'policy': {'input_shapes': policy_input_shapes, 'output_slices': policy_output_slices},
+  }
+
+
+def _make_multi_policy_metadata(vision_input_shapes, policy_input_shapes,
+                                vision_output_slices, policy_output_slices):
+  return {
+    'vision': {'input_shapes': vision_input_shapes, 'output_slices': vision_output_slices},
+    'offPolicy': {'input_shapes': policy_input_shapes, 'output_slices': policy_output_slices},
+  }
+
+
+def _make_tri_policy_metadata(vision_input_shapes, policy_input_shapes,
+                              vision_output_slices, policy_output_slices):
+  return {
+    'vision': {'input_shapes': vision_input_shapes, 'output_slices': vision_output_slices},
+    'onPolicy': {'input_shapes': policy_input_shapes, 'output_slices': policy_output_slices},
+    'offPolicy': {'input_shapes': policy_input_shapes, 'output_slices': policy_output_slices},
+  }
+
+
+def _make_supercombo_metadata(input_shapes, output_slices):
+  return {'model': {'input_shapes': input_shapes, 'output_slices': output_slices}}
+
+
+SPLIT_VISION_INPUT_SHAPES = {'img': (1, 12, 128, 256), 'big_img': (1, 12, 128, 256)}
+SPLIT_POLICY_INPUT_SHAPES = {'features_buffer': (1, 25, 512), 'desire_pulse': (1, 25, 8), 'traffic_convention': (1, 2)}
+SPLIT_VISION_SLICES = {'hidden_state': slice(0, 512), 'pose': slice(512, 524)}
+SPLIT_POLICY_SLICES = {'plan': slice(0, 495), 'meta': slice(495, 550)}
+
+SUPERCOMBO_INPUT_SHAPES = {
+  'img': (1, 12, 128, 256), 'big_img': (1, 12, 128, 256),
+  'desire': (1, 100, 8), 'features_buffer': (1, 99, 512),
+  'lateral_control_params': (1, 2), 'prev_desired_curv': (1, 100, 1),
+  'traffic_convention': (1, 2),
+}
+SUPERCOMBO_SLICES = {'plan': slice(0, 495), 'hidden_state': slice(495, 1007), 'meta': slice(1007, 1062)}
+
+CAM_W, CAM_H = 1928, 1208
+
+ARCHETYPES = {
+  'vision_policy_split': Archetype(
+    name='vision_policy_split',
+    metadata_structure=_make_vision_policy_metadata(
+      SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES,
+      SPLIT_VISION_SLICES, SPLIT_POLICY_SLICES),
+    model_stubs=[DummyModel('vision', 'driving_test_tinygrad.pkl'),
+                 DummyModel('policy', 'driving_test_tinygrad.pkl')],
+    is_20hz=True,
+    expected_model_type='split',
+    expected_constants_class=SplitModelConstants,
+    expected_parser_module='parse_model_outputs_split',
+    expected_desire_key='desire',
+  ),
+  'vision_multi_policy': Archetype(
+    name='vision_multi_policy',
+    metadata_structure=_make_multi_policy_metadata(
+      SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES,
+      SPLIT_VISION_SLICES, SPLIT_POLICY_SLICES),
+    model_stubs=[DummyModel('vision', 'driving_test_tinygrad.pkl'),
+                 DummyModel('offPolicy', 'driving_test_tinygrad.pkl')],
+    is_20hz=True,
+    expected_model_type='multi_policy',
+    expected_constants_class=SplitModelConstants,
+    expected_parser_module='parse_model_outputs_split',
+    expected_desire_key='desire',
+  ),
+  'tri_policy': Archetype(
+    name='tri_policy',
+    metadata_structure=_make_tri_policy_metadata(
+      SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES,
+      SPLIT_VISION_SLICES, SPLIT_POLICY_SLICES),
+    model_stubs=[DummyModel('vision', 'driving_test_tinygrad.pkl'),
+                 DummyModel('onPolicy', 'driving_test_tinygrad.pkl'),
+                 DummyModel('offPolicy', 'driving_test_tinygrad.pkl')],
+    is_20hz=True,
+    expected_model_type='multi_policy',
+    expected_constants_class=SplitModelConstants,
+    expected_parser_module='parse_model_outputs_split',
+    expected_desire_key='desire',
+  ),
+  'supercombo_non20hz': Archetype(
+    name='supercombo_non20hz',
+    metadata_structure=_make_supercombo_metadata(SUPERCOMBO_INPUT_SHAPES, SUPERCOMBO_SLICES),
+    model_stubs=[DummyModel('supercombo', 'driving_test_tinygrad.pkl')],
+    is_20hz=False,
+    expected_model_type='supercombo',
+    expected_constants_class=ModelConstants,
+    expected_parser_module='parse_model_outputs',
+    expected_desire_key='desire',
+  ),
+}
+
+
+def make_pkl_data(archetype):
+  return {
+    'metadata': archetype.metadata_structure,
+    (CAM_W, CAM_H): {'run_policy': _noop_jit, 'warp_enqueue': _noop_jit},
+  }
+
+
+def write_pkl(tmp_path, archetype):
+  pkl_path = tmp_path / 'driving_test_tinygrad.pkl'
+  with open(pkl_path, 'wb') as f:
+    pickle.dump(make_pkl_data(archetype), f)
+  return pkl_path
+
+
+def make_bundle(archetype):
+  return DummyBundle(
+    models=archetype.model_stubs,
+    is_20hz=archetype.is_20hz,
+  )
+
+
+@pytest.fixture
+def patch_modeld(monkeypatch):
+  def _patch(bundle):
+    monkeypatch.setattr(helpers, 'get_active_bundle', lambda params=None: bundle, raising=False)
+    monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None: bundle, raising=False)
+
+  return _patch
+
+
+@pytest.fixture
+def model_state_factory(tmp_path, monkeypatch, patch_modeld):
+  from openpilot.common.hardware import hw
+
+  def _create(archetype):
+    write_pkl(tmp_path, archetype)
+    bundle = make_bundle(archetype)
+    patch_modeld(bundle)
+    monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path)))
+    return modeld_module.ModelState(cam_w=CAM_W, cam_h=CAM_H)
+
+  return _create

openpilot/sunnypilot/modeld_v2/tests/test_camera_offset_helper.py

@@ -0,0 +1,84 @@
+"""
+Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
+
+This file is part of sunnypilot and is licensed under the MIT License.
+See the LICENSE.md file in the root directory for more details.
+"""
+import numpy as np
+
+from openpilot.common.transformations.camera import DEVICE_CAMERAS
+from openpilot.common.transformations.model import get_warp_matrix
+from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelper
+
+
+class MockStruct:
+  def __init__(self, **kwargs):
+    for k, v in kwargs.items():
+      setattr(self, k, v)
+
+  def __getitem__(self, item):
+    return getattr(self, item)
+
+
+class TestCameraOffset:
+  def setup_method(self):
+    self.camera_offset = CameraOffsetHelper()
+    self.dc = DEVICE_CAMERAS[('mici', 'os04c10')]
+
+  def test_smoothing(self):
+    self.camera_offset.set_offset(0.2)
+
+    sm = MockStruct(
+      deviceState=MockStruct(deviceType='mici'),
+      roadCameraState=MockStruct(sensor='os04c10'),
+      liveCalibration=MockStruct(rpyCalib=[0.0, 0.0, 0.0], height=[1.22])
+    )
+
+    intrinsics_main = self.dc.fcam.intrinsics
+    intrinsics_extra = self.dc.ecam.intrinsics
+    device_from_calib_euler = np.array([0.0, 0.0, 0.0], dtype=np.float32)
+    main_transform = get_warp_matrix(device_from_calib_euler, intrinsics_main, False).astype(np.float32)
+    extra_transform = get_warp_matrix(device_from_calib_euler, intrinsics_extra, True).astype(np.float32)
+
+    self.camera_offset.update(main_transform, extra_transform, sm, False)
+    np.testing.assert_almost_equal(self.camera_offset.actual_camera_offset, 0.02)
+    self.camera_offset.update(main_transform, extra_transform, sm, False)
+    np.testing.assert_almost_equal(self.camera_offset.actual_camera_offset, 0.038)
+
+  def test_camera_offset_(self):
+    intrinsics = self.dc.fcam.intrinsics
+    transform = np.eye(3, dtype=np.float32)
+    height = 1.22
+    offset = 0.1
+
+    cy = intrinsics[1, 2]
+    expected_shear = np.eye(3, dtype=np.float32)
+    expected_shear[0, 1] = offset / height
+    expected_shear[0, 2] = -offset / height * cy
+
+    result = CameraOffsetHelper.apply_camera_offset(transform, intrinsics, height, offset)
+    np.testing.assert_array_almost_equal(result, expected_shear)
+
+  def test_update(self):
+    sm = MockStruct(
+      deviceState=MockStruct(deviceType='mici'),
+      roadCameraState=MockStruct(sensor='os04c10'),
+      liveCalibration=MockStruct(rpyCalib=[0.0, 0.0, 0.0], height=[1.22])
+    )
+    intrinsics_main = self.dc.fcam.intrinsics
+    intrinsics_extra = self.dc.ecam.intrinsics
+    device_from_calib_euler = np.array([0.0, 0.0, 0.0], dtype=np.float32)
+    main_transform = get_warp_matrix(device_from_calib_euler, intrinsics_main, False).astype(np.float32)
+    extra_transform = get_warp_matrix(device_from_calib_euler, intrinsics_extra, True).astype(np.float32)
+
+    self.camera_offset.set_offset(0.0)  # test default offset doesn't change transformation
+    main_out, extra_out = self.camera_offset.update(main_transform, extra_transform, sm, False)
+    np.testing.assert_array_equal(main_out, main_transform)
+    np.testing.assert_array_equal(extra_out, extra_transform)
+
+    self.camera_offset.set_offset(0.2)  # test valid offset changes transformation
+    main_out, extra_out = self.camera_offset.update(main_transform, extra_transform, sm, False)
+    assert not np.array_equal(main_out, main_transform)
+    assert not np.array_equal(extra_out, extra_transform)
+    assert main_out[0, 1] != 0.0
+    assert main_out[0, 2] != 0.0

openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py

@@ -0,0 +1,260 @@
+"""
+Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
+
+This file is part of sunnypilot and is licensed under the MIT License.
+See the LICENSE.md file in the root directory for more details.
+"""
+
+import pytest
+
+import openpilot.sunnypilot.models.helpers as helpers
+import openpilot.sunnypilot.modeld_v2.modeld as modeld_module
+from openpilot.sunnypilot.modeld_v2.modeld import _find_driving_pkl
+from openpilot.sunnypilot.modeld_v2.tests.conftest import DummyModel, DummyBundle, ARCHETYPES, CAM_W, CAM_H, \
+  SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES
+
+ModelState = modeld_module.ModelState
+
+
+# Pkl discovery
+
+class TestFindDrivingPkl:
+  def test_returns_none_when_no_bundle(self):
+    assert _find_driving_pkl(None) is None
+
+  def test_returns_none_when_no_models(self):
+    bundle = DummyBundle(models=[])
+    assert _find_driving_pkl(bundle) is None
+
+  def test_returns_none_when_pkl_not_on_disk(self):
+    bundle = DummyBundle(models=[
+      DummyModel('vision', 'driving_fof_tinygrad.pkl'),
+      DummyModel('policy', 'driving_fof_tinygrad.pkl'),
+    ])
+    assert _find_driving_pkl(bundle) is None
+
+  def test_finds_pkl_by_artifact_name(self, tmp_path, monkeypatch):
+    (tmp_path / 'driving_fof_tinygrad.pkl').write_bytes(b'fake')
+    from openpilot.common.hardware import hw
+    monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path)))
+
+    bundle = DummyBundle(models=[
+      DummyModel('vision', 'driving_fof_tinygrad.pkl'),
+      DummyModel('policy', 'driving_fof_tinygrad.pkl'),
+    ])
+    result = _find_driving_pkl(bundle)
+    assert result is not None
+    assert 'driving_fof_tinygrad.pkl' in result
+
+
+# Init — assertion guard
+
+class TestModelStateCombinedInit:
+  def test_asserts_when_no_pkl(self, monkeypatch):
+    bundle = DummyBundle(models=[], is_20hz=True)
+    monkeypatch.setattr(helpers, 'get_active_bundle', lambda params=None: bundle, raising=False)
+    monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None: bundle, raising=False)
+    with pytest.raises(AssertionError, match="No driving pkl found"):
+      ModelState(cam_w=CAM_W, cam_h=CAM_H)
+
+
+class TestStockEquivalence:
+
+  def test_split_queue_keys_match_stock(self, model_state_factory):
+    from openpilot.selfdrive.modeld.compile_modeld import make_input_queues
+    from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip
+
+    state = model_state_factory(ARCHETYPES['vision_policy_split'])
+
+    frame_skip = derive_frame_skip(SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES)
+    # action_t is a deep-model prerequisite the SP loader doesn't provide yet; see skip_keys below
+    stock_shapes = {**SPLIT_VISION_INPUT_SHAPES, **SPLIT_POLICY_INPUT_SHAPES, 'action_t': (1, 2)}
+    stock_queues, stock_npy = make_input_queues(stock_shapes, frame_skip, device='NPY')
+
+    # TODO-SP: remove action_t skip once SP adds prerequisite for deep models (action_t input queue)
+    # prev_feat is a stock QCOM corruption workaround handled inside the SP loader's JIT path
+    skip_keys = {'action_t', 'prev_feat'}
+    # stock packs the per-key policy inputs into packed_npy_inputs; the npy views carry the individual keys
+    stock_queue_keys = set(stock_queues.keys())
+    if 'packed_npy_inputs' in stock_queue_keys:
+      stock_queue_keys.remove('packed_npy_inputs')
+      stock_queue_keys |= set(stock_npy.keys())
+    assert set(state.input_queues.keys()) == stock_queue_keys - skip_keys, \
+      f"Queue keys differ: v2={set(state.input_queues.keys())}, stock={stock_queue_keys}"
+    assert set(state.numpy_inputs.keys()) == set(stock_npy.keys()) - skip_keys, \
+      f"Npy keys differ: v2={set(state.numpy_inputs.keys())}, stock={set(stock_npy.keys())}"
+
+  def test_split_queue_keys_work_with_desire_key(self, model_state_factory):
+    from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues
+
+    policy_shapes_desire = {'features_buffer': (1, 25, 512), 'desire': (1, 25, 8), 'traffic_convention': (1, 2)}
+    frame_skip = derive_frame_skip(SPLIT_VISION_INPUT_SHAPES, policy_shapes_desire)
+    queues, npy = make_split_input_queues(SPLIT_VISION_INPUT_SHAPES, policy_shapes_desire, frame_skip, device='NPY')
+
+    assert 'desire_q' in queues
+    assert 'desire' in npy
+    assert 'img_q' in queues
+    assert 'feat_q' in queues
+
+  def test_split_vision_input_names_match_stock(self, model_state_factory):
+    state = model_state_factory(ARCHETYPES['vision_policy_split'])
+    assert state.vision_input_names == ['img', 'big_img']
+
+  def test_split_output_slices_preserved(self, model_state_factory):
+    arch = ARCHETYPES['vision_policy_split']
+    state = model_state_factory(arch)
+    assert state.vision_output_slices == arch.metadata_structure['vision']['output_slices']
+    assert state.policy_output_slices == arch.metadata_structure['policy']['output_slices']
+
+
+ARCHETYPE_NAMES = list(ARCHETYPES.keys())
+
+
+class TestModelTypeDetection:
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_combined_model_type(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    assert state._combined_model_type == arch.expected_model_type, \
+      f"{arch.name}: got {state._combined_model_type}, expected {arch.expected_model_type}"
+
+
+class TestConstantsSelection:
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_constants_class(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    assert type(state.constants) is arch.expected_constants_class, \
+      f"{arch.name}: got {type(state.constants).__name__}, expected {arch.expected_constants_class.__name__}"
+
+
+class TestParserSelection:
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_parser_module(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    parser_module = type(state.parser).__module__
+    assert parser_module.endswith(arch.expected_parser_module), \
+      f"{arch.name}: parser from {parser_module}, expected module ending with {arch.expected_parser_module}"
+
+
+class TestDesireKeyDetection:
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_desire_key(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    assert state.desire_key == arch.expected_desire_key, \
+      f"{arch.name}: got {state.desire_key}, expected {arch.expected_desire_key}"
+
+
+class TestVisionInputNames:
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_vision_names_contain_img(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    assert len(state.vision_input_names) >= 1
+    for name in state.vision_input_names:
+      assert 'img' in name, f"{arch.name}: vision input name '{name}' missing 'img'"
+
+
+class TestOutputSlices:
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_vision_slices_populated(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    assert len(state.vision_output_slices) > 0, f"{arch.name}: vision_output_slices empty"
+
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_policy_slices_match_type(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    if arch.expected_model_type == 'supercombo':
+      assert state.policy_output_slices == {}, f"{arch.name}: supercombo should have empty policy slices"
+    else:
+      assert len(state.policy_output_slices) > 0, f"{arch.name}: split/multi should have policy slices"
+
+
+class TestInputQueueCreation:
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_queues_not_empty(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    assert len(state.input_queues) > 0, f"{arch.name}: input_queues empty"
+
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_npy_contains_transforms(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    assert 'tfm' in state.numpy_inputs, f"{arch.name}: 'tfm' missing from npy"
+    assert 'big_tfm' in state.numpy_inputs, f"{arch.name}: 'big_tfm' missing from npy"
+    assert state.numpy_inputs['tfm'].shape == (3, 3)
+    assert state.numpy_inputs['big_tfm'].shape == (3, 3)
+
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_npy_contains_desire(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    assert arch.expected_desire_key in state.numpy_inputs, \
+      f"{arch.name}: '{arch.expected_desire_key}' missing from npy"
+
+
+class TestFrameBufferParams:
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_frame_buf_params_per_vision_input(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    for name in state.vision_input_names:
+      assert name in state.frame_buf_params, f"{arch.name}: frame_buf_params missing '{name}'"
+      nv12_info = state.frame_buf_params[name]
+      assert len(nv12_info) >= 4, f"{arch.name}: nv12_info for '{name}' too short"
+
+
+class TestBundleOverrides:
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_smoothing_params_from_overrides(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    assert state.LAT_SMOOTH_SECONDS == 0.1
+    assert state.LONG_SMOOTH_SECONDS == 0.3
+
+  @pytest.mark.parametrize("archetype_name", ARCHETYPE_NAMES)
+  def test_generation_from_bundle(self, archetype_name, model_state_factory):
+    arch = ARCHETYPES[archetype_name]
+    state = model_state_factory(arch)
+    assert state.generation == 10
+
+
+class TestMlsimProperty:
+  def test_mlsim_false_for_gen10(self, model_state_factory):
+    state = model_state_factory(ARCHETYPES['supercombo_non20hz'])
+    assert state.mlsim is False
+
+  def test_mlsim_true_for_gen11(self, tmp_path, monkeypatch, patch_modeld):
+    from openpilot.common.hardware import hw
+    from openpilot.sunnypilot.modeld_v2.tests.conftest import write_pkl, ARCHETYPES as A
+
+    arch = A['supercombo_non20hz']
+    write_pkl(tmp_path, arch)
+    bundle = DummyBundle(models=arch.model_stubs, is_20hz=arch.is_20hz, generation=11)
+    patch_modeld(bundle)
+    monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path)))
+
+    state = ModelState(cam_w=CAM_W, cam_h=CAM_H)
+    assert state.mlsim is True
+
+
+class TestCrossArchetypeMismatch:
+  def test_wrong_is_20hz_changes_constants(self, tmp_path, monkeypatch, patch_modeld):
+    from openpilot.common.hardware import hw
+    from openpilot.sunnypilot.modeld_v2.tests.conftest import write_pkl
+    from openpilot.sunnypilot.modeld_v2.constants import ModelConstants
+
+    arch = ARCHETYPES['vision_policy_split']
+    write_pkl(tmp_path, arch)
+    bundle = DummyBundle(models=arch.model_stubs, is_20hz=False)
+    patch_modeld(bundle)
+    monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path)))
+
+    state = ModelState(cam_w=CAM_W, cam_h=CAM_H)
+    assert type(state.constants) is ModelConstants, \
+      "Wrong is_20hz should produce wrong constants class"

openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py

@@ -0,0 +1,161 @@
+"""
+Copyright (c) 2021-, Haibin Wen, sunnypilot, and a number of other contributors.
+
+This file is part of sunnypilot and is licensed under the MIT License.
+See the LICENSE.md file in the root directory for more details.
+"""
+
+import numpy as np
+import pytest
+
+from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, _detect_desire_key
+
+
+class TestDeriveFrameSkip:
+  def test_non20hz_supercombo(self):
+    vision = {}
+    policy = {'features_buffer': (1, 99, 512), 'desire': (1, 100, 8)}
+    assert derive_frame_skip(vision, policy) == 1
+
+  def test_20hz_supercombo(self):
+    vision = {}
+    policy = {'features_buffer': (1, 24, 512), 'desire': (1, 25, 8)}
+    assert derive_frame_skip(vision, policy) == 4
+
+  def test_split_vision_policy(self):
+    vision = {'img': (1, 12, 128, 256)}
+    policy = {'features_buffer': (1, 25, 512), 'desire_pulse': (1, 25, 8)}
+    assert derive_frame_skip(vision, policy) == 4
+
+  def test_no_features_buffer(self):
+    assert derive_frame_skip({}, {}) == 1
+
+
+class TestFrameSkipBufferLengthEquivalence:
+  @pytest.mark.parametrize("frame_skip,expected_buffer_length", [
+    (1, 2),
+    (4, 5),
+  ])
+  def test_img_buffer_size_matches_warp_buffer_length(self, frame_skip, expected_buffer_length):
+    n_frames = 2
+    img_buf_dim0 = frame_skip * (n_frames - 1) + 1
+    assert img_buf_dim0 == expected_buffer_length, \
+      f"frame_skip={frame_skip}: img_buf[0]={img_buf_dim0}, expected {expected_buffer_length}"
+
+  @pytest.mark.parametrize("is_20hz,expected_frame_skip,expected_buffer_length", [
+    (False, 1, 2),
+    (True, 4, 5),
+  ])
+  def test_is_20hz_to_frame_skip_to_buffer_length(self, is_20hz, expected_frame_skip, expected_buffer_length):
+    if is_20hz:
+      policy_shapes = {'features_buffer': (1, 24, 512)}
+    else:
+      policy_shapes = {'features_buffer': (1, 99, 512)}
+    frame_skip = derive_frame_skip({}, policy_shapes)
+    assert frame_skip == expected_frame_skip
+
+    n_frames = 2
+    img_buf_dim0 = frame_skip * (n_frames - 1) + 1
+    assert img_buf_dim0 == expected_buffer_length
+
+
+class TestTemporalSamplingEquivalence:
+  def test_non20hz_desire_sampling_identity(self):
+    buf = np.random.randn(100, 1, 8).astype(np.float32)
+    frame_skip = 1
+    sampled = buf[::frame_skip].reshape(-1, 8)
+    assert sampled.shape == (100, 8)
+    np.testing.assert_array_equal(sampled, buf[:, 0, :])
+
+  def test_20hz_desire_sampling_max(self):
+    buf = np.zeros((100, 1, 8), dtype=np.float32)
+    buf[99, 0, 3] = 1.0
+    frame_skip = 4
+    reshaped = buf.reshape(-1, frame_skip, 1, 8).max(axis=1)
+    sampled = reshaped.reshape(-1, 8)
+    assert sampled.shape == (25, 8)
+    assert sampled[24, 3] == 1.0
+    assert sampled[23, 3] == 0.0
+
+  def test_split_features_buffer_sampling_skip(self):
+    buf = np.arange(100 * 512, dtype=np.float32).reshape(100, 1, 512)
+    frame_skip = 4
+    sampled = buf[::frame_skip].reshape(-1, 512)
+    assert sampled.shape == (25, 512)
+    np.testing.assert_array_equal(sampled[0], buf[0, 0])
+    np.testing.assert_array_equal(sampled[1], buf[4, 0])
+    np.testing.assert_array_equal(sampled[24], buf[96, 0])
+
+  def test_non20hz_features_buffer_sampling_identity(self):
+    buf = np.arange(99 * 512, dtype=np.float32).reshape(99, 1, 512)
+    frame_skip = 1
+    sampled = buf[::frame_skip].reshape(-1, 512)
+    assert sampled.shape == (99, 512)
+    np.testing.assert_array_equal(sampled, buf[:, 0, :])
+
+
+class TestTemporalIdxEquivalence:
+  @pytest.mark.parametrize("mode,desire_shape,fb_shape,frame_skip", [
+    ('non20hz', (1, 100, 8), (1, 99, 512), 1),
+    ('20hz', (1, 25, 8), (1, 24, 512), 4),
+    ('split', (1, 25, 8), (1, 25, 512), 4),
+  ])
+  def test_features_buffer_idx_equivalence(self, mode, desire_shape, fb_shape, frame_skip):
+    history = fb_shape[1]
+
+    if mode == 'non20hz':
+      modelstate_idxs = np.arange(history)
+      buf_len = history
+    elif mode == '20hz':
+      buf_len = (history + 1) * 4
+      step = int(-buf_len / history)
+      modelstate_idxs = np.arange(step, step * (history + 1), step)[::-1]
+    elif mode == 'split':
+      buf_len = history * 4
+      skip = buf_len // history
+      modelstate_idxs = np.arange(buf_len)[-1 - (skip * (history - 1))::skip]
+
+    assert len(modelstate_idxs) == fb_shape[1], \
+      f"{mode}: ModelState idx count {len(modelstate_idxs)} != input shape {fb_shape[1]}"
+
+  @pytest.mark.parametrize("mode,desire_shape,fb_shape,frame_skip", [
+    ('non20hz', (1, 100, 8), (1, 99, 512), 1),
+    ('20hz', (1, 25, 8), (1, 24, 512), 4),
+    ('split', (1, 25, 8), (1, 25, 512), 4),
+  ])
+  def test_desire_idx_equivalence(self, mode, desire_shape, fb_shape, frame_skip):
+    history = desire_shape[1]
+
+    compile_desire_buf_len = frame_skip * history if mode != 'non20hz' else history
+    compile_sampled_count = compile_desire_buf_len // frame_skip if frame_skip > 1 else compile_desire_buf_len
+    assert compile_sampled_count == history, \
+      f"{mode}: compile desire samples {compile_sampled_count} != model input {history}"
+
+
+class TestDetectDesireKey:
+  def test_finds_desire(self):
+    shapes = {'features_buffer': (1, 99, 512), 'desire': (1, 100, 8), 'traffic_convention': (1, 2)}
+    assert _detect_desire_key(shapes) == 'desire'
+
+  def test_finds_desire_pulse(self):
+    shapes = {'features_buffer': (1, 25, 512), 'desire_pulse': (1, 25, 8), 'traffic_convention': (1, 2)}
+    assert _detect_desire_key(shapes) == 'desire_pulse'
+
+  def test_returns_none_when_no_desire(self):
+    shapes = {'features_buffer': (1, 99, 512), 'traffic_convention': (1, 2)}
+    assert _detect_desire_key(shapes) is None
+
+
+class TestOutputSlicePreservation:
+  def test_vision_hidden_state_slice_used_for_features(self):
+    mock_slices = {'hidden_state': slice(0, 512), 'plan': slice(512, 1024)}
+    features_slice = mock_slices['hidden_state']
+    fake_output = np.random.randn(1, 1024).astype(np.float32)
+    features = fake_output[:, features_slice]
+    assert features.shape == (1, 512)
+
+  def test_policy_output_slices_independent(self):
+    vision_slices = {'hidden_state': slice(0, 512)}
+    policy_slices = {'plan': slice(0, 495), 'meta': slice(495, 550)}
+    assert set(vision_slices.keys()) & set(policy_slices.keys()) == set(), \
+      "vision and policy slices should not overlap in keys"

openpilot/sunnypilot/modeld_v2/tests/test_recovery_power.py

@@ -0,0 +1,70 @@
+import numpy as np
+
+from openpilot.cereal import log
+
+from openpilot.sunnypilot.modeld_v2.constants import Plan
+from openpilot.sunnypilot.modeld_v2.modeld import ModelState
+import openpilot.sunnypilot.modeld_v2.modeld as modeld
+
+
+class MockStruct:
+  def __init__(self, **kwargs):
+    for k, v in kwargs.items():
+      setattr(self, k, v)
+
+
+def test_recovery_power_scaling():
+  state = MockStruct(
+    PLANPLUS_CONTROL=0.75,
+    LONG_SMOOTH_SECONDS=0.3,
+    LAT_SMOOTH_SECONDS=0.1,
+    MIN_LAT_CONTROL_SPEED=0.3,
+    mlsim=True,
+    generation=12,
+    constants=MockStruct(T_IDXS=np.arange(100), DESIRE_LEN=8)
+  )
+  prev_action = log.ModelDataV2.Action()
+  recorded_vel: list = []
+  recorded_curv_plans: list = []
+
+  def mock_accel(plan_vel, plan_accel, t_idxs, action_t=0.0):
+    recorded_vel.append(plan_vel.copy())
+    return 0.0, False
+
+  def mock_curvature(output, plan, vego, lat_action_t, mlsim):
+    recorded_curv_plans.append(plan.copy())
+    return 0.0
+
+  modeld.get_accel_from_plan = mock_accel
+  modeld.get_curvature_from_output = mock_curvature
+  plan = np.random.rand(1, 100, 15).astype(np.float32)
+  planplus = np.random.rand(1, 100, 15).astype(np.float32)
+  merged_plan = plan + planplus
+
+  model_output: dict = {
+    'plan': merged_plan.copy(),
+    'planplus': planplus.copy()
+  }
+
+  test_cases: list = [
+    # (control, v_ego)
+    (0.55, 20.0),
+    (1.0, 25.0),
+    (1.5, 25.1),
+    (2.0, 20.0),
+    (0.75, 19.0),
+    (0.8, 25.1),
+  ]
+
+  for control, v_ego in test_cases:
+    state.PLANPLUS_CONTROL = control
+    recorded_vel.clear()
+    recorded_curv_plans.clear()
+    ModelState.get_action_from_model(state, model_output, prev_action, 0.0, 0.0, v_ego)
+
+    expected_accel_plan_vel = plan[0, :, Plan.VELOCITY][:, 0] + planplus[0, :, Plan.VELOCITY][:, 0]
+    np.testing.assert_allclose(recorded_vel[0], expected_accel_plan_vel, rtol=1e-5, atol=1e-6)
+
+    # For the below, yes, I know this isn't the same slicing as fillmodlmsg. This is to show that the values are only scaled on curv
+    expected_curv_plan_vel = plan[0, :, Plan.VELOCITY][:, 0] + control * planplus[0, :, Plan.VELOCITY][:, 0]
+    np.testing.assert_allclose(recorded_curv_plans[0][:, Plan.VELOCITY][:, 0], expected_curv_plan_vel, rtol=1e-5, atol=1e-6)

openpilot/sunnypilot/modeld_v2/tests/test_warp.py

@@ -0,0 +1,103 @@
+import os
+os.environ['DEV'] = 'CPU'
+import pytest
+import numpy as np
+from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
+from openpilot.sunnypilot.modeld_v2.warp import CAMERA_CONFIGS
+from openpilot.sunnypilot.modeld_v2.warp import Warp, MODEL_W, MODEL_H
+
+VISION_NAME_PAIRS = [  # needed to account for supercombos input_imgs
+  ('img', 'big_img'),
+  ('input_imgs', 'big_input_imgs'),
+]
+
+
+class MockVisionBuf:
+  def __init__(self, w, h):
+    self.width = w
+    self.height = h
+    _, _, _, yuv_size = get_nv12_info(w, h)
+    self.data = np.zeros(yuv_size, dtype=np.uint8)
+
+
+@pytest.mark.parametrize("buffer_length", [2, 5])
+def test_warp_initialization(buffer_length):
+  warp = Warp(buffer_length)
+  assert warp.buffer_length == buffer_length
+  assert warp.img_buffer_shape == (buffer_length * 6, MODEL_H // 2, MODEL_W // 2)
+
+
+@pytest.mark.parametrize("buffer_length", [2, 5])
+@pytest.mark.parametrize("cam_w, cam_h", CAMERA_CONFIGS)
+@pytest.mark.parametrize("road, wide", VISION_NAME_PAIRS)
+def test_warp_process(buffer_length, cam_w, cam_h, road, wide):
+  warp = Warp(buffer_length)
+  mock_buf = MockVisionBuf(cam_w, cam_h)
+  transform = np.eye(3, dtype=np.float32).flatten()
+  bufs = {road: mock_buf, wide: mock_buf}
+  transforms = {road: transform, wide: transform}
+
+  out = warp.process(bufs, transforms)
+  assert isinstance(out, dict)
+  assert road in out and wide in out
+  assert out[road].shape == (1, 12, MODEL_H // 2, MODEL_W // 2)
+  assert out[wide].shape == (1, 12, MODEL_H // 2, MODEL_W // 2)
+
+  key = (cam_w, cam_h)
+  assert key in warp.jit_cache
+
+  out2 = warp.process(bufs, transforms)
+  assert out2[road].shape == out[road].shape
+
+
+@pytest.mark.parametrize("road, wide", VISION_NAME_PAIRS)
+def test_warp_buffer_shift(road, wide):
+  warp = Warp(2)
+  cam_w, cam_h = CAMERA_CONFIGS[1]
+  transform = np.eye(3, dtype=np.float32).flatten()
+
+  buf1 = MockVisionBuf(cam_w, cam_h)
+  buf1.data[0] = 255
+  bufs1 = {road: buf1, wide: buf1}
+  transforms = {road: transform, wide: transform}
+  out1 = warp.process(bufs1, transforms)
+  road1 = out1[road].numpy().copy()
+
+  buf2 = MockVisionBuf(cam_w, cam_h)
+  buf2.data[0] = 128
+  bufs2 = {road: buf2, wide: buf2}
+  out2 = warp.process(bufs2, transforms)
+  assert not np.array_equal(road1, out2[road].numpy())
+
+
+@pytest.mark.parametrize("buffer_length", [2, 5])
+@pytest.mark.parametrize("road, wide", VISION_NAME_PAIRS)
+def test_warp_buffer_accumulation(buffer_length, road, wide):
+  warp = Warp(buffer_length)
+  cam_w, cam_h = CAMERA_CONFIGS[0]
+  transform = np.eye(3, dtype=np.float32).flatten()
+  transforms = {road: transform, wide: transform}
+  outputs = []
+
+  for i in range(buffer_length + 1):
+    buf = MockVisionBuf(cam_w, cam_h)
+    buf.data[:] = i * 10
+    out = warp.process({road: buf, wide: buf}, transforms)
+    outputs.append(out[road].numpy().copy())
+
+  assert warp.full_buffers['img'].shape == (buffer_length * 6, MODEL_H // 2, MODEL_W // 2)
+  for i in range(1, len(outputs)):
+    assert not np.array_equal(outputs[i - 1], outputs[i])
+
+
+def test_warp_different_cameras_same_instance():
+  warp = Warp(2)
+  transform = np.eye(3, dtype=np.float32).flatten()
+
+  buf1 = MockVisionBuf(*CAMERA_CONFIGS[0])
+  warp.process({'img': buf1, 'big_img': buf1}, {'img': transform, 'big_img': transform})
+  assert len(warp.jit_cache) == 1
+
+  buf2 = MockVisionBuf(*CAMERA_CONFIGS[1])
+  warp.process({'img': buf2, 'big_img': buf2}, {'img': transform, 'big_img': transform})
+  assert len(warp.jit_cache) == 2

openpilot/sunnypilot/modeld_v2/warp.py

@@ -0,0 +1,171 @@
+import pickle
+import time
+import numpy as np
+from pathlib import Path
+from tinygrad.tensor import Tensor
+from tinygrad.engine.jit import TinyJit
+from tinygrad.device import Device
+
+from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
+from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE
+from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
+from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, make_frame_prepare as _make_frame_prepare
+
+CAMERA_CONFIGS = [
+  (_ar_ox_fisheye.width, _ar_ox_fisheye.height),
+  (_os_fisheye.width, _os_fisheye.height),
+]
+
+
+def make_frame_prepare(cam_w, cam_h, model_w, model_h):
+  nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
+  return _make_frame_prepare(nv12, model_w, model_h)
+
+
+def warp_pkl_path(w, h):
+  from openpilot.selfdrive.modeld.helpers import MODELS_DIR
+  return MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl'
+
+
+def make_update_img_input(frame_prepare, model_w, model_h):
+  def update_img_input_tinygrad(tensor, frame, M_inv):
+    M_inv = M_inv.to(Device.DEFAULT)
+    new_img = frame_prepare(frame, M_inv)
+    tensor.assign(tensor[6:].cat(new_img, dim=0).contiguous())
+    return Tensor.cat(tensor[:6], tensor[-6:], dim=0).contiguous().reshape(1, 12, model_h//2, model_w//2)
+  return update_img_input_tinygrad
+
+
+def make_update_both_imgs(frame_prepare, model_w, model_h):
+  update_img = make_update_img_input(frame_prepare, model_w, model_h)
+  def update_both_imgs_tinygrad(calib_img_buffer, new_img, M_inv,
+                                calib_big_img_buffer, new_big_img, M_inv_big):
+    calib_img_pair = update_img(calib_img_buffer, new_img, M_inv)
+    calib_big_img_pair = update_img(calib_big_img_buffer, new_big_img, M_inv_big)
+    return calib_img_pair, calib_big_img_pair
+  return update_both_imgs_tinygrad
+
+MODELS_DIR = Path(__file__).parent / 'models'
+MODEL_W, MODEL_H = MEDMODEL_INPUT_SIZE
+UPSTREAM_BUFFER_LENGTH = 5
+
+
+def v2_warp_pkl_path(cam_w, cam_h, buffer_length):
+  return MODELS_DIR / f'warp_{cam_w}x{cam_h}_b{buffer_length}_tinygrad.pkl'
+
+
+def compile_v2_warp(cam_w, cam_h, buffer_length):
+  _, _, _, yuv_size = get_nv12_info(cam_w, cam_h)
+  img_buffer_shape = (buffer_length * 6, MODEL_H // 2, MODEL_W // 2)
+
+  print(f"Compiling v2 warp for {cam_w}x{cam_h} buffer_length={buffer_length}...")
+
+  frame_prepare = make_frame_prepare(cam_w, cam_h, MODEL_W, MODEL_H)
+  update_both_imgs = make_update_both_imgs(frame_prepare, MODEL_W, MODEL_H)
+  update_img_jit = TinyJit(update_both_imgs, prune=True)
+
+  full_buffer = Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize()
+  big_full_buffer = Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize()
+  new_frame_np = np.random.randint(0, 256, yuv_size, dtype=np.uint8)
+  new_big_frame_np = np.random.randint(0, 256, yuv_size, dtype=np.uint8)
+  for i in range(10):
+    img_inputs = [full_buffer,
+                  Tensor.from_blob(new_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(),
+                  Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
+    big_img_inputs = [big_full_buffer,
+                      Tensor.from_blob(new_big_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(),
+                      Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
+    inputs = img_inputs + big_img_inputs
+    Device.default.synchronize()
+
+    st = time.perf_counter()
+    _ = update_img_jit(*inputs)
+    mt = time.perf_counter()
+    Device.default.synchronize()
+    et = time.perf_counter()
+    print(f"  [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms")
+
+  pkl_path = v2_warp_pkl_path(cam_w, cam_h, buffer_length)
+  with open(pkl_path, "wb") as f:
+    pickle.dump(update_img_jit, f)
+  print(f"  Saved to {pkl_path}")
+
+  jit = pickle.load(open(pkl_path, "rb"))
+  verify_frame = np.random.randint(0, 256, yuv_size, dtype=np.uint8)
+  verify_big_frame = np.random.randint(0, 256, yuv_size, dtype=np.uint8)
+  fresh_inputs = [
+    Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize(),
+    Tensor.from_blob(verify_frame.ctypes.data, (yuv_size,), dtype='uint8').realize(),
+    Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY'),
+    Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize(),
+    Tensor.from_blob(verify_big_frame.ctypes.data, (yuv_size,), dtype='uint8').realize(),
+    Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY'),
+  ]
+  jit(*fresh_inputs)
+
+
+class Warp:
+  def __init__(self, buffer_length=2):
+    self.buffer_length = buffer_length
+    self.img_buffer_shape = (buffer_length * 6, MODEL_H // 2, MODEL_W // 2)
+
+    self.jit_cache = {}
+    self.full_buffers = {k: Tensor.zeros(self.img_buffer_shape, dtype='uint8').contiguous().realize() for k in ['img', 'big_img']}
+    self._blob_cache: dict[int, Tensor] = {}
+    self._nv12_cache: dict[tuple[int, int], int] = {}
+    self.transforms_np = {k: np.zeros((3, 3), dtype=np.float32) for k in ['img', 'big_img']}
+    self.transforms = {k: Tensor(v, device='NPY').realize() for k, v in self.transforms_np.items()}
+
+  def process(self, bufs, transforms):
+    if not bufs:
+      return {}
+    road = next(n for n in bufs if 'big' not in n)
+    wide = next(n for n in bufs if 'big' in n)
+    cam_w, cam_h = bufs[road].width, bufs[road].height
+    key = (cam_w, cam_h)
+
+    if key not in self.jit_cache:
+      v2_pkl = v2_warp_pkl_path(cam_w, cam_h, self.buffer_length)
+      if v2_pkl.exists():
+        with open(v2_pkl, 'rb') as f:
+          self.jit_cache[key] = pickle.load(f)
+      elif self.buffer_length == UPSTREAM_BUFFER_LENGTH:
+        upstream_pkl = warp_pkl_path(cam_w, cam_h)
+        if upstream_pkl.exists():
+          with open(upstream_pkl, 'rb') as f:
+            self.jit_cache[key] = pickle.load(f)
+      if key not in self.jit_cache:
+        frame_prepare = make_frame_prepare(cam_w, cam_h, MODEL_W, MODEL_H)
+        update_both_imgs = make_update_both_imgs(frame_prepare, MODEL_W, MODEL_H)
+        self.jit_cache[key] = TinyJit(update_both_imgs, prune=True)
+
+    if key not in self._nv12_cache:
+      self._nv12_cache[key] = get_nv12_info(cam_w, cam_h)[3]
+    yuv_size = self._nv12_cache[key]
+
+    road_ptr = bufs[road].data.ctypes.data
+    wide_ptr = bufs[wide].data.ctypes.data
+    if road_ptr not in self._blob_cache:
+      self._blob_cache[road_ptr] = Tensor.from_blob(road_ptr, (yuv_size,), dtype='uint8')
+    if wide_ptr not in self._blob_cache:
+      self._blob_cache[wide_ptr] = Tensor.from_blob(wide_ptr, (yuv_size,), dtype='uint8')
+    road_blob = self._blob_cache[road_ptr]
+    wide_blob = self._blob_cache[wide_ptr] if wide_ptr != road_ptr else Tensor.from_blob(wide_ptr, (yuv_size,), dtype='uint8')
+    np.copyto(self.transforms_np['img'], transforms[road].reshape(3, 3))
+    np.copyto(self.transforms_np['big_img'], transforms[wide].reshape(3, 3))
+
+    Device.default.synchronize()
+    res = self.jit_cache[key](
+      self.full_buffers['img'], road_blob, self.transforms['img'],
+      self.full_buffers['big_img'], wide_blob, self.transforms['big_img'],
+    )
+    out_road = res[0].realize()
+    out_wide = res[1].realize()
+
+    return {road: out_road, wide: out_wide}
+
+
+if __name__ == "__main__":
+  for cam_w, cam_h in CAMERA_CONFIGS:
+    for bl in [2, 5]:
+      compile_v2_warp(cam_w, cam_h, bl)