IQ.Pilot Prebuilt Release @ 27f668a

iqpilot/selfdrive/iqmodeld/tools/install_models_pc.py

@@ -0,0 +1,128 @@
+"""
+Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
+"""
+
+from __future__ import annotations
+
+import base64
+import pickle
+import shutil
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import onnx
+
+from iqpilot.system.hardware.hw import Paths
+
+_MODEL_STEMS = ("driving_off_policy", "driving_on_policy", "driving_policy", "driving_vision")
+
+
+@dataclass(frozen=True)
+class _ModelBundle:
+  stem: str
+  onnx_path: Path
+  artifact_path: Path
+  metadata_path: Path
+
+
+def _tensor_shape(value_info) -> tuple[int, ...]:
+  return tuple(int(dim.dim_value) for dim in value_info.type.tensor_type.shape.dim)
+
+
+def _metadata_property(graph_model, key: str) -> str | None:
+  for property_item in graph_model.metadata_props:
+    if property_item.key == key:
+      return property_item.value
+  return None
+
+
+def _decode_output_slices(encoded_value: str):
+  return pickle.loads(base64.b64decode(encoded_value.encode()))
+
+
+def _metadata_record(graph_model) -> dict:
+  encoded_slices = _metadata_property(graph_model, "output_slices")
+  if encoded_slices is None:
+    raise ValueError("output_slices metadata missing")
+  return {
+    "model_checkpoint": _metadata_property(graph_model, "model_checkpoint"),
+    "output_slices": _decode_output_slices(encoded_slices),
+    "input_shapes": {item.name: _tensor_shape(item) for item in graph_model.graph.input},
+    "output_shapes": {item.name: _tensor_shape(item) for item in graph_model.graph.output},
+  }
+
+
+def generate_metadata_pkl(model_path, output_path):
+  try:
+    graph_model = onnx.load(str(model_path))
+    metadata = _metadata_record(graph_model)
+  except Exception:
+    return False
+
+  with open(output_path, "wb") as handle:
+    pickle.dump(metadata, handle)
+  return True
+
+
+def _discover_model_bundles(model_dir: Path) -> list[_ModelBundle]:
+  bundles: list[_ModelBundle] = []
+  for stem in _MODEL_STEMS:
+    onnx_path = model_dir / f"{stem}.onnx"
+    if not onnx_path.exists():
+      continue
+    bundles.append(_ModelBundle(
+      stem=stem,
+      onnx_path=onnx_path,
+      artifact_path=model_dir / f"{stem}_tinygrad.pkl",
+      metadata_path=model_dir / f"{stem}_metadata.pkl",
+    ))
+  return bundles
+
+
+def _prompt_short_name(found_stems: list[str]) -> str | None:
+  try:
+    response = input(f"Found models ({', '.join(found_stems)}). Enter model short name (e.g. wmiv4): ").strip()
+  except EOFError:
+    return None
+  return response or None
+
+
+def _ensure_metadata_file(bundle: _ModelBundle) -> None:
+  if bundle.metadata_path.exists():
+    return
+  generate_metadata_pkl(bundle.onnx_path, bundle.metadata_path)
+
+
+def _install_bundle(bundle: _ModelBundle, suffix: str, destination_root: Path) -> None:
+  _ensure_metadata_file(bundle)
+  renamed_artifact = destination_root / f"{bundle.stem}_{suffix}_tinygrad.pkl"
+  renamed_metadata = destination_root / f"{bundle.stem}_{suffix}_metadata.pkl"
+  if bundle.artifact_path.exists():
+    shutil.move(str(bundle.artifact_path), str(renamed_artifact))
+  if bundle.metadata_path.exists():
+    shutil.move(str(bundle.metadata_path), str(renamed_metadata))
+
+
+def install_models(model_dir):
+  source_root = Path(model_dir)
+  bundles = _discover_model_bundles(source_root)
+  if not bundles:
+    return
+
+  short_name = _prompt_short_name([bundle.stem for bundle in bundles])
+  if short_name is None:
+    print("No name provided, skipping installation.")
+    return
+
+  destination_root = Path(Paths.model_root())
+  destination_root.mkdir(parents=True, exist_ok=True)
+  for bundle in bundles:
+    _install_bundle(bundle, short_name, destination_root)
+
+
+if __name__ == "__main__":
+  if len(sys.argv) < 2:
+    print("Usage: install_models_pc.py <model_dir>")
+    sys.exit(1)
+  install_models(sys.argv[1])