sync: post-merge fixes for tinygrad bump

.github/workflows/build-all-tinygrad-models.yaml

@@ -120,6 +120,7 @@ jobs:
     with:
       upstream_branch: ${{ matrix.model.ref }}
       custom_name: ${{ matrix.model.display_name }}
+      is_20hz: ${{ matrix.model.is_20hz }}
       recompiled_dir: ${{ needs.setup.outputs.recompiled_dir }}
       json_version: ${{ needs.setup.outputs.json_version }}
     secrets: inherit
@@ -157,6 +158,7 @@ jobs:
     with:
       upstream_branch: ${{ matrix.model.ref }}
       custom_name: ${{ matrix.model.display_name }}
+      is_20hz: ${{ matrix.model.is_20hz }}
       recompiled_dir: ${{ needs.setup.outputs.recompiled_dir }}
       json_version: ${{ needs.setup.outputs.json_version }}
       artifact_suffix: -retry

.github/workflows/build-single-tinygrad-model.yaml

@@ -24,6 +24,11 @@ on:
         required: false
         type: string
         default: ''
+      is_20hz:
+        description: 'Is this a 20Hz model'
+        required: false
+        type: boolean
+        default: true
       bypass_push:
         description: 'Bypass pushing to GitLab for build-all'
         required: false
@@ -39,6 +44,11 @@ on:
         description: 'Custom name for the model (no date, only name)'
         required: false
         type: string
+      is_20hz:
+        description: 'Is this a 20Hz model'
+        required: false
+        type: boolean
+        default: true
       recompiled_dir:
         description: 'Existing recompiled directory number (e.g. 3 for recompiled3)'
         required: true
@@ -82,7 +92,7 @@ jobs:
     with:
       upstream_branch: ${{ inputs.upstream_branch }}
       custom_name: ${{ inputs.custom_name || inputs.upstream_branch }}
-      is_20hz: true
+      is_20hz: ${{ inputs.is_20hz }}
       artifact_suffix: ${{ inputs.artifact_suffix }}
     secrets: inherit
 

.github/workflows/sunnypilot-build-model.yaml

@@ -164,18 +164,54 @@ jobs:
           source /etc/profile
           export UV_PROJECT_ENVIRONMENT=${HOME}/venv
           export VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT
-          export PYTHONPATH="${PYTHONPATH}:${{ env.TINYGRAD_PATH }}"
+          export PYTHONPATH="${PYTHONPATH}:${{ env.TINYGRAD_PATH }}:${{ github.workspace }}"
 
-          # Loop through all .onnx files
+          COMPILE_MODELD="${{ github.workspace }}/sunnypilot/modeld_v2/compile_modeld.py"
+          MODEL_SIZE=$(python3 -c "from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE as s; print(f'{s[0]}x{s[1]}')")
+          CAMERA_RES=$(python3 -c "from openpilot.common.transformations.camera import _ar_ox_fisheye as a, _os_fisheye as o; print(f'{a.width}x{a.height} {o.width}x{o.height}')")
+          TG_FLAGS="DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1"
+
+          # Generate metadata for all ONNX files
           find "${{ env.MODELS_DIR }}" -maxdepth 1 -name '*.onnx' | while IFS= read -r onnx_file; do
-            base_name=$(basename "$onnx_file" .onnx)
-            output_file="${{ env.MODELS_DIR }}/${base_name}_tinygrad.pkl"
-
-            echo "Compiling: $onnx_file -> $output_file"
-            QCOM=1 python3 "${{ env.TINYGRAD_PATH }}/examples/openpilot/compile3.py" "$onnx_file" "$output_file"
-            DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 python3 "${{ env.MODELS_DIR }}/../get_model_metadata.py" "$onnx_file" || true
+            echo "Generating metadata: $onnx_file"
+            env ${TG_FLAGS} python3 "${{ env.MODELS_DIR }}/../get_model_metadata.py" "$onnx_file" || true
           done
 
+          # Detect model type and build compile args
+          VISION_ONNX="${{ env.MODELS_DIR }}/driving_vision.onnx"
+          POLICY_ONNX="${{ env.MODELS_DIR }}/driving_policy.onnx"
+          OFF_POLICY_ONNX="${{ env.MODELS_DIR }}/driving_off_policy.onnx"
+          ON_POLICY_ONNX="${{ env.MODELS_DIR }}/driving_on_policy.onnx"
+          SUPERCOMBO_ONNX="${{ env.MODELS_DIR }}/supercombo.onnx"
+
+          MODEL_TYPE="" ONNX_ARGS="" OUTPUT_NAME=""
+          if [ -f "$VISION_ONNX" ]; then
+            ONNX_ARGS="--vision-onnx $VISION_ONNX"
+            if [ -f "$ON_POLICY_ONNX" ] && [ -f "$OFF_POLICY_ONNX" ]; then
+              MODEL_TYPE=vision_multi_policy
+              ONNX_ARGS="$ONNX_ARGS --off-policy-onnx $OFF_POLICY_ONNX --on-policy-onnx $ON_POLICY_ONNX"
+            elif [ -f "$OFF_POLICY_ONNX" ] && [ -f "$POLICY_ONNX" ]; then
+              MODEL_TYPE=vision_multi_policy
+              ONNX_ARGS="$ONNX_ARGS --policy-onnx $POLICY_ONNX --off-policy-onnx $OFF_POLICY_ONNX"
+            elif [ -f "$POLICY_ONNX" ]; then
+              MODEL_TYPE=vision_policy
+              ONNX_ARGS="$ONNX_ARGS --policy-onnx $POLICY_ONNX"
+            fi
+          elif [ -f "$SUPERCOMBO_ONNX" ]; then
+            MODEL_TYPE=supercombo
+            ONNX_ARGS="--supercombo-onnx $SUPERCOMBO_ONNX"
+          fi
+
+          if [ -n "$MODEL_TYPE" ]; then
+            echo "Detected: $MODEL_TYPE -> driving_tinygrad.pkl"
+            env ${TG_FLAGS} python3 "$COMPILE_MODELD" \
+              --model-type $MODEL_TYPE \
+              --model-size $MODEL_SIZE \
+              --camera-resolutions $CAMERA_RES \
+              $ONNX_ARGS \
+              --output "${{ env.MODELS_DIR }}/driving_tinygrad.pkl"
+          fi
+
       - name: Validate Model Outputs
         run: |
           source /etc/profile
@@ -194,6 +230,8 @@ jobs:
           rsync -avm \
             --include='*.dlc' \
             --include='*.pkl' \
+            --include='*.chunk*' \
+            --include='*.chunkmanifest' \
             --include='*.onnx' \
             --exclude='*' \
             --delete-excluded \

SConstruct

@@ -190,7 +190,7 @@ else:
 np_version = SCons.Script.Value(np.__version__)
 Export('envCython', 'np_version')
 
-Export('env', 'arch', 'acados')
+Export('env', 'arch', 'acados', 'release')
 
 # Setup cache dir
 default_cache_dir = '/data/scons_cache' if arch == "larch64" else '/tmp/scons_cache'

