#!/usr/bin/env python3 import json import re import requests import shutil import time import urllib.parse from pathlib import Path from urllib.parse import quote_plus from openpilot.common.basedir import BASEDIR from openpilot.frogpilot.assets.download_functions import GITLAB_URL, download_file, get_remote_file_size, get_repository_url, handle_error, handle_request_error, verify_download from openpilot.frogpilot.common.frogpilot_utilities import delete_file, extract_tar, load_json_file, update_json_file from openpilot.frogpilot.common.frogpilot_variables import ( DEFAULT_MODEL, DEFAULT_MODEL_NAME, DEFAULT_MODEL_VERSION, MODELS_PATH, RESOURCES_REPO, TINYGRAD_FILES, params, params_default, params_memory, update_frogpilot_toggles ) VERSION = "v16" VERSION_PATH = MODELS_PATH / "model_version" CANCEL_DOWNLOAD_PARAM = "CancelModelDownload" DOWNLOAD_PROGRESS_PARAM = "ModelDownloadProgress" MODEL_DOWNLOAD_PARAM = "ModelToDownload" MODEL_DOWNLOAD_ALL_PARAM = "DownloadAllModels" UPDATE_TINYGRAD_PARAM = "UpdateTinygrad" DEFAULT_TINYGRAD_SIZE = 87746736 TAR_FILE_NAME = f"Tinygrad_{VERSION}.tar.gz" TINYGRAD_MODELD_PATH = Path(BASEDIR) / "frogpilot/tinygrad_modeld" TINYGRAD_REPO_PATH = Path(BASEDIR) / "tinygrad_repo" class ModelManager: def __init__(self, boot_run=False): self.downloading_model = False self.available_models = (params.get("AvailableModels", encoding="utf-8") or "").split(",") self.available_model_names = (params.get("AvailableModelNames", encoding="utf-8") or "").split(",") self.model_versions = (params.get("ModelVersions", encoding="utf-8") or "").split(",") self.model_sizes_path = MODELS_PATH / "model_sizes.json" self.tinygrad_sizes_path = MODELS_PATH / "tinygrad_sizes.json" self.model_sizes = load_json_file(self.model_sizes_path) self.tinygrad_sizes = load_json_file(self.tinygrad_sizes_path) self.session = requests.Session() self.session.headers.update({"Accept-Language": "en"}) self.session.headers.update({"User-Agent": "frogpilot-model-downloader/1.0 (https://github.com/FrogAi/FrogPilot)"}) if boot_run: self.copy_default_model() self.validate_models() def check_models(self, boot_run, repo_url): downloaded_models = [ model for model in MODELS_PATH.iterdir() if (MODELS_PATH / f"{model}.thneed").is_file() or all((MODELS_PATH / f"{model}_{filename}").is_file() for filename, _ in TINYGRAD_FILES) ] for model_file in downloaded_models: if not any(model in model_file.name for model in set(self.available_models)): 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").removesuffix("_default") not in self.available_models: params.put("Model", params_default.get("Model", encoding="utf-8")) if not (not boot_run and params.get_bool("AutomaticallyDownloadModels")): return model_sizes = self.fetch_all_model_sizes(repo_url) if not model_sizes: print("No model size data available. Skipping model checks...") return need_to_update_models = False for model in self.available_models: if self.is_tinygrad_model(model): model_file = MODELS_PATH / f"{model}.thneed" if not model_file.is_file(): need_to_update_models = True continue expected_size = model_sizes.get(model_file.name) local_size = self.model_sizes.get(model_file.name) if expected_size > 0 and local_size != expected_size: print(f"Model {model} is outdated. Deleting {model_file}...") delete_file(model_file) need_to_update_models = True else: model_missing = False model_outdated = False for filename, _ in TINYGRAD_FILES: expected_file = MODELS_PATH / f"{model}_{filename}" if not expected_file.is_file(): model_missing = True need_to_update_models = True break for filename, _ in TINYGRAD_FILES: model_file = f"{model}_{filename}" expected_size = model_sizes.get(model_file) local_size = self.model_sizes.get(model_file) if expected_size > 0 and local_size != expected_size: model_outdated = True need_to_update_models = True break if model_missing or model_outdated: print(f"Model {model} is either missing required files or outdated. Deleting...") for filename, _ in TINYGRAD_FILES: delete_file(MODELS_PATH / f"{model}_{filename}") if need_to_update_models: params_memory.put_bool(MODEL_DOWNLOAD_ALL_PARAM, True) def check_tinygrad(self, repo_url): tinygrad_url = f"{repo_url}/Tinygrad/{TAR_FILE_NAME}" expected_size = get_remote_file_size(tinygrad_url, self.session) local_size = int(self.tinygrad_sizes.get(TAR_FILE_NAME, 0)) if expected_size > 0 and local_size != expected_size: print(f"Tinygrad version {VERSION} is outdated, expected_size: {expected_size}, local_size: {local_size}, flagging for update...") params.put_bool("TinygradUpdateAvailable", True) def copy_default_model(self): classic_default_model_path = MODELS_PATH / "wd-40.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() or source_path.stat().st_size != classic_default_model_path.stat().st_size): shutil.copyfile(source_path, classic_default_model_path) print(f"Copied the classic default model from {source_path} to {classic_default_model_path}") self.update_model_size(classic_default_model_path) default_model_path = MODELS_PATH / "national-public-radio.thneed" source_path = Path(__file__).parents[2] / "selfdrive/modeld/models/supercombo.thneed" if source_path.is_file() and (not default_model_path.is_file() or source_path.stat().st_size != default_model_path.stat().st_size): shutil.copyfile(source_path, default_model_path) print(f"Copied the default model from {source_path} to {default_model_path}") self.update_model_size(default_model_path) for filename, description in TINYGRAD_FILES: source = TINYGRAD_MODELD_PATH / "models" / filename target = MODELS_PATH / f"{DEFAULT_MODEL}_{filename}" if source.is_file() and (not target.is_file() or source.stat().st_size != target.stat().st_size): shutil.copyfile(source, target) print(f"Copied the tinygrad {description} from {source} to {target}") def download_all_models(self): repo_url = get_repository_url(self.session) if not repo_url: handle_error(None, "GitHub and GitLab are offline...", "Repository unavailable", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM) return self.fetch_models(f"{repo_url}/Versions/model_names_{VERSION}.json", repo_url) for model in self.available_models: if params_memory.get_bool(CANCEL_DOWNLOAD_PARAM): handle_error(None, "Download cancelled...", "Download cancelled...", MODEL_DOWNLOAD_ALL_PARAM, DOWNLOAD_PROGRESS_PARAM) return if self.is_tinygrad_model(model): already_downloaded = (MODELS_PATH / f"{model}.thneed").is_file() else: already_downloaded = all((MODELS_PATH / f"{model}_{filename}").is_file() for filename, _ in TINYGRAD_FILES) if already_downloaded: continue print(f"Model {model} is not downloaded. Preparing to download...") params_memory.put(DOWNLOAD_PROGRESS_PARAM, f"Downloading \"{self.available_model_names[self.available_models.index(model)]}\"...") self.download_model(model) params_memory.put(DOWNLOAD_PROGRESS_PARAM, "All models downloaded!") params_memory.remove(MODEL_DOWNLOAD_ALL_PARAM) def download_model(self, model_to_download): self.downloading_model = True repo_url = get_repository_url(self.session) if not repo_url: handle_error(None, "GitHub and GitLab are offline...", "Repository unavailable", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM) self.downloading_model = False return if self.is_tinygrad_model(model_to_download): model_path = MODELS_PATH / f"{model_to_download}.thneed" model_url = f"{repo_url}/Models/{model_to_download}.thneed" print(f"Downloading model: {model_to_download}") download_file(CANCEL_DOWNLOAD_PARAM, model_path, DOWNLOAD_PROGRESS_PARAM, model_url, MODEL_DOWNLOAD_PARAM, self.session) if params_memory.get_bool(CANCEL_DOWNLOAD_PARAM): delete_file(model_path) handle_error(None, "Download cancelled...", "Download cancelled...", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM) self.downloading_model = False return if verify_download(model_path, model_url, self.session): print(f"Model {model_to_download} downloaded and verified successfully!") self.update_model_size(model_path) params_memory.put(DOWNLOAD_PROGRESS_PARAM, "Downloaded!") params_memory.remove(MODEL_DOWNLOAD_PARAM) self.downloading_model = False return print(f"Verification failed for model {model_to_download}. Retrying from GitLab...") fallback_url = f"{GITLAB_URL}/Models/{model_to_download}.thneed" download_file(CANCEL_DOWNLOAD_PARAM, model_path, DOWNLOAD_PROGRESS_PARAM, fallback_url, MODEL_DOWNLOAD_PARAM, self.session) if params_memory.get_bool(CANCEL_DOWNLOAD_PARAM): delete_file(model_path) handle_error(None, "Download cancelled...", "Download cancelled...", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM) self.downloading_model = False return if verify_download(model_path, fallback_url, self.session): print(f"Model {model_to_download} downloaded and verified successfully from GitLab!") self.update_model_size(model_path) 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) self.downloading_model = False else: all_model_sizes = self.fetch_all_model_sizes(repo_url) or {} tinygrad_filenames = [f"{model_to_download}_{file_key}" for file_key, _ in TINYGRAD_FILES] file_sizes = [] file_sources = [] missing = [name for name in tinygrad_filenames if int(all_model_sizes.get(name, 0)) <= 0] if missing: handle_error(None, "Missing size metadata...", f"Sizes not found for: {', '.join(missing)}...", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM) self.downloading_model = False return for filename in tinygrad_filenames: primary_url = f"{repo_url}/Models/compiled/{filename}" file_size = int(all_model_sizes.get(filename, 0)) file_sizes.append(file_size) file_sources.append((primary_url, None)) downloaded_offset_bytes = 0 known_file_sizes = [size for size in file_sizes if size > 0] total_model_bytes = sum(known_file_sizes) if len(known_file_sizes) == len(file_sizes) else 0 for (file_key, description), part_bytes, (primary_url, fallback_url) in zip(TINYGRAD_FILES, file_sizes, file_sources): filename = f"{model_to_download}_{file_key}" model_path = MODELS_PATH / filename print(f"Downloading {description} for model: {model_to_download}") download_file(CANCEL_DOWNLOAD_PARAM, model_path, DOWNLOAD_PROGRESS_PARAM, primary_url, MODEL_DOWNLOAD_PARAM, self.session, offset_bytes=downloaded_offset_bytes, total_bytes=total_model_bytes) if params_memory.get_bool(CANCEL_DOWNLOAD_PARAM): delete_file(model_path) handle_error(None, "Download cancelled...", "Download cancelled...", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM) self.downloading_model = False return if verify_download(model_path, primary_url, self.session): print(f"{description.capitalize()} for {model_to_download} downloaded and verified successfully!") if total_model_bytes: downloaded_offset_bytes += part_bytes continue print(f"Verification failed for {filename}. Retrying from GitLab...") fallback_url = f"{GITLAB_URL}/Models/compiled/{filename}" download_file(CANCEL_DOWNLOAD_PARAM, model_path, DOWNLOAD_PROGRESS_PARAM, fallback_url, MODEL_DOWNLOAD_PARAM, self.session, offset_bytes=downloaded_offset_bytes, total_bytes=total_model_bytes) if params_memory.get_bool(CANCEL_DOWNLOAD_PARAM): delete_file(model_path) handle_error(None, "Download cancelled...", "Download cancelled...", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM) self.downloading_model = False return if verify_download(model_path, fallback_url, self.session): print(f"{description.capitalize()} for {model_to_download} downloaded and verified successfully from GitLab!") if total_model_bytes: downloaded_offset_bytes += part_bytes else: handle_error(model_path, "Verification failed...", f"GitLab verification failed for {filename}", MODEL_DOWNLOAD_PARAM, DOWNLOAD_PROGRESS_PARAM) self.downloading_model = False return print(f"Updating model sizes for {model_to_download}...") for filename, _ in TINYGRAD_FILES: file_path = MODELS_PATH / f"{model_to_download}_{filename}" self.update_model_size(file_path) params_memory.put(DOWNLOAD_PROGRESS_PARAM, "Downloaded!") params_memory.remove(MODEL_DOWNLOAD_PARAM) self.downloading_model = False def fetch_all_model_sizes(self, repo_url): is_github = "github" in repo_url is_gitlab = "gitlab" in repo_url repo_encoded = quote_plus(RESOURCES_REPO) model_sizes = {} try: def fetch_dir_sizes(api_url): sizes = {} print(f"Fetching model metadata: {api_url}") response = self.session.get(api_url, timeout=10) response.raise_for_status() content = response.json() model_files = [file for file in content if "." in file["name"]] if is_github: for file in model_files: sizes[file["name"]] = file.get("size", 0) else: for file in model_files: file_path = quote_plus(file["path"]) metadata_url = f"https://gitlab.com/api/v4/projects/{repo_encoded}/repository/files/{file_path}/raw?ref=Models" head_response = self.session.head(metadata_url, timeout=10) if head_response.ok: sizes[file["name"]] = int(head_response.headers.get("content-length", 0)) return sizes if is_github: top_api_url = f"https://api.github.com/repos/{RESOURCES_REPO}/contents?ref=Models" version_api_url = f"https://api.github.com/repos/{RESOURCES_REPO}/contents/compiled?ref=Models" elif is_gitlab: top_api_url = f"https://gitlab.com/api/v4/projects/{repo_encoded}/repository/tree?ref=Models" version_api_url = f"https://gitlab.com/api/v4/projects/{repo_encoded}/repository/tree?path=compiled&ref=Models" else: print(f"Unsupported repository URL: {repo_url}") return model_sizes model_sizes.update(fetch_dir_sizes(top_api_url)) model_sizes.update(fetch_dir_sizes(version_api_url)) return model_sizes except requests.exceptions.RequestException as e: handle_request_error(f"Failed to fetch model sizes from {'GitHub' if is_github else 'GitLab'}: {e}", None, None, None) return {} def fetch_models(self, url, repo_url, boot_run=False): try: response = self.session.get(url, timeout=10) response.raise_for_status() model_info = response.json().get("models", []) if model_info: self.update_model_params(model_info) self.check_models(boot_run, repo_url) self.check_tinygrad(repo_url) except Exception as exception: handle_request_error(exception, None, None, None) return [] def is_tinygrad_model(self, model): return self.model_versions[self.available_models.index(model)] in {"v1", "v2", "v3", "v4", "v5", "v6"} def update_model_params(self, model_info): self.available_models = [model["id"] for model in model_info] self.available_model_names = [model["name"] for model in model_info] self.model_versions = [model["version"] for model in model_info] params.put("AvailableModels", ",".join(self.available_models)) params.put("AvailableModelNames", ",".join(self.available_model_names)) params.put("ModelVersions", ",".join(self.model_versions)) print("Models list updated successfully!") def update_models(self, boot_run): if self.downloading_model: return repo_url = get_repository_url(self.session) if repo_url is None: print("GitHub and GitLab are offline...") return self.fetch_models(f"{repo_url}/Versions/model_names_{VERSION}.json", repo_url, boot_run) def update_model_size(self, file_path): self.model_sizes[file_path.name] = file_path.stat().st_size update_json_file(self.model_sizes_path, self.model_sizes) print(f"Updated size for {file_path.name} in {self.model_sizes_path.name}") def update_tinygrad_size(self, file_path): self.tinygrad_sizes[TAR_FILE_NAME] = file_path.stat().st_size update_json_file(self.tinygrad_sizes_path, self.tinygrad_sizes) print(f"Updated size for {TAR_FILE_NAME} in {self.tinygrad_sizes_path.name}") def update_tinygrad(self): repo_url = get_repository_url(self.session) if not repo_url: handle_error(None, "GitHub and GitLab are offline...", "Repository unavailable", None, None) return primary_url = f"{repo_url}/Tinygrad/{TAR_FILE_NAME}" fallback_url = f"https://gitlab.com/{RESOURCES_REPO}/-/raw/Tinygrad/{TAR_FILE_NAME}" tinygrad_tar_path = Path("/data/tmp/tinygrad.tar.gz") try: print(f"Attempting to download tinygrad from {primary_url}...") download_file(CANCEL_DOWNLOAD_PARAM, tinygrad_tar_path, DOWNLOAD_PROGRESS_PARAM, primary_url, UPDATE_TINYGRAD_PARAM, self.session) if params_memory.get_bool(CANCEL_DOWNLOAD_PARAM): delete_file(tinygrad_tar_path) handle_error(None, "Tinygrad update cancelled...", "Tinygrad update cancelled...", UPDATE_TINYGRAD_PARAM, DOWNLOAD_PROGRESS_PARAM) params_memory.remove("CancelModelDownload") return if not verify_download(tinygrad_tar_path, primary_url, self.session): print(f"Verification failed for {primary_url}. Retrying from GitLab...") download_file(CANCEL_DOWNLOAD_PARAM, tinygrad_tar_path, DOWNLOAD_PROGRESS_PARAM, fallback_url, UPDATE_TINYGRAD_PARAM, self.session) if params_memory.get_bool(CANCEL_DOWNLOAD_PARAM): delete_file(tinygrad_tar_path) handle_error(None, "Tinygrad update cancelled...", "Tinygrad update cancelled...", UPDATE_TINYGRAD_PARAM, DOWNLOAD_PROGRESS_PARAM) params_memory.remove("CancelModelDownload") return if not verify_download(tinygrad_tar_path, fallback_url, self.session): handle_error(tinygrad_tar_path, "Verification Failed", "Tinygrad verification failed", UPDATE_TINYGRAD_PARAM, DOWNLOAD_PROGRESS_PARAM) return print("Tinygrad downloaded successfully! Proceeding with installation...") self.update_tinygrad_size(tinygrad_tar_path) params_memory.put(DOWNLOAD_PROGRESS_PARAM, "Installing...") print("Deleting old tinygrad directories...") delete_file(TINYGRAD_MODELD_PATH) print(f"Removed {TINYGRAD_MODELD_PATH}") delete_file(TINYGRAD_REPO_PATH) print(f"Removed {TINYGRAD_REPO_PATH}") extract_tar(tinygrad_tar_path, Path(BASEDIR)) print("Tinygrad update completed successfully!") params.put_bool("TinygradUpdateAvailable", False) params_memory.put(DOWNLOAD_PROGRESS_PARAM, "Updated!") params_memory.remove(UPDATE_TINYGRAD_PARAM) self.update_tinygrad_models(repo_url) except Exception as exception: handle_error(tinygrad_tar_path, "Update Failed", f"An unexpected error occurred: {exception}", UPDATE_TINYGRAD_PARAM, DOWNLOAD_PROGRESS_PARAM) def update_tinygrad_models(self, repo_url=None): print("Updating old Tinygrad models...") installed_tinygrad_models = set() for filename, _ in TINYGRAD_FILES: suffix = f"_{filename}" for file_path in MODELS_PATH.glob(f"*{suffix}"): model_name = file_path.name.rsplit(suffix, 1)[0] if model_name in set(self.available_models): installed_tinygrad_models.add(model_name) delete_file(file_path) self.copy_default_model() update_frogpilot_toggles() if repo_url is None: return current_model = params.get("Model", encoding="utf-8").removesuffix("_default") models_to_redownload = [current_model] models_to_redownload += [model for model in sorted(installed_tinygrad_models) if model != current_model] if DEFAULT_MODEL in models_to_redownload: models_to_redownload.remove(DEFAULT_MODEL) if models_to_redownload: print(f"Redownloading the following models: {', '.join(models_to_redownload)}") self.fetch_models(f"{repo_url}/Versions/model_names_{VERSION}.json", repo_url, boot_run=True) for model in models_to_redownload: if params_memory.get_bool(CANCEL_DOWNLOAD_PARAM): handle_error(None, "Download cancelled...", "Download cancelled...", MODEL_DOWNLOAD_ALL_PARAM, DOWNLOAD_PROGRESS_PARAM) return params_memory.put(DOWNLOAD_PROGRESS_PARAM, f"Downloading \"{self.available_model_names[self.available_models.index(model)]}\"...") self.download_model(model) else: print("No previously installed tinygrad models to redownload") update_frogpilot_toggles() def validate_models(self): current = params.get("Model", encoding="utf-8") default = params_default.get("Model", encoding="utf-8") if current.endswith("_default") and current != default: print(f"Model '{current}' does not match default '{default}', resetting...") params.put("Model", default) if VERSION_PATH.is_file(): version_name = VERSION_PATH.read_text().strip() if version_name != VERSION or int(self.tinygrad_sizes.get(TAR_FILE_NAME, 0)) == 0: self.update_tinygrad_models() self.tinygrad_sizes[TAR_FILE_NAME] = DEFAULT_TINYGRAD_SIZE update_json_file(self.tinygrad_sizes_path, self.tinygrad_sizes) print(f"Updated size for {TAR_FILE_NAME} in {self.tinygrad_sizes_path.name}") params.remove("TinygradUpdateAvailable") VERSION_PATH.write_text(VERSION) print(f"Updated {VERSION_PATH} to {VERSION}") if len(self.available_models) != len(self.available_model_names) or len(self.available_models) != len(self.model_versions): print("Model lists are out of sync. Resetting parameters...") self.available_models = DEFAULT_MODEL self.available_model_names = DEFAULT_MODEL_NAME self.model_versions = DEFAULT_MODEL_VERSION params.put("AvailableModels", self.available_models) params.put("AvailableModelNames", self.available_model_names) params.put("ModelVersions", self.model_versions)