101 lines
3.3 KiB
Python
101 lines
3.3 KiB
Python
"""
|
|
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
EGPU_MODELS: dict[str, dict] = {
|
|
"lebrowski": {
|
|
"model_name": "big_driving_supercombo",
|
|
"sha256": "a501760a9d1d5fef0eab2b8c5d122d06124fc26dc8e0782e0aa94b82a208f0ff",
|
|
"output_len": 2580,
|
|
"frame_skip": 4,
|
|
"download": {"kind": "comma_lfs", "size": 1757355221},
|
|
"input_shapes": {
|
|
"img": (1, 12, 128, 256),
|
|
"big_img": (1, 12, 128, 256),
|
|
"desire_pulse": (1, 25, 8),
|
|
"traffic_convention": (1, 2),
|
|
"action_t": (1, 2),
|
|
"features_buffer": (1, 24, 512),
|
|
},
|
|
"output_slices": {
|
|
"lane_lines": slice(0, 528),
|
|
"lane_lines_prob": slice(528, 536),
|
|
"road_edges": slice(536, 800),
|
|
"meta": slice(800, 855),
|
|
"desire_pred": slice(855, 887),
|
|
"pose": slice(887, 899),
|
|
"wide_from_device_euler": slice(899, 905),
|
|
"road_transform": slice(905, 917),
|
|
"plan": slice(917, 1907),
|
|
"lead": slice(1907, 2051),
|
|
"lead_prob": slice(2051, 2054),
|
|
"desire_state": slice(2054, 2062),
|
|
"action": slice(2062, 2066),
|
|
"hidden_state": slice(2066, 2578),
|
|
"pad": slice(-2, None),
|
|
},
|
|
},
|
|
"comma_small": {
|
|
"model_name": "driving_supercombo",
|
|
"sha256": "659727c4d4839adc4992a254409a54259a8756a743f2d567bf5fdc6579f8009b",
|
|
"output_len": 2576,
|
|
"frame_skip": 4,
|
|
"download": {"kind": "comma_lfs", "size": 60881999},
|
|
"output_slices": {
|
|
"meta": slice(0, 55),
|
|
"desire_pred": slice(55, 87),
|
|
"pose": slice(87, 99),
|
|
"wide_from_device_euler": slice(99, 105),
|
|
"road_transform": slice(105, 117),
|
|
"lane_lines": slice(117, 645),
|
|
"lane_lines_prob": slice(645, 653),
|
|
"road_edges": slice(653, 917),
|
|
"lead": slice(917, 1061),
|
|
"lead_prob": slice(1061, 1064),
|
|
"hidden_state": slice(1064, 1576),
|
|
"plan": slice(1576, 2566),
|
|
"desire_state": slice(2566, 2574),
|
|
"pad": slice(-2, None),
|
|
},
|
|
},
|
|
}
|
|
|
|
DEFAULT_EGPU_MODEL = "lebrowski"
|
|
|
|
BIG_MODEL_NAME = "big_driving_supercombo"
|
|
|
|
|
|
def get_egpu_model(key: str | None = None) -> dict:
|
|
name = key or DEFAULT_EGPU_MODEL
|
|
if name not in EGPU_MODELS:
|
|
raise KeyError(f"unknown eGPU model {name!r}; known: {sorted(EGPU_MODELS)}")
|
|
return {"key": name, **EGPU_MODELS[name]}
|
|
|
|
|
|
def resolve_egpu_model(params, selected: str | bytes | None = None, allow_refresh: bool = True) -> dict | None:
|
|
name = selected if selected is not None else (params.get("IQEmacModel") if params is not None else None)
|
|
name = name.decode() if isinstance(name, bytes) else (name or "")
|
|
if not name or name == DEFAULT_EGPU_MODEL:
|
|
return get_egpu_model()
|
|
from iqpilot.selfdrive.iqmodeld.big_catalog import cached_catalog, refresh_catalog
|
|
catalog = cached_catalog(params)
|
|
if name not in catalog and params is not None and allow_refresh:
|
|
refresh_catalog(params)
|
|
catalog = cached_catalog(params)
|
|
meta = catalog.get(name)
|
|
if meta is None or meta.get("model_name") != BIG_MODEL_NAME:
|
|
return None
|
|
return {"key": name, **meta}
|
|
|
|
|
|
def download_descriptor(meta: dict) -> tuple[str, int]:
|
|
dl = meta.get("download")
|
|
if not dl:
|
|
return "", 0
|
|
if dl["kind"] == "comma_lfs":
|
|
return f"commalfs:{meta['sha256']}", int(dl["size"])
|
|
if dl["kind"] == "url":
|
|
return str(dl["url"]), int(dl.get("size", 0))
|
|
return "", 0
|