106 lines
4.3 KiB
Python
106 lines
4.3 KiB
Python
"""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"
|