diff --git a/docs/v1-scope.md b/docs/v1-scope.md index 3c341b2..1123f2a 100644 --- a/docs/v1-scope.md +++ b/docs/v1-scope.md @@ -77,7 +77,8 @@ Nothing else gates V1. When that sentence is true end to end, V1 is done. - **Model Manager outdated-weight cleanup, 2026-09-04:** although the Model Manager UI remains deferred, it gained the outdated-weight notice and `Delete outdated` as the agreed handling for a rembg namespace move (issue - #22). Bulk re-download is deliberately not part of it. + #22). Bulk re-download returned in #29 with confirmation, cancellation, batch + progress and per-model deletion after verification. --- diff --git a/scripts/guardrails-baseline.json b/scripts/guardrails-baseline.json index e448e5b..c0f07f5 100644 --- a/scripts/guardrails-baseline.json +++ b/scripts/guardrails-baseline.json @@ -20,7 +20,7 @@ "long_functions": 2 }, "src/matteloop/jobs/models/session.py": { - "lines": 432, + "lines": 439, "long_functions": 1 }, "src/matteloop/jobs/render.py": { diff --git a/src/matteloop/jobs/models/session.py b/src/matteloop/jobs/models/session.py index 90874b9..37d777b 100644 --- a/src/matteloop/jobs/models/session.py +++ b/src/matteloop/jobs/models/session.py @@ -192,7 +192,12 @@ def prepare(self, model_id: str, extras: dict[str, object]) -> PreparationResult self._active_result = result return result - def fetch(self, model_id: str, progress: ProgressCallback | None = None) -> Path: + def fetch( + self, + model_id: str, + progress: ProgressCallback | None = None, + cancelled: CancellationCheck | None = None, + ) -> Path: """Download and verify one weight without starting a session for it. prepare() also launches the segmentation client, which is wrong for @@ -209,7 +214,7 @@ def fetch(self, model_id: str, progress: ProgressCallback | None = None) -> Path spec, self._cache_root, self._progress if progress is None else progress, - self._cancelled, + self._cancelled if cancelled is None else cancelled, ) def remove(self, model_id: str) -> bool: @@ -270,12 +275,14 @@ def remove(self, model_id: str) -> bool: ) raise mapped from cause - def remove_obsolete_versions(self) -> int: + def remove_obsolete_versions(self, model_id: str | None = None) -> int: with self._lock: self._require_open() removed = 0 for version in self._catalog.obsolete_rembg_versions: version_path = self._cache_root / version + if model_id is not None: + version_path = version_path / self._catalog.get(model_id).id if not version_path.is_dir() or version_path.is_symlink(): continue shutil.rmtree(version_path) diff --git a/src/matteloop/ui/model_manager.py b/src/matteloop/ui/model_manager.py deleted file mode 100644 index a064c58..0000000 --- a/src/matteloop/ui/model_manager.py +++ /dev/null @@ -1,753 +0,0 @@ -"""Model-cache dialog and controller for the existing session manager.""" - -from __future__ import annotations - -import os -import stat -import time -from collections.abc import Callable -from dataclasses import dataclass -from pathlib import Path -from typing import Protocol - -from PySide6.QtCore import QObject, QSize, Qt, QThread, QUrl, Signal, Slot -from PySide6.QtGui import QDesktopServices -from PySide6.QtWidgets import ( - QDialog, - QHBoxLayout, - QLabel, - QListWidget, - QListWidgetItem, - QMessageBox, - QPushButton, - QVBoxLayout, - QWidget, -) - -from matteloop.core.parameters import V1_MODEL_IDS -from matteloop.core.state import JobState, ModelAvailabilityChanged -from matteloop.jobs.models.catalog import ModelCatalog -from matteloop.ui.aligned_rows import ( - ROW_DATA_ROLE, - AlignedColumn, - AlignedRow, - AlignedRowDelegate, - install_aligned_row, -) -from matteloop.ui.ports import StateStore -from matteloop.ui.source_presentation import ( - DownloadRateEstimator, - format_model_download_detail, - format_source_file_size, -) - -MODEL_ENTRY_ROLE = ROW_DATA_ROLE + 11 -_DOWNLOAD_BLOCKED_MESSAGE = "Cannot download a model while a job is running." - - -class ModelRemovalService(Protocol): - @property - def active_id(self) -> str | None: ... - - def remove(self, model_id: str) -> bool: ... - - def remove_obsolete_versions(self) -> int: ... - - def fetch(self, model_id: str, progress: object = None) -> object: ... - - -@dataclass(frozen=True, slots=True) -class ModelEntry: - """One catalog model plus the cache facts shown by the manager.""" - - model_id: str - display_name: str - download_size_bytes: int - artifact_path: Path - disk_size_bytes: int | None - active: bool - outdated_size_bytes: int | None = None - outdated_rembg_version: str | None = None - - @property - def cached(self) -> bool: - return self.disk_size_bytes is not None - - -def present_model(entry: ModelEntry) -> AlignedRow: - """Present one model with aligned metadata and spoken status words.""" - if entry.cached: - cache_status = "cached locally" - cache_detail = cache_status - if entry.outdated_rembg_version is not None: - cache_detail += "; outdated copy from rembg " + entry.outdated_rembg_version - glyph = "◆" if entry.active else "✓" - elif entry.outdated_size_bytes is not None: - cache_status = "outdated weight" - version = entry.outdated_rembg_version or "obsolete rembg" - cache_detail = ( - f"{cache_status} from rembg {version}; " - f"{format_source_file_size(entry.outdated_size_bytes)} on disk" - ) - glyph = "⟳" - else: - cache_status = "not cached" - cache_detail = cache_status - glyph = "◆" if entry.active else "↓" - active_status = "active model" if entry.active else "not active" - size = format_source_file_size(entry.download_size_bytes) - detail = f"{entry.display_name}; {size}; {cache_detail}; {active_status}" - return AlignedRow( - glyph, - "cached" if entry.cached else "uncached", - ( - AlignedColumn(entry.display_name), - AlignedColumn(size, True), - AlignedColumn(cache_status), - AlignedColumn(active_status), - ), - detail, - ) - - -class ModelManagerDialog(QDialog): - """List the V1 model cache and expose actions for its selected row.""" - - download_requested = Signal(object) - remove_requested = Signal(object) - delete_outdated_requested = Signal() - show_cache_requested = Signal() - - def __init__( - self, - catalog: ModelCatalog, - cache_root: Path, - *, - active_model: Callable[[], str | None] | None = None, - parent: QWidget | None = None, - ) -> None: - super().__init__(parent) - if type(catalog) is not ModelCatalog: - raise TypeError("catalog must be an exact ModelCatalog") - if not isinstance(cache_root, Path): - raise TypeError("cache_root must be a Path") - self.setObjectName("model_manager") - self.setWindowTitle("Model manager") - self.setAccessibleName("Model manager") - self.setModal(True) - self.resize(1080, 460) - self._catalog = catalog - self._cache_root = cache_root - self._active_model = active_model or (lambda: None) - self._entries: tuple[ModelEntry, ...] = () - self._obsolete_size_bytes = 0 - self._busy = False - self._removal_guard: Callable[[ModelEntry], str | None] = ( - self._active_removal_guard - ) - self._build_widgets() - self._build_layout() - self._connect_signals() - self._update_actions() - - def _build_widgets(self) -> None: - self._message = QLabel() - self._message.setObjectName("model_manager_message") - self._message.setProperty("secondary", True) - self.total_size_label = QLabel() - self.total_size_label.setObjectName("model_cache_total") - self.total_size_label.setAccessibleName("Total model cache size") - self.cache_location_label = QLabel(str(self._cache_root)) - self.cache_location_label.setObjectName("model_cache_location") - self.cache_location_label.setAccessibleName("Model cache location") - self.cache_location_label.setToolTip(str(self._cache_root)) - self.cache_location_label.setAccessibleDescription(str(self._cache_root)) - self.model_list = QListWidget() - self.model_list.setObjectName("model_list") - self.model_list.setAccessibleName("V1 models") - self.model_list.setSelectionMode(QListWidget.SelectionMode.SingleSelection) - self.model_list.setItemDelegate(AlignedRowDelegate(self.model_list)) - self.model_list.setUniformItemSizes(True) - self.download_button = QPushButton("Download weight") - self.download_button.setObjectName("download_model") - self.download_button.setAccessibleName("Download selected model weight") - self.remove_button = QPushButton("Remove downloaded weight") - self.remove_button.setObjectName("remove_model") - self.remove_button.setAccessibleName("Remove selected model weight") - self.outdated_notice_label = QLabel() - self.outdated_notice_label.setObjectName("outdated_model_notice") - self.outdated_notice_label.setWordWrap(True) - self.outdated_notice_label.setVisible(False) - self.delete_outdated_button = QPushButton("Delete outdated") - self.delete_outdated_button.setObjectName("delete_outdated") - self.delete_outdated_button.setAccessibleName("Delete outdated model weights") - self.show_cache_button = QPushButton("Show cache location") - self.show_cache_button.setObjectName("show_model_cache") - self.show_cache_button.setAccessibleName("Show model cache location") - self.close_button = QPushButton("Close") - self.close_button.setAccessibleName("Close model manager") - - def _build_layout(self) -> None: - layout = QVBoxLayout(self) - layout.addWidget(QLabel("Manage the downloaded V1 model weights.")) - layout.addWidget(self._message) - layout.addWidget(self.total_size_label) - layout.addWidget(self.cache_location_label) - layout.addWidget(self.model_list, 1) - layout.addWidget(self.outdated_notice_label) - actions = QHBoxLayout() - actions.addWidget(self.download_button) - actions.addWidget(self.remove_button) - actions.addWidget(self.delete_outdated_button) - actions.addWidget(self.show_cache_button) - actions.addStretch(1) - actions.addWidget(self.close_button) - layout.addLayout(actions) - - def _connect_signals(self) -> None: - self.model_list.currentItemChanged.connect( - lambda _current, _previous: self._update_actions() - ) - self.download_button.clicked.connect(self._download_selected) - self.remove_button.clicked.connect(self._remove_selected) - self.delete_outdated_button.clicked.connect(self.delete_outdated_requested.emit) - self.show_cache_button.clicked.connect(self.show_cache_requested.emit) - self.close_button.clicked.connect(self.close) - - @property - def cache_root(self) -> Path: - return self._cache_root - - @property - def entries(self) -> tuple[ModelEntry, ...]: - return self._entries - - @property - def outdated_entries(self) -> tuple[ModelEntry, ...]: - return tuple( - entry for entry in self._entries if entry.outdated_size_bytes is not None - ) - - @property - def obsolete_size_bytes(self) -> int: - return self._obsolete_size_bytes - - @property - def obsolete_versions(self) -> tuple[str, ...]: - return self._catalog.obsolete_rembg_versions - - @property - def selected_entry(self) -> ModelEntry | None: - item = self.model_list.currentItem() - if item is None: - return None - value = item.data(MODEL_ENTRY_ROLE) - return value if isinstance(value, ModelEntry) else None - - def set_removal_guard(self, guard: Callable[[ModelEntry], str | None]) -> None: - self._removal_guard = guard - self._update_actions() - - def set_busy(self, busy: bool) -> None: - self._busy = busy - self.model_list.setEnabled(not busy) - self.download_button.setEnabled(not busy) - self.show_cache_button.setEnabled(not busy) - self.close_button.setEnabled(not busy) - self._update_actions() - - def set_message(self, message: str) -> None: - self._message.setText(message) - self._message.setAccessibleDescription(message) - - def _active_removal_guard(self, entry: ModelEntry) -> str | None: - if self._active_model() == entry.model_id: - return "Cannot remove the model used by the active session." - return None - - def refresh(self) -> None: - """Read only the thirteen catalog paths and rebuild the aligned rows.""" - selected_id = self.selected_entry.model_id if self.selected_entry else None - active_id = self._active_model() - active_id = active_id if isinstance(active_id, str) else None - self._entries = tuple( - _model_entry(self._catalog, self._cache_root, model_id, active_id) - for model_id in V1_MODEL_IDS - ) - self._obsolete_size_bytes = _obsolete_directory_size( - self._cache_root, self._catalog.obsolete_rembg_versions - ) - self.model_list.clear() - for entry in self._entries: - row = present_model(entry) - item = QListWidgetItem(row.display_text) - install_aligned_row(item, row) - item.setData(MODEL_ENTRY_ROLE, entry) - item.setSizeHint(QSize(0, 46)) - self.model_list.addItem(item) - if selected_id is not None: - selected_index = next( - ( - index - for index, entry in enumerate(self._entries) - if entry.model_id == selected_id - ), - -1, - ) - if selected_index >= 0: - self.model_list.setCurrentRow(selected_index) - if self.model_list.currentItem() is None and self.model_list.count(): - self.model_list.setCurrentRow(0) - total = sum(entry.disk_size_bytes or 0 for entry in self._entries) - total += self._obsolete_size_bytes - self.total_size_label.setText( - f"Total on disk: {format_source_file_size(total)}" - ) - self.total_size_label.setAccessibleDescription( - f"Total on disk: {format_source_file_size(total)}" - ) - self.set_message(f"{len(self._entries)} V1 model(s); cache: {self._cache_root}") - self._update_outdated_notice(self._obsolete_size_bytes) - self._update_actions() - - def _update_outdated_notice(self, size_bytes: int) -> None: - if size_bytes <= 0: - self.outdated_notice_label.clear() - self.outdated_notice_label.setVisible(False) - return - versions = ", ".join(self._catalog.obsolete_rembg_versions) - self.outdated_notice_label.setText( - f"Outdated weights from rembg {versions} occupy " - f"{format_source_file_size(size_bytes)} on disk and cannot be used by " - "this version." - ) - self.outdated_notice_label.setVisible(True) - - def _update_actions(self) -> None: - entry = self.selected_entry - blocked = self._removal_guard(entry) if entry is not None else None - enabled = ( - not self._busy and entry is not None and entry.cached and blocked is None - ) - self.remove_button.setEnabled(enabled) - if blocked is not None: - self.remove_button.setToolTip(blocked) - self.remove_button.setAccessibleDescription(blocked) - else: - self.remove_button.setToolTip( - "Remove the selected downloaded weight and free its disk space." - ) - self.remove_button.setAccessibleDescription("") - self.download_button.setEnabled( - not self._busy and entry is not None and not entry.cached - ) - self.download_button.setToolTip( - "Download the selected weight now instead of waiting for a preview." - ) - has_outdated = self._obsolete_size_bytes > 0 - self.delete_outdated_button.setVisible(has_outdated) - self.delete_outdated_button.setEnabled(not self._busy and has_outdated) - - def _download_selected(self) -> None: - entry = self.selected_entry - if entry is not None and not entry.cached: - self.download_requested.emit(entry) - - def _remove_selected(self) -> None: - entry = self.selected_entry - if entry is not None and self._removal_guard(entry) is None: - self.remove_requested.emit(entry) - - -class _ModelRemovalWorker(QObject): - removed = Signal(str, bool) - failed = Signal(str, object) - finished = Signal() - - def __init__(self, manager: ModelRemovalService, model_id: str) -> None: - super().__init__() - self._manager = manager - self._model_id = model_id - - @Slot() - def run(self) -> None: - try: - self.removed.emit(self._model_id, self._manager.remove(self._model_id)) - except Exception as error: - self.failed.emit(self._model_id, error) - finally: - self.finished.emit() - - -class _ModelDownloadWorker(QObject): - downloaded = Signal(str) - failed = Signal(str, object) - progressed = Signal(int, int) - finished = Signal() - - def __init__(self, manager: ModelRemovalService, model_id: str) -> None: - super().__init__() - self._manager = manager - self._model_id = model_id - - @Slot() - def run(self) -> None: - try: - self._manager.fetch(self._model_id, self._report) - except Exception as error: - self.failed.emit(self._model_id, error) - else: - self.downloaded.emit(self._model_id) - finally: - self.finished.emit() - - def _report(self, completed: int, total: int) -> None: - self.progressed.emit(completed, total) - - -class _ObsoleteRemovalWorker(QObject): - removed = Signal(int) - failed = Signal(object) - finished = Signal() - - def __init__(self, manager: ModelRemovalService) -> None: - super().__init__() - self._manager = manager - - @Slot() - def run(self) -> None: - try: - self.removed.emit(self._manager.remove_obsolete_versions()) - except Exception as error: - self.failed.emit(error) - finally: - self.finished.emit() - - -class ModelManagerController(QObject): - """Own model-manager actions while keeping widgets on the GUI thread.""" - - def __init__( - self, - store: StateStore, - *, - catalog: ModelCatalog, - cache_root: Path, - manager: ModelRemovalService | None, - dialog_parent: QWidget | None = None, - ) -> None: - super().__init__(dialog_parent) - self._store = store - self._manager = manager - active_model = ( - manager_active_id if manager is None else lambda: manager.active_id - ) - self.dialog = ModelManagerDialog( - catalog, - cache_root, - active_model=active_model, - parent=dialog_parent, - ) - self.dialog.set_removal_guard(self._removal_block_reason) - self.dialog.download_requested.connect(self._download_requested) - self.dialog.remove_requested.connect(self._remove_requested) - self.dialog.delete_outdated_requested.connect(self._delete_outdated_requested) - self.dialog.show_cache_requested.connect(self._show_cache) - self._remove_thread: QThread | None = None - self._remove_worker: ( - _ModelRemovalWorker | _ModelDownloadWorker | _ObsoleteRemovalWorker | None - ) = None - self._removal_entry: ModelEntry | None = None - - def set_dialog_parent(self, parent: QWidget) -> None: - self.dialog.setParent(parent, Qt.WindowType.Dialog) - - def open(self) -> None: - self.dialog.refresh() - self.dialog.open() - - def close(self) -> None: - thread = self._remove_thread - if thread is not None: - thread.quit() - thread.wait(5000) - self._remove_thread = None - self.dialog.close() - - def _removal_block_reason(self, entry: ModelEntry) -> str | None: - if self._store.state.job.phase is not JobState.IDLE: - return "Cannot remove a model while a job is running." - if self._manager is None: - return "Model removal is unavailable in this runtime." - if self._manager.active_id == entry.model_id: - return "Cannot remove the model used by the active session." - return None - - def _download_requested(self, value: object) -> None: - manager = self._manager - if not isinstance(value, ModelEntry) or value.cached or manager is None: - return - if self._remove_thread or self._store.state.job.phase is not JobState.IDLE: - if self._remove_thread is None: - self.dialog.set_message(_DOWNLOAD_BLOCKED_MESSAGE) - return - worker = _ModelDownloadWorker(manager, value.model_id) - thread = QThread(self) - worker.moveToThread(thread) - self._download_rate = DownloadRateEstimator() - worker.progressed.connect(self._download_progressed) - worker.downloaded.connect(self._download_succeeded) - worker.failed.connect(self._download_failed) - worker.finished.connect(thread.quit) - worker.finished.connect(worker.deleteLater) - thread.finished.connect(self._removal_thread_finished) - thread.finished.connect(thread.deleteLater) - self._remove_thread = thread - self._remove_worker = worker - self._removal_entry = value - self.dialog.set_busy(True) - self.dialog.set_message(f"Downloading {value.display_name}…") - thread.started.connect(worker.run) - thread.start() - - def _delete_outdated_requested(self) -> None: - manager = self._manager - if self.dialog.obsolete_size_bytes <= 0: - return - if self._store.state.job.phase is not JobState.IDLE: - self.dialog.set_message("Cannot remove a model while a job is running.") - return - if manager is None: - self.dialog.set_message("Model removal is unavailable in this runtime.") - return - versions = ", ".join(self.dialog.obsolete_versions) - directory_word = ( - "directory" if len(self.dialog.obsolete_versions) == 1 else "directories" - ) - size = format_source_file_size(self.dialog.obsolete_size_bytes) - answer = QMessageBox.question( - self.dialog, - f"Delete all weights from rembg {versions}?", - f"This removes the whole {versions} {directory_word}: {size} on disk. " - "Weights this version needs are downloaded again on demand.", - QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, - QMessageBox.StandardButton.No, - ) - if answer is not QMessageBox.StandardButton.Yes: - return - self._start_obsolete_removal(manager) - - def _start_obsolete_removal(self, manager: ModelRemovalService) -> None: - worker = _ObsoleteRemovalWorker(manager) - thread = QThread(self) - worker.moveToThread(thread) - worker.removed.connect(self._obsolete_removal_succeeded) - worker.failed.connect(self._obsolete_removal_failed) - worker.finished.connect(thread.quit) - worker.finished.connect(worker.deleteLater) - thread.finished.connect(self._removal_thread_finished) - thread.finished.connect(thread.deleteLater) - self._remove_thread = thread - self._remove_worker = worker - self._removal_entry = None - self.dialog.set_busy(True) - self.dialog.set_message("Removing outdated model weights…") - thread.started.connect(worker.run) - thread.start() - - @Slot(int, int) - def _download_progressed(self, completed: int, total: int) -> None: - entry = self._removal_entry - if entry is None: - return - speed = self._download_rate.update(completed, time.monotonic()) - self.dialog.set_message( - format_model_download_detail(entry.display_name, completed, total, speed) - ) - - @Slot(str) - def _download_succeeded(self, model_id: str) -> None: - entry = self._removal_entry - if entry is None or entry.model_id != model_id: - return - self.dialog.refresh() - self.dialog.set_message(f"Downloaded {entry.display_name}.") - if self._store.state.parameters.model_id == model_id: - self._store.dispatch(ModelAvailabilityChanged(True)) - - @Slot(str, object) - def _download_failed(self, model_id: str, error: object) -> None: - entry = self._removal_entry - if entry is None or entry.model_id != model_id: - return - self.dialog.set_message(f"Could not download the weight: {error}") - QMessageBox.warning(self.dialog, "Could not download model", str(error)) - - @Slot(int) - def _obsolete_removal_succeeded(self, removed: int) -> None: - self.dialog.refresh() - self.dialog.set_message( - f"Removed {removed} outdated model version(s) from the cache." - ) - - @Slot(object) - def _obsolete_removal_failed(self, error: object) -> None: - self.dialog.refresh() - self.dialog.set_message( - f"Could not remove outdated weights: {type(error).__name__}: {error}" - ) - - def _remove_requested(self, value: object) -> None: - if not isinstance(value, ModelEntry) or not value.cached: - return - if self._removal_block_reason(value) is not None: - return - if self._remove_thread is not None: - return - answer = QMessageBox.question( - self.dialog, - "Remove model weight?", - f"Remove {value.display_name}'s downloaded weight?\n" - f"This frees {format_source_file_size(value.disk_size_bytes or 0)}.", - QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, - QMessageBox.StandardButton.No, - ) - if answer is not QMessageBox.StandardButton.Yes: - return - manager = self._manager - if manager is None: - return - self._start_removal(manager, value) - - def _start_removal(self, manager: ModelRemovalService, entry: ModelEntry) -> None: - worker = _ModelRemovalWorker(manager, entry.model_id) - thread = QThread(self) - worker.moveToThread(thread) - worker.removed.connect(self._removal_succeeded) - worker.failed.connect(self._removal_failed) - worker.finished.connect(thread.quit) - worker.finished.connect(worker.deleteLater) - thread.finished.connect(self._removal_thread_finished) - thread.finished.connect(thread.deleteLater) - self._remove_thread = thread - self._remove_worker = worker - self._removal_entry = entry - self.dialog.set_busy(True) - thread.started.connect(worker.run) - thread.start() - - @Slot(str, bool) - def _removal_succeeded(self, model_id: str, removed: bool) -> None: - entry = self._removal_entry - if entry is None: - return - if entry.model_id != model_id: - return - if not removed: - self.dialog.refresh() - self.dialog.set_message( - f"No weight could be removed at {entry.artifact_path}." - ) - return - freed = format_source_file_size(entry.disk_size_bytes or 0) - self.dialog.refresh() - self.dialog.set_message(f"Removed {entry.display_name}; freed {freed}.") - if self._store.state.parameters.model_id == model_id: - self._store.dispatch(ModelAvailabilityChanged(False)) - QMessageBox.information( - self.dialog, - "Model removed", - f"Removed {entry.display_name}'s downloaded weight.\nFreed {freed}.", - ) - - @Slot(str, object) - def _removal_failed(self, model_id: str, error: object) -> None: - entry = self._removal_entry - if entry is None or entry.model_id != model_id: - return - self.dialog.set_message(f"Could not remove the selected weight: {error}") - QMessageBox.warning(self.dialog, "Could not remove model", str(error)) - - @Slot() - def _removal_thread_finished(self) -> None: - self.dialog.set_busy(False) - self._remove_worker = None - self._remove_thread = None - - @Slot() - def _show_cache(self) -> None: - # Nothing has created the cache directory before the first download, - # and openUrl on a missing directory fails silently. - root = self.dialog.cache_root - try: - root.mkdir(parents=True, exist_ok=True) - except OSError as error: - self.dialog.set_message(f"Could not open the model folder: {error}") - return - if not QDesktopServices.openUrl(QUrl.fromLocalFile(str(root))): - self.dialog.set_message(f"Model folder: {root}") - - -def manager_active_id() -> str | None: - """Default active-model callback used by a read-only fallback dialog.""" - return None - - -def _model_entry( - catalog: ModelCatalog, - cache_root: Path, - model_id: str, - active_id: str | None, -) -> ModelEntry: - spec = catalog.get(model_id) - artifact = spec.artifact - if artifact is None: - raise ValueError(f"V1 model {model_id!r} has no artifact") - artifact_path = ( - cache_root / catalog.rembg_version / model_id / artifact.runtime_filename - ) - disk_size_bytes = _regular_file_size(artifact_path) - outdated_size_bytes: int | None = None - outdated_rembg_version: str | None = None - for version in catalog.obsolete_rembg_versions: - outdated_path = cache_root / version / model_id / artifact.runtime_filename - outdated_size_bytes = _regular_file_size(outdated_path) - if outdated_size_bytes is not None: - outdated_rembg_version = version - break - return ModelEntry( - model_id=model_id, - display_name=spec.display_name, - download_size_bytes=artifact.size_bytes, - artifact_path=artifact_path, - disk_size_bytes=disk_size_bytes, - active=model_id == active_id, - outdated_size_bytes=outdated_size_bytes, - outdated_rembg_version=outdated_rembg_version, - ) - - -def _obsolete_directory_size(cache_root: Path, versions: tuple[str, ...]) -> int: - return sum(_directory_size(cache_root / version) for version in versions) - - -def _directory_size(path: Path) -> int: - try: - info = path.lstat() - except OSError: - return 0 - if not stat.S_ISDIR(info.st_mode): - return 0 - total = 0 - for root, _directories, filenames in os.walk(path, followlinks=False): - for filename in filenames: - total += _regular_file_size(Path(root) / filename) or 0 - return total - - -def _regular_file_size(path: Path) -> int | None: - try: - info = path.lstat() - except OSError: - return None - return info.st_size if stat.S_ISREG(info.st_mode) else None diff --git a/src/matteloop/ui/model_manager/__init__.py b/src/matteloop/ui/model_manager/__init__.py new file mode 100644 index 0000000..e335e03 --- /dev/null +++ b/src/matteloop/ui/model_manager/__init__.py @@ -0,0 +1,12 @@ +"""Model-cache dialog and controller for the existing session manager.""" + +from __future__ import annotations + +from ._controller import ModelManagerController, ModelRemovalService +from ._dialog import ModelManagerDialog + +__all__ = ( + "ModelManagerController", + "ModelManagerDialog", + "ModelRemovalService", +) diff --git a/src/matteloop/ui/model_manager/_controller.py b/src/matteloop/ui/model_manager/_controller.py new file mode 100644 index 0000000..15bd60b --- /dev/null +++ b/src/matteloop/ui/model_manager/_controller.py @@ -0,0 +1,465 @@ +from __future__ import annotations + +import time +from pathlib import Path +from typing import Protocol + +from PySide6.QtCore import QObject, Qt, QThread, QUrl, Slot +from PySide6.QtGui import QDesktopServices +from PySide6.QtWidgets import QMessageBox, QWidget + +from matteloop.core.errors import AppError, ErrorCode +from matteloop.core.state import JobState, ModelAvailabilityChanged +from matteloop.jobs.context import CancellationState +from matteloop.jobs.models.catalog import ModelCatalog +from matteloop.ui.ports import StateStore +from matteloop.ui.source_presentation import ( + DownloadRateEstimator, + format_model_download_detail, + format_source_file_size, +) + +from ._dialog import ModelManagerDialog +from ._entries import ModelEntry, manager_active_id +from ._workers import ( + _ModelDownloadWorker, + _ModelRemovalWorker, + _ObsoleteRemovalWorker, + _OutdatedRedownloadWorker, +) + +_DOWNLOAD_BLOCKED_MESSAGE = "Cannot download a model while a job is running." + + +def _confirm_redownload( + dialog: ModelManagerDialog, entries: tuple[ModelEntry, ...] +) -> QMessageBox.StandardButton: + count = len(entries) + size = format_source_file_size(sum(entry.download_size_bytes for entry in entries)) + versions = ", ".join(dialog.obsolete_versions) + directory_word = ( + "directory" if len(dialog.obsolete_versions) == 1 else "directories" + ) + return QMessageBox.question( + dialog, + "Re-download outdated model weights?", + f"Download {count} model weight(s) ({size})? " + f"Each outdated copy from rembg {versions} is deleted once its replacement " + "has been verified. " + f"Other files in the {versions} {directory_word} stay until you use " + "Delete outdated.", + QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, + QMessageBox.StandardButton.No, + ) + + +class ModelRemovalService(Protocol): + @property + def active_id(self) -> str | None: ... + + def remove(self, model_id: str) -> bool: ... + + def remove_obsolete_versions(self, model_id: str | None = None) -> int: ... + + def fetch( + self, model_id: str, progress: object = None, cancelled: object = None + ) -> object: ... + + +class ModelManagerController(QObject): + """Own model-manager actions while keeping widgets on the GUI thread.""" + + def __init__( + self, + store: StateStore, + *, + catalog: ModelCatalog, + cache_root: Path, + manager: ModelRemovalService | None, + dialog_parent: QWidget | None = None, + ) -> None: + super().__init__(dialog_parent) + self._store = store + self._manager = manager + active_model = ( + manager_active_id if manager is None else lambda: manager.active_id + ) + self.dialog = ModelManagerDialog( + catalog, + cache_root, + active_model=active_model, + parent=dialog_parent, + ) + self.dialog.set_removal_guard(self._removal_block_reason) + self.dialog.download_requested.connect(self._download_requested) + self.dialog.remove_requested.connect(self._remove_requested) + self.dialog.redownload_outdated_requested.connect( + self._redownload_outdated_requested + ) + self.dialog.cancel_requested.connect(self._cancel_requested) + self.dialog.delete_outdated_requested.connect(self._delete_outdated_requested) + self.dialog.show_cache_requested.connect(self._show_cache) + self._remove_thread: QThread | None = None + self._remove_worker: ( + _ModelRemovalWorker + | _ModelDownloadWorker + | _ObsoleteRemovalWorker + | _OutdatedRedownloadWorker + | None + ) = None + self._removal_entry: ModelEntry | None = None + self._cancel: CancellationState | None = None + self._batch: dict[str, ModelEntry] = {} + self._batch_done = 0 + self._batch_failure: str | None = None + self._batch_position = "" + + def set_dialog_parent(self, parent: QWidget) -> None: + self.dialog.setParent(parent, Qt.WindowType.Dialog) + + def open(self) -> None: + self.dialog.refresh() + self.dialog.open() + + def close(self) -> None: + cancel = self._cancel + if cancel is not None: + cancel.request() + thread = self._remove_thread + if thread is not None: + thread.quit() + thread.wait() + self._remove_thread = None + self._cancel = None + self.dialog.set_busy(False) + self.dialog.close() + + def _removal_block_reason(self, entry: ModelEntry) -> str | None: + if self._store.state.job.phase is not JobState.IDLE: + return "Cannot remove a model while a job is running." + if self._manager is None: + return "Model removal is unavailable in this runtime." + if self._manager.active_id == entry.model_id: + return "Cannot remove the model used by the active session." + return None + + def _download_requested(self, value: object) -> None: + manager = self._manager + if not isinstance(value, ModelEntry) or value.cached or manager is None: + return + if self._remove_thread or self._store.state.job.phase is not JobState.IDLE: + if self._remove_thread is None: + self.dialog.set_message(_DOWNLOAD_BLOCKED_MESSAGE) + return + cancel = CancellationState() + worker = _ModelDownloadWorker(manager, value.model_id, cancel) + thread = QThread(self) + worker.moveToThread(thread) + self._download_rate = DownloadRateEstimator() + worker.progressed.connect(self._download_progressed) + worker.downloaded.connect(self._download_succeeded) + worker.failed.connect(self._download_failed) + worker.finished.connect(thread.quit) + worker.finished.connect(worker.deleteLater) + thread.finished.connect(self._removal_thread_finished) + thread.finished.connect(thread.deleteLater) + self._remove_thread = thread + self._remove_worker = worker + self._removal_entry = value + self._cancel = cancel + self._batch_position = "" + self.dialog.set_busy(True, cancellable=True) + self.dialog.set_message(f"Downloading {value.display_name}…") + thread.started.connect(worker.run) + thread.start() + + def _delete_outdated_requested(self) -> None: + manager = self._manager + if self.dialog.obsolete_size_bytes <= 0: + return + if self._store.state.job.phase is not JobState.IDLE: + self.dialog.set_message("Cannot remove a model while a job is running.") + return + if manager is None: + self.dialog.set_message("Model removal is unavailable in this runtime.") + return + versions = ", ".join(self.dialog.obsolete_versions) + directory_word = ( + "directory" if len(self.dialog.obsolete_versions) == 1 else "directories" + ) + size = format_source_file_size(self.dialog.obsolete_size_bytes) + answer = QMessageBox.question( + self.dialog, + f"Delete all weights from rembg {versions}?", + f"This removes the whole {versions} {directory_word}: {size} on disk. " + "Weights this version needs are downloaded again on demand.", + QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, + QMessageBox.StandardButton.No, + ) + if answer is not QMessageBox.StandardButton.Yes: + return + self._start_obsolete_removal(manager) + + def _redownload_outdated_requested(self) -> None: + entries = self.dialog.redownload_entries + if not entries: + return + if self._store.state.job.phase is not JobState.IDLE: + self.dialog.set_message(_DOWNLOAD_BLOCKED_MESSAGE) + return + manager = self._manager + if manager is None: + self.dialog.set_message("Model removal is unavailable in this runtime.") + return + if self._remove_thread is not None: + return + if ( + _confirm_redownload(self.dialog, entries) + is not QMessageBox.StandardButton.Yes + ): + return + cancel = CancellationState() + self._batch = {entry.model_id: entry for entry in entries} + worker = _OutdatedRedownloadWorker(manager, tuple(self._batch), cancel) + thread = QThread(self) + worker.moveToThread(thread) + worker.started.connect(self._batch_model_started) + worker.progressed.connect(self._download_progressed) + worker.downloaded.connect(self._batch_model_downloaded) + worker.failed.connect(self._batch_model_failed) + worker.finished.connect(self._batch_finished) + worker.finished.connect(thread.quit) + worker.finished.connect(worker.deleteLater) + thread.finished.connect(self._removal_thread_finished) + thread.finished.connect(thread.deleteLater) + self._remove_thread = thread + self._remove_worker = worker + self._cancel = cancel + self._removal_entry = None + self._batch_done = 0 + self._batch_failure = None + self._batch_position = "" + self.dialog.set_busy(True, cancellable=True) + self.dialog.set_message( + f"Re-downloading {len(self._batch)} outdated model weight(s)…" + ) + thread.started.connect(worker.run) + thread.start() + + @Slot() + def _cancel_requested(self) -> None: + cancel = self._cancel + if cancel is None or not cancel.request(): + return + self.dialog.set_cancelling() + + def _start_obsolete_removal(self, manager: ModelRemovalService) -> None: + worker = _ObsoleteRemovalWorker(manager) + thread = QThread(self) + worker.moveToThread(thread) + worker.removed.connect(self._obsolete_removal_succeeded) + worker.failed.connect(self._obsolete_removal_failed) + worker.finished.connect(thread.quit) + worker.finished.connect(worker.deleteLater) + thread.finished.connect(self._removal_thread_finished) + thread.finished.connect(thread.deleteLater) + self._remove_thread = thread + self._remove_worker = worker + self._removal_entry = None + self.dialog.set_busy(True) + self.dialog.set_message("Removing outdated model weights…") + thread.started.connect(worker.run) + thread.start() + + @Slot(int, int) + def _download_progressed(self, completed: int, total: int) -> None: + if self._cancel is not None and self._cancel.requested: + return + entry = self._removal_entry + if entry is None: + return + speed = self._download_rate.update(completed, time.monotonic()) + self.dialog.set_message( + format_model_download_detail( + entry.display_name + self._batch_position, + completed, + total, + speed, + ) + ) + + @Slot(str) + def _batch_model_started(self, model_id: str) -> None: + entry = self._batch.get(model_id) + if entry is None: + return + self._removal_entry = entry + self._batch_position = f" ({self._batch_done + 1} of {len(self._batch)})" + self._download_rate = DownloadRateEstimator() + self.dialog.set_message( + f"Downloading {entry.display_name}{self._batch_position}…" + ) + + @Slot(str) + def _batch_model_downloaded(self, model_id: str) -> None: + entry = self._batch.get(model_id) + if entry is None: + return + self.dialog.set_message( + f"Downloaded {entry.display_name}{self._batch_position}." + ) + self._batch_done += 1 + if self._store.state.parameters.model_id == model_id: + self._store.dispatch(ModelAvailabilityChanged(True)) + + @Slot(str, object) + def _batch_model_failed(self, model_id: str, error: object) -> None: + entry = self._batch.get(model_id) + if entry is not None: + self._batch_failure = f"Could not download {entry.display_name}: {error}." + + @Slot() + def _batch_finished(self) -> None: + self.dialog.refresh() + if self._batch_failure is not None: + message = ( + f"{self._batch_failure} Re-downloaded {self._batch_done} of " + f"{len(self._batch)} outdated model weight(s)." + ) + elif self._cancel is not None and self._cancel.requested: + message = ( + f"Download cancelled. Re-downloaded {self._batch_done} of " + f"{len(self._batch)} outdated model weight(s)." + ) + else: + message = f"Re-downloaded {len(self._batch)} outdated model weight(s)." + self.dialog.set_message(message) + + @Slot(str) + def _download_succeeded(self, model_id: str) -> None: + entry = self._removal_entry + if entry is None or entry.model_id != model_id: + return + self.dialog.refresh() + self.dialog.set_message(f"Downloaded {entry.display_name}.") + if self._store.state.parameters.model_id == model_id: + self._store.dispatch(ModelAvailabilityChanged(True)) + + @Slot(str, object) + def _download_failed(self, model_id: str, error: object) -> None: + entry = self._removal_entry + if entry is None or entry.model_id != model_id: + return + if isinstance(error, AppError) and error.code is ErrorCode.JOB_CANCELLED: + self.dialog.set_message("Download cancelled.") + return + self.dialog.set_message(f"Could not download the weight: {error}") + QMessageBox.warning(self.dialog, "Could not download model", str(error)) + + @Slot(int) + def _obsolete_removal_succeeded(self, removed: int) -> None: + self.dialog.refresh() + self.dialog.set_message( + f"Removed {removed} outdated model version(s) from the cache." + ) + + @Slot(object) + def _obsolete_removal_failed(self, error: object) -> None: + self.dialog.refresh() + self.dialog.set_message( + f"Could not remove outdated weights: {type(error).__name__}: {error}" + ) + + def _remove_requested(self, value: object) -> None: + if not isinstance(value, ModelEntry) or not value.cached: + return + if self._removal_block_reason(value) is not None: + return + if self._remove_thread is not None: + return + answer = QMessageBox.question( + self.dialog, + "Remove model weight?", + f"Remove {value.display_name}'s downloaded weight?\n" + f"This frees {format_source_file_size(value.disk_size_bytes or 0)}.", + QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, + QMessageBox.StandardButton.No, + ) + if answer is not QMessageBox.StandardButton.Yes: + return + manager = self._manager + if manager is None: + return + self._start_removal(manager, value) + + def _start_removal(self, manager: ModelRemovalService, entry: ModelEntry) -> None: + worker = _ModelRemovalWorker(manager, entry.model_id) + thread = QThread(self) + worker.moveToThread(thread) + worker.removed.connect(self._removal_succeeded) + worker.failed.connect(self._removal_failed) + worker.finished.connect(thread.quit) + worker.finished.connect(worker.deleteLater) + thread.finished.connect(self._removal_thread_finished) + thread.finished.connect(thread.deleteLater) + self._remove_thread = thread + self._remove_worker = worker + self._removal_entry = entry + self.dialog.set_busy(True) + thread.started.connect(worker.run) + thread.start() + + @Slot(str, bool) + def _removal_succeeded(self, model_id: str, removed: bool) -> None: + entry = self._removal_entry + if entry is None: + return + if entry.model_id != model_id: + return + if not removed: + self.dialog.refresh() + self.dialog.set_message( + f"No weight could be removed at {entry.artifact_path}." + ) + return + freed = format_source_file_size(entry.disk_size_bytes or 0) + self.dialog.refresh() + self.dialog.set_message(f"Removed {entry.display_name}; freed {freed}.") + if self._store.state.parameters.model_id == model_id: + self._store.dispatch(ModelAvailabilityChanged(False)) + QMessageBox.information( + self.dialog, + "Model removed", + f"Removed {entry.display_name}'s downloaded weight.\nFreed {freed}.", + ) + + @Slot(str, object) + def _removal_failed(self, model_id: str, error: object) -> None: + entry = self._removal_entry + if entry is None or entry.model_id != model_id: + return + self.dialog.set_message(f"Could not remove the selected weight: {error}") + QMessageBox.warning(self.dialog, "Could not remove model", str(error)) + + @Slot() + def _removal_thread_finished(self) -> None: + self.dialog.set_busy(False) + self._remove_worker = None + self._remove_thread = None + self._cancel = None + self._batch = {} + self._batch_position = "" + self._removal_entry = None + + @Slot() + def _show_cache(self) -> None: + # Nothing has created the cache directory before the first download, + # and openUrl on a missing directory fails silently. + root = self.dialog.cache_root + try: + root.mkdir(parents=True, exist_ok=True) + except OSError as error: + self.dialog.set_message(f"Could not open the model folder: {error}") + return + if not QDesktopServices.openUrl(QUrl.fromLocalFile(str(root))): + self.dialog.set_message(f"Model folder: {root}") diff --git a/src/matteloop/ui/model_manager/_dialog.py b/src/matteloop/ui/model_manager/_dialog.py new file mode 100644 index 0000000..6c43110 --- /dev/null +++ b/src/matteloop/ui/model_manager/_dialog.py @@ -0,0 +1,327 @@ +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path + +from PySide6.QtCore import QSize, Signal +from PySide6.QtWidgets import ( + QDialog, + QHBoxLayout, + QLabel, + QListWidget, + QListWidgetItem, + QPushButton, + QVBoxLayout, + QWidget, +) + +from matteloop.core.parameters import V1_MODEL_IDS +from matteloop.jobs.models.catalog import ModelCatalog +from matteloop.ui.aligned_rows import AlignedRowDelegate, install_aligned_row +from matteloop.ui.source_presentation import format_source_file_size + +from ._entries import ( + MODEL_ENTRY_ROLE, + ModelEntry, + _model_entry, + _obsolete_directory_size, + present_model, +) + + +class ModelManagerDialog(QDialog): + """List the V1 model cache and expose actions for its selected row.""" + + download_requested = Signal(object) + remove_requested = Signal(object) + delete_outdated_requested = Signal() + redownload_outdated_requested = Signal() + cancel_requested = Signal() + show_cache_requested = Signal() + + def __init__( + self, + catalog: ModelCatalog, + cache_root: Path, + *, + active_model: Callable[[], str | None] | None = None, + parent: QWidget | None = None, + ) -> None: + super().__init__(parent) + if type(catalog) is not ModelCatalog: + raise TypeError("catalog must be an exact ModelCatalog") + if not isinstance(cache_root, Path): + raise TypeError("cache_root must be a Path") + self.setObjectName("model_manager") + self.setWindowTitle("Model manager") + self.setAccessibleName("Model manager") + self.setModal(True) + self.resize(1080, 460) + self._catalog = catalog + self._cache_root = cache_root + self._active_model = active_model or (lambda: None) + self._entries: tuple[ModelEntry, ...] = () + self._obsolete_size_bytes = 0 + self._busy = False + self._removal_guard: Callable[[ModelEntry], str | None] = ( + self._active_removal_guard + ) + self._build_widgets() + self._build_layout() + self._connect_signals() + self._update_actions() + + def _build_widgets(self) -> None: + self._message = QLabel() + self._message.setObjectName("model_manager_message") + self._message.setProperty("secondary", True) + self.total_size_label = QLabel() + self.total_size_label.setObjectName("model_cache_total") + self.total_size_label.setAccessibleName("Total model cache size") + self.cache_location_label = QLabel(str(self._cache_root)) + self.cache_location_label.setObjectName("model_cache_location") + self.cache_location_label.setAccessibleName("Model cache location") + self.cache_location_label.setToolTip(str(self._cache_root)) + self.cache_location_label.setAccessibleDescription(str(self._cache_root)) + self.model_list = QListWidget() + self.model_list.setObjectName("model_list") + self.model_list.setAccessibleName("V1 models") + self.model_list.setSelectionMode(QListWidget.SelectionMode.SingleSelection) + self.model_list.setItemDelegate(AlignedRowDelegate(self.model_list)) + self.model_list.setUniformItemSizes(True) + self.download_button = QPushButton("Download weight") + self.download_button.setObjectName("download_model") + self.download_button.setAccessibleName("Download selected model weight") + self.remove_button = QPushButton("Remove downloaded weight") + self.remove_button.setObjectName("remove_model") + self.remove_button.setAccessibleName("Remove selected model weight") + self.redownload_outdated_button = QPushButton("Re-download outdated") + self.redownload_outdated_button.setObjectName("redownload_outdated") + self.redownload_outdated_button.setAccessibleName( + "Re-download outdated model weights" + ) + self.outdated_notice_label = QLabel() + self.outdated_notice_label.setObjectName("outdated_model_notice") + self.outdated_notice_label.setWordWrap(True) + self.outdated_notice_label.setVisible(False) + self.delete_outdated_button = QPushButton("Delete outdated") + self.delete_outdated_button.setObjectName("delete_outdated") + self.delete_outdated_button.setAccessibleName("Delete outdated model weights") + self.show_cache_button = QPushButton("Show cache location") + self.show_cache_button.setObjectName("show_model_cache") + self.show_cache_button.setAccessibleName("Show model cache location") + self.close_button = QPushButton("Close") + self.close_button.setAccessibleName("Close model manager") + self.cancel_button = QPushButton("Cancel") + self.cancel_button.setObjectName("cancel_model_download") + self.cancel_button.setAccessibleName("Cancel the running model download") + self.cancel_button.setVisible(False) + + def _build_layout(self) -> None: + layout = QVBoxLayout(self) + layout.addWidget(QLabel("Manage the downloaded V1 model weights.")) + layout.addWidget(self._message) + layout.addWidget(self.total_size_label) + layout.addWidget(self.cache_location_label) + layout.addWidget(self.model_list, 1) + layout.addWidget(self.outdated_notice_label) + actions = QHBoxLayout() + actions.addWidget(self.download_button) + actions.addWidget(self.remove_button) + actions.addWidget(self.redownload_outdated_button) + actions.addWidget(self.delete_outdated_button) + actions.addWidget(self.show_cache_button) + actions.addStretch(1) + actions.addWidget(self.cancel_button) + actions.addWidget(self.close_button) + layout.addLayout(actions) + + def _connect_signals(self) -> None: + self.model_list.currentItemChanged.connect( + lambda _current, _previous: self._update_actions() + ) + self.download_button.clicked.connect(self._download_selected) + self.remove_button.clicked.connect(self._remove_selected) + self.redownload_outdated_button.clicked.connect( + self.redownload_outdated_requested.emit + ) + self.delete_outdated_button.clicked.connect(self.delete_outdated_requested.emit) + self.cancel_button.clicked.connect(self.cancel_requested.emit) + self.show_cache_button.clicked.connect(self.show_cache_requested.emit) + self.close_button.clicked.connect(self.close) + + @property + def cache_root(self) -> Path: + return self._cache_root + + @property + def entries(self) -> tuple[ModelEntry, ...]: + return self._entries + + @property + def outdated_entries(self) -> tuple[ModelEntry, ...]: + return tuple( + entry for entry in self._entries if entry.outdated_size_bytes is not None + ) + + @property + def redownload_entries(self) -> tuple[ModelEntry, ...]: + return tuple( + entry + for entry in self._entries + if entry.outdated_size_bytes is not None and not entry.cached + ) + + @property + def obsolete_size_bytes(self) -> int: + return self._obsolete_size_bytes + + @property + def obsolete_versions(self) -> tuple[str, ...]: + return self._catalog.obsolete_rembg_versions + + @property + def selected_entry(self) -> ModelEntry | None: + item = self.model_list.currentItem() + if item is None: + return None + value = item.data(MODEL_ENTRY_ROLE) + return value if isinstance(value, ModelEntry) else None + + def set_removal_guard(self, guard: Callable[[ModelEntry], str | None]) -> None: + self._removal_guard = guard + self._update_actions() + + def set_busy(self, busy: bool, *, cancellable: bool = False) -> None: + self._busy = busy + self.model_list.setEnabled(not busy) + self.download_button.setEnabled(not busy) + self.show_cache_button.setEnabled(not busy) + self.close_button.setEnabled(not busy) + if busy and cancellable: + self.cancel_button.setText("Cancel") + self.cancel_button.setVisible(True) + self.cancel_button.setEnabled(True) + else: + self.cancel_button.setText("Cancel") + self.cancel_button.setVisible(False) + self.cancel_button.setEnabled(True) + self._update_actions() + + def set_cancelling(self) -> None: + self.cancel_button.setText("Cancelling…") + self.cancel_button.setEnabled(False) + self.set_message("Cancelling…") + + def reject(self) -> None: + if self._busy: + self.cancel_requested.emit() + return + super().reject() + + def set_message(self, message: str) -> None: + self._message.setText(message) + self._message.setAccessibleDescription(message) + + def _active_removal_guard(self, entry: ModelEntry) -> str | None: + if self._active_model() == entry.model_id: + return "Cannot remove the model used by the active session." + return None + + def refresh(self) -> None: + """Read only the thirteen catalog paths and rebuild the aligned rows.""" + selected_id = self.selected_entry.model_id if self.selected_entry else None + active_id = self._active_model() + active_id = active_id if isinstance(active_id, str) else None + self._entries = tuple( + _model_entry(self._catalog, self._cache_root, model_id, active_id) + for model_id in V1_MODEL_IDS + ) + self._obsolete_size_bytes = _obsolete_directory_size( + self._cache_root, self._catalog.obsolete_rembg_versions + ) + self.model_list.clear() + for entry in self._entries: + row = present_model(entry) + item = QListWidgetItem(row.display_text) + install_aligned_row(item, row) + item.setData(MODEL_ENTRY_ROLE, entry) + item.setSizeHint(QSize(0, 46)) + self.model_list.addItem(item) + if selected_id is not None: + selected_index = next( + ( + index + for index, entry in enumerate(self._entries) + if entry.model_id == selected_id + ), + -1, + ) + if selected_index >= 0: + self.model_list.setCurrentRow(selected_index) + if self.model_list.currentItem() is None and self.model_list.count(): + self.model_list.setCurrentRow(0) + total = sum(entry.disk_size_bytes or 0 for entry in self._entries) + total += self._obsolete_size_bytes + self.total_size_label.setText( + f"Total on disk: {format_source_file_size(total)}" + ) + self.total_size_label.setAccessibleDescription( + f"Total on disk: {format_source_file_size(total)}" + ) + self.set_message(f"{len(self._entries)} V1 model(s); cache: {self._cache_root}") + self._update_outdated_notice(self._obsolete_size_bytes) + self._update_actions() + + def _update_outdated_notice(self, size_bytes: int) -> None: + if size_bytes <= 0: + self.outdated_notice_label.clear() + self.outdated_notice_label.setVisible(False) + return + versions = ", ".join(self._catalog.obsolete_rembg_versions) + self.outdated_notice_label.setText( + f"Outdated weights from rembg {versions} occupy " + f"{format_source_file_size(size_bytes)} on disk and cannot be used by " + "this version." + ) + self.outdated_notice_label.setVisible(True) + + def _update_actions(self) -> None: + entry = self.selected_entry + blocked = self._removal_guard(entry) if entry is not None else None + enabled = ( + not self._busy and entry is not None and entry.cached and blocked is None + ) + self.remove_button.setEnabled(enabled) + if blocked is not None: + self.remove_button.setToolTip(blocked) + self.remove_button.setAccessibleDescription(blocked) + else: + self.remove_button.setToolTip( + "Remove the selected downloaded weight and free its disk space." + ) + self.remove_button.setAccessibleDescription("") + self.download_button.setEnabled( + not self._busy and entry is not None and not entry.cached + ) + self.download_button.setToolTip( + "Download the selected weight now instead of waiting for a preview." + ) + has_outdated = self._obsolete_size_bytes > 0 + self.delete_outdated_button.setVisible(has_outdated) + self.delete_outdated_button.setEnabled(not self._busy and has_outdated) + has_redownloadable = bool(self.redownload_entries) + self.redownload_outdated_button.setVisible(has_redownloadable) + self.redownload_outdated_button.setEnabled( + not self._busy and has_redownloadable + ) + + def _download_selected(self) -> None: + entry = self.selected_entry + if entry is not None and not entry.cached: + self.download_requested.emit(entry) + + def _remove_selected(self) -> None: + entry = self.selected_entry + if entry is not None and self._removal_guard(entry) is None: + self.remove_requested.emit(entry) diff --git a/src/matteloop/ui/model_manager/_entries.py b/src/matteloop/ui/model_manager/_entries.py new file mode 100644 index 0000000..0b8be15 --- /dev/null +++ b/src/matteloop/ui/model_manager/_entries.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import os +import stat +from dataclasses import dataclass +from pathlib import Path + +from matteloop.jobs.models.catalog import ModelCatalog +from matteloop.ui.aligned_rows import ROW_DATA_ROLE, AlignedColumn, AlignedRow +from matteloop.ui.source_presentation import format_source_file_size + +MODEL_ENTRY_ROLE = ROW_DATA_ROLE + 11 + + +@dataclass(frozen=True, slots=True) +class ModelEntry: + """One catalog model plus the cache facts shown by the manager.""" + + model_id: str + display_name: str + download_size_bytes: int + artifact_path: Path + disk_size_bytes: int | None + active: bool + outdated_size_bytes: int | None = None + outdated_rembg_version: str | None = None + + @property + def cached(self) -> bool: + return self.disk_size_bytes is not None + + +def present_model(entry: ModelEntry) -> AlignedRow: + """Present one model with aligned metadata and spoken status words.""" + if entry.cached: + cache_status = "cached locally" + cache_detail = cache_status + if entry.outdated_rembg_version is not None: + cache_detail += "; outdated copy from rembg " + entry.outdated_rembg_version + glyph = "◆" if entry.active else "✓" + elif entry.outdated_size_bytes is not None: + cache_status = "outdated weight" + version = entry.outdated_rembg_version or "obsolete rembg" + cache_detail = ( + f"{cache_status} from rembg {version}; " + f"{format_source_file_size(entry.outdated_size_bytes)} on disk" + ) + glyph = "⟳" + else: + cache_status = "not cached" + cache_detail = cache_status + glyph = "◆" if entry.active else "↓" + active_status = "active model" if entry.active else "not active" + size = format_source_file_size(entry.download_size_bytes) + detail = f"{entry.display_name}; {size}; {cache_detail}; {active_status}" + return AlignedRow( + glyph, + "cached" if entry.cached else "uncached", + ( + AlignedColumn(entry.display_name), + AlignedColumn(size, True), + AlignedColumn(cache_status), + AlignedColumn(active_status), + ), + detail, + ) + + +def manager_active_id() -> str | None: + """Default active-model callback used by a read-only fallback dialog.""" + return None + + +def _model_entry( + catalog: ModelCatalog, + cache_root: Path, + model_id: str, + active_id: str | None, +) -> ModelEntry: + spec = catalog.get(model_id) + artifact = spec.artifact + if artifact is None: + raise ValueError(f"V1 model {model_id!r} has no artifact") + artifact_path = ( + cache_root / catalog.rembg_version / model_id / artifact.runtime_filename + ) + disk_size_bytes = _regular_file_size(artifact_path) + outdated_size_bytes: int | None = None + outdated_rembg_version: str | None = None + for version in catalog.obsolete_rembg_versions: + outdated_path = cache_root / version / model_id / artifact.runtime_filename + outdated_size_bytes = _regular_file_size(outdated_path) + if outdated_size_bytes is not None: + outdated_rembg_version = version + break + return ModelEntry( + model_id=model_id, + display_name=spec.display_name, + download_size_bytes=artifact.size_bytes, + artifact_path=artifact_path, + disk_size_bytes=disk_size_bytes, + active=model_id == active_id, + outdated_size_bytes=outdated_size_bytes, + outdated_rembg_version=outdated_rembg_version, + ) + + +def _obsolete_directory_size(cache_root: Path, versions: tuple[str, ...]) -> int: + return sum(_directory_size(cache_root / version) for version in versions) + + +def _directory_size(path: Path) -> int: + try: + info = path.lstat() + except OSError: + return 0 + if not stat.S_ISDIR(info.st_mode): + return 0 + total = 0 + for root, _directories, filenames in os.walk(path, followlinks=False): + for filename in filenames: + total += _regular_file_size(Path(root) / filename) or 0 + return total + + +def _regular_file_size(path: Path) -> int | None: + try: + info = path.lstat() + except OSError: + return None + return info.st_size if stat.S_ISREG(info.st_mode) else None diff --git a/src/matteloop/ui/model_manager/_workers.py b/src/matteloop/ui/model_manager/_workers.py new file mode 100644 index 0000000..e3cd483 --- /dev/null +++ b/src/matteloop/ui/model_manager/_workers.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from PySide6.QtCore import QObject, Signal, Slot + +from matteloop.core.errors import AppError, ErrorCode +from matteloop.jobs.context import CancellationState + +if TYPE_CHECKING: + from ._controller import ModelRemovalService + + +class _ModelRemovalWorker(QObject): + removed = Signal(str, bool) + failed = Signal(str, object) + finished = Signal() + + def __init__(self, manager: ModelRemovalService, model_id: str) -> None: + super().__init__() + self._manager = manager + self._model_id = model_id + + @Slot() + def run(self) -> None: + try: + self.removed.emit(self._model_id, self._manager.remove(self._model_id)) + except Exception as error: + self.failed.emit(self._model_id, error) + finally: + self.finished.emit() + + +class _ModelDownloadWorker(QObject): + downloaded = Signal(str) + failed = Signal(str, object) + progressed = Signal(int, int) + finished = Signal() + + def __init__( + self, manager: ModelRemovalService, model_id: str, cancel: CancellationState + ) -> None: + super().__init__() + self._manager = manager + self._model_id = model_id + self._cancel = cancel + + @Slot() + def run(self) -> None: + try: + self._manager.fetch(self._model_id, self._report, self._cancelled) + except Exception as error: + self.failed.emit(self._model_id, error) + else: + self.downloaded.emit(self._model_id) + finally: + self.finished.emit() + + def _report(self, completed: int, total: int) -> None: + self.progressed.emit(completed, total) + + def _cancelled(self) -> bool: + return self._cancel.requested + + +class _OutdatedRedownloadWorker(QObject): + started = Signal(str) + progressed = Signal(int, int) + downloaded = Signal(str) + failed = Signal(str, object) + finished = Signal() + + def __init__( + self, + manager: ModelRemovalService, + model_ids: tuple[str, ...], + cancel: CancellationState, + ) -> None: + super().__init__() + self._manager = manager + self._model_ids = model_ids + self._cancel = cancel + + @Slot() + def run(self) -> None: + try: + for model_id in self._model_ids: + if self._cancel.requested: + return + self.started.emit(model_id) + try: + self._manager.fetch(model_id, self._report, self._cancelled) + self._manager.remove_obsolete_versions(model_id) + except AppError as error: + if error.code is ErrorCode.JOB_CANCELLED: + return + self.failed.emit(model_id, error) + return + except Exception as error: + self.failed.emit(model_id, error) + return + self.downloaded.emit(model_id) + finally: + self.finished.emit() + + def _report(self, completed: int, total: int) -> None: + self.progressed.emit(completed, total) + + def _cancelled(self) -> bool: + return self._cancel.requested + + +class _ObsoleteRemovalWorker(QObject): + removed = Signal(int) + failed = Signal(object) + finished = Signal() + + def __init__(self, manager: ModelRemovalService) -> None: + super().__init__() + self._manager = manager + + @Slot() + def run(self) -> None: + try: + self.removed.emit(self._manager.remove_obsolete_versions()) + except Exception as error: + self.failed.emit(error) + finally: + self.finished.emit() diff --git a/tests/jobs/models/test_session.py b/tests/jobs/models/test_session.py index d49d6af..3ae7c66 100644 --- a/tests/jobs/models/test_session.py +++ b/tests/jobs/models/test_session.py @@ -36,6 +36,7 @@ def __init__(self, cache_root: Path, events: list[str]) -> None: self.events = events self.fail: BaseException | None = None self.calls: list[str] = [] + self.cancelled_checks: list[Callable[[], bool]] = [] def download( self, @@ -46,6 +47,7 @@ def download( ) -> Path: self.events.append(f"download:{spec.id}") self.calls.append(spec.id) + self.cancelled_checks.append(cancelled) if self.fail is not None: raise self.fail assert destination == self.cache_root @@ -446,6 +448,52 @@ def test_remove_obsolete_versions_removes_only_catalog_listed_directories( assert unrelated.is_dir() +def test_fetch_passes_the_callers_cancellation_check_to_the_downloader( + tmp_path: Path, +) -> None: + manager, downloader, _clients, _events = _manager(tmp_path) + + def own() -> bool: + return False + + manager.fetch("u2netp", cancelled=own) + manager.fetch("u2netp") + + assert downloader.cancelled_checks[0] is own + assert downloader.cancelled_checks[1] is manager._cancelled + + +def test_remove_obsolete_versions_removes_only_the_named_model_directory( + tmp_path: Path, +) -> None: + manager, _downloader, _clients, _events = _manager(tmp_path) + catalog = ModelCatalog.load_resource() + obsolete = tmp_path / catalog.obsolete_rembg_versions[0] + named = obsolete / "u2netp" + other = obsolete / "u2net" + named.mkdir(parents=True) + (named / "u2netp.onnx").write_bytes(b"old") + other.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + symlinked = obsolete / "u2net_human_seg" + symlinked.symlink_to(outside, target_is_directory=True) + + assert manager.remove_obsolete_versions("u2netp") == 1 + assert not named.exists() + assert other.is_dir() + assert obsolete.is_dir() + + assert manager.remove_obsolete_versions("u2net_human_seg") == 0 + assert symlinked.is_symlink() + assert outside.is_dir() + + with pytest.raises(AppError) as unknown: + manager.remove_obsolete_versions("not-a-model") + assert unknown.value.code is ErrorCode.MODEL_NOT_FOUND + assert other.is_dir() + + def test_remove_rejects_symlink_traversal_and_unknown_ids(tmp_path: Path) -> None: manager, _downloader, _clients, _events = _manager(tmp_path) outside = tmp_path / "outside.onnx" diff --git a/tests/ui/test_model_manager.py b/tests/ui/test_model_manager.py index df8b18e..4863c81 100644 --- a/tests/ui/test_model_manager.py +++ b/tests/ui/test_model_manager.py @@ -2,12 +2,15 @@ import shutil from pathlib import Path -from threading import get_ident +from threading import Event, get_ident +import pytest from PySide6.QtCore import QUrl from PySide6.QtGui import QDesktopServices from PySide6.QtWidgets import QMessageBox +from matteloop.core.errors import AppError, ErrorCode +from matteloop.core.parameters import ParameterState from matteloop.core.state import ( ActiveJob, AppState, @@ -28,7 +31,7 @@ def __init__( active_id: str | None = None, targets: dict[str, Path] | None = None, ) -> None: - self.active_id = active_id + self._active_id = active_id self.targets = targets or {} self.removed: list[str] = [] self.fetched: list[str] = [] @@ -36,27 +39,70 @@ def __init__( self.obsolete_roots: tuple[Path, ...] = () self.obsolete_calls = 0 self.operations: list[str] = [] - - def fetch(self, model_id: str, progress=None) -> Path: - self.thread_ids.append(get_ident()) - self.fetched.append(model_id) - self.operations.append(f"fetch:{model_id}") - if progress is not None: - progress(2_000_000, 8_000_000) - target = self.targets.get(model_id) - assert target is not None - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(b"downloaded-weight") - return target - - def remove_obsolete_versions(self) -> int: + self.fetch_cancel_checks: list[bool] = [] + self.in_flight = False + self.block: Event | None = None + self.block_ids: set[str] = set() + self.cancel_observed = Event() + self.fail_ids: dict[str, BaseException] = {} + self.obsolete_block: Event | None = None + + @property + def active_id(self) -> str | None: + """Fail if the GUI reads manager state while fetch holds its lock.""" + if self.in_flight: + raise AssertionError("manager read while a fetch holds its lock") + return self._active_id + + def fetch(self, model_id: str, progress=None, cancelled=None) -> Path: + self.in_flight = True + try: + self.thread_ids.append(get_ident()) + self.fetched.append(model_id) + self.operations.append(f"fetch:{model_id}") + self.fetch_cancel_checks.append(cancelled is not None) + if self.block is not None and ( + not self.block_ids or model_id in self.block_ids + ): + assert cancelled is not None + while not cancelled() and not self.block.is_set(): + self.block.wait(0.01) + if cancelled(): + self.cancel_observed.set() + raise AppError( + ErrorCode.JOB_CANCELLED, + "model-download", + "error.job.cancelled", + "download cancelled", + "retry-model-download", + ) + error = self.fail_ids.get(model_id) + if error is not None: + raise error + if progress is not None: + progress(2_000_000, 8_000_000) + target = self.targets.get(model_id) + assert target is not None + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(b"downloaded-weight") + return target + finally: + self.in_flight = False + + def remove_obsolete_versions(self, model_id: str | None = None) -> int: self.thread_ids.append(get_ident()) self.obsolete_calls += 1 - self.operations.append("remove-obsolete") + self.operations.append( + "remove-obsolete" if model_id is None else f"remove-obsolete:{model_id}" + ) + if self.obsolete_block is not None: + while not self.obsolete_block.wait(0.01): + pass removed = 0 for root in self.obsolete_roots: - if root.is_dir(): - shutil.rmtree(root) + target = root if model_id is None else root / model_id + if target.is_dir(): + shutil.rmtree(target) removed += 1 return removed @@ -330,7 +376,9 @@ def record(message: str) -> None: controller.dialog.download_button.click() + assert controller.dialog.cancel_button.isVisible() qtbot.waitUntil(lambda: manager.fetched == ["u2netp"], timeout=5000) + assert manager.fetch_cancel_checks == [True] assert manager.thread_ids and manager.thread_ids[0] != get_ident() # The dialog reports the same size detail the job popup would, so a # multi-hundred-megabyte download is not a frozen-looking button. @@ -349,6 +397,7 @@ def record(message: str) -> None: timeout=5000, ) assert not controller.dialog.download_button.isEnabled() + assert controller.dialog.cancel_button.isHidden() controller.close() @@ -391,6 +440,9 @@ def test_model_manager_refuses_to_download_while_a_job_runs( controller.dialog.download_button.click() assert manager.fetched == [] + assert controller.dialog._message.text() == ( + "Cannot download a model while a job is running." + ) controller.close() @@ -426,6 +478,7 @@ def test_model_manager_lists_an_outdated_weight_with_size_and_rembg_version( "be used by this version." ) assert dialog.delete_outdated_button.isHidden() is False + assert dialog.redownload_outdated_button.isHidden() is False def test_model_manager_reports_and_deletes_outdated_copy_alongside_current_weight( @@ -529,6 +582,525 @@ def test_model_manager_deletes_outdated_directory_without_touching_other_cache( controller.close() +def test_model_manager_shows_re_download_only_when_an_outdated_row_lacks_a_current_weight( # noqa: E501 + tmp_path: Path, qtbot +) -> None: + outdated = _outdated_model_path(tmp_path, "u2netp") + outdated.parent.mkdir(parents=True) + outdated.write_bytes(b"old-weight") + controller = ModelManagerController( + ReducerStore(), + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=FakeModelRemovalService(), + ) + qtbot.addWidget(controller.dialog) + controller.open() + + assert controller.dialog.redownload_outdated_button.isVisible() + assert controller.dialog.delete_outdated_button.isVisible() + controller.close() + + current = _model_path(tmp_path, "u2netp") + current.parent.mkdir(parents=True) + current.write_bytes(b"current-weight") + cached_controller = ModelManagerController( + ReducerStore(), + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=FakeModelRemovalService(), + ) + qtbot.addWidget(cached_controller.dialog) + cached_controller.open() + + assert cached_controller.dialog.redownload_outdated_button.isHidden() + assert cached_controller.dialog.delete_outdated_button.isVisible() + cached_controller.close() + + +def test_model_manager_confirms_the_download_size_before_re_downloading_outdated_weights( # noqa: E501 + tmp_path: Path, monkeypatch, qtbot +) -> None: + for model_id in ("u2netp", "silueta"): + outdated = _outdated_model_path(tmp_path, model_id) + outdated.parent.mkdir(parents=True) + outdated.write_bytes(b"old-weight") + manager = FakeModelRemovalService( + targets={ + model_id: _model_path(tmp_path, model_id) + for model_id in ("u2netp", "silueta") + } + ) + controller = ModelManagerController( + ReducerStore(), + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=manager, + ) + qtbot.addWidget(controller.dialog) + controller.open() + confirmation: list[tuple[str, str]] = [] + monkeypatch.setattr( + QMessageBox, + "question", + lambda _parent, title, body, *_args: ( + confirmation.append((title, body)) or QMessageBox.StandardButton.No + ), + ) + + controller.dialog.redownload_outdated_button.click() + + assert confirmation == [ + ( + "Re-download outdated model weights?", + "Download 2 model weight(s) (46.5 MiB)? " + "Each outdated copy from rembg 2.0.72 is deleted once its replacement " + "has been verified. Other files in the 2.0.72 directory stay until you " + "use Delete outdated.", + ) + ] + assert manager.fetched == [] + controller.close() + + +def test_model_manager_re_downloads_each_outdated_weight_then_deletes_its_old_copy( + tmp_path: Path, monkeypatch, qtbot +) -> None: + model_ids = ("u2netp", "silueta") + for model_id in model_ids: + outdated = _outdated_model_path(tmp_path, model_id) + outdated.parent.mkdir(parents=True) + outdated.write_bytes(b"old-weight") + unlisted = tmp_path / "2.0.72" / "bria-rmbg" / "model.onnx" + unlisted.parent.mkdir(parents=True) + unlisted.write_bytes(b"unlisted-weight") + manager = FakeModelRemovalService( + targets={model_id: _model_path(tmp_path, model_id) for model_id in model_ids}, + ) + manager.block = Event() + manager.block_ids = {"silueta"} + manager.obsolete_roots = (tmp_path / "2.0.72",) + controller = ModelManagerController( + ReducerStore(AppState(parameters=ParameterState(model_id="u2netp"))), + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=manager, + ) + qtbot.addWidget(controller.dialog) + controller.open() + monkeypatch.setattr( + QMessageBox, + "question", + lambda *_args, **_kwargs: QMessageBox.StandardButton.Yes, + ) + + controller.dialog.redownload_outdated_button.click() + + qtbot.waitUntil( + lambda: "Silueta (2 of 2)" in controller.dialog._message.text(), timeout=5000 + ) + manager.block.set() + qtbot.waitUntil(lambda: controller._remove_thread is None, timeout=5000) + assert manager.operations == [ + "fetch:u2netp", + "remove-obsolete:u2netp", + "fetch:silueta", + "remove-obsolete:silueta", + ] + assert manager.fetch_cancel_checks == [True, True] + assert manager.thread_ids + assert all(thread_id != get_ident() for thread_id in manager.thread_ids) + assert not _outdated_model_path(tmp_path, "u2netp").exists() + assert not _outdated_model_path(tmp_path, "silueta").exists() + assert unlisted.read_bytes() == b"unlisted-weight" + assert controller.dialog._message.text() == ( + "Re-downloaded 2 outdated model weight(s)." + ) + assert controller.dialog.outdated_notice_label.text() == ( + "Outdated weights from rembg 2.0.72 occupy 15.0 B on disk and cannot " + "be used by this version." + ) + assert controller.dialog.redownload_outdated_button.isHidden() + assert controller.dialog.delete_outdated_button.isVisible() + assert controller._store.state.model_available is True + controller.close() + + +def test_model_manager_shows_the_batch_position_while_downloading( + tmp_path: Path, monkeypatch, qtbot +) -> None: + model_ids = ("u2netp", "silueta") + for model_id in model_ids: + outdated = _outdated_model_path(tmp_path, model_id) + outdated.parent.mkdir(parents=True) + outdated.write_bytes(b"old-weight") + manager = FakeModelRemovalService( + targets={model_id: _model_path(tmp_path, model_id) for model_id in model_ids} + ) + manager.obsolete_roots = (tmp_path / "2.0.72",) + controller = ModelManagerController( + ReducerStore(), + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=manager, + ) + qtbot.addWidget(controller.dialog) + controller.open() + messages: list[str] = [] + original = controller.dialog.set_message + + def record(message: str) -> None: + messages.append(message) + original(message) + + controller.dialog.set_message = record # type: ignore[method-assign] + monkeypatch.setattr( + QMessageBox, + "question", + lambda *_args, **_kwargs: QMessageBox.StandardButton.Yes, + ) + + controller.dialog.redownload_outdated_button.click() + + qtbot.waitUntil(lambda: controller._remove_thread is None, timeout=5000) + assert any("Silueta (2 of 2)" in message for message in messages) + assert any("Downloaded U²-Net P (1 of 2)." in message for message in messages) + controller.close() + + +def test_model_manager_cancel_stops_the_batch_and_keeps_the_rest_outdated( + tmp_path: Path, monkeypatch, qtbot +) -> None: + model_ids = ("u2netp", "silueta") + for model_id in model_ids: + outdated = _outdated_model_path(tmp_path, model_id) + outdated.parent.mkdir(parents=True) + outdated.write_bytes(b"old-weight") + manager = FakeModelRemovalService( + targets={model_id: _model_path(tmp_path, model_id) for model_id in model_ids} + ) + manager.block = Event() + manager.block_ids = {"u2netp"} + controller = ModelManagerController( + ReducerStore(), + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=manager, + ) + qtbot.addWidget(controller.dialog) + controller.open() + monkeypatch.setattr( + QMessageBox, + "question", + lambda *_args, **_kwargs: QMessageBox.StandardButton.Yes, + ) + + controller.dialog.redownload_outdated_button.click() + qtbot.waitUntil(lambda: manager.in_flight, timeout=5000) + controller.dialog.cancel_button.click() + + assert controller.dialog.cancel_button.text() == "Cancelling…" + assert controller.dialog.cancel_button.isEnabled() is False + qtbot.waitUntil(manager.cancel_observed.is_set, timeout=5000) + qtbot.waitUntil(lambda: controller._remove_thread is None, timeout=5000) + assert manager.fetched == ["u2netp"] + assert not _model_path(tmp_path, "u2netp").exists() + assert _outdated_model_path(tmp_path, "u2netp").exists() + assert _outdated_model_path(tmp_path, "silueta").exists() + assert controller.dialog._message.text() == ( + "Download cancelled. Re-downloaded 0 of 2 outdated model weight(s)." + ) + assert controller.dialog.outdated_notice_label.isVisible() + assert controller.dialog.redownload_outdated_button.isEnabled() + assert controller.dialog.cancel_button.isHidden() + assert controller.dialog.close_button.isEnabled() + controller.close() + + +@pytest.mark.parametrize("dismissal", ("reject", "close")) +def test_model_manager_escape_cancels_a_running_download_instead_of_hiding_the_dialog( + tmp_path: Path, monkeypatch, qtbot, dismissal: str +) -> None: + outdated = _outdated_model_path(tmp_path, "u2netp") + outdated.parent.mkdir(parents=True) + outdated.write_bytes(b"old-weight") + manager = FakeModelRemovalService( + targets={"u2netp": _model_path(tmp_path, "u2netp")} + ) + manager.block = Event() + manager.block_ids = {"u2netp"} + controller = ModelManagerController( + ReducerStore(), + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=manager, + ) + qtbot.addWidget(controller.dialog) + controller.open() + monkeypatch.setattr( + QMessageBox, + "question", + lambda *_args, **_kwargs: QMessageBox.StandardButton.Yes, + ) + controller.dialog.redownload_outdated_button.click() + qtbot.waitUntil(lambda: manager.in_flight, timeout=5000) + + if dismissal == "reject": + controller.dialog.reject() + else: + controller.dialog.close() + + assert controller.dialog.isVisible() + qtbot.waitUntil(manager.cancel_observed.is_set, timeout=5000) + qtbot.waitUntil(lambda: controller._remove_thread is None, timeout=5000) + assert controller.dialog._message.text().startswith("Download cancelled.") + controller.close() + + +def test_model_manager_ignores_dismissal_while_a_removal_runs( + tmp_path: Path, monkeypatch, qtbot +) -> None: + outdated = _outdated_model_path(tmp_path, "u2netp") + outdated.parent.mkdir(parents=True) + outdated.write_bytes(b"old-weight") + manager = FakeModelRemovalService() + manager.obsolete_roots = (tmp_path / "2.0.72",) + manager.obsolete_block = Event() + controller = ModelManagerController( + ReducerStore(), + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=manager, + ) + qtbot.addWidget(controller.dialog) + controller.open() + monkeypatch.setattr( + QMessageBox, + "question", + lambda *_args, **_kwargs: QMessageBox.StandardButton.Yes, + ) + controller.dialog.delete_outdated_button.click() + qtbot.waitUntil(lambda: manager.obsolete_calls == 1, timeout=5000) + + controller.dialog.reject() + + assert controller.dialog.isVisible() + assert controller._cancel is None + manager.obsolete_block.set() + qtbot.waitUntil(lambda: controller._remove_thread is None, timeout=5000) + controller.dialog.reject() + assert controller.dialog.isVisible() is False + controller.close() + + +def test_model_manager_stops_the_batch_at_the_first_failed_download( + tmp_path: Path, monkeypatch, qtbot +) -> None: + model_ids = ("u2netp", "silueta", "birefnet-dis") + for model_id in model_ids: + outdated = _outdated_model_path(tmp_path, model_id) + outdated.parent.mkdir(parents=True) + outdated.write_bytes(b"old-weight") + manager = FakeModelRemovalService( + targets={model_id: _model_path(tmp_path, model_id) for model_id in model_ids} + ) + manager.fail_ids["silueta"] = OSError("no route") + manager.obsolete_roots = (tmp_path / "2.0.72",) + controller = ModelManagerController( + ReducerStore(), + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=manager, + ) + qtbot.addWidget(controller.dialog) + controller.open() + warnings: list[str] = [] + monkeypatch.setattr( + QMessageBox, + "warning", + lambda _parent, _title, message: warnings.append(message), + ) + monkeypatch.setattr( + QMessageBox, + "question", + lambda *_args, **_kwargs: QMessageBox.StandardButton.Yes, + ) + + controller.dialog.redownload_outdated_button.click() + + qtbot.waitUntil(lambda: controller._remove_thread is None, timeout=5000) + assert manager.operations == [ + "fetch:u2netp", + "remove-obsolete:u2netp", + "fetch:silueta", + ] + assert _model_path(tmp_path, "u2netp").exists() + assert _outdated_model_path(tmp_path, "u2netp").exists() is False + assert _outdated_model_path(tmp_path, "silueta").exists() + assert _outdated_model_path(tmp_path, "birefnet-dis").exists() + failure_name = next( + entry.display_name + for entry in controller.dialog.entries + if entry.model_id == "silueta" + ) + assert controller.dialog._message.text() == ( + f"Could not download {failure_name}: no route. " + "Re-downloaded 1 of 3 outdated model weight(s)." + ) + assert warnings == [] + assert controller.dialog.redownload_outdated_button.isVisible() + controller.close() + + +def test_model_manager_close_cancels_the_running_download_before_waiting( + tmp_path: Path, monkeypatch, qtbot +) -> None: + outdated = _outdated_model_path(tmp_path, "u2netp") + outdated.parent.mkdir(parents=True) + outdated.write_bytes(b"old-weight") + manager = FakeModelRemovalService( + targets={"u2netp": _model_path(tmp_path, "u2netp")} + ) + manager.block = Event() + manager.block_ids = {"u2netp"} + controller = ModelManagerController( + ReducerStore(), + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=manager, + ) + qtbot.addWidget(controller.dialog) + controller.open() + monkeypatch.setattr( + QMessageBox, + "question", + lambda *_args, **_kwargs: QMessageBox.StandardButton.Yes, + ) + controller.dialog.redownload_outdated_button.click() + qtbot.waitUntil(lambda: manager.in_flight, timeout=5000) + + controller.close() + + assert manager.cancel_observed.is_set() + assert manager.in_flight is False + assert controller.dialog.isVisible() is False + + +def test_model_manager_close_cancels_a_single_download_before_waiting( + tmp_path: Path, qtbot +) -> None: + target = _model_path(tmp_path, "u2netp") + manager = FakeModelRemovalService(targets={"u2netp": target}) + manager.block = Event() + manager.block_ids = {"u2netp"} + controller = ModelManagerController( + ReducerStore(), + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=manager, + ) + qtbot.addWidget(controller.dialog) + controller.open() + index = next( + index + for index, entry in enumerate(controller.dialog.entries) + if entry.model_id == "u2netp" + ) + controller.dialog.model_list.setCurrentRow(index) + controller.dialog.download_button.click() + qtbot.waitUntil(lambda: manager.in_flight, timeout=5000) + + controller.close() + + assert manager.cancel_observed.is_set() + assert manager.in_flight is False + assert controller.dialog.isVisible() is False + + +def test_model_manager_single_download_can_be_cancelled_without_a_warning_box( + tmp_path: Path, monkeypatch, qtbot +) -> None: + target = _model_path(tmp_path, "u2netp") + manager = FakeModelRemovalService(targets={"u2netp": target}) + manager.block = Event() + manager.block_ids = {"u2netp"} + controller = ModelManagerController( + ReducerStore(), + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=manager, + ) + qtbot.addWidget(controller.dialog) + controller.open() + index = next( + index + for index, entry in enumerate(controller.dialog.entries) + if entry.model_id == "u2netp" + ) + controller.dialog.model_list.setCurrentRow(index) + warnings: list[str] = [] + monkeypatch.setattr( + QMessageBox, + "warning", + lambda _parent, _title, message: warnings.append(message), + ) + controller.dialog.download_button.click() + qtbot.waitUntil(lambda: manager.in_flight, timeout=5000) + + controller.dialog.cancel_button.click() + + qtbot.waitUntil(manager.cancel_observed.is_set, timeout=5000) + qtbot.waitUntil(lambda: controller._remove_thread is None, timeout=5000) + assert controller.dialog._message.text() == "Download cancelled." + assert warnings == [] + controller.close() + + +def test_model_manager_refuses_re_download_while_a_job_runs( + tmp_path: Path, monkeypatch, qtbot +) -> None: + outdated = _outdated_model_path(tmp_path, "u2netp") + outdated.parent.mkdir(parents=True) + outdated.write_bytes(b"old-weight") + manager = FakeModelRemovalService() + running = ReducerStore( + AppState( + source=SourceState.READY, + source_id="source-id", + source_value=object(), + job_request_id="request-id", + job=ActiveJob( + "job-id", + JobKind.RENDER, + JobState.RENDERING, + "Encode", + FocusTarget.NONE, + ), + ) + ) + controller = ModelManagerController( + running, + catalog=ModelCatalog.load_resource(), + cache_root=tmp_path, + manager=manager, + ) + qtbot.addWidget(controller.dialog) + controller.open() + monkeypatch.setattr( + QMessageBox, + "question", + lambda *_args, **_kwargs: QMessageBox.StandardButton.Yes, + ) + + controller.dialog.redownload_outdated_button.click() + + assert controller.dialog._message.text() == ( + "Cannot download a model while a job is running." + ) + assert manager.fetched == [] + controller.close() + + def test_model_manager_refreshes_rows_after_outdated_removal_fails( tmp_path: Path, monkeypatch, qtbot ) -> None: