Add main window: open, suggest, upscale, save as, overwrite

One QProcess drives both tools; Open and Upscale are disabled while it
runs so a result never lands on a replaced job. Images load with EXIF
orientation applied and go to sd-cli and imajev as an upright PNG.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

2 files changed, 284 insertions, 3 deletions

diff --git a/test_upscale_qt.py b/test_upscale_qt.py
index 439807c..1fbfe2c 100644
--- a/test_upscale_qt.py
+++ b/test_upscale_qt.py
@@ -96,6 +96,22 @@ def test_compare_view():
     assert shot.pixelColor(300, 200) == QColor('blue')
 
 
+def test_main_window():
+    w = uq.Main()
+    w.cfg.setValue('imajev_model', '')       # never run the real VLM in tests
+    w.cfg.setValue('upscalers_dir', os.path.join(TMP, 'models'))
+    w.reload_models()
+    assert not w.a_up.isEnabled()            # no image yet
+    path = os.path.join(TMP, 'in.png')
+    img = QImage(8, 6, QImage.Format.Format_RGB32)
+    img.fill(QColor('green'))
+    img.save(path)
+    w.load(path)
+    assert len(w.jobs) == 1 and w.jobs[0].image.size() == img.size()
+    assert os.path.exists(w.jobs[0].input)
+    assert w.a_up.isEnabled() and not w.a_save.isEnabled()
+
+
 if __name__ == '__main__':
     for name, fn in list(globals().items()):
         if name.startswith('test_'):
diff --git a/upscale-qt b/upscale-qt
index 41e3f5a..b60bd8d 100755
--- a/upscale-qt
+++ b/upscale-qt
@@ -16,15 +16,25 @@
 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 
 import os
+import shutil
 import sys
+from dataclasses import dataclass
 from pathlib import Path
 
-from PyQt6.QtCore import QPointF, QRectF, Qt
-from PyQt6.QtGui import QColor, QImage, QImageReader, QPainter, QPen
-from PyQt6.QtWidgets import QWidget
+from PyQt6.QtCore import QPointF, QProcess, QRectF, QSettings, QTemporaryDir, Qt
+from PyQt6.QtGui import (QAction, QColor, QImage, QImageReader, QKeySequence,
+                         QPainter, QPen)
+from PyQt6.QtWidgets import (QApplication, QComboBox, QDialog,
+                             QDialogButtonBox, QFileDialog, QFormLayout,
+                             QLineEdit, QMainWindow, QMessageBox, QWidget)
 
 MODEL_EXTS = ('.pth', '.safetensors', '.gguf')
 JPEG_QUALITY = 95
