From cb6cce379d0fb53fbed740ca488e1fa7f0e94b19 Mon Sep 17 00:00:00 2001 From: smb Date: Fri, 4 Sep 2026 11:58:26 +0200 Subject: [PATCH 1/3] refactor: split the model manager into a package model_manager.py stood at 753 lines against the 800-line module budget, so the next feature would have had to be squeezed in or would have pushed a compliant module over the ceiling. The bulk re-download of outdated weights (#29) needs the room. The module becomes a package with the layering it already had implicitly: _entries holds the row model and the cache lookups, _dialog the widget, _workers the three background workers, and _controller the protocol and the controller that wires them. The largest unit is 303 lines. __init__ re-exports the three names the two importers in the repository use, so no call site changes and no compatibility shim is needed: the UI tests monkeypatch Qt classes, not module globals. Code moved, nothing rewritten; a sorted-line comparison of the old module against the concatenated units differs only in import statements. Co-Authored-By: Claude Opus 5 (1M context) --- src/matteloop/ui/model_manager.py | 753 ------------------ src/matteloop/ui/model_manager/__init__.py | 12 + src/matteloop/ui/model_manager/_controller.py | 303 +++++++ src/matteloop/ui/model_manager/_dialog.py | 278 +++++++ src/matteloop/ui/model_manager/_entries.py | 131 +++ src/matteloop/ui/model_manager/_workers.py | 73 ++ 6 files changed, 797 insertions(+), 753 deletions(-) delete mode 100644 src/matteloop/ui/model_manager.py create mode 100644 src/matteloop/ui/model_manager/__init__.py create mode 100644 src/matteloop/ui/model_manager/_controller.py create mode 100644 src/matteloop/ui/model_manager/_dialog.py create mode 100644 src/matteloop/ui/model_manager/_entries.py create mode 100644 src/matteloop/ui/model_manager/_workers.py 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..28da9db --- /dev/null +++ b/src/matteloop/ui/model_manager/_controller.py @@ -0,0 +1,303 @@ +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.state import JobState, ModelAvailabilityChanged +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, +) + +_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: ... + + +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}") diff --git a/src/matteloop/ui/model_manager/_dialog.py b/src/matteloop/ui/model_manager/_dialog.py new file mode 100644 index 0000000..49126ce --- /dev/null +++ b/src/matteloop/ui/model_manager/_dialog.py @@ -0,0 +1,278 @@ +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() + 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) 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..1cf3362 --- /dev/null +++ b/src/matteloop/ui/model_manager/_workers.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from PySide6.QtCore import QObject, Signal, Slot + +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) -> 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() From e877891d459672ee404ee01c3132ec5e82cc7944 Mon Sep 17 00:00:00 2001 From: smb Date: Fri, 4 Sep 2026 12:26:31 +0200 Subject: [PATCH 2/3] feat: let a caller pass its own cancellation check to fetch fetch() hard-coded the manager's own cancellation state, which in production is the preview job's and is false whenever no job runs. A dialog that downloads a weight could therefore pass a progress callback but had no way to stop what it started. It now takes an optional cancellation check and keeps the manager's when none is given, so every existing caller is unchanged. remove_obsolete_versions() takes an optional model id so one outdated weight can be removed once its replacement has been verified, instead of clearing the whole namespace before the first download. The catalog resolves the id, so the extra path segment is always a known model, and the symlink and directory guards are untouched. Growth of the frozen jobs/models/session.py from 432 to 439 lines was agreed with the maintainer before the baseline was updated, as the guardrails require. Co-Authored-By: Claude Opus 5 (1M context) --- scripts/guardrails-baseline.json | 2 +- src/matteloop/jobs/models/session.py | 13 ++++++-- tests/jobs/models/test_session.py | 48 ++++++++++++++++++++++++++++ 3 files changed, 59 insertions(+), 4 deletions(-) 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/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" From b7c1d432954b7b91f360b7e8f5f83c5a7ad2e7e9 Mon Sep 17 00:00:00 2001 From: smb Date: Fri, 4 Sep 2026 12:26:32 +0200 Subject: [PATCH 3/3] feat: re-download outdated model weights as a cancellable batch The bulk action removed before #24 landed started up to 6.35 GiB of downloads from one unconfirmed click, survived closing the dialog and left shutdown blocked on the manager lock. It returns with the four things that were missing. A confirmation names the count and the size in the units the manager already prints, and says both that each old copy is deleted only once its replacement verifies and that the other files in the obsolete directory stay until Delete outdated. The batch owns a cancellation state that reaches the downloader, so Cancel, Esc and the title-bar X all stop it within one chunk; the dialog stays open and reports what it managed. Controller close() cancels, joins, releases the busy state and closes the dialog, so quitting the application no longer blocks behind a running fetch. The status line carries the model, its position in the batch and the transfer rate. The first failure stops the batch. Because each old weight is deleted only after its replacement is verified, what remains is truthful: the notice, the rows and the button describe exactly the models still to fetch, and one click resumes. That is why there is no failure aggregation and no second summary format. While a worker thread lives, the GUI thread reads nothing that takes the manager lock; rows refresh once at the end. The test fake raises when its active id is read during a fetch, and the batch test blocks on its second model so the slots that run between models are covered by that guard rather than racing past it. Closes #29 Co-Authored-By: Claude Opus 5 (1M context) --- docs/v1-scope.md | 3 +- src/matteloop/ui/model_manager/_controller.py | 176 ++++- src/matteloop/ui/model_manager/_dialog.py | 51 +- src/matteloop/ui/model_manager/_workers.py | 60 +- tests/ui/test_model_manager.py | 610 +++++++++++++++++- 5 files changed, 870 insertions(+), 30 deletions(-) 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/src/matteloop/ui/model_manager/_controller.py b/src/matteloop/ui/model_manager/_controller.py index 28da9db..15bd60b 100644 --- a/src/matteloop/ui/model_manager/_controller.py +++ b/src/matteloop/ui/model_manager/_controller.py @@ -8,7 +8,9 @@ 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 ( @@ -23,20 +25,45 @@ _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) -> int: ... + def remove_obsolete_versions(self, model_id: str | None = None) -> int: ... - def fetch(self, model_id: str, progress: object = None) -> object: ... + def fetch( + self, model_id: str, progress: object = None, cancelled: object = None + ) -> object: ... class ModelManagerController(QObject): @@ -66,13 +93,26 @@ def __init__( 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 | None + _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) @@ -82,11 +122,16 @@ def open(self) -> None: 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(5000) + 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: @@ -106,7 +151,8 @@ def _download_requested(self, value: object) -> None: if self._remove_thread is None: self.dialog.set_message(_DOWNLOAD_BLOCKED_MESSAGE) return - worker = _ModelDownloadWorker(manager, value.model_id) + cancel = CancellationState() + worker = _ModelDownloadWorker(manager, value.model_id, cancel) thread = QThread(self) worker.moveToThread(thread) self._download_rate = DownloadRateEstimator() @@ -120,7 +166,9 @@ def _download_requested(self, value: object) -> None: self._remove_thread = thread self._remove_worker = worker self._removal_entry = value - self.dialog.set_busy(True) + 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() @@ -152,6 +200,59 @@ def _delete_outdated_requested(self) -> None: 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) @@ -172,14 +273,68 @@ def _start_obsolete_removal(self, manager: ModelRemovalService) -> None: @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, completed, total, speed) + 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 @@ -195,6 +350,9 @@ 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)) @@ -288,6 +446,10 @@ 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: diff --git a/src/matteloop/ui/model_manager/_dialog.py b/src/matteloop/ui/model_manager/_dialog.py index 49126ce..6c43110 100644 --- a/src/matteloop/ui/model_manager/_dialog.py +++ b/src/matteloop/ui/model_manager/_dialog.py @@ -35,6 +35,8 @@ class ModelManagerDialog(QDialog): 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__( @@ -93,6 +95,11 @@ def _build_widgets(self) -> None: 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) @@ -105,6 +112,10 @@ def _build_widgets(self) -> None: 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) @@ -117,9 +128,11 @@ def _build_layout(self) -> None: 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) @@ -129,7 +142,11 @@ def _connect_signals(self) -> None: ) 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) @@ -147,6 +164,14 @@ def outdated_entries(self) -> tuple[ModelEntry, ...]: 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 @@ -167,14 +192,33 @@ 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: + 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) @@ -266,6 +310,11 @@ def _update_actions(self) -> None: 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 diff --git a/src/matteloop/ui/model_manager/_workers.py b/src/matteloop/ui/model_manager/_workers.py index 1cf3362..e3cd483 100644 --- a/src/matteloop/ui/model_manager/_workers.py +++ b/src/matteloop/ui/model_manager/_workers.py @@ -4,6 +4,9 @@ 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 @@ -34,15 +37,18 @@ class _ModelDownloadWorker(QObject): progressed = Signal(int, int) finished = Signal() - def __init__(self, manager: ModelRemovalService, model_id: str) -> None: + 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._manager.fetch(self._model_id, self._report, self._cancelled) except Exception as error: self.failed.emit(self._model_id, error) else: @@ -53,6 +59,56 @@ def run(self) -> None: 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) 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: