| from __future__ import annotations |
|
|
| from dataclasses import asdict |
| from datetime import datetime |
| import json |
| from pathlib import Path |
| from uuid import uuid4 |
|
|
| from PySide6.QtCore import QSize, Qt, QThread, QTimer, QUrl, Signal |
| from PySide6.QtGui import QBrush, QDesktopServices, QIcon, QImage, QImageReader, QPixmap |
| from PySide6.QtWidgets import ( |
| QAbstractItemView, |
| QComboBox, |
| QDialog, |
| QDialogButtonBox, |
| QFileDialog, |
| QFrame, |
| QGridLayout, |
| QHBoxLayout, |
| QHeaderView, |
| QLabel, |
| QLineEdit, |
| QListWidget, |
| QListWidgetItem, |
| QMessageBox, |
| QPlainTextEdit, |
| QProgressBar, |
| QPushButton, |
| QSpinBox, |
| QTabWidget, |
| QTableWidget, |
| QTableWidgetItem, |
| QVBoxLayout, |
| QWidget, |
| ) |
|
|
| from adam.assets import Asset, AssetRegistry |
| from adam.config import ConfigManager |
| from adam.job_manager import JobManager |
| from adam.models import Job |
| from adam.eve import EveResult, EveVisionModel, save_eve_results |
| from adam.studio import ( |
| PreviewEvaluation, |
| StudioStore, |
| TrainingRecipe, |
| caption_path, |
| checkpoint_files, |
| exact_duplicate_groups, |
| ) |
| from adam.ui.theme import COLORS |
|
|
|
|
| def _card() -> QFrame: |
| frame = QFrame() |
| frame.setProperty("card", True) |
| return frame |
|
|
|
|
| def _title(text: str) -> QLabel: |
| label = QLabel(text) |
| label.setObjectName("CardTitle") |
| return label |
|
|
|
|
| def _header(title: str, subtitle: str) -> QWidget: |
| widget = QWidget() |
| layout = QVBoxLayout(widget) |
| layout.setContentsMargins(0, 0, 0, 10) |
| heading = QLabel(title) |
| heading.setObjectName("PageTitle") |
| detail = QLabel(subtitle) |
| detail.setProperty("muted", True) |
| detail.setWordWrap(True) |
| layout.addWidget(heading) |
| layout.addWidget(detail) |
| return widget |
|
|
|
|
| def _thumbnail(path: Path, width: int, height: int) -> QPixmap: |
| """Decode close to display size instead of loading a full-resolution image.""" |
| reader = QImageReader(str(path)) |
| reader.setAutoTransform(True) |
| source_size = reader.size() |
| if source_size.isValid(): |
| source_size.scale(QSize(width, height), Qt.KeepAspectRatio) |
| reader.setScaledSize(source_size) |
| image = reader.read() |
| return QPixmap.fromImage(image) if not image.isNull() else QPixmap() |
|
|
|
|
| class ImageScanWorker(QThread): |
| scanned = Signal(object, int) |
|
|
| def __init__(self, folder: str, token: int, limit: int = 2500) -> None: |
| super().__init__() |
| self.folder = folder |
| self.token = token |
| self.limit = limit |
|
|
| def run(self) -> None: |
| root = Path(self.folder).expanduser() |
| paths: list[Path] = [] |
| if root.is_dir(): |
| for path in root.rglob("*"): |
| if self.isInterruptionRequested(): |
| return |
| if path.is_file() and path.suffix.casefold() in { |
| ".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif" |
| }: |
| paths.append(path) |
| if len(paths) >= self.limit: |
| break |
| if not self.isInterruptionRequested(): |
| self.scanned.emit(sorted(paths), self.token) |
|
|
|
|
| class EveReviewWorker(QThread): |
| progress = Signal(int, str) |
| completed = Signal(object) |
| failed = Signal(str) |
|
|
| def __init__( |
| self, |
| paths: list[Path], |
| good_references: list[Path], |
| bad_references: list[Path], |
| keep_threshold: float, |
| reject_threshold: float, |
| ) -> None: |
| super().__init__() |
| self.paths = paths |
| self.good_references = good_references |
| self.bad_references = bad_references |
| self.keep_threshold = keep_threshold |
| self.reject_threshold = reject_threshold |
|
|
| def run(self) -> None: |
| try: |
| self.progress.emit(1, "Loading EVE's local vision model…") |
|
|
| def on_progress(done: int, total: int) -> None: |
| if self.isInterruptionRequested(): |
| raise RuntimeError("EVE review cancelled.") |
| percent = 5 + int(done / max(1, total) * 94) |
| self.progress.emit(percent, f"EVE analyzed {done} of {total} images…") |
|
|
| results = EveVisionModel().review( |
| self.paths, |
| self.good_references, |
| self.bad_references, |
| keep_threshold=self.keep_threshold, |
| reject_threshold=self.reject_threshold, |
| progress=on_progress, |
| ) |
| self.progress.emit(100, "EVE finished sorting the dataset.") |
| self.completed.emit(results) |
| except Exception as exc: |
| self.failed.emit(str(exc)) |
|
|
|
|
| class EveReviewDialog(QDialog): |
| applied = Signal(object) |
|
|
| def __init__( |
| self, |
| root_path: Path, |
| dataset_path: str, |
| paths: list[Path], |
| store: StudioStore, |
| parent: QWidget | None = None, |
| ) -> None: |
| super().__init__(parent) |
| self.root_path = root_path |
| self.dataset_path = dataset_path |
| self.paths = paths |
| self.store = store |
| self.results: list[EveResult] = [] |
| self.worker: EveReviewWorker | None = None |
| self.setWindowTitle("EVE · AI Dataset Review") |
| self.setMinimumSize(980, 700) |
| root = QVBoxLayout(self) |
| root.addWidget(_header( |
| "EVE · AI Dataset Review", |
| "Give EVE a few good examples and optional bad examples. EVE proposes Keep, Reject, or Uncertain; nothing changes until you apply the review.", |
| )) |
|
|
| references = QHBoxLayout() |
| self.good_references = self._reference_panel( |
| references, "GOOD REFERENCES", "Add good images…", self._add_good_references |
| ) |
| self.bad_references = self._reference_panel( |
| references, "BAD REFERENCES (OPTIONAL)", "Add bad images…", self._add_bad_references |
| ) |
| root.addLayout(references) |
|
|
| controls = QHBoxLayout() |
| self.keep_threshold = QSpinBox(); self.keep_threshold.setRange(51, 99); self.keep_threshold.setValue(75); self.keep_threshold.setSuffix("%") |
| self.reject_threshold = QSpinBox(); self.reject_threshold.setRange(1, 49); self.reject_threshold.setValue(25); self.reject_threshold.setSuffix("%") |
| self.analyze_button = QPushButton("Analyze dataset with EVE") |
| self.analyze_button.setProperty("primary", True) |
| self.analyze_button.clicked.connect(self._analyze) |
| controls.addWidget(QLabel("Keep at or above")); controls.addWidget(self.keep_threshold) |
| controls.addWidget(QLabel("Reject at or below")); controls.addWidget(self.reject_threshold) |
| controls.addStretch(); controls.addWidget(self.analyze_button) |
| root.addLayout(controls) |
| self.progress = QProgressBar(); self.progress.setRange(0, 100); self.progress.setValue(0) |
| self.status = QLabel( |
| "Tip: 3–10 varied good references work best. Bad references help EVE distinguish visually similar mistakes." |
| ) |
| self.status.setProperty("muted", True); self.status.setWordWrap(True) |
| root.addWidget(self.progress); root.addWidget(self.status) |
|
|
| self.tabs = QTabWidget() |
| self.result_lists: dict[str, QListWidget] = {} |
| for decision, label in (("keep", "KEEP"), ("reject", "REJECT"), ("unreviewed", "UNCERTAIN")): |
| gallery = QListWidget() |
| gallery.setViewMode(QListWidget.IconMode) |
| gallery.setIconSize(QSize(120, 90)); gallery.setGridSize(QSize(165, 145)) |
| gallery.setResizeMode(QListWidget.Adjust) |
| gallery.setSelectionMode(QAbstractItemView.ExtendedSelection) |
| gallery.itemDoubleClicked.connect(lambda item: QDesktopServices.openUrl(QUrl.fromLocalFile(str(item.data(Qt.UserRole))))) |
| self.result_lists[decision] = gallery |
| self.tabs.addTab(gallery, label) |
| root.addWidget(self.tabs, 1) |
|
|
| moves = QHBoxLayout() |
| select_all = QPushButton("Select all in current group") |
| clear_selection = QPushButton("Clear selection") |
| to_keep = QPushButton("Move selected to Keep") |
| to_reject = QPushButton("Move selected to Reject") |
| to_uncertain = QPushButton("Move selected to Uncertain") |
| select_all.clicked.connect(self._select_all_current) |
| clear_selection.clicked.connect(self._clear_current_selection) |
| to_keep.clicked.connect(lambda: self._move_selected("keep")) |
| to_reject.clicked.connect(lambda: self._move_selected("reject")) |
| to_uncertain.clicked.connect(lambda: self._move_selected("unreviewed")) |
| moves.addWidget(select_all); moves.addWidget(clear_selection) |
| moves.addWidget(to_keep); moves.addWidget(to_reject); moves.addWidget(to_uncertain); moves.addStretch() |
| root.addLayout(moves) |
|
|
| buttons = QDialogButtonBox(QDialogButtonBox.Close) |
| self.close_button = buttons.button(QDialogButtonBox.Close) |
| self.apply_button = QPushButton("Apply EVE review") |
| self.apply_button.setProperty("primary", True); self.apply_button.setEnabled(False) |
| self.apply_button.clicked.connect(self._apply) |
| buttons.addButton(self.apply_button, QDialogButtonBox.AcceptRole) |
| buttons.rejected.connect(self.reject) |
| root.addWidget(buttons) |
|
|
| def _reference_panel(self, row: QHBoxLayout, title: str, button_text: str, callback) -> QListWidget: |
| frame = _card(); layout = QVBoxLayout(frame); layout.addWidget(_title(title)) |
| listing = QListWidget(); listing.setMaximumHeight(115) |
| button = QPushButton(button_text); button.clicked.connect(callback) |
| clear = QPushButton("Clear"); clear.clicked.connect(listing.clear) |
| actions = QHBoxLayout(); actions.addWidget(button); actions.addWidget(clear) |
| layout.addWidget(listing); layout.addLayout(actions); row.addWidget(frame, 1) |
| return listing |
|
|
| @staticmethod |
| def _reference_paths(listing: QListWidget) -> list[Path]: |
| return [Path(str(listing.item(index).data(Qt.UserRole))) for index in range(listing.count())] |
|
|
| def _add_references(self, listing: QListWidget) -> None: |
| selected, _ = QFileDialog.getOpenFileNames( |
| self, "Choose EVE reference images", self.dataset_path, |
| "Images (*.png *.jpg *.jpeg *.webp *.bmp *.gif)", |
| ) |
| existing = {str(path) for path in self._reference_paths(listing)} |
| for raw_path in selected: |
| path = str(Path(raw_path).resolve()) |
| if path in existing: |
| continue |
| item = QListWidgetItem(Path(path).name); item.setData(Qt.UserRole, path) |
| listing.addItem(item); existing.add(path) |
|
|
| def _add_good_references(self) -> None: |
| self._add_references(self.good_references) |
|
|
| def _add_bad_references(self) -> None: |
| self._add_references(self.bad_references) |
|
|
| def _analyze(self) -> None: |
| good = self._reference_paths(self.good_references) |
| if not good: |
| QMessageBox.information(self, "Good references required", "Add at least one good reference image for EVE.") |
| return |
| if self.reject_threshold.value() >= self.keep_threshold.value(): |
| QMessageBox.warning(self, "Check thresholds", "Reject confidence must be lower than Keep confidence.") |
| return |
| self.analyze_button.setEnabled(False); self.apply_button.setEnabled(False) |
| self.close_button.setEnabled(False) |
| self.progress.setValue(0); self.status.setText("EVE is starting. The first run may download its vision model once.") |
| self.worker = EveReviewWorker( |
| self.paths, good, self._reference_paths(self.bad_references), |
| self.keep_threshold.value() / 100, self.reject_threshold.value() / 100, |
| ) |
| self.worker.progress.connect(self._progress) |
| self.worker.completed.connect(self._completed) |
| self.worker.failed.connect(self._failed) |
| self.worker.finished.connect(self._worker_finished) |
| self.worker.start() |
|
|
| def _progress(self, percent: int, message: str) -> None: |
| self.progress.setValue(percent); self.status.setText(message) |
|
|
| def _completed(self, results: object) -> None: |
| if not isinstance(results, list): |
| self._failed("EVE returned an invalid review.") |
| return |
| self.results = results |
| save_eve_results(self.root_path, self.dataset_path, self.results) |
| self._rebuild_results() |
| self.apply_button.setEnabled(True) |
| counts = {key: sum(result.suggestion == key for result in self.results) for key in self.result_lists} |
| self.status.setText( |
| f"EVE proposes {counts['keep']} Keep, {counts['reject']} Reject, and {counts['unreviewed']} Uncertain. Review both sides before applying." |
| ) |
|
|
| def _failed(self, message: str) -> None: |
| self.status.setText(f"EVE could not finish: {message}") |
| QMessageBox.warning(self, "EVE review stopped", message) |
|
|
| def _worker_finished(self) -> None: |
| self.analyze_button.setEnabled(True) |
| self.close_button.setEnabled(True) |
| if self.worker: |
| self.worker.deleteLater() |
| self.worker = None |
|
|
| def _rebuild_results(self) -> None: |
| for listing in self.result_lists.values(): |
| listing.clear() |
| for result in self.results: |
| path = Path(result.path) |
| decision_score = result.match_score if result.suggestion == "keep" else 1.0 - result.match_score if result.suggestion == "reject" else result.match_score |
| text = f"{path.name}\n{decision_score * 100:.0f}% " + ("match" if result.suggestion != "reject" else "reject confidence") |
| item = QListWidgetItem(QIcon(_thumbnail(path, 120, 90)), text) |
| item.setData(Qt.UserRole, result.path) |
| self.result_lists[result.suggestion].addItem(item) |
| self._update_tab_labels() |
|
|
| def _update_tab_labels(self) -> None: |
| labels = {"keep": "KEEP", "reject": "REJECT", "unreviewed": "UNCERTAIN"} |
| for index, key in enumerate(("keep", "reject", "unreviewed")): |
| self.tabs.setTabText(index, f"{labels[key]} ({self.result_lists[key].count()})") |
|
|
| def _current_result_list(self) -> QListWidget: |
| return self.tabs.currentWidget() |
|
|
| def _select_all_current(self) -> None: |
| self._current_result_list().selectAll() |
|
|
| def _clear_current_selection(self) -> None: |
| self._current_result_list().clearSelection() |
|
|
| def _move_selected(self, decision: str) -> None: |
| source = self._current_result_list() |
| selected_items = source.selectedItems() |
| if not selected_items: |
| return |
| destination = self.result_lists[decision] |
| if source is destination: |
| return |
| selected_paths = {str(item.data(Qt.UserRole)) for item in selected_items} |
| for result in self.results: |
| if result.path in selected_paths: |
| result.suggestion = decision |
| result.decision_confidence = 1.0 |
| |
| |
| selected_rows = sorted((source.row(item) for item in selected_items), reverse=True) |
| moved_items = [source.takeItem(row) for row in selected_rows] |
| for item in reversed(moved_items): |
| destination.addItem(item) |
| item.setSelected(True) |
| self._update_tab_labels() |
|
|
| def _apply(self) -> None: |
| counts = {key: sum(result.suggestion == key for result in self.results) for key in self.result_lists} |
| answer = QMessageBox.question( |
| self, "Apply EVE review", |
| f"Apply {counts['keep']} Keep and {counts['reject']} Reject decisions?\n\n" |
| f"The {counts['unreviewed']} uncertain images will remain unreviewed. Rejected files are not moved until you choose Exclude rejected.", |
| ) |
| if answer != QMessageBox.Yes: |
| return |
| decisions = {result.path: result.suggestion for result in self.results} |
| changed = self.store.apply_decisions(self.dataset_path, decisions) |
| save_eve_results(self.root_path, self.dataset_path, self.results) |
| self.applied.emit(decisions) |
| self.status.setText(f"Applied EVE review ({changed} decisions changed). You can continue reviewing manually.") |
|
|
| def closeEvent(self, event) -> None: |
| if self.worker and self.worker.isRunning(): |
| self.worker.requestInterruption() |
| self.status.setText("EVE is stopping after the current image. The window will be safe to close when analysis ends.") |
| event.ignore() |
| return |
| super().closeEvent(event) |
|
|
|
|
| class DatasetReviewTab(QWidget): |
| def __init__(self, assets: AssetRegistry, store: StudioStore) -> None: |
| super().__init__() |
| self.assets = assets |
| self.store = store |
| self.paths: list[Path] = [] |
| self.dataset_path = "" |
| self._load_index = 0 |
| self._load_token = 0 |
| self._requested_row = 0 |
| self._scan_workers: set[ImageScanWorker] = set() |
|
|
| root = QVBoxLayout(self) |
| root.setContentsMargins(10, 14, 10, 10) |
| top = QVBoxLayout() |
| dataset_row = QHBoxLayout() |
| review_actions = QHBoxLayout() |
| self.dataset = QComboBox() |
| self.dataset.setMinimumWidth(300) |
| browse = QPushButton("Open another dataset…") |
| browse.clicked.connect(self._browse) |
| refresh = QPushButton("Refresh") |
| refresh.clicked.connect(self.refresh) |
| duplicates = QPushButton("Check duplicates") |
| duplicates.clicked.connect(self._duplicates) |
| keep_all = QPushButton("Keep all images") |
| keep_all.setProperty("primary", True) |
| keep_all.setToolTip( |
| "Accept every image in this dataset, then reject only the individual images you do not want." |
| ) |
| keep_all.clicked.connect(self._keep_all_images) |
| eve_review = QPushButton("EVE AI Review…") |
| eve_review.setToolTip( |
| "Sort this dataset from visual reference images, then review EVE's Keep, Reject, and Uncertain groups." |
| ) |
| eve_review.clicked.connect(self._open_eve_review) |
| apply_rejected = QPushButton("Exclude rejected") |
| apply_rejected.setToolTip( |
| "Move rejected images out of the training dataset into ADAM's recoverable quarantine." |
| ) |
| apply_rejected.clicked.connect(self._apply_rejected) |
| restore_rejected = QPushButton("Restore excluded") |
| restore_rejected.clicked.connect(self._restore_rejected) |
| dataset_row.addWidget(QLabel("Dataset")) |
| dataset_row.addWidget(self.dataset, 1) |
| dataset_row.addWidget(browse) |
| dataset_row.addWidget(refresh) |
| review_actions.addWidget(duplicates) |
| review_actions.addWidget(keep_all) |
| review_actions.addWidget(eve_review) |
| review_actions.addWidget(apply_rejected) |
| review_actions.addWidget(restore_rejected) |
| review_actions.addStretch() |
| top.addLayout(dataset_row) |
| top.addLayout(review_actions) |
| root.addLayout(top) |
|
|
| self.summary = QLabel("Choose a dataset to begin reviewing it.") |
| self.summary.setProperty("muted", True) |
| root.addWidget(self.summary) |
|
|
| body = QHBoxLayout() |
| self.gallery = QListWidget() |
| self.gallery.setViewMode(QListWidget.IconMode) |
| self.gallery.setIconSize(QPixmap(150, 110).size()) |
| self.gallery.setGridSize(QPixmap(178, 158).size()) |
| self.gallery.setResizeMode(QListWidget.Adjust) |
| self.gallery.setSelectionMode(QAbstractItemView.SingleSelection) |
| self.gallery.currentRowChanged.connect(self._selected) |
| body.addWidget(self.gallery, 3) |
|
|
| detail = _card() |
| detail.setMinimumWidth(330) |
| detail_layout = QVBoxLayout(detail) |
| detail_layout.addWidget(_title("IMAGE REVIEW")) |
| self.preview = QLabel("Select an image") |
| self.preview.setAlignment(Qt.AlignCenter) |
| self.preview.setMinimumHeight(230) |
| self.preview.setStyleSheet( |
| f"background: #050d14; border: 1px solid {COLORS['border']}; border-radius: 8px;" |
| ) |
| self.file_label = QLabel() |
| self.file_label.setWordWrap(True) |
| self.file_label.setProperty("muted", True) |
| self.caption = QPlainTextEdit() |
| self.caption.setPlaceholderText("Caption text stored beside the image…") |
| self.caption.setMaximumHeight(120) |
| detail_layout.addWidget(self.preview) |
| detail_layout.addWidget(self.file_label) |
| detail_layout.addWidget(QLabel("Caption")) |
| detail_layout.addWidget(self.caption) |
| buttons = QGridLayout() |
| keep = QPushButton("Keep") |
| keep.setProperty("primary", True) |
| reject = QPushButton("Reject") |
| reject.setProperty("danger", True) |
| restore = QPushButton("Mark unreviewed") |
| save_caption = QPushButton("Save caption") |
| keep.clicked.connect(lambda: self._decide("keep")) |
| reject.clicked.connect(lambda: self._decide("reject")) |
| restore.clicked.connect(lambda: self._decide("unreviewed")) |
| save_caption.clicked.connect(self._save_caption) |
| buttons.addWidget(keep, 0, 0) |
| buttons.addWidget(reject, 0, 1) |
| buttons.addWidget(restore, 1, 0) |
| buttons.addWidget(save_caption, 1, 1) |
| detail_layout.addLayout(buttons) |
| detail_layout.addStretch() |
| body.addWidget(detail, 2) |
| root.addLayout(body, 1) |
|
|
| self.dataset.currentIndexChanged.connect(self.refresh) |
| self.reload_assets() |
|
|
| def reload_assets(self) -> None: |
| current = self.dataset.currentData() |
| self.dataset.blockSignals(True) |
| self.dataset.clear() |
| for asset in self.assets.assets: |
| if asset.kind == "dataset" and Path(asset.path).is_dir(): |
| self.dataset.addItem(asset.name, asset.path) |
| self.dataset.blockSignals(False) |
| index = self.dataset.findData(current) |
| if index >= 0: |
| self.dataset.setCurrentIndex(index) |
| self.refresh() |
|
|
| def _browse(self) -> None: |
| selected = QFileDialog.getExistingDirectory( |
| self, "Choose an image dataset", self.dataset_path or str(Path.home()) |
| ) |
| if not selected: |
| return |
| index = self.dataset.findData(selected) |
| if index < 0: |
| self.dataset.addItem(Path(selected).name, selected) |
| index = self.dataset.count() - 1 |
| self.dataset.setCurrentIndex(index) |
|
|
| def refresh(self) -> None: |
| previous_path = self.dataset_path |
| self.dataset_path = str(self.dataset.currentData() or "") |
| if self.dataset_path != previous_path: |
| self._requested_row = 0 |
| self.gallery.clear() |
| self._load_token += 1 |
| token = self._load_token |
| self._load_index = 0 |
| for worker in self._scan_workers: |
| worker.requestInterruption() |
| if not self.dataset_path: |
| self.paths = [] |
| self.summary.setText("Choose a dataset to begin reviewing it.") |
| self.preview.setText("No supported images found") |
| return |
| self.summary.setText("Scanning dataset…") |
| worker = ImageScanWorker(self.dataset_path, token) |
| self._scan_workers.add(worker) |
| worker.scanned.connect(self._scan_finished) |
| worker.finished.connect( |
| lambda worker=worker: self._scan_workers.discard(worker) |
| ) |
| worker.finished.connect(worker.deleteLater) |
| worker.start() |
|
|
| def _scan_finished(self, paths: object, token: int) -> None: |
| if token != self._load_token or not isinstance(paths, list): |
| return |
| self.paths = paths |
| self._update_review_summary(loading=True) |
| if not self.paths: |
| self.preview.setText("No supported images found") |
| return |
| QTimer.singleShot(0, lambda: self._load_next_thumbnail(token)) |
|
|
| def _update_review_summary(self, *, loading: bool = False) -> None: |
| """Update counts without rescanning files or rebuilding thumbnails.""" |
| review = self.store.review(self.dataset_path) |
| kept = rejected = 0 |
| for path in self.paths: |
| decision = review.decisions.get(str(path.resolve()), "unreviewed") if review else "unreviewed" |
| kept += decision == "keep" |
| rejected += decision == "reject" |
| reviewed = kept + rejected |
| captions = sum(caption_path(path).is_file() for path in self.paths) |
| self.summary.setText( |
| f"{len(self.paths)} images · {reviewed} reviewed · {kept} kept · " |
| f"{rejected} rejected · {captions} captions" |
| + (" · Loading thumbnails…" if loading else "") |
| ) |
|
|
| @staticmethod |
| def _style_review_item(item: QListWidgetItem, path: Path, decision: str) -> None: |
| prefix = {"keep": "✓ ", "reject": "× ", "unreviewed": ""}[decision] |
| item.setText(prefix + path.name) |
| if decision == "reject": |
| item.setForeground(Qt.red) |
| elif decision == "keep": |
| item.setForeground(Qt.green) |
| else: |
| item.setForeground(QBrush()) |
|
|
| def _load_next_thumbnail(self, token: int) -> None: |
| """Decode one image per event-loop turn so large datasets stay responsive.""" |
| if token != self._load_token or self._load_index >= len(self.paths): |
| if token == self._load_token: |
| self.summary.setText(self.summary.text().replace(" · Loading thumbnails…", "")) |
| return |
| path = self.paths[self._load_index] |
| review = self.store.review(self.dataset_path) if self.dataset_path else None |
| decision = review.decisions.get(str(path.resolve()), "unreviewed") if review else "unreviewed" |
| item = QListWidgetItem(path.name) |
| item.setData(Qt.UserRole, str(path)) |
| pixmap = _thumbnail(path, 150, 110) |
| if not pixmap.isNull(): |
| item.setIcon(QIcon(pixmap)) |
| self._style_review_item(item, path, decision) |
| self.gallery.addItem(item) |
| target_row = min(self._requested_row, len(self.paths) - 1) |
| if self._load_index == target_row: |
| self.gallery.setCurrentRow(target_row) |
| self._load_index += 1 |
| QTimer.singleShot(0, lambda: self._load_next_thumbnail(token)) |
|
|
| def _current_path(self) -> Path | None: |
| item = self.gallery.currentItem() |
| return Path(str(item.data(Qt.UserRole))) if item else None |
|
|
| def _selected(self, _row: int) -> None: |
| path = self._current_path() |
| if not path: |
| return |
| pixmap = _thumbnail(path, 310, 260) |
| self.preview.setPixmap( |
| pixmap |
| ) |
| self.file_label.setText(str(path)) |
| try: |
| text = caption_path(path).read_text(encoding="utf-8") |
| except OSError: |
| text = "" |
| self.caption.setPlainText(text) |
|
|
| def _decide(self, decision: str) -> None: |
| path = self._current_path() |
| if not path or not self.dataset_path: |
| return |
| row = self.gallery.currentRow() |
| self.store.set_decision(self.dataset_path, str(path), decision) |
| item = self.gallery.item(row) |
| if item is not None: |
| self._style_review_item(item, path, decision) |
| self._update_review_summary() |
| if decision in {"keep", "reject"} and row + 1 < self.gallery.count(): |
| self.gallery.setCurrentRow(row + 1) |
| self.gallery.scrollToItem(self.gallery.currentItem()) |
|
|
| def _save_caption(self) -> None: |
| path = self._current_path() |
| if not path: |
| return |
| try: |
| caption_path(path).write_text( |
| self.caption.toPlainText().strip() + "\n", encoding="utf-8" |
| ) |
| except OSError as exc: |
| QMessageBox.warning(self, "Caption not saved", str(exc)) |
| return |
| self.summary.setText(self.summary.text() + " · Caption saved") |
|
|
| def _apply_rejected(self) -> None: |
| if not self.dataset_path: |
| return |
| review = self.store.review(self.dataset_path) |
| count = sum(value == "reject" for value in review.decisions.values()) |
| if not count: |
| QMessageBox.information( |
| self, "No rejected images", "Mark images as rejected before excluding them." |
| ) |
| return |
| answer = QMessageBox.question( |
| self, |
| "Exclude rejected images", |
| f"Move {count} rejected image(s) and their captions out of this training dataset?\n\n" |
| "They remain recoverable with Restore excluded.", |
| ) |
| if answer != QMessageBox.Yes: |
| return |
| moved = self.store.apply_rejections(self.dataset_path) |
| self._requested_row = 0 |
| self.refresh() |
| self.summary.setText( |
| f"Excluded {moved} rejected image(s) from training. They remain recoverable." |
| ) |
|
|
| def _keep_all_images(self) -> None: |
| if not self.dataset_path or not self.paths: |
| QMessageBox.information( |
| self, "No images", "Choose a dataset and wait for its images to finish loading." |
| ) |
| return |
| answer = QMessageBox.question( |
| self, |
| "Keep all images", |
| f"Mark all {len(self.paths)} images in this dataset as kept?\n\n" |
| "You can still reject individual images afterward.", |
| ) |
| if answer != QMessageBox.Yes: |
| return |
| changed = self.store.set_all_decisions( |
| self.dataset_path, self.paths, "keep" |
| ) |
| for row in range(self.gallery.count()): |
| item = self.gallery.item(row) |
| self._style_review_item(item, Path(str(item.data(Qt.UserRole))), "keep") |
| self._update_review_summary() |
| self.summary.setText( |
| self.summary.text() |
| + (f" · All images kept ({changed} changed)" if changed else " · All images already kept") |
| ) |
|
|
| def _open_eve_review(self) -> None: |
| if not self.dataset_path or not self.paths: |
| QMessageBox.information( |
| self, "No dataset ready", "Choose a dataset and wait for its image scan to finish." |
| ) |
| return |
| root_path = self.store.path.parent.parent |
| dialog = EveReviewDialog(root_path, self.dataset_path, self.paths, self.store, self) |
| dialog.applied.connect(self._eve_decisions_applied) |
| dialog.exec() |
|
|
| def _eve_decisions_applied(self, decisions: object) -> None: |
| if not isinstance(decisions, dict): |
| return |
| for row in range(self.gallery.count()): |
| item = self.gallery.item(row) |
| path = Path(str(item.data(Qt.UserRole))) |
| decision = str(decisions.get(str(path.resolve()), "unreviewed")) |
| self._style_review_item(item, path, decision) |
| self._update_review_summary() |
| self.summary.setText(self.summary.text() + " · EVE review applied") |
|
|
| def _restore_rejected(self) -> None: |
| if not self.dataset_path: |
| return |
| restored = self.store.restore_rejections(self.dataset_path) |
| self._requested_row = 0 |
| self.refresh() |
| self.summary.setText( |
| f"Restored {restored} excluded image(s)." |
| if restored |
| else "No excluded images were available to restore." |
| ) |
|
|
| def _duplicates(self) -> None: |
| groups = exact_duplicate_groups(self.paths) |
| exact_members = {value for group in groups for value in group} |
| hashes: list[tuple[Path, int]] = [] |
| for path in self.paths[:500]: |
| image = QImage(str(path)) |
| if image.isNull(): |
| continue |
| sample = image.convertToFormat(QImage.Format_Grayscale8).scaled( |
| 8, 8, Qt.IgnoreAspectRatio, Qt.SmoothTransformation |
| ) |
| values = [ |
| sample.pixelColor(x, y).red() for y in range(8) for x in range(8) |
| ] |
| average = sum(values) / len(values) |
| bits = 0 |
| for index, value in enumerate(values): |
| if value >= average: |
| bits |= 1 << index |
| hashes.append((path, bits)) |
| near: list[tuple[str, str]] = [] |
| for index, (first_path, first_hash) in enumerate(hashes): |
| for second_path, second_hash in hashes[index + 1 :]: |
| if str(first_path) in exact_members and str(second_path) in exact_members: |
| continue |
| if (first_hash ^ second_hash).bit_count() <= 5: |
| near.append((first_path.name, second_path.name)) |
| if len(near) >= 20: |
| break |
| if len(near) >= 20: |
| break |
| if not groups and not near: |
| QMessageBox.information( |
| self, |
| "Duplicate check", |
| "No exact or visually similar duplicate candidates were found.", |
| ) |
| return |
| lines = [ |
| " = ".join(Path(value).name for value in group) for group in groups[:20] |
| ] |
| lines.extend(f"≈ {first} / {second}" for first, second in near) |
| scope = ( |
| " Visual similarity checked the first 500 images." |
| if len(self.paths) > 500 |
| else "" |
| ) |
| QMessageBox.warning( |
| self, |
| "Duplicate candidates found", |
| f"{len(groups)} exact group(s) and {len(near)} visually similar " |
| f"candidate pair(s).{scope}\n\n" + "\n".join(lines), |
| ) |
|
|
|
|
| class ExperimentsTab(QWidget): |
| def __init__( |
| self, jobs: JobManager, assets: AssetRegistry, store: StudioStore |
| ) -> None: |
| super().__init__() |
| self.jobs = jobs |
| self.assets = assets |
| self.store = store |
| root = QVBoxLayout(self) |
| root.setContentsMargins(10, 14, 10, 10) |
| hint = QLabel( |
| "Select one run for details or two runs to compare their recipes and outcomes." |
| ) |
| hint.setProperty("muted", True) |
| root.addWidget(hint) |
| self.table = QTableWidget(0, 7) |
| self.table.setHorizontalHeaderLabels( |
| ["RUN", "PROJECT", "TRAINER", "EPOCHS", "STATUS", "CREATED", "OUTPUT"] |
| ) |
| self.table.setSelectionBehavior(QAbstractItemView.SelectRows) |
| self.table.setSelectionMode(QAbstractItemView.ExtendedSelection) |
| self.table.setEditTriggers(QAbstractItemView.NoEditTriggers) |
| self.table.verticalHeader().hide() |
| header = self.table.horizontalHeader() |
| header.setSectionResizeMode(1, QHeaderView.Stretch) |
| for column in (0, 2, 3, 4, 5, 6): |
| header.setSectionResizeMode(column, QHeaderView.ResizeToContents) |
| self.table.itemSelectionChanged.connect(self._show_selection) |
| root.addWidget(self.table, 3) |
| detail = _card() |
| detail_layout = QVBoxLayout(detail) |
| detail_layout.addWidget(_title("RUN COMPARISON")) |
| self.comparison = QPlainTextEdit() |
| self.comparison.setReadOnly(True) |
| self.comparison.setMaximumHeight(180) |
| detail_layout.addWidget(self.comparison) |
| actions = QHBoxLayout() |
| self.best = QPushButton("Mark model as best") |
| self.open = QPushButton("Open output") |
| self.recipe = QPushButton("Save as reusable recipe") |
| self.best.clicked.connect(self._toggle_best) |
| self.open.clicked.connect(self._open_output) |
| self.recipe.clicked.connect(self._save_recipe) |
| actions.addWidget(self.best) |
| actions.addWidget(self.recipe) |
| actions.addWidget(self.open) |
| detail_layout.addLayout(actions) |
| root.addWidget(detail, 1) |
| jobs.job_created.connect(lambda _job: self.refresh()) |
| jobs.job_updated.connect(lambda _job: self.refresh()) |
| self.refresh() |
|
|
| @staticmethod |
| def _training(job: Job) -> tuple[str, int]: |
| for step in job.plan.steps: |
| if step.tool_id.endswith("_trainer"): |
| return step.tool_id.removesuffix("_trainer").upper(), int( |
| step.arguments.get("epochs", 0) or 0 |
| ) |
| return "—", 0 |
|
|
| def refresh(self) -> None: |
| selected = {job.id for job in self._selected_jobs()} |
| self.table.setRowCount(len(self.jobs.jobs)) |
| for row, job in enumerate(self.jobs.jobs): |
| trainer, epochs = self._training(job) |
| values = [ |
| job.id, |
| job.plan.project_name, |
| trainer, |
| str(epochs or "—"), |
| job.status.value, |
| self._date(job.created_at), |
| "Ready" if job.output_folder else "—", |
| ] |
| for column, value in enumerate(values): |
| item = QTableWidgetItem(value) |
| if column != 1: |
| item.setTextAlignment(Qt.AlignCenter) |
| self.table.setItem(row, column, item) |
| if job.id in selected: |
| self.table.selectRow(row) |
| self._show_selection() |
|
|
| def _selected_jobs(self) -> list[Job]: |
| rows = sorted({index.row() for index in self.table.selectionModel().selectedRows()}) |
| result = [] |
| for row in rows: |
| item = self.table.item(row, 0) |
| if item: |
| try: |
| result.append(self.jobs.get(item.text())) |
| except KeyError: |
| pass |
| return result |
|
|
| def _show_selection(self) -> None: |
| jobs = self._selected_jobs() |
| if not jobs: |
| self.comparison.setPlainText("Select a run to inspect it.") |
| else: |
| blocks = [] |
| for job in jobs[:2]: |
| trainer, epochs = self._training(job) |
| duration = self._duration(job) |
| blocks.append( |
| f"{job.plan.project_name} ({job.id})\n" |
| f" Trainer: {trainer} · Epochs: {epochs or '—'}\n" |
| f" Status: {job.status.value} · Progress: {job.progress}%\n" |
| f" Duration: {duration} · Output: {job.output_folder or 'Not recorded'}" |
| ) |
| self.comparison.setPlainText("\n\n".join(blocks)) |
| one = len(jobs) == 1 |
| self.open.setEnabled(one and bool(jobs[0].output_folder)) |
| self.recipe.setEnabled(one and self._training(jobs[0])[0] != "—") |
| model = self._model_for_job(jobs[0]) if one else None |
| self.best.setEnabled(model is not None) |
| if model: |
| self.best.setText( |
| "Unmark best" if model.id in self.store.best_models else "Mark model as best" |
| ) |
|
|
| def _model_for_job(self, job: Job) -> Asset | None: |
| if not job.output_folder: |
| return None |
| output = Path(job.output_folder) |
| return next( |
| ( |
| asset |
| for asset in self.assets.assets |
| if asset.kind == "model" |
| and ( |
| Path(asset.path) == output |
| or output in Path(asset.path).parents |
| or Path(asset.path) in output.parents |
| ) |
| ), |
| None, |
| ) |
|
|
| def _toggle_best(self) -> None: |
| jobs = self._selected_jobs() |
| model = self._model_for_job(jobs[0]) if len(jobs) == 1 else None |
| if model: |
| self.store.toggle_best(model.id) |
| self._show_selection() |
|
|
| def _open_output(self) -> None: |
| jobs = self._selected_jobs() |
| if len(jobs) == 1 and jobs[0].output_folder: |
| QDesktopServices.openUrl(QUrl.fromLocalFile(jobs[0].output_folder)) |
|
|
| def _save_recipe(self) -> None: |
| jobs = self._selected_jobs() |
| if len(jobs) != 1: |
| return |
| job = jobs[0] |
| for step in job.plan.steps: |
| if step.tool_id.endswith("_trainer"): |
| self.store.add_recipe( |
| TrainingRecipe( |
| name=f"{job.plan.project_name} recipe", |
| trainer=step.tool_id.removesuffix("_trainer"), |
| epochs=int(step.arguments.get("epochs", 100) or 100), |
| base_model=str(step.arguments.get("base_model", "")), |
| notes=f"Created from job {job.id}", |
| ) |
| ) |
| QMessageBox.information( |
| self, "Recipe saved", "The run settings are now reusable in Recipes." |
| ) |
| return |
|
|
| @staticmethod |
| def _date(value: str) -> str: |
| try: |
| return datetime.fromisoformat(value).astimezone().strftime("%b %d · %H:%M") |
| except ValueError: |
| return value[:16] |
|
|
| @staticmethod |
| def _duration(job: Job) -> str: |
| if not job.started_at or not job.ended_at: |
| return "In progress" if job.started_at else "Not started" |
| try: |
| seconds = int( |
| ( |
| datetime.fromisoformat(job.ended_at) |
| - datetime.fromisoformat(job.started_at) |
| ).total_seconds() |
| ) |
| return f"{seconds // 3600}h {(seconds % 3600) // 60}m" if seconds >= 3600 else f"{seconds // 60}m {seconds % 60}s" |
| except ValueError: |
| return "Unknown" |
|
|
|
|
| class PreviewLabTab(QWidget): |
| plan_requested = Signal(str) |
|
|
| def __init__(self, assets: AssetRegistry, store: StudioStore) -> None: |
| super().__init__() |
| self.assets = assets |
| self.store = store |
| self._output_paths: list[Path] = [] |
| self._output_index = 0 |
| self._output_token = 0 |
| self._output_scan_workers: set[ImageScanWorker] = set() |
| root = QHBoxLayout(self) |
| root.setContentsMargins(10, 14, 10, 10) |
| form = _card() |
| form_layout = QGridLayout(form) |
| form_layout.addWidget(_title("PROMPT & CHECKPOINT EVALUATION"), 0, 0, 1, 2) |
| self.model = QComboBox() |
| self.checkpoint = QComboBox() |
| self.prompt = QPlainTextEdit() |
| self.prompt.setPlaceholderText( |
| "A consistent evaluation prompt for comparing model checkpoints…" |
| ) |
| self.prompt.setMaximumHeight(110) |
| self.seed = QSpinBox() |
| self.seed.setRange(0, 2_147_483_647) |
| self.count = QSpinBox() |
| self.count.setRange(1, 16) |
| self.count.setValue(4) |
| self.rating = QSpinBox() |
| self.rating.setRange(0, 5) |
| self.rating.setSuffix(" / 5") |
| self.notes = QPlainTextEdit() |
| self.notes.setPlaceholderText("What worked, what drifted, and what to try next…") |
| self.notes.setMaximumHeight(100) |
| rows = ( |
| ("Model", self.model), |
| ("Checkpoint", self.checkpoint), |
| ("Evaluation prompt", self.prompt), |
| ("Seed", self.seed), |
| ("Preview count", self.count), |
| ("Rating", self.rating), |
| ("Notes", self.notes), |
| ) |
| for row, (label, widget) in enumerate(rows, 1): |
| form_layout.addWidget(QLabel(label), row, 0, Qt.AlignTop) |
| form_layout.addWidget(widget, row, 1) |
| save = QPushButton("Save evaluation") |
| save.clicked.connect(self._save) |
| request = QPushButton("Plan preview job") |
| request.setProperty("primary", True) |
| request.clicked.connect(self._request) |
| form_layout.addWidget(save, len(rows) + 1, 0) |
| form_layout.addWidget(request, len(rows) + 1, 1) |
| root.addWidget(form, 2) |
|
|
| history = _card() |
| history_layout = QVBoxLayout(history) |
| history_layout.addWidget(_title("MODEL OUTPUTS")) |
| self.outputs = QListWidget() |
| self.outputs.setViewMode(QListWidget.IconMode) |
| self.outputs.setIconSize(QPixmap(110, 80).size()) |
| self.outputs.setMaximumHeight(210) |
| self.outputs.itemDoubleClicked.connect( |
| lambda item: QDesktopServices.openUrl( |
| QUrl.fromLocalFile(str(item.data(Qt.UserRole))) |
| ) |
| ) |
| history_layout.addWidget(self.outputs) |
| history_layout.addWidget(_title("EVALUATION HISTORY")) |
| self.history = QListWidget() |
| history_layout.addWidget(self.history) |
| root.addWidget(history, 1) |
| self.model.currentIndexChanged.connect(self._reload_checkpoints) |
| self.reload_assets() |
|
|
| def reload_assets(self) -> None: |
| current = self.model.currentData() |
| self.model.blockSignals(True) |
| self.model.clear() |
| for asset in self.assets.assets: |
| if asset.kind == "model": |
| star = "★ " if asset.id in self.store.best_models else "" |
| self.model.addItem(star + asset.name, asset.id) |
| self.model.blockSignals(False) |
| index = self.model.findData(current) |
| self.model.setCurrentIndex(index if index >= 0 else 0) |
| self._reload_checkpoints() |
| self._reload_history() |
|
|
| def _asset(self) -> Asset | None: |
| model_id = str(self.model.currentData() or "") |
| return next((asset for asset in self.assets.assets if asset.id == model_id), None) |
|
|
| def _reload_checkpoints(self) -> None: |
| self.checkpoint.clear() |
| self.outputs.clear() |
| self._output_token += 1 |
| token = self._output_token |
| self._output_index = 0 |
| for worker in self._output_scan_workers: |
| worker.requestInterruption() |
| asset = self._asset() |
| if not asset: |
| return |
| found = checkpoint_files(asset.path) |
| if asset.checkpoint and Path(asset.checkpoint).exists(): |
| found = [Path(asset.checkpoint), *[path for path in found if str(path) != asset.checkpoint]] |
| if not found: |
| self.checkpoint.addItem("Latest model output", asset.path) |
| else: |
| for path in found: |
| self.checkpoint.addItem(path.name, str(path)) |
| worker = ImageScanWorker(asset.path, token, limit=80) |
| self._output_scan_workers.add(worker) |
| worker.scanned.connect(self._preview_scan_finished) |
| worker.finished.connect( |
| lambda worker=worker: self._output_scan_workers.discard(worker) |
| ) |
| worker.finished.connect(worker.deleteLater) |
| worker.start() |
|
|
| def _preview_scan_finished(self, paths: object, token: int) -> None: |
| if token != self._output_token or not isinstance(paths, list): |
| return |
| self._output_paths = paths |
| QTimer.singleShot(0, lambda: self._load_next_preview(token)) |
|
|
| def _load_next_preview(self, token: int) -> None: |
| if token != self._output_token or self._output_index >= len(self._output_paths): |
| return |
| path = self._output_paths[self._output_index] |
| item = QListWidgetItem(path.name) |
| item.setData(Qt.UserRole, str(path)) |
| pixmap = _thumbnail(path, 110, 80) |
| if not pixmap.isNull(): |
| item.setIcon(QIcon(pixmap)) |
| self.outputs.addItem(item) |
| self._output_index += 1 |
| QTimer.singleShot(0, lambda: self._load_next_preview(token)) |
|
|
| def _save(self) -> None: |
| asset = self._asset() |
| if not asset: |
| QMessageBox.warning(self, "No model", "Register or finish a model first.") |
| return |
| self.store.add_evaluation( |
| PreviewEvaluation( |
| model_id=asset.id, |
| checkpoint=str(self.checkpoint.currentData() or ""), |
| prompt=self.prompt.toPlainText().strip(), |
| seed=self.seed.value(), |
| rating=self.rating.value(), |
| notes=self.notes.toPlainText().strip(), |
| ) |
| ) |
| self._reload_history() |
|
|
| def _reload_history(self) -> None: |
| self.history.clear() |
| names = {asset.id: asset.name for asset in self.assets.assets} |
| for evaluation in reversed(self.store.evaluations[-100:]): |
| self.history.addItem( |
| f"{'★' * evaluation.rating or 'Unrated'} · " |
| f"{names.get(evaluation.model_id, 'Unknown model')}\n" |
| f"{evaluation.prompt or 'No prompt recorded'}" |
| ) |
|
|
| def _request(self) -> None: |
| asset = self._asset() |
| if not asset: |
| return |
| prompt = self.prompt.toPlainText().strip() |
| checkpoint = str(self.checkpoint.currentData() or "") |
| request = ( |
| f"Generate {self.count.value()} previews for the {asset.name} model" |
| + (f" from checkpoint {checkpoint}" if checkpoint else "") |
| + (f" using this evaluation prompt: {prompt}" if prompt else "") |
| + f". Use seed {self.seed.value()}." |
| ) |
| self.plan_requested.emit(request) |
|
|
|
|
| class RecipesTab(QWidget): |
| plan_requested = Signal(str) |
|
|
| def __init__(self, store: StudioStore) -> None: |
| super().__init__() |
| self.store = store |
| root = QHBoxLayout(self) |
| root.setContentsMargins(10, 14, 10, 10) |
| self.list = QListWidget() |
| self.list.currentRowChanged.connect(self._selected) |
| root.addWidget(self.list, 1) |
| form = _card() |
| layout = QGridLayout(form) |
| layout.addWidget(_title("REPRODUCIBLE TRAINING RECIPE"), 0, 0, 1, 2) |
| self.name = QLineEdit() |
| self.trainer = QComboBox() |
| self.trainer.addItem("LoRA", "lora") |
| self.trainer.addItem("DDPM", "ddpm") |
| self.trainer.addItem("Flow Matching", "flow") |
| self.epochs = QSpinBox() |
| self.epochs.setRange(1, 100_000) |
| self.epochs.setValue(100) |
| self.images = QSpinBox() |
| self.images.setRange(10, 100_000) |
| self.images.setValue(60) |
| self.base_model = QLineEdit() |
| self.preview_prompt = QLineEdit() |
| self.notes = QPlainTextEdit() |
| self.notes.setMaximumHeight(100) |
| for row, (label, widget) in enumerate( |
| ( |
| ("Name", self.name), |
| ("Trainer", self.trainer), |
| ("Epochs", self.epochs), |
| ("Dataset target", self.images), |
| ("Base model", self.base_model), |
| ("Preview prompt", self.preview_prompt), |
| ("Notes", self.notes), |
| ), |
| 1, |
| ): |
| layout.addWidget(QLabel(label), row, 0) |
| layout.addWidget(widget, row, 1) |
| save = QPushButton("Save recipe") |
| save.clicked.connect(self._save) |
| use = QPushButton("Create model from recipe") |
| use.setProperty("primary", True) |
| use.clicked.connect(self._use) |
| layout.addWidget(save, 8, 0) |
| layout.addWidget(use, 8, 1) |
| transfer = QHBoxLayout() |
| export = QPushButton("Export recipe…") |
| import_button = QPushButton("Import recipe…") |
| export.clicked.connect(self._export) |
| import_button.clicked.connect(self._import) |
| transfer.addWidget(import_button) |
| transfer.addWidget(export) |
| layout.addLayout(transfer, 9, 0, 1, 2) |
| root.addWidget(form, 2) |
| self.refresh() |
|
|
| def refresh(self) -> None: |
| row = self.list.currentRow() |
| self.list.clear() |
| for recipe in self.store.recipes: |
| item = QListWidgetItem( |
| f"{recipe.name}\n{recipe.trainer.upper()} · {recipe.epochs} epochs" |
| ) |
| item.setData(Qt.UserRole, recipe.id) |
| self.list.addItem(item) |
| if self.list.count(): |
| self.list.setCurrentRow(max(0, min(row, self.list.count() - 1))) |
|
|
| def _current(self) -> TrainingRecipe | None: |
| item = self.list.currentItem() |
| recipe_id = str(item.data(Qt.UserRole)) if item else "" |
| return next((recipe for recipe in self.store.recipes if recipe.id == recipe_id), None) |
|
|
| def _selected(self, _row: int) -> None: |
| recipe = self._current() |
| if not recipe: |
| return |
| self.name.setText(recipe.name) |
| self.trainer.setCurrentIndex(max(0, self.trainer.findData(recipe.trainer))) |
| self.epochs.setValue(recipe.epochs) |
| self.images.setValue(recipe.image_count) |
| self.base_model.setText(recipe.base_model) |
| self.preview_prompt.setText(recipe.preview_prompt) |
| self.notes.setPlainText(recipe.notes) |
|
|
| def _save(self) -> None: |
| current = self._current() |
| recipe = TrainingRecipe( |
| id=current.id if current else uuid4().hex[:10], |
| created_at=current.created_at if current else datetime.now().astimezone().isoformat(), |
| name=self.name.text().strip() or "Untitled recipe", |
| trainer=str(self.trainer.currentData()), |
| epochs=self.epochs.value(), |
| image_count=self.images.value(), |
| base_model=self.base_model.text().strip(), |
| preview_prompt=self.preview_prompt.text().strip(), |
| notes=self.notes.toPlainText().strip(), |
| ) |
| self.store.add_recipe(recipe) |
| self.refresh() |
|
|
| def _use(self) -> None: |
| trainer = str(self.trainer.currentData()).upper() |
| self.plan_requested.emit( |
| f"Create a {trainer} model using {self.images.value()} images and train " |
| f"for {self.epochs.value()} epochs. Use the recipe named " |
| f"{self.name.text().strip() or 'Untitled recipe'}." |
| ) |
|
|
| def _export(self) -> None: |
| recipe = self._current() |
| if not recipe: |
| QMessageBox.information(self, "No recipe", "Select or save a recipe first.") |
| return |
| selected, _filter = QFileDialog.getSaveFileName( |
| self, |
| "Export ADAM recipe", |
| f"{recipe.name}.adam-recipe.json", |
| "ADAM recipes (*.json)", |
| ) |
| if selected: |
| try: |
| Path(selected).write_text( |
| json.dumps( |
| {"format": "adam-training-recipe-v1", "recipe": asdict(recipe)}, |
| indent=2, |
| ), |
| encoding="utf-8", |
| ) |
| except OSError as exc: |
| QMessageBox.warning(self, "Recipe not exported", str(exc)) |
|
|
| def _import(self) -> None: |
| selected, _filter = QFileDialog.getOpenFileName( |
| self, "Import ADAM recipe", "", "ADAM recipes (*.json)" |
| ) |
| if not selected: |
| return |
| try: |
| payload = json.loads(Path(selected).read_text(encoding="utf-8")) |
| if payload.get("format") != "adam-training-recipe-v1": |
| raise ValueError("This is not an ADAM training recipe.") |
| recipe = TrainingRecipe.from_dict(dict(payload["recipe"])) |
| recipe.id = uuid4().hex[:10] |
| self.store.add_recipe(recipe) |
| except (OSError, ValueError, TypeError, KeyError, json.JSONDecodeError) as exc: |
| QMessageBox.warning(self, "Recipe not imported", str(exc)) |
| return |
| self.refresh() |
|
|
|
|
| class StudioPage(QWidget): |
| plan_requested = Signal(str) |
|
|
| def __init__( |
| self, |
| root_path: Path, |
| jobs: JobManager, |
| assets: AssetRegistry, |
| config: ConfigManager, |
| ) -> None: |
| super().__init__() |
| del config |
| self.store = StudioStore(root_path) |
| root = QVBoxLayout(self) |
| root.setContentsMargins(24, 20, 24, 17) |
| root.setSpacing(8) |
| root.addWidget( |
| _header( |
| "Training studio", |
| "Review datasets, compare experiments, evaluate checkpoints, and preserve reproducible recipes.", |
| ) |
| ) |
| self.tabs = QTabWidget() |
| self.datasets = DatasetReviewTab(assets, self.store) |
| self.experiments = ExperimentsTab(jobs, assets, self.store) |
| self.previews = PreviewLabTab(assets, self.store) |
| self.recipes = RecipesTab(self.store) |
| self.tabs.addTab(self.datasets, "Datasets") |
| self.tabs.addTab(self.experiments, "Experiments") |
| self.tabs.addTab(self.previews, "Checkpoint Lab") |
| self.tabs.addTab(self.recipes, "Recipes") |
| root.addWidget(self.tabs, 1) |
| self.previews.plan_requested.connect(self.plan_requested) |
| self.recipes.plan_requested.connect(self.plan_requested) |
|
|
| def refresh(self) -> None: |
| self.datasets.assets.load() |
| self.datasets.reload_assets() |
| self.experiments.assets.load() |
| self.experiments.refresh() |
| self.previews.assets.load() |
| self.previews.reload_assets() |
| self.recipes.refresh() |
|
|
| def shutdown(self) -> None: |
| workers = [ |
| *self.datasets._scan_workers, |
| *self.previews._output_scan_workers, |
| ] |
| for worker in workers: |
| worker.requestInterruption() |
| for worker in workers: |
| worker.wait(1000) |
|
|
