test and compile

expect tests to fail as they arent updated

openpilot/sunnypilot/modeld_v2/compile_modeld.py

@@ -32,7 +32,7 @@ def _patch_tinygrad_fetch_fw():
   helpers.fetch_fw = fetch_fw
 _patch_tinygrad_fetch_fw()
 
-from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, make_frame_prepare, sample_desire, sample_skip, shift_and_sample
+import openpilot.selfdrive.modeld.compile_modeld as stock
 from tinygrad import dtypes
 from tinygrad.device import Device
 from tinygrad.engine.jit import TinyJit
@@ -41,7 +41,7 @@ from tinygrad.tensor import Tensor
 MODEL_TYPES = ('vision_policy', 'supercombo', 'vision_multi_policy')
 WARP_INPUTS = ['tfm', 'big_tfm']
 POLICY_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs']
-
+nv12_copy_size = stock.nv12_copy_size
 
 def _detect_desire_key(shapes: dict) -> str | None:
   return next((key for key in shapes if key.startswith('desire')), None)
@@ -138,7 +138,9 @@ def make_supercombo_input_queues(input_shapes: dict, frame_skip: int,
   return generate_queues_and_npy(input_shapes, frame_skip, device, is_supercombo=True)
 
 
-def make_random_images(keys, shape, device):
+def make_random_images(keys, shape, device, rng):
+  if device == 'NPY':
+    return {k: Tensor(rng.integers(0, 256, size=shape, dtype=np.uint8), device='NPY').realize() for k in keys}
   return {k: Tensor.randint(shape, low=0, high=256, dtype=dtypes.uint8, device=device).realize() for k in keys}
 
 
@@ -151,14 +153,15 @@ def make_warp_queues(device=Device.DEFAULT):
   return queues, npy
 
 
-def make_warp(nv12: NV12Frame, model_w: int, model_h: int):
-  frame_prepare = make_frame_prepare(nv12, model_w, model_h)
+def make_warp(nv12: stock.NV12Frame, model_w: int, model_h: int):
+  frame_prepare = stock.make_frame_prepare(nv12, model_w, model_h)
 
   def warp(tfm, big_tfm, frame, big_frame):
     tfm = tfm.to(Device.DEFAULT)
     big_tfm = big_tfm.to(Device.DEFAULT)
-    frame = frame.to(Device.DEFAULT)
-    big_frame = big_frame.to(Device.DEFAULT)
+    if Device.DEFAULT == 'AMD':
+      frame = frame.to(Device.DEFAULT)
+      big_frame = big_frame.to(Device.DEFAULT)
     Tensor.realize(tfm, big_tfm, frame, big_frame)
 
     warped_frame = frame_prepare(frame, tfm).unsqueeze(0)
@@ -168,8 +171,8 @@ def make_warp(nv12: NV12Frame, model_w: int, model_h: int):
 
 
 def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, frame_skip: int, input_shapes: dict):
-  sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
-  sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
+  sample_skip_fn = partial(stock.sample_skip, frame_skip=frame_skip)
+  sample_desire_fn = partial(stock.sample_desire, frame_skip=frame_skip)
 
   desire_key = _detect_desire_key(input_shapes)
   road_key, wide_key = _detect_vision_keys(input_shapes)
@@ -186,14 +189,14 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice,
     warped_dev = warped.to(Device.DEFAULT)
     Tensor.realize(packed_npy_inputs_dev, warped_dev)
 
-    img = shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn)
-    big_img = shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn)
+    img = stock.shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn)
+    big_img = stock.shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn)
 
     unpacked_tensors = [tensor.reshape(shape) for tensor, shape in zip(packed_npy_inputs_dev.split(npy_sizes), npy_shapes.values(), strict=True)]
     unpacked_dict = dict(zip(npy_shapes.keys(), unpacked_tensors, strict=True))
 
     desire_dev = unpacked_dict['desire']
-    desire_buf = shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn)
+    desire_buf = stock.shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn)
 
     inputs = {desire_key: desire_buf}
     for key, tensor_val in unpacked_dict.items():
@@ -202,13 +205,13 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice,
 
     if 'prev_feat' in unpacked_dict:
       prev_feat_dev = unpacked_dict['prev_feat']
-      inputs['features_buffer'] = shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn).reshape(input_shapes['features_buffer'])
+      inputs['features_buffer'] = stock.shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn).reshape(input_shapes['features_buffer'])
 
     if vision_runner:
       vision_out_cast = next(iter(vision_runner({road_key: img, wide_key: big_img}).values())).cast('float32').realize()
       if 'features_buffer' not in inputs:
         new_feat = vision_out_cast[:, features_slice].reshape(1, -1).unsqueeze(0)
-        inputs['features_buffer'] = shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
+        inputs['features_buffer'] = stock.shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
       policy_outs = [next(iter(pol_runner(inputs).values())).cast('float32').realize() for pol_runner in policy_runners]
       return (vision_out_cast, *policy_outs) if len(policy_outs) > 1 else (vision_out_cast, policy_outs[0])
 
@@ -219,27 +222,28 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice,
     policy_out = next(iter(policy_runners[0](inputs).values())).cast('float32').realize()
     if 'features_buffer' not in inputs and features_slice is not None:
       new_feat = policy_out[:, features_slice].reshape(1, -1).unsqueeze(0)
-      shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
+      stock.shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
     return policy_out
 
   return run_policy
 
 
-def compile_jit(jit, make_random_inputs, input_keys, make_queues):
+def compile_jit(jit, input_keys, make_queues, make_random_inputs=None, benchmark_runs: int = 1):
   SEED = 42
-  def random_inputs_run(fn, seed, test_val=None, test_buffers=None, expect_match=True):
-    input_queues, npy = make_queues(Device.DEFAULT)
+  def random_inputs_run(fn, seed, n_runs, test_val=None, test_buffers=None, expect_match=True):
+    queues_res = make_queues(Device.DEFAULT)
+    input_queues, npy = queues_res[0], queues_res[1]
+    frame_views = queues_res[2] if len(queues_res) > 2 else {}
     rng = np.random.default_rng(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):
       for v in npy.values():
         v[:] = rng.standard_normal(v.shape).astype(v.dtype)
+      for v in frame_views.values():
+        v[:] = rng.integers(0, 256, size=v.shape, dtype=np.uint8)
       Device.default.synchronize()
-      random_inputs = make_random_inputs()
+      random_inputs = make_random_inputs(rng=rng) if make_random_inputs is not None else {}
       st = time.perf_counter()
       outs = fn(**{k: input_queues[k] for k in input_keys if k in input_queues}, **random_inputs)
       mt = time.perf_counter()
@@ -260,14 +264,15 @@ def compile_jit(jit, make_random_inputs, input_keys, make_queues):
     return val, buffers
 
   print('capture + replay')
-  test_val, test_buffers = random_inputs_run(jit, SEED)
-  print('pickle round trip')
+  test_val, test_buffers = random_inputs_run(jit, SEED, 3)
+  print(f'pickle round trip ({benchmark_runs} runs per seed)')
   with tempfile.TemporaryFile(dir=".") as f:
     dump_oob(jit, f)
     f.seek(0)
-    deserialized_jit = load_oob(f)
-  random_inputs_run(deserialized_jit, SEED, test_val=test_val, test_buffers=test_buffers)
-  return deserialized_jit
+    loaded_jit = load_oob(f)
+  random_inputs_run(loaded_jit, SEED, benchmark_runs, test_val, test_buffers, expect_match=True)
+  random_inputs_run(loaded_jit, SEED+1, benchmark_runs, test_val, test_buffers, expect_match=False)
+  return jit
 
 
 def _parse_size(size_str: str) -> tuple[int, int]:
@@ -317,6 +322,7 @@ if __name__ == "__main__":
   parser.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH')
   parser.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True)
   parser.add_argument('--frame-skip', type=int, default=None, help='frame skip value (auto-derived if not provided)')
+  parser.add_argument('--benchmark-runs', type=int, default=1, help='benchmark runs')
   parser.add_argument('--output', required=True)
 
   parser.add_argument('--vision-onnx', help='vision ONNX (for split models)')
@@ -335,50 +341,65 @@ if __name__ == "__main__":
   args.on_policy_onnx = read_file_chunked_to_disk(args.on_policy_onnx)
   args.supercombo_onnx = read_file_chunked_to_disk(args.supercombo_onnx)
 
-  vision_runner = OnnxRunner(args.vision_onnx) if args.vision_onnx else None
-
-  if args.model_type == 'vision_policy':
-    assert vision_runner and args.policy_onnx
-    policy_runners = [OnnxRunner(args.policy_onnx)]
-    output_data['metadata'] = {'vision': make_metadata_dict(args.vision_onnx), 'policy': make_metadata_dict(args.policy_onnx)}
-  elif args.model_type == 'supercombo':
+  if args.model_type == 'supercombo':
     assert args.supercombo_onnx
-    policy_runners = [OnnxRunner(args.supercombo_onnx)]
-    output_data['metadata'] = {'model': make_metadata_dict(args.supercombo_onnx)}
-  elif args.model_type == 'vision_multi_policy':
-    assert vision_runner
-    policy_runners, policy_names = _load_policy_runners(args)
-    output_data['metadata'] = {'vision': make_metadata_dict(args.vision_onnx)}
-    for name in policy_names:
-      runner_arg = getattr(args, f"{name}_onnx")
-      output_data['metadata'][name] = make_metadata_dict(runner_arg)
+    model_metadata = make_metadata_dict(args.supercombo_onnx)
+    output_data['metadata'] = {'model': model_metadata, **model_metadata}
+    output_data['input_devices'] = {'model': Device.DEFAULT}
+    output_data['run_model'] = {}
+    derived_frame_skip = args.frame_skip or derive_frame_skip({}, model_metadata['input_shapes'])
+    model_runner = OnnxRunner(args.supercombo_onnx)
+    run_policy = stock.make_run_policy(model_runner, model_metadata, derived_frame_skip)
+    for cam_w, cam_h in args.camera_resolutions:
+      print(f"Compiling unified run_model JIT for {cam_w}x{cam_h}...")
+      nv12 = stock.NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
+      frame_copy_size = stock.nv12_copy_size(nv12.stride, nv12.y_height, nv12.uv_height)
+      make_model_queues = partial(stock.make_input_queues, model_metadata['input_shapes'], derived_frame_skip,
+                                  frame_copy_size=frame_copy_size)
+      warp = stock.make_warp(nv12, model_w, model_h)
+      run_model_jit = TinyJit(stock.make_run_model(warp, run_policy, model_metadata, frame_copy_size), prune=True)
+      output_data['run_model'][(cam_w, cam_h)] = compile_jit(run_model_jit, stock.MODELD_INPUTS, make_model_queues, benchmark_runs=args.benchmark_runs)
+  else:
+    vision_runner = OnnxRunner(args.vision_onnx) if args.vision_onnx else None
+    if args.model_type == 'vision_policy':
+      assert vision_runner and args.policy_onnx
+      policy_runners = [OnnxRunner(args.policy_onnx)]
+      output_data['metadata'] = {'vision': make_metadata_dict(args.vision_onnx), 'policy': make_metadata_dict(args.policy_onnx)}
+    elif args.model_type == 'vision_multi_policy':
+      assert vision_runner
+      policy_runners, policy_names = _load_policy_runners(args)
+      output_data['metadata'] = {'vision': make_metadata_dict(args.vision_onnx)}
+      for name in policy_names:
+        runner_arg = getattr(args, f"{name}_onnx")
+        output_data['metadata'][name] = make_metadata_dict(runner_arg)
 
-  policy_keys = [key for key in output_data['metadata'].keys() if key != 'vision']
-  first_policy_meta = output_data['metadata'][policy_keys[0]] if policy_keys else {}
-  vision_meta = output_data['metadata'].get('vision', {})
+    policy_keys = [key for key in output_data['metadata'].keys() if key != 'vision']
+    first_policy_meta = output_data['metadata'][policy_keys[0]] if policy_keys else {}
+    vision_meta = output_data['metadata'].get('vision', {})
 
-  derived_frame_skip = args.frame_skip or derive_frame_skip(vision_meta.get('input_shapes', {}), first_policy_meta.get('input_shapes', {}))
-  all_shapes = {key: value for meta in output_data['metadata'].values() for key, value in meta['input_shapes'].items()}
-  feat_meta = output_data['metadata'].get('vision') or output_data['metadata'].get('model') or output_data['metadata'].get('policy')
-  assert feat_meta is not None
-  features_slice = feat_meta['output_slices']['hidden_state']
-  is_supercombo = vision_runner is None
+    derived_frame_skip = args.frame_skip or derive_frame_skip(vision_meta.get('input_shapes', {}), first_policy_meta.get('input_shapes', {}))
+    all_shapes = {key: value for meta in output_data['metadata'].values() for key, value in meta['input_shapes'].items()}
+    feat_meta = output_data['metadata'].get('vision') or output_data['metadata'].get('policy')
+    assert feat_meta is not None
+    features_slice = feat_meta['output_slices']['hidden_state']
 
-  print(f"Compiling run_policy JIT (model_size={model_w}x{model_h}, frame_skip={derived_frame_skip})...")
-  run_policy_func = make_run_policy(vision_runner, policy_runners, features_slice, derived_frame_skip, all_shapes)
-  run_policy_jit = TinyJit(run_policy_func, prune=True)
-  make_policy_queues = partial(generate_queues_and_npy, all_shapes, derived_frame_skip, is_supercombo=is_supercombo)
-  make_random_model_inputs = partial(make_random_images, keys=['warped'], shape=(2, 6, model_h // 2, model_w // 2), device=Device.DEFAULT)
-  output_data['run_policy'] = compile_jit(run_policy_jit, make_random_model_inputs, POLICY_INPUTS, make_policy_queues)
+    print(f"Compiling run_policy JIT (model_size={model_w}x{model_h}, frame_skip={derived_frame_skip})...")
+    run_policy_func = make_run_policy(vision_runner, policy_runners, features_slice, derived_frame_skip, all_shapes)
+    run_policy_jit = TinyJit(run_policy_func, prune=True)
+    make_policy_queues = partial(generate_queues_and_npy, all_shapes, derived_frame_skip, is_supercombo=False)
+    make_random_model_inputs = partial(make_random_images, keys=['warped'], shape=(2, 6, model_h // 2, model_w // 2), device=Device.DEFAULT)
+    output_data['run_policy'] = compile_jit(run_policy_jit, POLICY_INPUTS, make_policy_queues, make_random_inputs=make_random_model_inputs)
 
-  for cam_w, cam_h in args.camera_resolutions:
-    print(f"Compiling warp JIT for {cam_w}x{cam_h}...")
-    nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
-    make_random_warp_inputs = partial(make_random_images, keys=['frame', 'big_frame'], shape=nv12.size, device=Device.DEFAULT)
-    warp = TinyJit(make_warp(nv12, model_w, model_h), prune=True)
-    output_data[(cam_w, cam_h)] = compile_jit(warp, make_random_warp_inputs, WARP_INPUTS, make_warp_queues)
+    for cam_w, cam_h in args.camera_resolutions:
+      print(f"Compiling warp JIT for {cam_w}x{cam_h}...")
+      nv12 = stock.NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
+      frame_copy_size = stock.nv12_copy_size(nv12.stride, nv12.y_height, nv12.uv_height)
+      warp_input_dev = 'NPY' if Device.DEFAULT == 'AMD' else Device.DEFAULT
+      make_random_warp_inputs = partial(make_random_images, keys=['frame', 'big_frame'], shape=frame_copy_size, device=warp_input_dev)
+      warp = TinyJit(make_warp(nv12, model_w, model_h), prune=True)
+      output_data[(cam_w, cam_h)] = compile_jit(warp, WARP_INPUTS, make_warp_queues, make_random_inputs=make_random_warp_inputs)
 
-  output_data['metadata']['warp_dev'] = Device.DEFAULT
+    output_data['metadata']['warp_dev'] = Device.DEFAULT
 
   with open(args.output, "wb") as file:
     dump_oob(output_data, file)

openpilot/sunnypilot/modeld_v2/constants.py

@@ -14,6 +14,8 @@ class ModelConstants:
 
   # model inputs constants
   MODEL_FREQ = 20
+  MODEL_RUN_FREQ = 20
+  MODEL_CONTEXT_FREQ = 5
   FEATURE_LEN = 512
   FULL_HISTORY_BUFFER_LEN = 99
   DESIRE_LEN = 8

openpilot/sunnypilot/modeld_v2/modeld.py

@@ -38,11 +38,18 @@ 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.selfdrive.modeld.modeld import ChestnutState
 
+from openpilot.selfdrive.modeld.compile_modeld import (
+  MODELD_INPUTS,
+  make_input_queues as make_stock_input_queues,
+)
 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.parse_model_outputs import Parser as CombinedParser
+from openpilot.sunnypilot.modeld_v2.constants import ModelConstants, 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.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues, make_supercombo_input_queues, WARP_INPUTS, POLICY_INPUTS
+from openpilot.sunnypilot.modeld_v2.compile_modeld import (derive_frame_skip, make_split_input_queues,
+                                                           make_supercombo_input_queues, nv12_copy_size,
+                                                           WARP_INPUTS, POLICY_INPUTS)
 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
@@ -112,29 +119,40 @@ class ModelState(ModelStateBase):
     jits = load_oob(open_file_chunked(pkl_path))
 
     metadata = jits['metadata']
-    self.WARP_DEV = metadata.get('warp_dev', 'QCOM' if COMMA_HARDWARE else 'CPU')
-    self.DEV = 'AMD' if self.chestnut else ('QCOM' if COMMA_HARDWARE else 'CPU')
+    self.use_frame_buffers = metadata.get('warp_dev') == 'AMD'
+    self.WARP_DEV = metadata.get('warp_dev', 'QCOM') if COMMA_HARDWARE else 'CPU'
+    self.DEV = ('AMD' if self.chestnut else 'QCOM') if COMMA_HARDWARE else 'CPU'
     self.QUEUE_DEV = self.DEV
-    self.run_policy = jits['run_policy']
-    self.warp = jits[(cam_w, cam_h)]
+    self.is_run_model = 'run_model' in jits
 
-    if 'model' in metadata:
-      model_metadata = metadata['model']
+    nv12_info = get_nv12_info(cam_w, cam_h)
+    self.frame_copy_size = nv12_copy_size(*nv12_info[:3])
+    self.full_frames: dict = {}
+    self._blob_cache: dict = {}
+    self.frame_buffers: dict = {}
+
+    if self.is_run_model or 'model' in metadata:
+      model_metadata = metadata.get('model', metadata)
+      self.input_shapes = model_metadata['input_shapes']
       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 = [key for key in model_metadata['input_shapes'] if 'img' in key]
-      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.QUEUE_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'
+      self._vision_input_names = [key for key in self.input_shapes if 'img' in key]
+      self.frame_skip = derive_frame_skip({}, self.input_shapes)
+      if self.is_run_model:
+        self.input_queues, self.numpy_inputs, self.frame_buffers = make_stock_input_queues(
+          self.input_shapes, self.frame_skip, device=self.DEV, frame_copy_size=self.frame_copy_size)
+        self.frame_views, self.npy = self.frame_buffers, self.numpy_inputs
+        self.run_model, self.run_policy, self.warp = jits['run_model'][(cam_w, cam_h)], None, None
       else:
-        self._combined_model_type = 'multi_policy'
+        self.input_queues, self.numpy_inputs = make_supercombo_input_queues(self.input_shapes, self.frame_skip, device=self.QUEUE_DEV)
+        self.run_model, self.run_policy, self.warp = None, jits['run_policy'], jits[(cam_w, cam_h)]
+    else:
+      self.run_model, self.run_policy, self.warp = None, jits['run_policy'], jits[(cam_w, cam_h)]
+      vision_metadata = metadata['vision']
+      policy_keys = [k for k in metadata if k not in ('vision', 'warp_dev')]
+      self._combined_model_type = 'split' if policy_keys == ['policy'] else '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]
@@ -150,50 +168,49 @@ class ModelState(ModelStateBase):
     self._desire_key = next(key for key in self.numpy_inputs if key.startswith('desire'))
     self._road_key = next(key for key in self._vision_input_names if 'big' not in key)
     self._wide_key = next(key for key in self._vision_input_names if 'big' in key)
+    self.frame_buf_params = dict.fromkeys(self._vision_input_names, nv12_info)
 
     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()
 
     if self._combined_model_type != 'supercombo':
       from openpilot.sunnypilot.modeld_v2.parse_model_outputs_split import Parser as SplitParser
       self.parser = SplitParser()
     else:
-      from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser as CombinedParser
       self.parser = CombinedParser()
 
     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)
 
-    yuv_size = self.frame_buf_params[self._road_key][3]
-    frame_tensor = Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize()
-    big_frame_tensor = Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize()
-    self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=frame_tensor, big_frame=big_frame_tensor)
+    if not self.is_run_model:
+      if self.use_frame_buffers:
+        self.frame_buffers = {k: np.zeros(self.frame_copy_size, dtype=np.uint8) for k in self._vision_input_names}
+        self.full_frames = {k: Tensor(self.frame_buffers[k], device='NPY').realize() for k in self._vision_input_names}
+      else:
+        self.full_frames = {k: Tensor(np.zeros(nv12_info[3], dtype=np.uint8), device=self.WARP_DEV).contiguous().realize() for k in self._vision_input_names}
+      self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[self._road_key], big_frame=self.full_frames[self._wide_key])
 
   def warmup(self) -> None:
-    dummy_frames = {k: np.zeros(self.frame_buf_params[k][3], dtype=np.uint8) for k in self._vision_input_names}
+    dummy_size = self.frame_copy_size if (self.is_run_model or self.use_frame_buffers) else self.frame_buf_params[self._road_key][3]
+    dummy_frames = {k: np.zeros(dummy_size, dtype=np.uint8) for k in self._vision_input_names}
     transforms = {k: np.eye(3, dtype=np.float32) for k in [self._road_key, self._wide_key] if k}
-
-    dummy_inputs = {}
-    for k, v in self.numpy_inputs.items():
-      if k not in ['tfm', 'big_tfm', 'prev_feat']:
-        dummy_inputs[k] = np.zeros(v.shape, dtype=v.dtype)
-
-    self.run(dummy_frames, transforms, dummy_inputs, prepare_only=False)
-
-    for v in self.numpy_inputs.values():
-      v[:] = 0
+    dummy_inputs = {k: np.zeros(v.shape, dtype=v.dtype) for k, v in self.numpy_inputs.items() if k not in ['tfm', 'big_tfm', 'prev_feat']}
+    self.run(dummy_frames, transforms, dummy_inputs)
+    if self.is_run_model:
+      self.input_queues, self.numpy_inputs, self.frame_buffers = make_stock_input_queues(
+        self.input_shapes, self.frame_skip, device=self.DEV, frame_copy_size=self.frame_copy_size)
+      self.frame_views = self.frame_buffers
+      self.npy = self.numpy_inputs
+    else:
+      for v in self.numpy_inputs.values():
+        v[:] = 0
+      if not self.use_frame_buffers:
+        self.full_frames.clear()
+        self._blob_cache.clear()
     self.prev_desire[:] = 0
-    self.full_frames.clear()
-    self._blob_cache.clear()
-
 
   @property
   def mlsim(self) -> bool:
@@ -208,34 +225,39 @@ class ModelState(ModelStateBase):
     return self._desire_key
 
   def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray],
-          inputs: dict[str, np.ndarray], prepare_only: bool,
+          inputs: dict[str, np.ndarray],
           after_enqueue: Callable[[], None] | None = None) -> dict[str, np.ndarray] | None:
-    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.WARP_DEV)
-      self.full_frames[key] = self._blob_cache[cache_key]
+    if self.is_run_model or self.use_frame_buffers:
+      for key, buf in bufs.items():
+        data = buf.data if hasattr(buf, 'data') else buf
+        np.copyto(self.frame_buffers[key], np.frombuffer(data, dtype=np.uint8, count=self.frame_copy_size))
+    else:
+      for key, buf in bufs.items():
+        ptr = np.frombuffer(buf.data, dtype=np.uint8).ctypes.data
+        cache_key = (key, ptr)
+        if cache_key not in self._blob_cache:
+          self._blob_cache[cache_key] = Tensor.from_blob(ptr, (self.frame_buf_params[key][3],), dtype='uint8', device=self.WARP_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', 'action_t'):
       if key in self.numpy_inputs and key in inputs:
         self.numpy_inputs[key][:] = inputs[key]
 
-    road_key = self._road_key
-    wide_key = self._wide_key
-    self.numpy_inputs['tfm'][:, :] = transforms[road_key].reshape(3, 3)
-    self.numpy_inputs['big_tfm'][:, :] = transforms[wide_key].reshape(3, 3)
+    self.numpy_inputs['tfm'][:, :] = transforms[self._road_key].reshape(3, 3)
+    self.numpy_inputs['big_tfm'][:, :] = transforms[self._wide_key].reshape(3, 3)
+
+    if self.is_run_model:
+      outs, = self.run_model(**{k: self.input_queues[k] for k in MODELD_INPUTS})
+      raw_outputs = outs
+    else:
+      warped = self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[self._road_key], big_frame=self.full_frames[self._wide_key])
+      raw_outputs = self.run_policy(**{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, warped=warped)
 
-    if prepare_only:
-      self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key])
-      return None
-    warped = self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key])
-    raw_outputs = self.run_policy(**{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, warped=warped)
     if after_enqueue is not None:
       after_enqueue()
 
@@ -245,7 +267,7 @@ class ModelState(ModelStateBase):
         raise RuntimeError("model output not finite")
       sliced = {k: model_output[np.newaxis, v] for k, v in self.vision_output_slices.items()}
       outputs = self.parser.parse_outputs(sliced)
-      if 'prev_feat' in self.numpy_inputs:
+      if 'prev_feat' in self.numpy_inputs and 'hidden_state' in self.vision_output_slices:
         self.numpy_inputs['prev_feat'][:] = model_output[self.vision_output_slices['hidden_state']]
     else:
       vision_output = raw_outputs[0].numpy().flatten()
@@ -360,7 +382,11 @@ def main(demo=False):
     loader.start()
     loader.join(BIG_MODEL_TIMEOUT)
     model = big_model
+    if model is None:
+      params.put_bool("ChestnutModelError", True)
     params.put_bool("ChestnutActive", model is not None)
+    if model is not None:
+      params.remove("ChestnutModelError")
 
   small_model = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height, chestnut=False) if model is None or CHESTNUT else None
   if model is None:
@@ -474,9 +500,6 @@ def main(demo=False):
     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}
@@ -501,11 +524,12 @@ def main(demo=False):
     try:
       send_chestnut = (chestnut_state is not None and
                        run_count % round(model.constants.MODEL_FREQ / SERVICE_LIST['chestnutState'].frequency) == 0)
-      model_output = model.run(bufs, transforms, inputs, prepare_only, chestnut_state.send if send_chestnut else None)
+      model_output = model.run(bufs, transforms, inputs, chestnut_state.send if send_chestnut else None)
     except Exception:
       if not params.get_bool("ChestnutActive"):
         raise
       cloudlog.exception("chestnut failed, falling back to small")
+      params.put_bool("ChestnutModelError", True)
       params.put_bool("ChestnutActive", False)
       assert small_model is not None
       model = small_model

openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py

@@ -103,6 +103,23 @@ class TestStockEquivalence(OpenpilotTestCase):
     assert state.vision_output_slices == arch.metadata_structure['vision']['output_slices']
     assert state.policy_output_slices == arch.metadata_structure['policy']['output_slices']
 
+  def test_unified_run_model(self, tmp_path, monkeypatch, patch_modeld):
+    from openpilot.common.hardware import hw
+    from openpilot.selfdrive.modeld.helpers import dump_oob
+    shapes = {'img': (1, 12, 128, 256), 'big_img': (1, 12, 128, 256), 'features_buffer': (1, 24, 32, 512),
+              'desire_pulse': (1, 25, 8), 'traffic_convention': (1, 2), 'action_t': (1, 2)}
+    pkl_data = {'metadata': {'model': {'input_shapes': shapes, 'output_slices': {}}},
+                'run_model': {(CAM_W, CAM_H): tests_helpers._noop_jit}}
+    with open(tmp_path / 'driving_test_tinygrad.pkl', 'wb') as f:
+      dump_oob(pkl_data, f)
+    bundle = DummyBundle(models=[DummyModel('supercombo', 'driving_test_tinygrad.pkl')])
+    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.is_run_model and state.run_model is not None
+    assert state.run_policy is None and state.warp is None
+    assert 'img' in state.frame_views and 'big_img' in state.frame_views
+
 
 ARCHETYPE_NAMES = list(ARCHETYPES.keys())
 

openpilot/sunnypilot/models/fetcher.py

@@ -139,7 +139,7 @@ class ModelCache:
 class ModelFetcher:
   """Handles fetching and caching of model data from remote source"""
   MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v22.json"
-  MODEL_URL_CHESTNUT = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_chestnut_v24.json"
+  MODEL_URL_CHESTNUT = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_chestnut_v25.json"
 
   MODEL_SOURCES = {
     "qcom": (MODEL_URL, ""),