From 37d743122edfe7f417c73d900fce6c70eaceff43 Mon Sep 17 00:00:00 2001 From: discountchubbs Date: Fri, 14 Aug 2026 14:35:28 -0700 Subject: [PATCH] this made me cry --- openpilot/sunnypilot/SConscript | 1 - openpilot/sunnypilot/modeld_v2/SConscript | 105 ----------- .../sunnypilot/modeld_v2/compile_modeld.py | 173 ++++++++++-------- openpilot/sunnypilot/modeld_v2/modeld.py | 30 ++- .../sunnypilot/modeld_v2/tests/helpers.py | 3 +- 5 files changed, 117 insertions(+), 195 deletions(-) delete mode 100644 openpilot/sunnypilot/modeld_v2/SConscript diff --git a/openpilot/sunnypilot/SConscript b/openpilot/sunnypilot/SConscript index 09ad39ab43..587deea5ff 100644 --- a/openpilot/sunnypilot/SConscript +++ b/openpilot/sunnypilot/SConscript @@ -1,3 +1,2 @@ SConscript(['common/transformations/SConscript']) -SConscript(['modeld_v2/SConscript']) SConscript(['selfdrive/locationd/SConscript']) diff --git a/openpilot/sunnypilot/modeld_v2/SConscript b/openpilot/sunnypilot/modeld_v2/SConscript deleted file mode 100644 index ded9f8d851..0000000000 --- a/openpilot/sunnypilot/modeld_v2/SConscript +++ /dev/null @@ -1,105 +0,0 @@ -import os -import glob -import sys -import subprocess - -from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye -from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE -from openpilot.common.hardware import HARDWARE, PC -from openpilot.selfdrive.modeld.helpers import usbgpu_present - -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] - - -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() - -# remove me after sync -def probe_devices(): - return set(subprocess.run( - [sys.executable, '-c', 'from tinygrad import Device\nprint("\\n".join(Device.get_available_devices()))'], - capture_output=True, text=True, check=True).stdout.strip().splitlines()) - -available = probe_devices() #remove me after sync -if 'QCOM' in available: # change to this after sync. if arch == 'comma_arm64': - tg_backend = 'QCOM' - tg_flags = f'DEV={tg_backend} IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1' -else: - tg_backend = 'CPU' - tg_flags = f'DEV=CPU HOME={os.path.expanduser("~")}' if arch == 'Darwin' else 'DEV=CPU:LLVM' - -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("#openpilot/selfdrive/modeld").File("compile_modeld.py").abspath) -script_deps = [File("compile_modeld.py"), upstream_compile_script] - -USBGPU = usbgpu_present() -if USBGPU: - usbgpu_tg_flags = f'DEBUG=2 DEV=USB+AMD:LLVM WARP_DEV={tg_backend} FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2' - usbgpu_lock = File("models/.usb_gpu.lock").abspath - -def compile_combined(model_type, onnx_args, output_name): - for usbgpu in ([False, True] if USBGPU else [False]): - prefix = 'big_' if usbgpu else ('big_' if os.getenv('BIG_INTO_SMALL') else '') - final_output_name = prefix + output_name - output_pkl = File(f"models/{final_output_name}").abspath - - active_tg_flags = usbgpu_tg_flags if usbgpu else tg_flags - - cmd = (f'{pythonpath_string} {active_tg_flags} 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')] - node = lenv.Command(output_pkl, tinygrad_files + script_deps + [File(f) for f in onnx_files if os.path.isfile(f)], cmd) - if usbgpu: - lenv.SideEffect(usbgpu_lock, node) - -# 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') diff --git a/openpilot/sunnypilot/modeld_v2/compile_modeld.py b/openpilot/sunnypilot/modeld_v2/compile_modeld.py index 32f4260ce2..4397209005 100755 --- a/openpilot/sunnypilot/modeld_v2/compile_modeld.py +++ b/openpilot/sunnypilot/modeld_v2/compile_modeld.py @@ -9,7 +9,7 @@ See the LICENSE.md file in the root directory for more details. import argparse import os import tempfile -from collections import defaultdict +import time from functools import partial from openpilot.selfdrive.modeld.helpers import dump_oob, load_oob import numpy as np @@ -38,6 +38,9 @@ from tinygrad.engine.jit import TinyJit from tinygrad.tensor import Tensor MODEL_TYPES = ('vision_policy', 'supercombo', 'vision_multi_policy') +WARP_INPUTS = ['tfm', 'big_tfm'] +POLICY_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs'] +WARP_DEV = os.getenv('WARP_DEV') def _detect_desire_key(shapes: dict) -> str | None: @@ -132,9 +135,35 @@ def make_supercombo_input_queues(input_shapes: dict, frame_skip: int, return generate_queues_and_npy(input_shapes, frame_skip, device, is_supercombo=True) -def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, model_size: tuple[int, int], - features_slice: slice, frame_skip: int, input_shapes: dict, prepare_only: bool): - frame_prepare = make_frame_prepare(nv12, *model_size) +def make_random_images(keys, shape, device): + return {k: Tensor.randint(shape, low=0, high=256, dtype=dtypes.uint8, device=device).realize() for k in keys} + + +def make_warp_queues(device=Device.DEFAULT): + npy = { + 'tfm': np.zeros((3, 3), dtype=np.float32), + 'big_tfm': np.zeros((3, 3), dtype=np.float32), + } + queues = {k: Tensor(v, device='NPY').realize() for k, v in npy.items()} + return queues, npy + + +def make_warp(nv12: NV12Frame, model_w: int, model_h: int): + frame_prepare = make_frame_prepare(nv12, model_w, model_h) + WARP_DEV = os.getenv('WARP_DEV', Device.DEFAULT) + + def warp(tfm, big_tfm, frame, big_frame): + tfm = tfm.to(WARP_DEV) + big_tfm = big_tfm.to(WARP_DEV) + Tensor.realize(tfm, big_tfm) + + warped_frame = frame_prepare(frame, tfm).unsqueeze(0) + warped_big_frame = frame_prepare(big_frame, big_tfm).unsqueeze(0) + return Tensor.cat(warped_frame, warped_big_frame) + return warp + + +def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, frame_skip: int, input_shapes: dict): sample_skip_fn = partial(sample_skip, frame_skip=frame_skip) sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) @@ -147,20 +176,14 @@ def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, mode is_supercombo = vision_runner is None npy_shapes, npy_sizes = get_policy_npy_shapes(input_shapes, is_supercombo=is_supercombo) - def runner(img_q, big_img_q, feat_q, packed_npy_inputs, frame, big_frame, tfm, big_tfm, **kwargs): + def run_policy(warped, img_q, big_img_q, feat_q, packed_npy_inputs, **kwargs): desire_q = kwargs['desire_q'] - packed_npy_inputs_dev = packed_npy_inputs.to(Device.DEFAULT) - tfm_dev = tfm.to(Device.DEFAULT) - big_tfm_dev = big_tfm.to(Device.DEFAULT) + warped_dev = warped.to(Device.DEFAULT) + Tensor.realize(packed_npy_inputs_dev, warped_dev) - Tensor.realize(packed_npy_inputs_dev, tfm_dev, big_tfm_dev) - - img = shift_and_sample(img_q, frame_prepare(frame, tfm_dev).unsqueeze(0), sample_skip_fn).realize() - big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm_dev).unsqueeze(0), sample_skip_fn).realize() - - if prepare_only: - return img, big_img + img = shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn).realize() + big_img = shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn).realize() unpacked_tensors = [tensor.reshape(shape) for tensor, shape in zip(packed_npy_inputs_dev.split(npy_sizes), npy_shapes.values(), strict=True)] unpacked_dict = dict(zip(npy_shapes.keys(), unpacked_tensors, strict=True)) @@ -195,50 +218,52 @@ def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, mode shift_and_sample(feat_q, new_feat, sample_skip_fn).realize() return policy_out - return runner + return run_policy -def compile_and_warmup(nv12: NV12Frame, model_size: tuple[int, int], prepare_only: bool, frame_skip: int, vision_runner, policy_runners: list, metadata: dict): - print(f"Compiling combined JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...") - - all_shapes = {key: value for meta in metadata.values() for key, value in meta['input_shapes'].items()} - - feat_meta = metadata.get('vision') or metadata.get('model') or metadata.get('policy') - if not feat_meta: - raise ValueError("Could not find vision, model, or policy metadata.") - - features_slice = feat_meta['output_slices']['hidden_state'] - WARP_DEV = os.getenv('WARP_DEV', Device.DEFAULT) - - is_supercombo = vision_runner is None - run_func = create_jit_runner(vision_runner, policy_runners, nv12, model_size, features_slice, frame_skip, all_shapes, prepare_only) - run_jit = TinyJit(run_func, prune=True) - def run_once(seed, queues=None, npy=None): - if queues is None or npy is None: - queues, npy = generate_queues_and_npy(all_shapes, frame_skip, Device.DEFAULT, is_supercombo=is_supercombo) +def compile_jit(jit, make_random_inputs, input_keys, make_queues): + SEED = 42 + def random_inputs_run(fn, seed, test_val=None, test_buffers=None, expect_match=True): + input_queues, npy = make_queues(Device.DEFAULT) rng = np.random.default_rng(seed) Tensor.manual_seed(seed) - frame = Tensor.randint(nv12.size, low=0, high=256, dtype=dtypes.uint8, device=WARP_DEV).realize() - big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype=dtypes.uint8, device=WARP_DEV).realize() - for value in npy.values(): - value[:] = rng.standard_normal(value.shape).astype(value.dtype) - Device.default.synchronize() - outs = run_jit(**queues, frame=frame, big_frame=big_frame) - Device.default.synchronize() - return [np.copy(value.numpy()) for value in (outs if isinstance(outs, tuple) else [outs])] if outs is not None else [] - warmup_queues, warmup_npy = generate_queues_and_npy(all_shapes, frame_skip, Device.DEFAULT, is_supercombo=is_supercombo) - for i in range(3): - run_once(42 + i, warmup_queues, warmup_npy) + testing = test_val is not None or test_buffers is not None + n_runs = 1 if testing else 3 - if not prepare_only: - baseline = run_once(42) - with tempfile.TemporaryFile(dir=".") as f: - dump_oob(run_jit, f) - f.seek(0) - run_jit = load_oob(f) - assert all(np.array_equal(baseline, deserialized) for baseline, deserialized in zip(baseline, run_once(42), strict=True)), "OOB pickling regression" - return run_jit + for i in range(n_runs): + for v in npy.values(): + v[:] = rng.standard_normal(v.shape).astype(v.dtype) + Device.default.synchronize() + random_inputs = make_random_inputs() + st = time.perf_counter() + outs = fn(**{k: input_queues[k] for k in input_keys if k in input_queues}, **random_inputs) + mt = time.perf_counter() + 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 if isinstance(outs, tuple) else [outs])] if outs is not None else [] + 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 val, buffers + + print('capture + replay') + test_val, test_buffers = random_inputs_run(jit, SEED) + print('pickle round trip') + with tempfile.TemporaryFile(dir=".") as f: + dump_oob(jit, f) + f.seek(0) + deserialized_jit = load_oob(f) + random_inputs_run(deserialized_jit, SEED, test_val=test_val, test_buffers=test_buffers) + return deserialized_jit def _parse_size(size_str: str) -> tuple[int, int]: @@ -260,19 +285,6 @@ def read_file_chunked_to_shm(path): return shm_path -def _compile_for_resolutions(camera_resolutions: list, model_size: tuple[int, int], frame_skip: int, - vision_runner, policy_runners: list, metadata: dict) -> dict: - from openpilot.system.camerad.cameras.nv12_info import get_nv12_info - return { - (cam_w, cam_h): { - name: compile_and_warmup(NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)), model_size, prepare_only, - frame_skip, vision_runner, policy_runners, metadata) - for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)] - } - for cam_w, cam_h in camera_resolutions - } - - def _load_policy_runners(args: argparse.Namespace) -> tuple[list, list]: runners, keys = [], [] for name, onnx_arg in [('policy', args.policy_onnx), ('off_policy', args.off_policy_onnx), ('on_policy', args.on_policy_onnx)]: @@ -284,7 +296,6 @@ def _load_policy_runners(args: argparse.Namespace) -> tuple[list, list]: if __name__ == "__main__": if 'USB' in os.getenv('DEV', '') or os.getenv('USBGPU'): - import time from openpilot.system.hardware.chestnut.flash import link_up for _ in range(10): if link_up(): @@ -293,7 +304,9 @@ if __name__ == "__main__": else: raise RuntimeError("Chestnut not ready, skipping big model build") + from openpilot.common.file_chunker import chunk_file, get_chunk_targets from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict + from openpilot.system.camerad.cameras.nv12_info import get_nv12_info from tinygrad.nn.onnx import OnnxRunner parser = argparse.ArgumentParser(description="Compile combined JIT pkl for sunnypilot modeld_v2") @@ -310,7 +323,8 @@ if __name__ == "__main__": parser.add_argument('--supercombo-onnx', help='supercombo ONNX (for supercombo)') args = parser.parse_args() - output_data = defaultdict(dict) + model_w, model_h = args.model_size + output_data = {} args.vision_onnx = read_file_chunked_to_shm(args.vision_onnx) args.policy_onnx = read_file_chunked_to_shm(args.policy_onnx) @@ -341,16 +355,31 @@ if __name__ == "__main__": vision_meta = output_data['metadata'].get('vision', {}) derived_frame_skip = args.frame_skip or derive_frame_skip(vision_meta.get('input_shapes', {}), first_policy_meta.get('input_shapes', {})) - output_data.update(_compile_for_resolutions(args.camera_resolutions, args.model_size, derived_frame_skip, - vision_runner, policy_runners, output_data['metadata'])) + all_shapes = {key: value for meta in output_data['metadata'].values() for key, value in meta['input_shapes'].items()} + feat_meta = output_data['metadata'].get('vision') or output_data['metadata'].get('model') or output_data['metadata'].get('policy') + assert feat_meta is not None + features_slice = feat_meta['output_slices']['hidden_state'] + is_supercombo = vision_runner is None + + print(f"Compiling run_policy JIT (model_size={model_w}x{model_h}, frame_skip={derived_frame_skip})...") + run_policy_func = make_run_policy(vision_runner, policy_runners, features_slice, derived_frame_skip, all_shapes) + run_policy_jit = TinyJit(run_policy_func, prune=True) + make_policy_queues = partial(generate_queues_and_npy, all_shapes, derived_frame_skip, is_supercombo=is_supercombo) + make_random_model_inputs = partial(make_random_images, keys=['warped'], shape=(2, 6, model_h // 2, model_w // 2), device=WARP_DEV) + output_data['run_policy'] = compile_jit(run_policy_jit, make_random_model_inputs, POLICY_INPUTS, make_policy_queues) + + for cam_w, cam_h in args.camera_resolutions: + print(f"Compiling warp JIT for {cam_w}x{cam_h}...") + nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) + make_random_warp_inputs = partial(make_random_images, keys=['frame', 'big_frame'], shape=nv12.size, device=WARP_DEV) + warp = TinyJit(make_warp(nv12, model_w, model_h), prune=True) + output_data[(cam_w, cam_h)] = compile_jit(warp, make_random_warp_inputs, WARP_INPUTS, make_warp_queues) with open(args.output, "wb") as file: dump_oob(output_data, file) 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) print(f"Chunked into {len(chunk_targets) - 1} file(s)") diff --git a/openpilot/sunnypilot/modeld_v2/modeld.py b/openpilot/sunnypilot/modeld_v2/modeld.py index b55687c7b9..edec917c75 100755 --- a/openpilot/sunnypilot/modeld_v2/modeld.py +++ b/openpilot/sunnypilot/modeld_v2/modeld.py @@ -41,7 +41,7 @@ from openpilot.sunnypilot.modeld_v2.fill_model_msg import fill_model_msg, fill_p from openpilot.sunnypilot.modeld_v2.constants import Plan from openpilot.sunnypilot.modeld_v2.meta_helper import load_meta_constants from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelper -from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues +from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues, make_supercombo_input_queues, WARP_INPUTS, POLICY_INPUTS from openpilot.sunnypilot.livedelay.helpers import get_lat_delay from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase @@ -112,11 +112,10 @@ class ModelState(ModelStateBase): self.WARP_DEV = 'QCOM' if COMMA_HARDWARE else 'CPU' self.DEV = 'AMD' if self.usbgpu else self.WARP_DEV self.QUEUE_DEV = self.DEV - metadata = jits['metadata'] - self._run_policy = jits[(cam_w, cam_h)]['run_policy'] - self._warp_enqueue = jits[(cam_w, cam_h)]['warp_enqueue'] + self.run_policy = jits['run_policy'] + self.warp = jits[(cam_w, cam_h)] if 'model' in metadata: model_metadata = metadata['model'] @@ -125,7 +124,6 @@ class ModelState(ModelStateBase): self._policy_slices_list = [] self._combined_model_type = 'supercombo' self._vision_input_names = [key for key in model_metadata['input_shapes'] if 'img' in key] - from openpilot.sunnypilot.modeld_v2.compile_modeld import make_supercombo_input_queues frame_skip = derive_frame_skip({}, model_metadata['input_shapes']) self.input_queues, self.numpy_inputs = make_supercombo_input_queues(model_metadata['input_shapes'], frame_skip, device=self.QUEUE_DEV) @@ -141,12 +139,11 @@ class ModelState(ModelStateBase): self._policy_slices_list = [metadata[k]['output_slices'] for k in policy_keys] self.policy_output_slices = self._policy_slices_list[0] self._has_on_policy = any('on' in k.lower() for k in policy_keys) - first_policy_metadata = metadata[policy_keys[0]] - vision_input_shapes = vision_metadata['input_shapes'] - policy_input_shapes = first_policy_metadata['input_shapes'] - self._vision_input_names = [k for k in vision_input_shapes if 'img' in k] - frame_skip = derive_frame_skip(vision_input_shapes, policy_input_shapes) - self.input_queues, self.numpy_inputs = make_split_input_queues(vision_input_shapes, policy_input_shapes, + self._vision_input_names = [key for key in vision_metadata['input_shapes'] if 'img' in key] + first_policy_meta = metadata[policy_keys[0]] + frame_skip = derive_frame_skip(vision_metadata['input_shapes'], first_policy_meta['input_shapes']) + self.input_queues, self.numpy_inputs = make_split_input_queues(vision_metadata['input_shapes'], + first_policy_meta['input_shapes'], frame_skip, device=self.QUEUE_DEV) self._desire_key = next(key for key in self.numpy_inputs if key.startswith('desire')) @@ -175,8 +172,9 @@ class ModelState(ModelStateBase): self.frame_buf_params = dict.fromkeys(self._vision_input_names, nv12_info) yuv_size = self.frame_buf_params[self._road_key][3] - self._warp_enqueue( - **self.input_queues, + self.warp( + tfm=self.input_queues['tfm'], + big_tfm=self.input_queues['big_tfm'], frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize(), big_frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize()) @@ -237,10 +235,10 @@ class ModelState(ModelStateBase): self.numpy_inputs['big_tfm'][:, :] = transforms[wide_key].reshape(3, 3) if prepare_only: - self._warp_enqueue(**self.input_queues, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key]) + self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key]) return None - - raw_outputs = self._run_policy(**self.input_queues, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key]) + warped = self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key]) + raw_outputs = self.run_policy(**{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, warped=warped) if self._combined_model_type == 'supercombo': model_output = raw_outputs.numpy().flatten() diff --git a/openpilot/sunnypilot/modeld_v2/tests/helpers.py b/openpilot/sunnypilot/modeld_v2/tests/helpers.py index 3c37d38256..ee59e82785 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/helpers.py +++ b/openpilot/sunnypilot/modeld_v2/tests/helpers.py @@ -163,7 +163,8 @@ ARCHETYPES = { def make_pkl_data(archetype): return { 'metadata': archetype.metadata_structure, - (CAM_W, CAM_H): {'run_policy': _noop_jit, 'warp_enqueue': _noop_jit}, + 'run_policy': _noop_jit, + (CAM_W, CAM_H): _noop_jit, }