+DEFAULTS = {
+    'upscalers_dir': '/data/LLM-models/SD/upscalers',
+    'imajev_model': '/data/LLM-models/SD/imajev/imajev-4b-Q4_K_M.gguf',
+    'imajev_mmproj': '/data/LLM-models/SD/imajev/imajev-4b-mmproj-f16.gguf',
+}
 PROMPT = (
     'You pick an image upscaler model. Judge the subject (photo of people, '
     'landscape, anime, illustration) and also the source quality: noise, '
@@ -74,6 +84,21 @@ def save_over(img, path):
     return True
 
 
+def tail(data, lines=15):
+    text = bytes(data).decode(errors='replace').replace('\r', '\n')
+    return '\n'.join([l for l in text.splitlines() if l.strip()][-lines:])
+
+
+@dataclass
+class Job:
+    path: str          # file the user opened, overwrite target
+    image: QImage      # original, EXIF orientation applied
+    input: str         # oriented PNG fed to sd-cli and imajev
+    model: str = ''
+    out: str = ''
+    result: QImage | None = None
+
+
 class CompareView(QWidget):
     """Before left of a draggable divider, after right, one shared transform.
 
@@ -167,3 +192,243 @@ class CompareView(QWidget):
 
     def resizeEvent(self, e):
         self.fit()  # ponytail: resize resets zoom, keep the view center if it annoys
+
+
+class SettingsDialog(QDialog):
+    def __init__(self, parent, cfg):
+        super().__init__(parent)
+        self.setWindowTitle('Settings')
+        self.cfg = cfg
+        self.edits = {k: QLineEdit(str(cfg.value(k, v))) for k, v in DEFAULTS.items()}
+        form = QFormLayout(self)
+        form.addRow('Upscalers folder', self.edits['upscalers_dir'])
+        form.addRow('imajev model (.gguf)', self.edits['imajev_model'])
+        form.addRow('imajev mmproj (.gguf)', self.edits['imajev_mmproj'])
+        bb = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok
+                              | QDialogButtonBox.StandardButton.Cancel)
+        bb.accepted.connect(self.accept)
+        bb.rejected.connect(self.reject)
+        form.addRow(bb)
+
+    def accept(self):
+        for k, e in self.edits.items():
+            self.cfg.setValue(k, e.text().strip())
+        super().accept()
+
+
+class Main(QMainWindow):
+    def __init__(self, path=None):
+        super().__init__()
+        self.setWindowTitle('upscale-qt')
+        self.resize(1200, 800)
+        self.cfg = QSettings('upscale-qt', 'upscale-qt')
+        self.tmp = QTemporaryDir()
+        self.jobs = []      # one job for now, batch mode adds more
+        self.cur = None
+        self.runs = 0       # unique temp file names
+        self.what = None    # 'suggest' or 'upscale' while a process runs
+        self.err = b''
+
+        self.view = CompareView()
+        self.setCentralWidget(self.view)
+        self.models = QComboBox()
+        self.models.setSizeAdjustPolicy(QComboBox.SizeAdjustPolicy.AdjustToContents)
+        self.models.setToolTip('Upscaler model')
+
+        tb = self.addToolBar('main')
+        tb.setMovable(False)
+
+        def act(text, slot, key=None):
+            a = QAction(text, self)
+            a.triggered.connect(slot)
+            if key is not None:
+                a.setShortcut(key)
+            tb.addAction(a)
+            return a
+
+        self.a_open = act('Open', self.open_dialog, QKeySequence.StandardKey.Open)
+        tb.addWidget(self.models)
+        self.a_up = act('Upscale', self.upscale, 'Ctrl+U')
+        self.a_save = act('Save As', self.save_as, QKeySequence.StandardKey.SaveAs)
+        self.a_over = act('Overwrite', self.overwrite)
+        act('Settings', self.settings)
+
+        self.proc = QProcess(self)
+        self.proc.finished.connect(self.proc_done)
+        self.proc.errorOccurred.connect(self.proc_error)
+        self.proc.readyReadStandardError.connect(self.proc_stderr)
+
+        self.reload_models()
+        if path:
+            self.load(path)
+
+    def get(self, key):
+        return str(self.cfg.value(key, DEFAULTS[key]))
+
+    def job(self):
+        return self.jobs[self.cur] if self.cur is not None else None
+
+    def names(self):
+        return [self.models.itemText(i) for i in range(self.models.count())]
+
+    def reload_models(self):
+        keep = self.models.currentText()
+        self.models.clear()
+        self.models.addItems(list_models(self.get('upscalers_dir')))
+        if keep:
+            self.models.setCurrentText(keep)
+        self.refresh()
+
+    def refresh(self):
+        busy = self.what is not None
+        job = self.job()
+        done = job is not None and job.result is not None
+        self.a_open.setEnabled(not busy)
+        self.models.setEnabled(not busy)
+        self.a_up.setEnabled(job is not None and not busy and self.models.count() > 0)
+        self.a_up.setToolTip('Upscale' if self.models.count() else
+                             f'No models in {self.get("upscalers_dir")}, see Settings')
+        self.a_save.setEnabled(done and not busy)
+        self.a_over.setEnabled(done and not busy)
+
+    def open_dialog(self):
+        path, _ = QFileDialog.getOpenFileName(
+            self, 'Open image', str(self.cfg.value('last_dir', str(Path.home()))),
+            'Images (*.png *.jpg *.jpeg *.webp *.bmp *.tif *.tiff);;All files (*)')
+        if path:
+            self.load(path)
+
+    def load(self, path):
+        r = QImageReader(path)
+        r.setAutoTransform(True)
+        img = r.read()
+        if img.isNull():
+            QMessageBox.warning(self, 'upscale-qt', f'Cannot read {path}:\n{r.errorString()}')
+            return
+        self.runs += 1
+        src = self.tmp.filePath(f'in{self.runs}.png')
+        if not img.save(src):
+            QMessageBox.warning(self, 'upscale-qt', f'Cannot write temp file {src}')
+            return
+        self.jobs = [Job(path, img, src)]
+        self.cur = 0
+        self.cfg.setValue('last_dir', str(Path(path).parent))
+        self.setWindowTitle(f'upscale-qt - {Path(path).name}')
+        self.view.set_images(img)
+        self.statusBar().showMessage(f'{img.width()}x{img.height()}')
+        self.refresh()
+        self.suggest()
+
+    def run(self, what, prog, args, msg):
+        self.what, self.err = what, b''
+        self.statusBar().showMessage(msg)
+        self.refresh()
+        self.proc.start(prog, args)
+
+    def suggest(self):
+        model, mmproj = self.get('imajev_model'), self.get('imajev_mmproj')
+        if not (self.names() and os.path.isfile(model) and os.path.isfile(mmproj)
+                and shutil.which('llama-mtmd-cli')):
+            return  # suggestion is optional, skip quietly
+        self.run('suggest', 'llama-mtmd-cli',
+                 llama_args(model, mmproj, self.job().input, self.names()),
+                 'imajev is looking at the image...')
+
+    def upscale(self):
+        job = self.job()
+        job.model = self.models.currentText()
+        self.runs += 1
+        job.out = self.tmp.filePath(f'out{self.runs}.png')
+        self.run('upscale', 'sd-cli',
+                 sd_args(os.path.join(self.get('upscalers_dir'), job.model),
+                         job.input, job.out),
+                 f'Upscaling with {job.model}...')
+
+    def proc_stderr(self):
+        self.err = (self.err + bytes(self.proc.readAllStandardError()))[-8000:]
+        if self.what == 'upscale':
+            last = tail(self.err, 1)
+            if last:
+                self.statusBar().showMessage(last)
+
+    def proc_done(self, code, status):
+        what, self.what = self.what, None
+        ok = code == 0 and status == QProcess.ExitStatus.NormalExit
+        job = self.job()
+        if what == 'suggest':
+            answer = bytes(self.proc.readAllStandardOutput()).decode(errors='replace')
+            pick = match_model(answer, self.names()) if ok else None
+            if pick:
+                self.models.setCurrentText(pick)
+                self.statusBar().showMessage(f'imajev suggests {pick}')
+            else:
+                self.statusBar().showMessage('imajev gave no usable suggestion')
+        elif what == 'upscale':
+            img = QImage(job.out) if ok else QImage()
+            if img.isNull():
+                self.statusBar().showMessage('Upscale failed')
+                QMessageBox.warning(self, 'Upscale failed', tail(self.err) or f'exit code {code}')
+            else:
+                job.result = img
+                self.view.set_images(job.image, img)
+                self.statusBar().showMessage(
+                    f'{job.image.width()}x{job.image.height()} -> '
+                    f'{img.width()}x{img.height()} with {job.model}')
+        self.refresh()
+
+    def proc_error(self, err):
+        if err != QProcess.ProcessError.FailedToStart:
+            return  # crashes also emit finished, handled there
+        what, self.what = self.what, None
+        prog = self.proc.program()
+        if what == 'upscale':
+            QMessageBox.warning(self, 'upscale-qt', f'Cannot start {prog}: is it installed?')
+        self.statusBar().showMessage(f'Cannot start {prog}')
+        self.refresh()
+
+    def save_as(self):
+        job = self.job()
+        src = Path(job.path)
+        path, _ = QFileDialog.getSaveFileName(
+            self, 'Save upscaled image', str(src.with_name(f'{src.stem}_x4{src.suffix}')),
+            'Images (*.png *.jpg *.jpeg *.webp *.tif *.tiff)')
+        if not path:
+            return
+        if job.result.save(path, None, JPEG_QUALITY):
+            self.statusBar().showMessage(f'Saved {path}')
+        else:
+            QMessageBox.warning(self, 'upscale-qt', f'Cannot save {path}')
+
+    def overwrite(self):
+        job = self.job()
+        r = job.result
+        if QMessageBox.question(
+                self, 'Overwrite',
+                f'Replace {job.path}\nwith the {r.width()}x{r.height()} upscaled image?\n'
+                'This cannot be undone.') != QMessageBox.StandardButton.Yes:
+            return
+        if save_over(r, job.path):
+            self.statusBar().showMessage(f'Overwrote {job.path}')
+        else:
+            QMessageBox.warning(self, 'upscale-qt', f'Cannot write {job.path}, original untouched')
+
+    def settings(self):
+        if SettingsDialog(self, self.cfg).exec():
+            self.reload_models()
+
+    def closeEvent(self, e):
+        if self.proc.state() != QProcess.ProcessState.NotRunning:
+            self.proc.kill()
+            self.proc.waitForFinished(2000)
+        super().closeEvent(e)
+
+
+def main():
+    app = QApplication(sys.argv)
+    w = Main(sys.argv[1] if len(sys.argv) > 1 else None)
+    w.show()
+    sys.exit(app.exec())
+
+
+if __name__ == '__main__':
+    main()