release/ci/model_generator.py

@@ -1,3 +1,10 @@
+"""
+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
 import pickle
 import sys
@@ -32,6 +39,9 @@ OPTIONAL_OUTPUT_KEYS = frozenset({
 def validate_model_outputs(metadata_paths: list[Path]) -> None:
   combined_keys: set[str] = set()
   for path in metadata_paths:
+    if path.stat().st_size == 0:
+      print(f"skipping empty metadata: {path}")
+      continue
     with open(path, "rb") as f:
       metadata = pickle.load(f)
     combined_keys.update(metadata.get("output_slices", {}).keys())
@@ -78,38 +88,65 @@ def create_short_name(full_name):
   return result[:8]
 
 
-def generate_metadata(model_path: Path, output_dir: Path, short_name: str):
-  model_path = model_path
-  output_path = output_dir
+def _read_pkl_bytes(pkl_path: Path) -> bytes:
+  manifest = Path(f"{pkl_path}.chunkmanifest")
+  if manifest.exists():
+    num_chunks = int(manifest.read_text().strip())
+    parts = []
+    for i in range(num_chunks):
+      chunk = Path(f"{pkl_path}.chunk{i + 1:02d}of{num_chunks:02d}")
+      parts.append(chunk.read_bytes())
+    return b''.join(parts)
+  return pkl_path.read_bytes()
+
+
+def _find_driving_pkl(output_path: Path) -> Path | None:
+  for pattern in ('driving_tinygrad.pkl', 'driving_*_tinygrad.pkl'):
+    matches = sorted(output_path.glob(pattern))
+    if matches:
+      return matches[0]
+  for pattern in ('driving_tinygrad.pkl.chunkmanifest', 'driving_*_tinygrad.pkl.chunkmanifest'):
+    matches = sorted(output_path.glob(pattern))
+    if matches:
+      return Path(str(matches[0]).removesuffix('.chunkmanifest'))
+  return None
+
+
+def _rename_pkl_with_chunks(old_pkl: Path, new_pkl: Path) -> Path:
+  manifest = Path(f"{old_pkl}.chunkmanifest")
+  if manifest.exists():
+    for f in sorted(old_pkl.parent.glob(f"{old_pkl.name}.chunk*")):
+      f.rename(old_pkl.parent / f.name.replace(old_pkl.name, new_pkl.name, 1))
+    return new_pkl
+  return old_pkl.rename(new_pkl)
+
+
+def generate_metadata(model_path: Path, output_dir: Path, short_name: str, driving_pkl: Path):
   base = model_path.stem
+  metadata_file = output_dir / f"{base}_metadata.pkl"
 
-  # Define output files for tinygrad and metadata
-  tinygrad_file = output_path / f"{base}_tinygrad.pkl"
-  metadata_file = output_path / f"{base}_metadata.pkl"
+  if short_name:
+    renamed_meta = output_dir / f"{base}_{short_name.lower()}_metadata.pkl"
+    if metadata_file.exists() and not renamed_meta.exists():
+      metadata_file = metadata_file.rename(renamed_meta)
+    elif renamed_meta.exists():
+      metadata_file = renamed_meta
 
-  if not tinygrad_file.exists() or not metadata_file.exists():
-    print(f"Error: Missing files for model {base} ({tinygrad_file} or {metadata_file})", file=sys.stderr)
+  if not metadata_file.exists():
+    print(f"Warning: Missing metadata for {base} ({metadata_file}), skipping", file=sys.stderr)
     return
 
-  # Calculate the sha256 hashes
-  with open(tinygrad_file, 'rb') as f:
-    tinygrad_hash = hashlib.sha256(f.read()).hexdigest()
+  tinygrad_hash = hashlib.sha256(_read_pkl_bytes(driving_pkl)).hexdigest()
 
   with open(metadata_file, 'rb') as f:
     metadata_hash = hashlib.sha256(f.read()).hexdigest()
 
-  # Rename the files if a custom file name is provided
-  if short_name:
-    tinygrad_file = tinygrad_file.rename(output_path / f"{base}_{short_name.lower()}_tinygrad.pkl")
-    metadata_file = metadata_file.rename(output_path / f"{base}_{short_name.lower()}_metadata.pkl")
-
-  # Build the metadata structure
   model_type = "offPolicy" if "off_policy" in base else "onPolicy" if "on_policy" in base else base.split("_")[-1]
 
-  model_metadata = {
+  return {
     "type": model_type,
     "artifact": {
-      "file_name": tinygrad_file.name,
+      "file_name": driving_pkl.name,
       "download_uri": {
         "url": "https://gitlab.com/sunnypilot/public/docs.sunnypilot.ai/-/raw/main/",
         "sha256": tinygrad_hash
@@ -124,9 +161,6 @@ def generate_metadata(model_path: Path, output_dir: Path, short_name: str):
     }
   }
 
-  # Return model metadata
-  return model_metadata
-
 
 def create_metadata_json(models: list, output_dir: Path, custom_name=None, short_name=None, is_20hz=False, upstream_branch="unknown"):
   metadata_json = {
@@ -181,14 +215,28 @@ if __name__ == "__main__":
 
   _output_dir = Path(args.output_dir)
   _output_dir.mkdir(exist_ok=True, parents=True)
+  _short_name = create_short_name(args.custom_name) if args.custom_name else None
+
+  _driving_pkl = _find_driving_pkl(_output_dir)
+  if not _driving_pkl:
+    print(f"No driving_tinygrad.pkl found in {_output_dir}", file=sys.stderr)
+    sys.exit(1)
+
+  if _short_name:
+    new_pkl = _output_dir / f"driving_{_short_name.lower()}_tinygrad.pkl"
+    if not new_pkl.exists():
+      _driving_pkl = _rename_pkl_with_chunks(_driving_pkl, new_pkl)
+    else:
+      _driving_pkl = new_pkl
+
   _models = []
 
   for _model_path in model_paths:
-    _model_metadata = generate_metadata(Path(_model_path), _output_dir, create_short_name(args.custom_name))
+    _model_metadata = generate_metadata(Path(_model_path), _output_dir, _short_name, _driving_pkl)
     if _model_metadata:
       _models.append(_model_metadata)
 
   if _models:
-    create_metadata_json(_models, _output_dir, args.custom_name, create_short_name(args.custom_name), args.is_20hz, args.upstream_branch)
+    create_metadata_json(_models, _output_dir, args.custom_name, _short_name, args.is_20hz, args.upstream_branch)
   else:
     print("No models processed.", file=sys.stderr)

sunnypilot/modeld_v2/SConscript

@@ -1,12 +1,88 @@
 import os
 import glob
 
-Import('env', 'arch')
+from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
+from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE
+from openpilot.system.hardware import HARDWARE, PC
+
+Import('env', 'arch', 'release')
 lenv = env.Clone()
 tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=env.Dir("#").abspath) if 'pycache' not in x]
 
-# Get model metadata
-PC = not os.path.isfile('/TICI')
+
+def get_camera_configs():
+  DEVICE_RESOLUTIONS = {
+    "tici": (_ar_ox_fisheye.width, _ar_ox_fisheye.height),
+    "tizi": (_ar_ox_fisheye.width, _ar_ox_fisheye.height),
+    "mici": (_os_fisheye.width, _os_fisheye.height),
+  }
+  if release or PC or 'CI' in os.environ:
+    return set(DEVICE_RESOLUTIONS.values())
+  return [DEVICE_RESOLUTIONS[HARDWARE.get_device_type()]]
+
+CAMERA_CONFIGS = get_camera_configs()
+
+tg_flags = {
+  'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0',
+  'Darwin': f'DEV=CPU HOME={os.path.expanduser("~")}',
+}.get(arch, 'DEV=CPU:LLVM')
+
+image_flag = {
+  'larch64': 'IMAGE=2',
+}.get(arch, 'IMAGE=0')
+
+model_w, model_h = MEDMODEL_INPUT_SIZE
+from openpilot.selfdrive.modeld.constants import ModelConstants
+frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
+camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in CAMERA_CONFIGS)
+
+pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + ':' + env.Dir("#").abspath + '"'
+compile_modeld_script = File("compile_modeld.py").abspath
+upstream_compile_script = File(Dir("#selfdrive/modeld").File("compile_modeld.py").abspath)
+script_deps = [File("compile_modeld.py"), upstream_compile_script]
+
+def compile_combined(model_type, onnx_args, output_name):
+  output_pkl = File(f"models/{output_name}").abspath
+  cmd = (f'{pythonpath_string} {tg_flags} {image_flag} python3 {compile_modeld_script} '
+         f'--model-type {model_type} '
+         f'--model-size {model_w}x{model_h} '
+         f'--camera-resolutions {camera_res_args} '
+         f'{onnx_args} '
+         f'--frame-skip {frame_skip} '
+         f'--output {output_pkl}')
+  onnx_files = [f for f in onnx_args.split() if f.endswith('.onnx')]
+  return lenv.Command(output_pkl, tinygrad_files + script_deps + [File(f) for f in onnx_files if os.path.isfile(f)], cmd)
+
+# Vision + Policy (stock default model)
+vision_onnx = File("models/driving_vision.onnx").abspath
+policy_onnx = File("models/driving_policy.onnx").abspath
+if os.path.isfile(vision_onnx) and os.path.isfile(policy_onnx):
+  compile_combined('vision_policy',
+                   f'--vision-onnx {vision_onnx} --policy-onnx {policy_onnx}',
+                   'driving_combined_tinygrad.pkl')
+
+# Vision + Off-Policy
+off_policy_onnx = File("models/driving_off_policy.onnx").abspath
+if os.path.isfile(vision_onnx) and os.path.isfile(off_policy_onnx):
+  policy_arg = f'--policy-onnx {policy_onnx}' if os.path.isfile(policy_onnx) else ''
+  compile_combined('vision_multi_policy',
+                   f'--vision-onnx {vision_onnx} {policy_arg} --off-policy-onnx {off_policy_onnx}',
+                   'driving_combined_multi_tinygrad.pkl')
+
+# Vision + On-Policy + Off-Policy
+on_policy_onnx = File("models/driving_on_policy.onnx").abspath
+if os.path.isfile(vision_onnx) and os.path.isfile(on_policy_onnx) and os.path.isfile(off_policy_onnx):
+  compile_combined('vision_multi_policy',
+                   f'--vision-onnx {vision_onnx} --off-policy-onnx {off_policy_onnx} --on-policy-onnx {on_policy_onnx}',
+                   'driving_combined_tri_tinygrad.pkl')
+
+# Supercombo
+supercombo_onnx = File("models/supercombo.onnx").abspath
+if os.path.isfile(supercombo_onnx):
+  compile_combined('supercombo',
+                   f'--supercombo-onnx {supercombo_onnx}',
+                   'driving_combined_supercombo_tinygrad.pkl')
+
 if PC:
   inputs = tinygrad_files + [File(Dir("#sunnypilot/modeld_v2").File("install_models_pc.py").abspath)]
   outputs = []
@@ -21,39 +97,3 @@ if PC:
   if outputs:
     lenv.Command(outputs, inputs, cmd)
 
-tg_flags = {
-  'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0',
-  'Darwin': f'DEV=CPU THREADS=0 HOME={os.path.expanduser("~")}',
-}.get(arch, 'DEV=CPU CPU_LLVM=1 THREADS=0')
-
-image_flag = {
-  'larch64': 'IMAGE=2',
-}.get(arch, 'IMAGE=0')
-
-def tg_compile(flags, model_name):
-  pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
-  fn = File(f"models/{model_name}").abspath
-  out = fn + "_tinygrad.pkl"
-
-  return lenv.Command(
-    out,
-    [fn + ".onnx"] + tinygrad_files,
-    f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {out}'
-  )
-
-# Compile models
-for model_name in ['supercombo', 'driving_vision', 'driving_off_policy', 'driving_on_policy', 'driving_policy']:
-  if File(f"models/{model_name}.onnx").exists():
-    tg_compile(tg_flags, model_name)
-
-script_files = [File("warp.py"), File(Dir("#selfdrive/modeld").File("compile_warp.py").abspath)]
-pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + ':' + env.Dir("#").abspath + '"'
-compile_warp_cmd = f'{pythonpath_string} {tg_flags} python3 -m sunnypilot.modeld_v2.warp'
-
-from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
-warp_targets = []
-for cam in [_ar_ox_fisheye, _os_fisheye]:
-  w, h = cam.width, cam.height
-  for bl in [2, 5]:
-    warp_targets.append(File(f"models/warp_{w}x{h}_b{bl}_tinygrad.pkl").abspath)
-lenv.Command(warp_targets, tinygrad_files + script_files, compile_warp_cmd)

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])
+
+  npy_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'):
+        npy_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:
+        npy_keys[key] = np.zeros(shape, dtype=np.float32)
+    elif len(shape) == 2:
+      npy_keys[key] = np.zeros(shape, dtype=np.float32)
+
+  if 'traffic_convention' not in npy_keys:
+    tc_shape = input_shapes.get('traffic_convention', (1, 2))
+    npy_keys['traffic_convention'] = np.zeros(tc_shape, dtype=np.float32)
+
+  npy_keys['tfm'] = np.zeros((3, 3), dtype=np.float32)
+  npy_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 npy_keys.items()},
+  }
+  return input_queues, npy_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)")

sunnypilot/modeld_v2/modeld.py

@@ -1,4 +1,11 @@
 #!/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.system.hardware import TICI
 os.environ['DEV'] = 'QCOM' if TICI else 'CPU'
@@ -6,6 +13,7 @@ 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 cereal.messaging as messaging
@@ -26,18 +34,40 @@ from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan,
 
 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.warp import Warp
 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
-from openpilot.sunnypilot.models.runners.helpers import get_model_runner
 
 PROCESS_NAME = "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.system.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
+
+  fallback = os.path.join(model_root, 'driving_tinygrad.pkl')
+  if _pkl_exists(fallback):
+    return fallback
+  return None
+
+
 class FrameMeta:
   frame_id: int = 0
   timestamp_sof: int = 0
@@ -49,117 +79,170 @@ class FrameMeta:
 
 
 class ModelState(ModelStateBase):
-  frames: dict[str, Warp]
   inputs: dict[str, np.ndarray]
-  prev_desire: np.ndarray  # for tracking the rising edge of the pulse
-  temporal_idxs: slice | np.ndarray
+  prev_desire: np.ndarray
 
-  def __init__(self):
+  def __init__(self, cam_w: int, cam_h: int):
     ModelStateBase.__init__(self)
-    try:
-      self.model_runner = get_model_runner()
-      self.constants = self.model_runner.constants
-    except Exception as e:
-      cloudlog.exception(f"Failed to initialize model runner: {str(e)}")
-      raise
 
-    model_bundle = get_active_bundle()
+    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}
+    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
 
-    buffer_length = 5 if self.model_runner.is_20hz else 2
-    self.warp = Warp(buffer_length)
+    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 read_file_chunked
+
+    cloudlog.warning(f"loading combined pkl: {pkl_path}")
+    jits = pickle.loads(read_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.npy = 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.npy = 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.numpy_inputs = {}
-    self.temporal_buffers = {}
-    self.temporal_idxs_map = {}
+    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)
 
-    for key, shape in self.model_runner.input_shapes.items():
-      if key not in self.model_runner.vision_input_names: # Policy inputs
-        self.numpy_inputs[key] = np.zeros(shape, dtype=np.float32)
+    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())
 
-        # Temporal input: shape is [batch, history, features]
-        if len(shape) == 3 and shape[1] > 1:
-          buffer_history_len = shape[1] * 4 if shape[1] < 99 else shape[1]  # Allow for higher history buffers in the future
-          feature_len = shape[2]
-          features_buffer_shape = self.model_runner.input_shapes.get('features_buffer')
-          if shape[1] in (24, 25) and features_buffer_shape is not None and features_buffer_shape[1] == 24:  # 20Hz
-            buffer_history_len = (features_buffer_shape[1] + 1) * 4
-            step = int(-buffer_history_len / shape[1])
-            self.temporal_idxs_map[key] = np.arange(step, step * (shape[1] + 1), step)[::-1]
-          elif shape[1] == 25:  # Split
-            skip = buffer_history_len // shape[1]
-            self.temporal_idxs_map[key] = np.arange(buffer_history_len)[-1 - (skip * (shape[1] - 1))::skip]
-          elif shape[1] >= 99:  # non20hz
-            self.temporal_idxs_map[key] = np.arange(shape[1])
-          self.temporal_buffers[key] = np.zeros((1, buffer_history_len, feature_len), dtype=np.float32)
 
   @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(key for key in self.numpy_inputs if key.startswith('desire'))
+    return next(k for k in self.npy 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:
-    # Model decides when action is completed, so desire input is just a pulse triggered on rising edge
-    inputs[self.desire_key][0] = 0
-    new_desire = np.where(inputs[self.desire_key] - self.prev_desire > .99, inputs[self.desire_key], 0)
-    self.prev_desire[:] = inputs[self.desire_key]
-    self.temporal_buffers[self.desire_key][0,:-1] = self.temporal_buffers[self.desire_key][0,1:]
-    self.temporal_buffers[self.desire_key][0,-1] = new_desire
+    from tinygrad.tensor import Tensor
 
-    # Roll buffer and assign based on desire.shape[1] value
-    if self.temporal_buffers[self.desire_key].shape[1] > self.numpy_inputs[self.desire_key].shape[1]:
-      skip = self.temporal_buffers[self.desire_key].shape[1] // self.numpy_inputs[self.desire_key].shape[1]
-      self.numpy_inputs[self.desire_key][:] = (self.temporal_buffers[self.desire_key][0].reshape(
-                                               self.numpy_inputs[self.desire_key].shape[0], self.numpy_inputs[self.desire_key].shape[1], skip, -1).max(axis=2))
-    else:
-      self.numpy_inputs[self.desire_key][:] = self.temporal_buffers[self.desire_key][0, self.temporal_idxs_map[self.desire_key]]
+    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]
 
-    for key in self.numpy_inputs:
-      if key in inputs and key not in [self.desire_key]:
-        self.numpy_inputs[key][:] = inputs[key]
+    desire_key = self.desire_key
+    inputs[desire_key][0] = 0
+    self.npy[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.npy and key in inputs:
+        self.npy[key][:] = inputs[key]
 
-    imgs_tensors = self.warp.process(bufs, transforms)
-    for name, tensor in imgs_tensors.items():
-      self.model_runner.inputs[name] = tensor
-    self.model_runner.prepare_inputs(self.numpy_inputs)
+    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.npy['tfm'][:, :] = transforms[road_key].reshape(3, 3)
+    self.npy['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
 
-    # Run model inference
-    outputs = self.model_runner.run_model()
+    raw_outputs = self._run_policy(**self.input_queues, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key])
 
-    # Update features_buffer
-    self.temporal_buffers['features_buffer'][0, :-1] = self.temporal_buffers['features_buffer'][0, 1:]
-    self.temporal_buffers['features_buffer'][0, -1] = outputs['hidden_state'][0, :]
-    self.numpy_inputs['features_buffer'][:] = self.temporal_buffers['features_buffer'][0, self.temporal_idxs_map['features_buffer']]
+    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)
 
-    if "desired_curvature" in outputs:
-      input_name_prev = None
-      if "prev_desired_curv" in self.numpy_inputs.keys():
-        input_name_prev = 'prev_desired_curv'
-      if input_name_prev and input_name_prev in self.temporal_buffers:
-        self.process_desired_curvature(outputs, input_name_prev)
+      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.npy:
+      buf = self.npy['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 process_desired_curvature(self, outputs, input_name_prev):
-    self.temporal_buffers[input_name_prev][0,:-1] = self.temporal_buffers[input_name_prev][0,1:]
-    self.temporal_buffers[input_name_prev][0,-1,:] = outputs['desired_curvature'][0, :]
-    self.numpy_inputs[input_name_prev][:] = self.temporal_buffers[input_name_prev][0, self.temporal_idxs_map[input_name_prev]]
-    if self.mlsim:
-      self.numpy_inputs[input_name_prev][:] = 0*self.temporal_buffers[input_name_prev][0, self.temporal_idxs_map[input_name_prev]]
-
   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]
@@ -186,10 +269,6 @@ def main(demo=False):
   setproctitle(PROCESS_NAME)
   config_realtime_process(7, 54)
 
-  cloudlog.warning("loading model")
-  model = ModelState()
-  cloudlog.warning("models loaded, modeld starting")
-
   # visionipc clients
   while True:
     available_streams = VisionIpcClient.available_streams("camerad", block=False)
@@ -213,6 +292,10 @@ def main(demo=False):
   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"])
@@ -246,6 +329,7 @@ def main(demo=False):
   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
@@ -318,14 +402,14 @@ def main(demo=False):
     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.model_runner.vision_input_names}
-    transforms = {name: model_transform_extra if 'big' in name else model_transform_main for name in model.model_runner.vision_input_names}
+    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.keys():
+    if 'lateral_control_params' in model.npy:
       inputs['lateral_control_params'] = np.array([v_ego, lat_delay], dtype=np.float32)
 
     mt1 = time.perf_counter()
@@ -343,7 +427,7 @@ def main(demo=False):
       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, load_meta_constants())
+                     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]

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.system.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

sunnypilot/modeld_v2/tests/test_buffer_logic_inspect.py

@@ -1,263 +0,0 @@
-import numpy as np
-import pytest
-from typing import Any
-
-import openpilot.sunnypilot.models.helpers as helpers
-import openpilot.sunnypilot.models.runners.helpers as runner_helpers
-import openpilot.sunnypilot.modeld_v2.modeld as modeld_module
-
-ModelState = modeld_module.ModelState
-
-# These are the shapes extracted/loaded from the model onnx
-SHAPE_MODE_PARAMS = [
-  ({'desire': (1, 100, 8), 'features_buffer': (1, 99, 512), "nav_features": (1, 256), "nav_instructions": (1, 150)}, 'non20hz'), # Optimus Prime
-  ({'desire': (1, 100, 8), 'features_buffer': (1, 99, 512), "lat_planner_state": (1, 4),}, 'non20hz'), # farmville
-  ({'desire': (1, 100, 8), 'features_buffer': (1, 99, 512), "lateral_control_params": (1, 2), "prev_desired_curv": (1, 100, 1)}, 'non20hz'), # wd40
-  ({'desire': (1, 100, 8), 'features_buffer': (1, 99, 512), 'prev_desired_curv': (1, 100, 1), "lateral_control_params": (1, 2),}, 'non20hz'), # NTS
-  ({'desire': (1, 25, 8), 'features_buffer': (1, 24, 512)}, '20hz'), # NPR
-  ({'desire': (1, 100, 8), 'features_buffer': (1, 99, 512), 'prev_desired_curv': (1, 100, 1), "lateral_control_params": (1, 2),}, 'non20hz'), # NTS
-  ({'desire': (1, 25, 8), 'features_buffer': (1, 25, 512)}, 'split'), # Steam Powered v2
-  ({'desire_pulse': (1, 25, 8), 'features_buffer': (1, 25, 512)}, 'split'), # desire rename
-]
-
-
-# This creates a dummy runner, override, and bundle instance for the tests to run, without actually trying to load a physical model.
-class DummyOverride:
-  def __init__(self, key: str, value: str) -> None:
-    self.key = key
-    self.value = value
-
-
-class DummyBundle:
-  def __init__(self) -> None:
-    self.overrides = [DummyOverride('lat', '.1'), DummyOverride('long', '.3')]
-    self.generation = 10  # default to non-mlsim for buffer-update tests, as raising to 11 here will zero curvature buffer
-
-
-class DummyModelRunner:
-  def __init__(self, input_shapes: dict[str, tuple[int, int, int]], constants: Any = None) -> None:
-    self.input_shapes = input_shapes
-    self.constants = constants or type('C', (), {
-      'FULL_HISTORY_BUFFER_LEN': 100,
-      'FEATURE_LEN': 512,
-      'DESIRE_LEN': 8,
-      'PREV_DESIRED_CURV_LEN': 1,
-      'INPUT_HISTORY_BUFFER_LEN': 25,
-      'TEMPORAL_SKIP': 4,
-    })()
-    self.vision_input_names: list[str] = []
-    shape = input_shapes.get('desire', (1, 0, 0)) # [batch, history, features]
-    if shape[1] == 25:
-      self.is_20hz = True
-    else:
-      self.is_20hz = False
-
-  # Minimal prepare/run methods so ModelState can be run without actually running the model
-  def prepare_inputs(self, numpy_inputs):
-    return None
-
-  def run_model(self):
-    return {
-      'hidden_state': np.zeros((1, self.constants.FEATURE_LEN), dtype=np.float32),
-      'desired_curvature': np.zeros((1, 1), dtype=np.float32),
-    }
-
-
-@pytest.fixture
-def shapes(request):
-  return request.param
-
-
-@pytest.fixture
-def bundle() -> DummyBundle:
-  return DummyBundle()
-
-
-@pytest.fixture
-def runner(shapes) -> DummyModelRunner:
-  return DummyModelRunner(shapes)
-
-
-@pytest.fixture
-def apply_patches(monkeypatch: pytest.MonkeyPatch, bundle: DummyBundle, runner: DummyModelRunner):
-  monkeypatch.setattr(helpers, 'get_active_bundle', lambda params=None: bundle, raising=False)
-  monkeypatch.setattr(runner_helpers, 'get_model_runner', lambda: runner, raising=False)
-  monkeypatch.setattr(modeld_module, 'get_model_runner', lambda: runner, raising=False)
-  monkeypatch.setattr(modeld_module, 'get_active_bundle', lambda params=None: bundle, raising=False)
-
-
-# These are expected shapes and indices based on the time the model was presented
-def get_expected_indices(shape, constants, mode, key=None):
-  if mode == 'split':
-    start = -1 - (constants.TEMPORAL_SKIP * (constants.INPUT_HISTORY_BUFFER_LEN - 1))
-    arr = np.arange(constants.FULL_HISTORY_BUFFER_LEN)
-    idxs = arr[start::constants.TEMPORAL_SKIP]
-    return idxs
-  elif mode == '20hz':
-    num_elements = shape[1]
-    step_size = int(-100 / num_elements)
-    idxs = np.arange(step_size, step_size * (num_elements + 1), step_size)[::-1]
-    return idxs
-  elif mode == 'non20hz':
-    return np.arange(shape[1])
-  return None
-
-
-@pytest.mark.parametrize("shapes,mode", SHAPE_MODE_PARAMS, indirect=["shapes"])
-def test_buffer_shapes_and_indices(shapes, mode, apply_patches):
-  state = ModelState()
-  constants = DummyModelRunner(shapes).constants
-  for key in shapes:
-    buf = state.temporal_buffers.get(key, None)
-    idxs = state.temporal_idxs_map.get(key, None)
-    if buf is None:
-      continue    # not all shapes are 3D, and the non-3D ones are not buffered
-    # Buffer shape logic
-    if mode == 'split':
-      expected_shape = (1, constants.FULL_HISTORY_BUFFER_LEN, shapes[key][2])
-      expected_idxs = get_expected_indices(shapes[key], constants, 'split', key)
-    elif mode == '20hz':
-      expected_shape = (1, constants.FULL_HISTORY_BUFFER_LEN, shapes[key][2])
-      expected_idxs = get_expected_indices(shapes[key], constants, '20hz', key)
-    elif mode == 'non20hz':
-      expected_shape = (1, shapes[key][1], shapes[key][2])
-      expected_idxs = get_expected_indices(shapes[key], constants, 'non20hz', key)
-
-    assert buf is not None, f"{key}: buffer not found"
-    assert buf.shape == expected_shape, f"{key}: buffer shape {buf.shape} != expected {expected_shape}"
-    if expected_idxs is not None:
-      assert np.all(idxs == expected_idxs), f"{key}: buffer idxs {idxs} != expected {expected_idxs}"
-    else:
-      assert idxs is None or idxs.size == 0, f"{key}: buffer idxs should be None or empty"
-
-
-def legacy_buffer_update(buf, new_val, mode, key, constants, idxs, input_shape, prev_desire=None):
-  # This is what we compare the new dynamic logic to, to ensure it does the same thing
-  if mode == 'split':
-    if key == 'desire' or key.startswith('desire'):
-      buf[0,:-1] = buf[0,1:]
-      buf[0,-1] = new_val
-      return buf.reshape((1, constants.INPUT_HISTORY_BUFFER_LEN, constants.TEMPORAL_SKIP, -1)).max(axis=2)
-    elif key == 'features_buffer':
-      buf[0,:-1] = buf[0,1:]
-      buf[0,-1] = new_val
-      return buf[0, idxs]
-    elif key == 'prev_desired_curv':
-      buf[0,:-1] = buf[0,1:]
-      buf[0,-1,:] = new_val
-      return buf[0, idxs]
-  elif mode == '20hz':
-    if key == 'desire':
-      buf[:-1] = buf[1:]
-      buf[-1] = new_val
-      reshape_dims = (1, buf.shape[1], -1, buf.shape[2])
-      reshaped = buf.reshape(reshape_dims).max(axis=2)
-      # Slice to last shape[1] elements to match model input shape
-      input_len = reshaped.shape[1]
-      model_input_len = 25  # For 20hz mode, desire shape[1] is 25
-      if input_len > model_input_len:
-        reshaped = reshaped[:, -model_input_len:, :]
-      return reshaped
-    elif key == 'features_buffer':
-      buffer_history_len = buf.shape[1]
-      legacy_buf = np.zeros((buffer_history_len, buf.shape[2]), dtype=np.float32)
-      legacy_buf[:] = buf[0]
-      legacy_buf[:-1] = legacy_buf[1:]
-      legacy_buf[-1] = new_val
-      return legacy_buf[idxs]
-    elif key == 'prev_desired_curv':
-      buffer_history_len = buf.shape[1]
-      legacy_buf = np.zeros((buffer_history_len, buf.shape[2]), dtype=np.float32)
-      legacy_buf[:] = buf[0]
-      legacy_buf[:-1] = legacy_buf[1:]
-      legacy_buf[-1,:] = new_val
-      return legacy_buf[idxs]
-  elif mode == 'non20hz':
-    if key == 'desire':
-      desire_len = constants.DESIRE_LEN
-      if prev_desire is None:
-        prev_desire = np.zeros(desire_len, dtype=np.float32)
-      # Set first element to zero
-      new_val = new_val.copy()
-      new_val[0] = 0
-      # Shift buffer by desire len
-      buf[0][:-desire_len] = buf[0][desire_len:]
-      # Only insert new desire if rising edge
-      buf[0][-desire_len:] = np.where(new_val - prev_desire > 0.99, new_val, 0)
-      prev_desire[:] = new_val
-      return buf[0]
-    elif key == 'features_buffer':
-      buf[0, :-1] = buf[0, 1:]
-      buf[0, -1] = new_val
-      return buf[0, -input_shape[1]:]  # (99, 512)
-    elif key == 'prev_desired_curv':
-      length = new_val.shape[0]
-      buf[0,:-length,0] = buf[0,length:,0]
-      buf[0,-length:,0] = new_val[:length]
-      return buf[0]
-  return None
-
-
-def dynamic_buffer_update(state, key, new_val, mode):
-  if key == 'desire' or key.startswith('desire'):
-    inputs = {k: np.zeros(v[2], dtype=np.float32) if len(v) == 3 else np.zeros(v[1], dtype=np.float32)
-              for k, v in state.model_runner.input_shapes.items() if k != key}
-    inputs[key] = new_val.copy()
-    # ModelState.run expects desire as a pulse, so we zero the first element.
-    inputs[key][0] = 0
-    state.run({}, {}, inputs, prepare_only=False)
-    return state.numpy_inputs[key]
-
-  if key == 'features_buffer':
-    inputs = {k: np.zeros(v[2], dtype=np.float32) if len(v) == 3 else np.zeros(v[1], dtype=np.float32)
-              for k, v in state.model_runner.input_shapes.items() if k != 'features_buffer'}
-    def run_model_stub():
-      return {
-        'hidden_state': np.asarray(new_val, dtype=np.float32).reshape(1, -1),
-      }
-    state.model_runner.run_model = run_model_stub
-    state.run({}, {}, inputs, prepare_only=False)
-    return state.numpy_inputs['features_buffer'][0]
-
-  if key == 'prev_desired_curv':
-    inputs = {k: np.zeros(v[2], dtype=np.float32) if len(v) == 3 else np.zeros(v[1], dtype=np.float32)
-              for k, v in state.model_runner.input_shapes.items() if k != 'prev_desired_curv'}
-    def run_model_stub():
-      return {
-        'hidden_state': np.zeros((1, state.constants.FEATURE_LEN), dtype=np.float32),
-        'desired_curvature': np.asarray(new_val, dtype=np.float32).reshape(1, -1),
-      }
-    state.model_runner.run_model = run_model_stub
-    state.run({}, {}, inputs, prepare_only=False)
-    return state.numpy_inputs['prev_desired_curv'][0]
-  return None
-
-
-@pytest.mark.parametrize("shapes,mode", SHAPE_MODE_PARAMS, indirect=["shapes"])
-@pytest.mark.parametrize("key", ["desire", "features_buffer", "prev_desired_curv"])
-def test_buffer_update_equivalence(shapes, mode, key, apply_patches):
-  state = ModelState()
-  if key == "desire":
-    desire_keys = [k for k in shapes.keys() if k.startswith('desire')]
-    if desire_keys:
-      actual_key = desire_keys[0]  # Use the first (and likely only) desire key
-  else:
-    actual_key = key
-
-  if actual_key not in state.numpy_inputs:
-    pytest.skip()
-
-  constants = DummyModelRunner(shapes).constants
-  buf = state.temporal_buffers.get(actual_key, None)
-  idxs = state.temporal_idxs_map.get(actual_key, None)
-  input_shape = shapes[actual_key]
-  prev_desire = np.zeros(constants.DESIRE_LEN, dtype=np.float32) if key == 'desire' else None
-
-  for step in range(20):    # multiple steps to ensure history is built up
-    new_val = np.full((input_shape[2],), step, dtype=np.float32)
-    expected = legacy_buffer_update(buf, new_val, mode, actual_key, constants, idxs, input_shape, prev_desire)
-    actual = dynamic_buffer_update(state, actual_key, new_val, mode)
-    if expected is not None and actual is not None and expected.shape != actual.shape:
-      if expected.ndim == 2 and actual.ndim == 2 and expected.shape[1] == actual.shape[1]:
-        expected = expected[-actual.shape[0]:]
-    assert np.allclose(actual, expected), f"{mode} {actual_key}: dynamic buffer update does not match legacy logic"

sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py

@@ -0,0 +1,263 @@
+"""
+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.system.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
+
+  def test_finds_fallback_driving_tinygrad(self, tmp_path, monkeypatch):
+    (tmp_path / 'driving_tinygrad.pkl').write_bytes(b'fake')
+    from openpilot.system.hardware import hw
+    monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path)))
+
+    bundle = DummyBundle(models=[DummyModel('vision', 'nonexistent.pkl')])
+    result = _find_driving_pkl(bundle)
+    assert result is not None
+    assert 'driving_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)
+    stock_queues, stock_npy = make_input_queues(SPLIT_VISION_INPUT_SHAPES, SPLIT_POLICY_INPUT_SHAPES, frame_skip,
+                                                device='NPY')
+
+    # TODO-SP: remove action_t skip once SP adds prerequisite for deep models (action_t input queue)
+    skip_keys = {'action_t'}
+    assert set(state.input_queues.keys()) == set(stock_queues.keys()) - skip_keys, \
+      f"Queue keys differ: v2={set(state.input_queues.keys())}, stock={set(stock_queues.keys())}"
+    assert set(state.npy.keys()) == set(stock_npy.keys()) - skip_keys, \
+      f"Npy keys differ: v2={set(state.npy.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.npy, f"{arch.name}: 'tfm' missing from npy"
+    assert 'big_tfm' in state.npy, f"{arch.name}: 'big_tfm' missing from npy"
+    assert state.npy['tfm'].shape == (3, 3)
+    assert state.npy['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.npy, \
+      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.system.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.system.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"

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"

sunnypilot/modeld_v2/tests/test_warp.py

@@ -2,7 +2,8 @@ import os
 os.environ['DEV'] = 'CPU'
 import pytest
 import numpy as np
-from openpilot.selfdrive.modeld.compile_warp import get_nv12_info, CAMERA_CONFIGS
+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

sunnypilot/modeld_v2/warp.py

@@ -7,10 +7,43 @@ from tinygrad.engine.jit import TinyJit
 from tinygrad.device import Device
 
 from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
-from openpilot.selfdrive.modeld.compile_warp import (
-  CAMERA_CONFIGS, MEDMODEL_INPUT_SIZE, make_frame_prepare, make_update_both_imgs,
-  warp_pkl_path,
-)
+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
@@ -58,7 +91,17 @@ def compile_v2_warp(cam_w, cam_h, buffer_length):
   print(f"  Saved to {pkl_path}")
 
   jit = pickle.load(open(pkl_path, "rb"))
-  jit(*inputs)
+  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:

sunnypilot/models/fetcher.py

@@ -110,13 +110,13 @@ class ModelCache:
 
   def set(self, data: dict) -> None:
     """Updates the cache with new model data"""
-    self.params.put(self._CACHE_KEY, data)
-    self.params.put(self._LAST_SYNC_KEY, int(time.monotonic() * 1e9))
+    self.params.put(self._CACHE_KEY, data, block=True)
+    self.params.put(self._LAST_SYNC_KEY, int(time.monotonic() * 1e9), block=True)
 
 
 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_v16.json"
+  MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v17.json"
 
   def __init__(self, params: Params):
     self.params = params

sunnypilot/models/helpers.py

@@ -6,7 +6,6 @@ See the LICENSE.md file in the root directory for more details.
 """
 
 import hashlib
-import os
 import pickle
 import numpy as np
 
@@ -28,16 +27,12 @@ ModelManager = custom.ModelManagerSP
 
 
 async def verify_file(file_path: str, expected_hash: str) -> bool:
-  """Verifies file hash against expected hash"""
-  if not os.path.exists(file_path):
+  from openpilot.common.file_chunker import read_file_chunked
+  try:
+    data = read_file_chunked(file_path)
+  except FileNotFoundError:
     return False
-
-  sha256_hash = hashlib.sha256()
-  with open(file_path, "rb") as file:
-    for chunk in iter(lambda: file.read(4096), b""):
-      sha256_hash.update(chunk)
-
-  return sha256_hash.hexdigest().lower() == expected_hash.lower()
+  return hashlib.sha256(data).hexdigest().lower() == expected_hash.lower()
 
 
 def is_bundle_version_compatible(bundle: dict) -> bool:
@@ -75,32 +70,13 @@ def get_active_bundle(params: Params = None) -> custom.ModelManagerSP.ModelBundl
   return None
 
 
-def get_active_model_runner(params: Params = None, force_check=False) -> custom.ModelManagerSP.Runner:
-  """
-  Determines and returns the active model runner type, based on provided parameters.
-  The function utilizes caching to prevent redundant calculations and checks.
-
-  If the cached "ModelRunnerTypeCache" exists in the provided parameters and `force_check`
-  is set to False, the cached value is directly returned. Otherwise, the function determines
-  the runner type based on the active model bundle. If a model bundle containing a drive
-  model exists, the runner type is derived based on the filename of the drive model.
-  Finally, it updates the cache with the determined runner type, if needed.
-
-  :param params: The parameter set used to retrieve caching and runner details. If `None`,
-      a default `Params` instance is created internally.
-  :type params: Params
-  :param force_check: A flag indicating whether to bypass cached results and always
-      re-determine the runner type. Defaults to `False`.
-  :type force_check: bool
-  :return: The determined or cached model runner type.
-  :rtype: custom.ModelManagerSP.Runner
-  """
+def get_active_model_runner(params: Params = None, force_check=False) -> int:
   if params is None:
     params = Params()
 
-  if (cached_runner_type := params.get("ModelRunnerTypeCache")) and not force_check:
-    if isinstance(cached_runner_type, str) and cached_runner_type.isdigit():
-      return int(cached_runner_type)
+  cached_runner_type = params.get("ModelRunnerTypeCache")
+  if cached_runner_type is not None and not force_check:
+    return cached_runner_type
 
   runner_type = custom.ModelManagerSP.Runner.stock
 
@@ -108,7 +84,7 @@ def get_active_model_runner(params: Params = None, force_check=False) -> custom.
     runner_type = active_bundle.runner.raw
 
   if cached_runner_type != runner_type:
-    params.put("ModelRunnerTypeCache", int(runner_type))
+    params.put("ModelRunnerTypeCache", int(runner_type), block=True)
 
   return runner_type
 

sunnypilot/models/manager.py

@@ -33,6 +33,17 @@ class ModelManagerSP:
     self._chunk_size = 128 * 1000  # 128 KB chunks
     self._download_start_times: dict[str, float] = {}  # Track start time per model
 
+  def _sync_artifact_progress(self, source_artifact) -> None:
+    """Mirror download progress to all artifacts sharing the same filename in the selected bundle."""
+    if not self.selected_bundle:
+      return
+    for model in self.selected_bundle.models:
+      for artifact in (model.artifact, model.metadata):
+        if artifact is not source_artifact and artifact.fileName == source_artifact.fileName:
+          artifact.downloadProgress.status = source_artifact.downloadProgress.status
+          artifact.downloadProgress.progress = source_artifact.downloadProgress.progress
+          artifact.downloadProgress.eta = source_artifact.downloadProgress.eta
+
   def _calculate_eta(self, filename: str, progress: float) -> int:
     """Calculate ETA based on elapsed time and current progress"""
     if filename not in self._download_start_times or progress <= 0:
@@ -63,7 +74,7 @@ class ModelManagerSP:
             f.write(chunk)
             bytes_downloaded += len(chunk)
 
-            if not self.params.get("ModelManager_DownloadIndex"):
+            if self.params.get("ModelManager_DownloadIndex") is None:
               raise Exception("Download cancelled")
 
             if total_size > 0:
@@ -71,13 +82,55 @@ class ModelManagerSP:
               model.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloading
               model.downloadProgress.progress = progress
               model.downloadProgress.eta = self._calculate_eta(model.fileName, progress)
+              self._sync_artifact_progress(model)
               self._report_status()
 
         # Clean up start time after download completes
         del self._download_start_times[model.fileName]
 
+  async def _download_chunked(self, base_url: str, base_path: str, artifact) -> None:
+    from openpilot.common.file_chunker import get_manifest_path, get_chunk_name
+    manifest_url = get_manifest_path(base_url)
+    manifest_path = get_manifest_path(base_path)
+
+    async with aiohttp.ClientSession() as session:
+      async with session.get(manifest_url) as resp:
+        if resp.status == 404:
+          raise FileNotFoundError
+        resp.raise_for_status()
+        num_chunks = int((await resp.read()).strip())
+
+    self._download_start_times[artifact.fileName] = time.monotonic()
+
+    for i in range(num_chunks):
+      chunk_url = get_chunk_name(base_url, i, num_chunks)
+      chunk_path = get_chunk_name(base_path, i, num_chunks)
+      chunk_downloaded = 0
+      async with aiohttp.ClientSession() as session:
+        async with session.get(chunk_url) as response:
+          response.raise_for_status()
+          chunk_size = int(response.headers.get("content-length", 0))
+          with open(chunk_path, 'wb') as f:
+            async for data in response.content.iter_chunked(self._chunk_size):
+              f.write(data)
+              chunk_downloaded += len(data)
+              if self.params.get("ModelManager_DownloadIndex") is None:
+                raise Exception("Download cancelled")
+              intra = chunk_downloaded / max(chunk_size, 1)
+              progress = min(99, (i + intra) / num_chunks * 100)
+              artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloading
+              artifact.downloadProgress.progress = progress
+              artifact.downloadProgress.eta = self._calculate_eta(artifact.fileName, progress)
+              self._sync_artifact_progress(artifact)
+              self._report_status()
+
+    with open(manifest_path, 'w') as f:
+      f.write(str(num_chunks))
+    if os.path.isfile(base_path):
+      os.remove(base_path)
+    del self._download_start_times[artifact.fileName]
+
   async def _process_artifact(self, artifact, destination_path: str) -> None:
-    """Processes a single model download including verification"""
     if not artifact.downloadUri.uri:
       return None
 
@@ -87,32 +140,38 @@ class ModelManagerSP:
     full_path = os.path.join(destination_path, filename)
 
     try:
-      # Check existing file
-      if os.path.exists(full_path) and await verify_file(full_path, expected_hash):
+      if await verify_file(full_path, expected_hash):
         artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.cached
         artifact.downloadProgress.progress = 100
         artifact.downloadProgress.eta = 0
+        self._sync_artifact_progress(artifact)
         self._report_status()
         return
 
-      # Download and verify
-      await self._download_file(url, full_path, artifact)
+      try:
+        await self._download_chunked(url, full_path, artifact)
+      except (FileNotFoundError, aiohttp.ClientResponseError):
+        await self._download_file(url, full_path, artifact)
+
       if not await verify_file(full_path, expected_hash):
         raise ValueError(f"Hash validation failed for {filename}")
 
       artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.downloaded
+      artifact.downloadProgress.progress = 100
       artifact.downloadProgress.eta = 0
+      self._sync_artifact_progress(artifact)
       self._report_status()
 
     except Exception as e:
       cloudlog.error(f"Error downloading {filename}: {str(e)}")
-      if os.path.exists(full_path):
-        os.remove(full_path)
+      for f in [full_path] + [p for p in (os.path.join(destination_path, f) for f in os.listdir(destination_path)) if filename in p]:
+        if os.path.isfile(f):
+          os.remove(f)
       artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.failed
       artifact.downloadProgress.eta = 0
+      self._sync_artifact_progress(artifact)
       self.selected_bundle.status = custom.ModelManagerSP.DownloadStatus.failed
       self._report_status()
-      # Clean up start time if it exists
       self._download_start_times.pop(artifact.fileName, None)
       raise
 
@@ -144,11 +203,22 @@ class ModelManagerSP:
     os.makedirs(destination_path, exist_ok=True)
 
     try:
-      tasks = [self._process_model(model, destination_path) for model in self.selected_bundle.models]
-      await asyncio.gather(*tasks)
+      seen_artifacts: set[str] = set()
+      for model in self.selected_bundle.models:
+        for artifact in (model.metadata, model.artifact):
+          if not artifact.fileName:
+            continue
+          if artifact.fileName in seen_artifacts:
+            artifact.downloadProgress.status = custom.ModelManagerSP.DownloadStatus.cached
+            artifact.downloadProgress.progress = 100
+            artifact.downloadProgress.eta = 0
+          else:
+            seen_artifacts.add(artifact.fileName)
+            await self._process_artifact(artifact, destination_path)
+
       self.active_bundle = self.selected_bundle
       self.active_bundle.status = custom.ModelManagerSP.DownloadStatus.downloaded
-      self.params.put("ModelManager_ActiveBundle", self.active_bundle.to_dict())
+      self.params.put("ModelManager_ActiveBundle", self.active_bundle.to_dict(), block=True)
       self.selected_bundle = None
 
     except Exception:
@@ -206,11 +276,12 @@ class ModelManagerSP:
         if hasattr(model, 'metadata') and model.metadata.fileName:
           active_files.append(model.metadata.fileName)
 
-    # Remove all files except active ones
+    # Remove all files except active ones (including their chunk files)
     model_dir = Paths.model_root()
     try:
       for filename in os.listdir(model_dir):
-        if filename not in active_files:
+        base = filename.split('.chunk')[0] if '.chunk' in filename else filename
+        if base not in active_files and filename not in active_files:
           file_path = os.path.join(model_dir, filename)
           if os.path.isfile(file_path):
             os.remove(file_path)

sunnypilot/models/model_name.py

@@ -1 +1 @@
-DEFAULT_MODEL = "POP model"
+DEFAULT_MODEL = "CD210"

sunnypilot/models/tests/model_hash

@@ -1 +1 @@
-5d4d21f1899de21137f69d74a4602c44cc5a6b04cf4e4aa9d0ec9206f8c30350
+32f57bdc91f910df1f48ddae7c59aaf6e751f9df6756da481a210577dbce8bcf