"""Configure the restoration ("расцензурить") engine. inpaint — no setup. deepmosaics — point at a user-installed DeepMosaics folder and its clean weights; a CUDA GPU is strongly recommended (set GPU id, -1 = CPU/slow). """ from __future__ import annotations from pathlib import Path from PySide6.QtWidgets import ( QComboBox, QDialog, QDialogButtonBox, QFileDialog, QFormLayout, QHBoxLayout, QLabel, QLineEdit, QPushButton, QWidget, ) from ..config import AppConfig class RestoreDialog(QDialog): def __init__(self, config: AppConfig, parent: QWidget | None = None) -> None: super().__init__(parent) self._cfg = config self.setWindowTitle("Движок восстановления") self.setMinimumWidth(520) self.engine = QComboBox() self.engine.addItem("Инпейнт (быстро, замазывает — без модели)", "inpaint") self.engine.addItem("DeepMosaics (реальное расцензуривание, нужна модель+GPU)", "deepmosaics") self.engine.setCurrentIndex(1 if config.restorer == "deepmosaics" else 0) self.engine.currentIndexChanged.connect(self._sync) self.dm_dir = QLineEdit(config.dm_dir or "") self.dm_model = QLineEdit(config.dm_model or "") self.dm_python = QLineEdit(config.dm_python or "") self.dm_python.setPlaceholderText("по умолчанию — python текущего venv") self.dm_gpu = QLineEdit(config.dm_gpu or "0") self.dm_gpu.setPlaceholderText("0 = первая CUDA-карта, -1 = CPU (медленно)") form = QFormLayout(self) form.addRow("Движок:", self.engine) form.addRow("Папка DeepMosaics:", self._with_browse(self.dm_dir, self._browse_dir)) form.addRow("Веса (clean_*.pth):", self._with_browse(self.dm_model, self._browse_model)) form.addRow("Python для DeepMosaics:", self._with_browse(self.dm_python, self._browse_python)) form.addRow("GPU id:", self.dm_gpu) hint = QLabel( "DeepMosaics ставится отдельно (git clone + зависимости). Рядом с весами " "clean_*.pth должен лежать mosaic_position.pth.\n" "Для покадрового режима берите clean_youknow_resnet_9blocks.pth — " "видеомодель clean_youknow_video.pth покадрово не работает. См. README." ) hint.setWordWrap(True) form.addRow(hint) buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) buttons.accepted.connect(self.accept) buttons.rejected.connect(self.reject) form.addRow(buttons) self._sync() def _with_browse(self, line: QLineEdit, slot) -> QWidget: w = QWidget() h = QHBoxLayout(w) h.setContentsMargins(0, 0, 0, 0) h.addWidget(line, 1) btn = QPushButton("…") btn.setMaximumWidth(32) btn.clicked.connect(slot) h.addWidget(btn) return w def _sync(self) -> None: is_dm = self.engine.currentData() == "deepmosaics" for w in (self.dm_dir, self.dm_model, self.dm_python, self.dm_gpu): w.setEnabled(is_dm) def _browse_dir(self) -> None: d = QFileDialog.getExistingDirectory(self, "Папка DeepMosaics", self.dm_dir.text()) if d: self.dm_dir.setText(d) def _browse_model(self) -> None: start = self.dm_model.text() or self.dm_dir.text() p, _ = QFileDialog.getOpenFileName(self, "Веса DeepMosaics", start, "Веса (*.pth);;Все файлы (*.*)") if p: self.dm_model.setText(p) def _browse_python(self) -> None: p, _ = QFileDialog.getOpenFileName(self, "Python для DeepMosaics", self.dm_python.text(), "python (*.exe);;Все файлы (*.*)") if p: self.dm_python.setText(p) def apply_to_config(self) -> None: self._cfg.restorer = self.engine.currentData() self._cfg.dm_dir = self.dm_dir.text().strip() or None self._cfg.dm_model = self.dm_model.text().strip() or None self._cfg.dm_python = self.dm_python.text().strip() or None self._cfg.dm_gpu = self.dm_gpu.text().strip() or "0"