mirror of
https://github.com/MoreTore/openpilot.git
synced 2026-08-03 16:49:33 +08:00
FrogAi
243 lines
11 KiB
Python
243 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
import json
|
|
import re
|
|
import requests
|
|
import shutil
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
|
|
from pathlib import Path
|
|
|
|
from openpilot.frogpilot.assets.download_functions import GITLAB_URL, download_file, get_repository_url, handle_error, handle_request_error, verify_download
|
|
from openpilot.frogpilot.common.frogpilot_utilities import delete_file
|
|
from openpilot.frogpilot.common.frogpilot_variables import DEFAULT_CLASSIC_MODEL, DEFAULT_MODEL, DEFAULT_TINYGRAD_MODEL, MODELS_PATH, params, params_default, params_memory
|
|
|
|
VERSION = "v15"
|
|
|
|
CANCEL_DOWNLOAD_PARAM = "CancelModelDownload"
|
|
DOWNLOAD_PROGRESS_PARAM = "ModelDownloadProgress"
|
|
MODEL_DOWNLOAD_PARAM = "ModelToDownload"
|
|
MODEL_DOWNLOAD_ALL_PARAM = "DownloadAllModels"
|
|
|
|
class ModelManager:
|
|
def __init__(self):
|
|
self.available_models = (params.get("AvailableModels", encoding="utf-8") or "").split(",")
|
|
self.model_versions = (params.get("ModelVersions", encoding="utf-8") or "").split(",")
|
|
|
|
self.downloading_model = False
|
|
|
|
@staticmethod
|
|
def fetch_models(url):
|
|
try:
|
|
with urllib.request.urlopen(url, timeout=10) as response:
|
|
return json.loads(response.read().decode("utf-8"))["models"]
|
|
except Exception as error:
|
|
handle_request_error(error, None, None, None, None)
|
|
return []
|
|
|
|
@staticmethod
|
|
def fetch_all_model_sizes(repo_url):
|
|
project_path = "FrogAi/FrogPilot-Resources"
|
|
branch = "Models"
|
|
|
|
if "github" in repo_url:
|
|
api_url = f"https://api.github.com/repos/{project_path}/contents?ref={branch}"
|
|
elif "gitlab" in repo_url:
|
|
api_url = f"https://gitlab.com/api/v4/projects/{urllib.parse.quote_plus(project_path)}/repository/tree?ref={branch}"
|
|
else:
|
|
return {}
|
|
|
|
try:
|
|
response = requests.get(api_url)
|
|
response.raise_for_status()
|
|
model_files = [file for file in response.json() if "." in file["name"]]
|
|
|
|
if "gitlab" in repo_url:
|
|
model_sizes = {}
|
|
for file in model_files:
|
|
file_path = file["path"]
|
|
metadata_url = f"https://gitlab.com/api/v4/projects/{urllib.parse.quote_plus(project_path)}/repository/files/{urllib.parse.quote_plus(file_path)}/raw?ref={branch}"
|
|
metadata_response = requests.head(metadata_url)
|
|
metadata_response.raise_for_status()
|
|
model_sizes[file["name"].rsplit(".", 1)[0]] = int(metadata_response.headers.get("content-length", 0))
|
|
return model_sizes
|
|
else:
|
|
return {file["name"].rsplit(".", 1)[0]: file["size"] for file in model_files if "size" in file}
|
|
|
|
except Exception as error:
|
|
handle_request_error(f"Failed to fetch model sizes from {'GitHub' if 'github' in repo_url else 'GitLab'}: {error}", None, None, None, None)
|
|
return {}
|
|
|
|
def handle_verification_failure(self, model, model_path, file_extension):
|
|
print(f"Verification failed for model {model}. Retrying from GitLab...")
|
|
model_url = f"{GITLAB_URL}/Models/{model}.{file_extension}"
|
|
download_file(CANCEL_DOWNLOAD_PARAM, model_path, DOWNLOAD_PROGRESS_PARAM, model_url, MODEL_DOWNLOAD_PARAM, params_memory)
|
|
|
|
if params_memory.get_bool(CANCEL_DOWNLOAD_PARAM):
|
|
handle_error(None, "Download cancelled...", "Download cancelled...", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM, params_memory)
|
|
self.downloading_model = False
|
|
return
|
|
|
|
if verify_download(model_path, model_url):
|
|
print(f"Model {model} downloaded and verified successfully!")
|
|
params_memory.put(DOWNLOAD_PROGRESS_PARAM, "Downloaded!")
|
|
params_memory.remove(MODEL_DOWNLOAD_PARAM)
|
|
self.downloading_model = False
|
|
else:
|
|
handle_error(model_path, "Verification failed...", "GitLab verification failed", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM, params_memory)
|
|
self.downloading_model = False
|
|
|
|
def download_model(self, model_to_download):
|
|
self.downloading_model = True
|
|
|
|
repo_url = get_repository_url()
|
|
if not repo_url:
|
|
handle_error(None, "GitHub and GitLab are offline...", "Repository unavailable", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM, params_memory)
|
|
self.downloading_model = False
|
|
return
|
|
|
|
file_extension = "thneed" if self.model_versions[self.available_models.index(model_to_download)] in {"v1", "v2", "v3", "v4", "v5", "v6"} else "pkl"
|
|
model_path = MODELS_PATH / f"{model_to_download}.{file_extension}"
|
|
model_url = f"{repo_url}/Models/{model_to_download}.{file_extension}"
|
|
print(f"Downloading model: {model_to_download}")
|
|
download_file(CANCEL_DOWNLOAD_PARAM, model_path, DOWNLOAD_PROGRESS_PARAM, model_url, MODEL_DOWNLOAD_PARAM, params_memory)
|
|
|
|
if params_memory.get_bool(CANCEL_DOWNLOAD_PARAM):
|
|
handle_error(None, "Download cancelled...", "Download cancelled...", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM, params_memory)
|
|
self.downloading_model = False
|
|
return
|
|
|
|
if verify_download(model_path, model_url):
|
|
print(f"Model {model_to_download} downloaded and verified successfully!")
|
|
params_memory.put(DOWNLOAD_PROGRESS_PARAM, "Downloaded!")
|
|
params_memory.remove(MODEL_DOWNLOAD_PARAM)
|
|
self.downloading_model = False
|
|
else:
|
|
self.handle_verification_failure(model_to_download, model_path, file_extension)
|
|
|
|
@staticmethod
|
|
def copy_default_model():
|
|
classic_default_model_path = MODELS_PATH / f"{DEFAULT_CLASSIC_MODEL}.thneed"
|
|
source_path = Path(__file__).parents[1] / "classic_modeld/models/supercombo.thneed"
|
|
if source_path.is_file() and not classic_default_model_path.is_file():
|
|
shutil.copyfile(source_path, classic_default_model_path)
|
|
print(f"Copied the classic default model from {source_path} to {classic_default_model_path}")
|
|
|
|
default_model_path = MODELS_PATH / f"{DEFAULT_MODEL}.thneed"
|
|
source_path = Path(__file__).parents[2] / "selfdrive/modeld/models/supercombo.thneed"
|
|
if source_path.is_file() and not default_model_path.is_file():
|
|
shutil.copyfile(source_path, default_model_path)
|
|
print(f"Copied the default model from {source_path} to {default_model_path}")
|
|
|
|
def check_models(self, boot_run, repo_url):
|
|
available_models = set(self.available_models) - {DEFAULT_MODEL, DEFAULT_CLASSIC_MODEL}
|
|
downloaded_models = {path.stem for path in MODELS_PATH.iterdir() if path.is_file()} - {DEFAULT_MODEL, DEFAULT_CLASSIC_MODEL}
|
|
|
|
outdated_models = downloaded_models - available_models
|
|
for model in outdated_models:
|
|
for model_file in MODELS_PATH.glob(f"{model}.*"):
|
|
print(f"Removing outdated model: {model_file}")
|
|
delete_file(model_file)
|
|
|
|
for tmp_file in MODELS_PATH.glob("tmp*"):
|
|
if tmp_file.is_file():
|
|
delete_file(tmp_file)
|
|
|
|
if params.get("Model", encoding="utf-8") not in self.available_models + [DEFAULT_TINYGRAD_MODEL]:
|
|
params.put("Model", params_default.get("Model", encoding="utf-8"))
|
|
|
|
automatically_download_models = not boot_run and params.get_bool("AutomaticallyDownloadModels")
|
|
if not automatically_download_models:
|
|
return
|
|
|
|
model_sizes = self.fetch_all_model_sizes(repo_url)
|
|
if not model_sizes:
|
|
print("No model size data available. Skipping model checks")
|
|
return
|
|
|
|
needs_download = False
|
|
|
|
for model in available_models:
|
|
expected_size = model_sizes.get(model)
|
|
if expected_size is None:
|
|
print(f"Size data for {model} not found in fetched metadata.")
|
|
continue
|
|
|
|
model_files = list(MODELS_PATH.glob(f"{model}.*"))
|
|
if not model_files:
|
|
needs_download = True
|
|
continue
|
|
|
|
for model_file in model_files:
|
|
if model_file.is_file():
|
|
local_size = model_file.stat().st_size
|
|
if local_size != expected_size:
|
|
print(f"Model {model} is outdated. Deleting {model_file}...")
|
|
delete_file(model_file)
|
|
needs_download = True
|
|
|
|
if needs_download:
|
|
self.download_all_models()
|
|
|
|
def update_model_params(self, model_info, repo_url):
|
|
self.available_models = [model["id"] for model in model_info if model["id"] != DEFAULT_TINYGRAD_MODEL]
|
|
self.model_versions = [model["version"] for model in model_info if model["id"] != DEFAULT_TINYGRAD_MODEL]
|
|
|
|
params.put("AvailableModels", ",".join(self.available_models))
|
|
params.put("AvailableModelNames", ",".join([model["name"] for model in model_info if model["id"] != DEFAULT_TINYGRAD_MODEL]))
|
|
params.put("ExperimentalModels", ",".join([model["id"] for model in model_info if model.get("experimental", False)]))
|
|
params.put("ModelVersions", ",".join(self.model_versions))
|
|
print("Models list updated successfully")
|
|
|
|
def update_models(self, boot_run=False):
|
|
if self.downloading_model:
|
|
return
|
|
|
|
repo_url = get_repository_url()
|
|
if repo_url is None:
|
|
print("GitHub and GitLab are offline...")
|
|
return
|
|
|
|
model_info = self.fetch_models(f"{repo_url}/Versions/model_names_{VERSION}.json")
|
|
if model_info:
|
|
self.update_model_params(model_info, repo_url)
|
|
self.check_models(boot_run, repo_url)
|
|
|
|
def queue_model_download(self, model, model_name=None):
|
|
while params_memory.get(MODEL_DOWNLOAD_PARAM, encoding="utf-8"):
|
|
time.sleep(1)
|
|
|
|
params_memory.put(MODEL_DOWNLOAD_PARAM, model)
|
|
if model_name:
|
|
params_memory.put(DOWNLOAD_PROGRESS_PARAM, f"Downloading \"{model_name}\"...")
|
|
|
|
def download_all_models(self):
|
|
repo_url = get_repository_url()
|
|
if not repo_url:
|
|
handle_error(None, "GitHub and GitLab are offline...", "Repository unavailable", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM, params_memory)
|
|
return
|
|
|
|
model_info = self.fetch_models(f"{repo_url}/Versions/model_names_{VERSION}.json")
|
|
if model_info:
|
|
available_models = [model["id"] for model in model_info if model["id"] != DEFAULT_TINYGRAD_MODEL]
|
|
available_model_names = [re.sub(r"[🗺️👀📡]", "", model["name"]).strip() for model in model_info if model["id"] != DEFAULT_TINYGRAD_MODEL]
|
|
|
|
for model_id, model_name in zip(available_models, available_model_names):
|
|
model_downloaded = list(MODELS_PATH.glob(f"{model_id}.*"))
|
|
if not model_downloaded:
|
|
print(f"Model {model_id} is not downloaded. Preparing to download...")
|
|
self.queue_model_download(model_id, model_name)
|
|
|
|
while not all(any(file.is_file() for file in MODELS_PATH.glob(f"{model}.*")) for model in available_models):
|
|
if params_memory.get_bool(CANCEL_DOWNLOAD_PARAM):
|
|
handle_error(None, "Download cancelled...", "Download cancelled...", MODEL_DOWNLOAD_ALL_PARAM, DOWNLOAD_PROGRESS_PARAM, params_memory)
|
|
handle_error(None, "Download cancelled...", "Download cancelled...", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM, params_memory)
|
|
return
|
|
time.sleep(1)
|
|
|
|
params_memory.put(DOWNLOAD_PROGRESS_PARAM, "All models downloaded!")
|
|
else:
|
|
handle_error(None, "Unable to fetch models...", "Model list unavailable", MODEL_DOWNLOAD_ALL_PARAM, DOWNLOAD_PROGRESS_PARAM, params_memory)
|
|
return
|