import json import threading from enum import IntEnum from pathlib import Path from openpilot.common.params import Params from openpilot.selfdrive.ui.ui_state import ui_state from openpilot.system.hardware import HARDWARE from openpilot.system.ui.lib.application import gui_app from openpilot.system.ui.widgets import Widget, DialogResult from openpilot.system.ui.widgets.confirm_dialog import ConfirmDialog, alert_dialog from openpilot.system.ui.widgets.list_view import ListItem, ToggleAction, TextAction, ITEM_TEXT_VALUE_COLOR from openpilot.system.ui.widgets.option_dialog import MultiOptionDialog from openpilot.system.ui.widgets.scroller_tici import Scroller from openpilot.frogpilot.common.frogpilot_utilities import clean_model_name from openpilot.frogpilot.common.frogpilot_variables import update_frogpilot_toggles from openpilot.frogpilot.system.ui.widgets.frogpilot_controls import ( FrogPilotButtonsControl, FrogPilotConfirmationDialog, ) MODEL_DIR = Path("/data/models/") TINYGRAD_SUFFIXES = [ "_driving_policy_metadata.pkl", "_driving_policy_tinygrad.pkl", "_driving_vision_metadata.pkl", "_driving_vision_tinygrad.pkl", ] class SubPanel(IntEnum): MAIN = 0 MODEL_LABELS = 1 def has_all_tinygrad_files(model_key: str) -> bool: """Check if a model has all required tinygrad files.""" for suffix in TINYGRAD_SUFFIXES: if not (MODEL_DIR / f"{model_key}{suffix}").exists(): return False return True class FrogPilotModelPanel(Widget): def __init__(self): super().__init__() self._current_panel = SubPanel.MAIN self._params = Params() self._params_memory = Params("", True) self._toggles = {} self._tuning_level = 0 # State tracking self._all_models_downloaded = False self._all_models_downloading = False self._cancelling_download = False self._current_model = "" self._default_model = "" self._finalizing_download = False self._model_downloading = False self._no_models_downloaded = False self._online = False self._parked = True self._started = False self._tinygrad_update = False self._updating_tinygrad = False # Model mappings self._available_model_names: list[str] = [] self._model_file_to_name: dict[str, str] = {} self._model_file_to_name_processed: dict[str, str] = {} # Get default model default_model_bytes = self._params.get_key_default_value("DrivingModel") self._default_model = default_model_bytes.decode() if default_model_bytes else "" self._build_main_panel() self._build_model_labels_panel() ui_state.add_offroad_transition_callback(self._update_toggles) def _build_main_panel(self): self._auto_download_item = ListItem( title="Automatically Download New Models", description="Automatically download new driving models as they become available.", action_item=ToggleAction( initial_state=self._params.get_bool("AutomaticallyDownloadModels"), callback=lambda state: self._simple_toggle("AutomaticallyDownloadModels", state), ), ) self._delete_model_control = FrogPilotButtonsControl( "Delete Driving Models", "Delete downloaded driving models to free up storage space.", "", button_texts=["DELETE", "DELETE ALL"], ) self._delete_model_control.set_click_callback(self._on_delete_model_click) self._download_model_control = FrogPilotButtonsControl( "Download Driving Models", "Manually download driving models to the device.", "", button_texts=["DOWNLOAD", "DOWNLOAD ALL"], ) self._download_model_control.set_click_callback(self._on_download_model_click) self._model_randomizer_item = ListItem( title="Model Randomizer", description="Select a random driving model each drive and use feedback prompts at the end of the drive to help find the model that best suits you!", action_item=ToggleAction( initial_state=self._params.get_bool("ModelRandomizer"), callback=self._on_model_randomizer_toggle, ), ) self._manage_blacklist_control = FrogPilotButtonsControl( "Manage Model Blacklist", "Add or remove driving models from the \"Model Randomizer\" blacklist.", "", button_texts=["ADD", "REMOVE", "REMOVE ALL"], ) self._manage_blacklist_control.set_click_callback(self._on_manage_blacklist_click) self._manage_scores_control = FrogPilotButtonsControl( "Manage Model Ratings", "View or reset saved model ratings used by the \"Model Randomizer\".", "", button_texts=["RESET", "VIEW"], ) self._manage_scores_control.set_click_callback(self._on_manage_scores_click) self._select_model_item = ListItem( title="Select Driving Model", description="Choose which driving model openpilot uses.", action_item=TextAction(lambda: self._get_current_model_display(), color=ITEM_TEXT_VALUE_COLOR), callback=self._on_select_model_click, ) self._update_tinygrad_control = FrogPilotButtonsControl( "Update Model Manager", "Update the \"Model Manager\" to support the latest models.", "", button_texts=["UPDATE"], ) self._update_tinygrad_control.set_click_callback(self._on_update_tinygrad_click) main_items = [ self._auto_download_item, self._delete_model_control, self._download_model_control, self._model_randomizer_item, self._manage_blacklist_control, self._manage_scores_control, self._select_model_item, self._update_tinygrad_control, ] self._toggles["AutomaticallyDownloadModels"] = self._auto_download_item self._toggles["DeleteModel"] = self._delete_model_control self._toggles["DownloadModel"] = self._download_model_control self._toggles["ModelRandomizer"] = self._model_randomizer_item self._toggles["ManageBlacklistedModels"] = self._manage_blacklist_control self._toggles["ManageScores"] = self._manage_scores_control self._toggles["SelectModel"] = self._select_model_item self._toggles["UpdateTinygrad"] = self._update_tinygrad_control self._main_scroller = Scroller(main_items, line_separator=True, spacing=0) def _build_model_labels_panel(self): self._model_labels_items: list[ListItem] = [] self._model_labels_scroller = Scroller(self._model_labels_items, line_separator=True, spacing=0) def _get_current_model_display(self) -> str: """Get the display string for the current model.""" display = self._current_model model_key = clean_model_name(self._params.get("DrivingModel", encoding="utf-8") or "") if model_key == self._default_model: display += " (Default)" return display def _simple_toggle(self, param: str, state: bool): self._params.put_bool(param, state) update_frogpilot_toggles() def _on_model_randomizer_toggle(self, state: bool): self._params.put_bool("ModelRandomizer", state) update_frogpilot_toggles() self._update_toggles() if state and not self._all_models_downloaded: gui_app.set_modal_overlay(ConfirmDialog( "The \"Model Randomizer\" works only with downloaded models. Download all models now?", "Yes", "No", )) def _on_delete_model_click(self, button_id: int): deletable_models = self._get_deletable_models() if not deletable_models: gui_app.set_modal_overlay(alert_dialog("No models available to delete.")) return if button_id == 0: # Delete single model gui_app.set_modal_overlay(MultiOptionDialog( "Select a driving model to delete", deletable_models, )) elif button_id == 1: # Delete all models gui_app.set_modal_overlay(ConfirmDialog( "Are you sure you want to delete all of your downloaded driving models?", "Delete", "Cancel", )) def _get_deletable_models(self) -> list[str]: """Get list of models that can be deleted (excludes current and default).""" deletable = [] if not MODEL_DIR.exists(): return deletable for file in MODEL_DIR.iterdir(): if not file.is_file(): continue base = file.stem for model_key in self._model_file_to_name_processed: if base.startswith(model_key): model_name = self._model_file_to_name_processed[model_key] if model_name not in deletable: deletable.append(model_name) break # Remove current model and default model from deletable list current_clean = clean_model_name(self._current_model) if current_clean in deletable: deletable.remove(current_clean) default_name = self._model_file_to_name_processed.get(clean_model_name(self._default_model), "") if default_name in deletable: deletable.remove(default_name) deletable.sort() return deletable def _delete_model(self, model_name: str): """Delete a specific model's files.""" model_file = None for key, name in self._model_file_to_name_processed.items(): if name == model_name: model_file = key break if not model_file or not MODEL_DIR.exists(): return for file in MODEL_DIR.iterdir(): if file.is_file() and file.stem.startswith(model_file): file.unlink() self._all_models_downloaded = False self._update_deletable_state() def _delete_all_models(self): """Delete all deletable models.""" deletable = self._get_deletable_models() if not MODEL_DIR.exists(): return for file in MODEL_DIR.iterdir(): if not file.is_file(): continue base = file.stem for model_key in self._model_file_to_name_processed: model_name = self._model_file_to_name_processed[model_key] if model_name in deletable and base.startswith(model_key): file.unlink() break self._all_models_downloaded = False self._no_models_downloaded = True self._update_deletable_state() def _update_deletable_state(self): """Update the enabled state of delete buttons.""" deletable = self._get_deletable_models() self._no_models_downloaded = len(deletable) == 0 can_delete = not (self._all_models_downloading or self._model_downloading or self._no_models_downloaded) self._delete_model_control.set_enabled(can_delete) def _on_download_model_click(self, button_id: int): if self._tinygrad_update: gui_app.set_modal_overlay(ConfirmDialog( "Tinygrad is out of date and must be updated before you can download new models. Update now?", "Yes", "No", )) return if button_id == 0: # Download single model or cancel if self._model_downloading: self._params_memory.put_bool("CancelModelDownload", True) self._cancelling_download = True else: downloadable = self._get_downloadable_models() if not downloadable: gui_app.set_modal_overlay(alert_dialog("All models are already downloaded.")) return gui_app.set_modal_overlay(MultiOptionDialog( "Select a driving model to download", downloadable, )) elif button_id == 1: # Download all or cancel if self._all_models_downloading: self._params_memory.put_bool("CancelModelDownload", True) self._cancelling_download = True else: self._params_memory.put_bool("DownloadAllModels", True) self._params_memory.put("ModelDownloadProgress", "Downloading...") self._download_model_control.set_text(1, "CANCEL") self._download_model_control.set_visible_button(0, False) self._all_models_downloading = True def _get_downloadable_models(self) -> list[str]: """Get list of models that can be downloaded.""" downloadable = list(self._available_model_names) for model_key in self._model_file_to_name: model_name = self._model_file_to_name[model_key] if has_all_tinygrad_files(model_key): if model_name in downloadable: downloadable.remove(model_name) downloadable.sort() return downloadable def _start_model_download(self, model_name: str): """Start downloading a specific model.""" model_key = None for key, name in self._model_file_to_name.items(): if name == model_name: model_key = key break if model_key: self._params_memory.put("ModelToDownload", model_key) self._params_memory.put("ModelDownloadProgress", "Downloading...") self._download_model_control.set_text(0, "CANCEL") self._download_model_control.set_visible_button(1, False) self._model_downloading = True def _on_manage_blacklist_click(self, button_id: int): blacklisted_str = self._params.get("BlacklistedModels", encoding="utf-8") or "" blacklisted = [m for m in blacklisted_str.split(",") if m] if button_id == 0: # Add to blacklist blacklistable = [] for model_key in self._model_file_to_name_processed: if model_key not in blacklisted: blacklistable.append(self._model_file_to_name_processed[model_key]) if len(blacklistable) <= 1: remaining = blacklistable[0] if blacklistable else "None" gui_app.set_modal_overlay(alert_dialog( f"There are no more driving models to blacklist. The only available model is \"{remaining}\"!" )) return blacklistable.sort() gui_app.set_modal_overlay(MultiOptionDialog( "Select a driving model to add to the blacklist", blacklistable, )) elif button_id == 1: # Remove from blacklist whitelistable = [] for model_key in blacklisted: model_name = self._model_file_to_name_processed.get(model_key, "") if model_name: whitelistable.append(model_name) if not whitelistable: gui_app.set_modal_overlay(alert_dialog("No models are currently blacklisted.")) return whitelistable.sort() gui_app.set_modal_overlay(MultiOptionDialog( "Select a driving model to remove from the blacklist", whitelistable, )) elif button_id == 2: # Remove all from blacklist if not blacklisted: gui_app.set_modal_overlay(alert_dialog("No models are currently blacklisted.")) return gui_app.set_modal_overlay(ConfirmDialog( "Are you sure you want to remove all of your blacklisted driving models?", "Yes", "No", )) def _add_to_blacklist(self, model_name: str): """Add a model to the blacklist.""" model_key = None for key, name in self._model_file_to_name_processed.items(): if name == model_name: model_key = key break if model_key: blacklisted_str = self._params.get("BlacklistedModels", encoding="utf-8") or "" blacklisted = [m for m in blacklisted_str.split(",") if m] if model_key not in blacklisted: blacklisted.append(model_key) self._params.put("BlacklistedModels", ",".join(blacklisted)) def _remove_from_blacklist(self, model_name: str): """Remove a model from the blacklist.""" model_key = None for key, name in self._model_file_to_name_processed.items(): if name == model_name: model_key = key break if model_key: blacklisted_str = self._params.get("BlacklistedModels", encoding="utf-8") or "" blacklisted = [m for m in blacklisted_str.split(",") if m] if model_key in blacklisted: blacklisted.remove(model_key) self._params.put("BlacklistedModels", ",".join(blacklisted)) def _clear_blacklist(self): """Clear all models from blacklist.""" self._params.remove("BlacklistedModels") def _on_manage_scores_click(self, button_id: int): if button_id == 0: # Reset scores gui_app.set_modal_overlay(ConfirmDialog( "Reset all model drives and ratings? This clears your drive history and collected feedback!", "Yes", "No", )) elif button_id == 1: # View scores self._update_model_labels() self._current_panel = SubPanel.MODEL_LABELS def _reset_model_scores(self): """Reset all model drives and scores.""" self._params.remove("ModelDrivesAndScores") def _update_model_labels(self): """Update the model labels panel with current ratings.""" self._model_labels_items.clear() scores_str = self._params.get("ModelDrivesAndScores", encoding="utf-8") or "{}" try: model_drives_and_scores = json.loads(scores_str) except json.JSONDecodeError: model_drives_and_scores = {} for model_name in sorted(self._available_model_names): clean_name = clean_model_name(model_name) model_data = model_drives_and_scores.get(clean_name, {}) drives = model_data.get("Drives", 0) score = model_data.get("Score", 0) if drives == 1: drives_display = f"{drives} Drive" elif drives > 0: drives_display = f"{drives} Drives" else: drives_display = "N/A" if drives > 0: score_display = f"Score: {score}%" else: score_display = "N/A" label_text = f"{score_display} ({drives_display})" item = ListItem( title=clean_name, action_item=TextAction(label_text, color=ITEM_TEXT_VALUE_COLOR), ) self._model_labels_items.append(item) self._model_labels_scroller = Scroller(self._model_labels_items, line_separator=True, spacing=0) def _on_select_model_click(self): selectable = [] for model_key in self._model_file_to_name: if model_key != clean_model_name(self._default_model) and has_all_tinygrad_files(model_key): selectable.append(self._model_file_to_name[model_key]) selectable.sort() # Add default model at the beginning default_name = self._model_file_to_name.get(clean_model_name(self._default_model), "") if default_name: selectable.insert(0, f"{default_name} (Default)") current_display = self._current_model model_key = clean_model_name(self._params.get("DrivingModel", encoding="utf-8") or "") if model_key == self._default_model: current_display += " (Default)" gui_app.set_modal_overlay(MultiOptionDialog( "Select a Model", selectable, current_display, )) def _select_model(self, model_name: str): """Select a driving model.""" model_name = model_name.replace(" (Default)", "") self._current_model = model_name model_key = None for key, name in self._model_file_to_name.items(): if name == model_name: model_key = key break if model_key: self._params.put("DrivingModel", model_key) update_frogpilot_toggles() if self._started: gui_app.set_modal_overlay(ConfirmDialog( "Reboot required to take effect.", "Reboot Now", "Reboot Later", )) self._update_deletable_state() def _on_update_tinygrad_click(self, button_id: int): if self._updating_tinygrad: self._params_memory.put_bool("CancelModelDownload", True) self._update_tinygrad_control.set_enabled(False) self._cancelling_download = True else: gui_app.set_modal_overlay(ConfirmDialog( "Updating Tinygrad will delete existing Tinygrad-based driving models and need to be re-downloaded. Proceed?", "Yes", "No", )) def _start_tinygrad_update(self): """Start the tinygrad update process.""" self._params_memory.put_bool("UpdateTinygrad", True) self._params_memory.put("ModelDownloadProgress", "Downloading...") self._update_tinygrad_control.set_text(0, "CANCEL") self._updating_tinygrad = True def _translate_progress(self, progress: str) -> str: """Translate download progress messages.""" translations = { "Downloading...": "Downloading...", "Downloaded!": "Downloaded!", "All models downloaded!": "All models downloaded!", "Repository unavailable": "Repository unavailable", } if progress in translations: return translations[progress] progress_lower = progress.lower() if "cancelled" in progress_lower: return "Download cancelled..." if "failed" in progress_lower: return "Download failed..." if "offline" in progress_lower: return "GitHub and GitLab are offline..." return progress def _update_download_state(self): """Update UI based on download progress.""" if self._finalizing_download: return progress = self._params_memory.get("ModelDownloadProgress", encoding="utf-8") or "" if self._all_models_downloading or self._model_downloading: import re download_failed = bool(re.search(r"cancelled|exists|failed|missing|offline", progress, re.IGNORECASE)) translated = self._translate_progress(progress) if progress in ("All models downloaded!", "Downloaded!") or download_failed: self._finalizing_download = True def finalize(): self._all_models_downloading = False self._cancelling_download = False self._finalizing_download = False self._model_downloading = False self._no_models_downloaded = False # Update all models downloaded state downloadable = self._get_downloadable_models() self._all_models_downloaded = len(downloadable) == 0 self._params_memory.remove("ModelDownloadProgress") self._download_model_control.set_enabled(True) self._download_model_control.set_text(0, "DOWNLOAD") self._download_model_control.set_text(1, "DOWNLOAD ALL") self._download_model_control.set_visible_button(0, True) self._download_model_control.set_visible_button(1, True) threading.Timer(2.5, finalize).start() if self._updating_tinygrad: import re download_failed = bool(re.search(r"cancelled|exists|failed|missing|offline", progress, re.IGNORECASE)) translated = self._translate_progress(progress) if progress == "Updated!" or download_failed: self._finalizing_download = True def finalize_tinygrad(): check_progress = self._params_memory.get("ModelDownloadProgress", encoding="utf-8") or "" self._model_downloading = bool(check_progress) if self._model_downloading: self._download_model_control.set_text(1, "CANCEL") self._download_model_control.set_visible_button(0, False) else: self._cancelling_download = False self._tinygrad_update = self._params.get_bool("TinygradUpdateAvailable") self._finalizing_download = False self._updating_tinygrad = False self._update_tinygrad_control.set_enabled(self._tinygrad_update) self._update_tinygrad_control.set_text(0, "UPDATE") threading.Timer(2.5, finalize_tinygrad).start() def _update_button_states(self): """Update button enabled/visible states.""" can_delete = not (self._all_models_downloading or self._model_downloading or self._no_models_downloaded) self._delete_model_control.set_enabled(can_delete) # Download buttons self._download_model_control.set_text(0, "CANCEL" if self._model_downloading else "DOWNLOAD") self._download_model_control.set_text(1, "CANCEL" if self._all_models_downloading else "DOWNLOAD ALL") can_download_single = (not self._all_models_downloaded and not self._all_models_downloading and not self._cancelling_download and not self._finalizing_download and not self._updating_tinygrad and self._online and self._parked) can_download_all = (not self._all_models_downloaded and not self._model_downloading and not self._cancelling_download and not self._finalizing_download and not self._updating_tinygrad and self._online and self._parked) self._download_model_control.set_enabled_buttons(0, can_download_single) self._download_model_control.set_enabled_buttons(1, can_download_all) self._download_model_control.set_visible_button(0, not self._all_models_downloading) self._download_model_control.set_visible_button(1, not self._model_downloading) # Tinygrad update button can_update = (not self._model_downloading and not self._cancelling_download and not self._finalizing_download and self._online and self._parked and self._tinygrad_update) self._update_tinygrad_control.set_enabled(can_update) def _update_toggles(self): self._tuning_level = self._params.get_int("TuningLevel") or 0 model_randomizer = self._params.get_bool("ModelRandomizer") # ManageBlacklistedModels and ManageScores only visible when ModelRandomizer enabled if hasattr(self._manage_blacklist_control, "set_visible"): self._manage_blacklist_control.set_visible(model_randomizer) if hasattr(self._manage_scores_control, "set_visible"): self._manage_scores_control.set_visible(model_randomizer) # SelectModel only visible when ModelRandomizer disabled if hasattr(self._select_model_item, "set_visible"): self._select_model_item.set_visible(not model_randomizer) def _load_model_data(self): """Load available models and current state.""" self._all_models_downloading = self._params_memory.get_bool("DownloadAllModels") progress = self._params_memory.get("ModelDownloadProgress", encoding="utf-8") or "" self._model_downloading = bool(progress) self._tinygrad_update = self._params.get_bool("TinygradUpdateAvailable") self._updating_tinygrad = self._params_memory.get_bool("UpdateTinygrad") self._model_downloading = self._model_downloading and not self._updating_tinygrad # Load available models available_models_str = self._params.get("AvailableModels", encoding="utf-8") or "" available_models = sorted([m for m in available_models_str.split(",") if m]) available_names_str = self._params.get("AvailableModelNames", encoding="utf-8") or "" self._available_model_names = sorted([m for m in available_names_str.split(",") if m]) # Build mappings self._model_file_to_name.clear() self._model_file_to_name_processed.clear() for i in range(min(len(available_models), len(self._available_model_names))): model_key = available_models[i] model_name = self._available_model_names[i] self._model_file_to_name[model_key] = model_name self._model_file_to_name_processed[model_key] = clean_model_name(model_name) # Check downloadable models downloadable = self._get_downloadable_models() self._all_models_downloaded = len(downloadable) == 0 # Check deletable models self._update_deletable_state() # Get current model model_key = clean_model_name(self._params.get("DrivingModel", encoding="utf-8") or "") if not has_all_tinygrad_files(model_key): model_key = self._default_model self._current_model = self._model_file_to_name.get(model_key, "") def _close_sub_panel(self): self._current_panel = SubPanel.MAIN def show_event(self): super().show_event() self._main_scroller.show_event() self._load_model_data() self._update_toggles() self._started = ui_state.started def hide_event(self): super().hide_event() self._current_panel = SubPanel.MAIN def _render(self, rect): # Update online/parked state self._started = ui_state.started self._parked = not self._started # Simplified - in real impl check frogpilot_scene.parked # Update download state self._update_download_state() self._update_button_states() if self._current_panel == SubPanel.MODEL_LABELS: self._model_labels_scroller.render(rect) else: self._main_scroller.render(rect)