github-actions[bot]
f9fcc7adab
date: 2026-06-28T09:48:35 master commit: da6313dbe95b3f24bb5d8018b0e5f950f5823ca7
76 lines
2.3 KiB
Python
Executable File
76 lines
2.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import sys
|
|
import shutil
|
|
import pickle
|
|
import codecs
|
|
from pathlib import Path
|
|
|
|
from openpilot.system.hardware.hw import Paths
|
|
from sunnypilot.modeld_v2.get_model_metadata import MetadataOnnxPBParser, get_name_and_shape, get_metadata_value_by_name
|
|
|
|
|
|
def generate_metadata_pkl(model_path, output_path):
|
|
try:
|
|
model = MetadataOnnxPBParser(model_path).parse()
|
|
output_slices = get_metadata_value_by_name(model, 'output_slices')
|
|
if not output_slices:
|
|
return False
|
|
metadata = {
|
|
'model_checkpoint': get_metadata_value_by_name(model, 'model_checkpoint'),
|
|
'output_slices': pickle.loads(codecs.decode(output_slices.encode(), "base64")),
|
|
'input_shapes': dict(get_name_and_shape(x) for x in model["graph"]["input"]),
|
|
'output_shapes': dict(get_name_and_shape(x) for x in model["graph"]["output"]),
|
|
}
|
|
with open(output_path, 'wb') as f:
|
|
pickle.dump(metadata, f)
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def install_models(model_dir):
|
|
model_dir = Path(model_dir)
|
|
models = ["driving_off_policy", "driving_on_policy", "driving_vision"]
|
|
found_models = []
|
|
|
|
for model in models:
|
|
if (model_dir / f"{model}.onnx").exists():
|
|
found_models.append(model)
|
|
|
|
if not found_models:
|
|
return
|
|
|
|
try:
|
|
custom_name = input(f"Found models ({', '.join(found_models)}). Enter model short name (e.g. wmiv4): ").strip()
|
|
except EOFError:
|
|
return
|
|
|
|
if not custom_name:
|
|
print("No name provided, skipping installation.")
|
|
return
|
|
|
|
dest_dir = Path(Paths.model_root())
|
|
dest_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
for model in found_models:
|
|
onnx_path = model_dir / f"{model}.onnx"
|
|
tinygrad_pkl = model_dir / f"{model}_tinygrad.pkl"
|
|
metadata_pkl = model_dir / f"{model}_metadata.pkl"
|
|
|
|
if not metadata_pkl.exists():
|
|
generate_metadata_pkl(onnx_path, metadata_pkl)
|
|
|
|
dest_tinygrad = dest_dir / f"{model}_{custom_name}_tinygrad.pkl"
|
|
dest_metadata = dest_dir / f"{model}_{custom_name}_metadata.pkl"
|
|
|
|
if tinygrad_pkl.exists():
|
|
shutil.move(str(tinygrad_pkl), str(dest_tinygrad))
|
|
if metadata_pkl.exists():
|
|
shutil.move(str(metadata_pkl), str(dest_metadata))
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2:
|
|
print("Usage: install_models_pc.py <model_dir>")
|
|
sys.exit(1)
|
|
install_models(sys.argv[1])
|