Add model matching, argument and overwrite helpers

match_model tries the longest stem first so an answer naming the anime
variant is not taken for its shorter prefix. save_over writes a .part file
in the original's format and renames it into place.

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

2 files changed, 156 insertions, 0 deletions

diff --git a/test_upscale_qt.py b/test_upscale_qt.py
new file mode 100644
index 0000000..526dc37
--- /dev/null
+++ b/test_upscale_qt.py
@@ -0,0 +1,84 @@
+# upscale-qt - tests for the non-GUI helpers and a GUI smoke test
+# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 2 as
+# published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+import os
+import tempfile
+from importlib.machinery import SourceFileLoader
+from importlib.util import module_from_spec, spec_from_loader
+
+TMP = tempfile.mkdtemp(prefix='upscale-qt-test-')
+os.environ['QT_QPA_PLATFORM'] = 'offscreen'
+os.environ['XDG_CONFIG_HOME'] = os.path.join(TMP, 'config')
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+_loader = SourceFileLoader('upscale_qt', os.path.join(HERE, 'upscale-qt'))
+uq = module_from_spec(spec_from_loader('upscale_qt', _loader))
+_loader.exec_module(uq)
+
+from PyQt6.QtGui import QColor, QImage, QImageReader  # noqa: E402
+
+NAMES = ['4x_detail.pth', 'RealESRGAN_x4plus.pth',
+         'RealESRGAN_x4plus_anime_6B.pth']
+
+
+def test_match_model():
+    m = uq.match_model
+    # the longer stem must win over its own prefix
+    assert m('RealESRGAN_x4plus_anime_6B.pth', NAMES) == 'RealESRGAN_x4plus_anime_6B.pth'
+    assert m('  realesrgan_x4plus\n', NAMES) == 'RealESRGAN_x4plus.pth'
+    assert m('I suggest 4x_detail.', NAMES) == '4x_detail.pth'
+    assert m('no idea', NAMES) is None
+    assert m('', []) is None
+
+
+def test_list_models():
+    d = os.path.join(TMP, 'models')
+    os.makedirs(d)
+    for f in ('b.safetensors', 'a.pth', 'notes.txt', 'c.GGUF'):
+        open(os.path.join(d, f), 'w').close()
+    assert uq.list_models(d) == ['a.pth', 'b.safetensors', 'c.GGUF']
+    assert uq.list_models(os.path.join(TMP, 'missing')) == []
+
+
+def test_args():
+    a = uq.sd_args('/m/x.pth', '/in.png', '/out.png')
+    assert a == ['-M', 'upscale', '--upscale-model', '/m/x.pth',
+                 '-i', '/in.png', '-o', '/out.png']
+    a = uq.llama_args('/m.gguf', '/mm.gguf', '/in.png', NAMES)
+    assert a[a.index('--image') + 1] == '/in.png'
+    prompt = a[a.index('-p') + 1]
+    assert all(n in prompt for n in NAMES)
+    assert 'noise' in prompt
+
+
+def test_save_over():
+    path = os.path.join(TMP, 'photo.JPG')   # uppercase suffix on purpose
+    small = QImage(4, 4, QImage.Format.Format_RGB32)
+    small.fill(QColor('red'))
+    assert small.save(path, 'jpeg')
+    big = QImage(16, 16, QImage.Format.Format_RGB32)
+    big.fill(QColor('blue'))
+    assert uq.save_over(big, path)
+    assert QImage(path).width() == 16
+    assert bytes(QImageReader.imageFormat(path)) == b'jpeg'
+    assert not os.path.exists(path + '.part')
+
+
+if __name__ == '__main__':
+    for name, fn in list(globals().items()):
+        if name.startswith('test_'):
+            fn()
+            print('ok', name)
diff --git a/upscale-qt b/upscale-qt
new file mode 100755
index 0000000..5baa355
--- /dev/null
+++ b/upscale-qt
@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+# upscale-qt - upscale an image with sd-cli and compare before/after
+# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 2 as
+# published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+
+import os
+import sys
+from pathlib import Path
+
+from PyQt6.QtGui import QImage, QImageReader
+
+MODEL_EXTS = ('.pth', '.safetensors', '.gguf')
+JPEG_QUALITY = 95
+PROMPT = (
+    'You pick an image upscaler model. Judge the subject (photo of people, '
+    'landscape, anime, illustration) and also the source quality: noise, '
+    'JPEG compression, blur. Noisy or compressed sources need a denoising '
+    'model, clean sources a detail-preserving one. Answer with exactly one '
+    'filename from this list and nothing else: {names}')
+
+
+def list_models(folder):
+    try:
+        return sorted(f for f in os.listdir(folder)
+                      if f.lower().endswith(MODEL_EXTS))
+    except OSError:
+        return []
+
+
+def match_model(answer, names):
+    """Map imajev's free-text answer to one of names, or None.
+
+    Longest stem first, so 'x4plus_anime_6B' is not taken for 'x4plus'."""
+    answer = answer.lower()
+    for name in sorted(names, key=lambda n: len(Path(n).stem), reverse=True):
+        if Path(name).stem.lower() in answer:
+            return name
+    return None
+
+
+def sd_args(model, src, out):
+    return ['-M', 'upscale', '--upscale-model', model, '-i', src, '-o', out]
+
+
+def llama_args(model, mmproj, src, names):
+    return ['-m', model, '--mmproj', mmproj, '--image', src,
+            '--temp', '0', '-n', '64',
+            '-p', PROMPT.format(names=' '.join(names))]
+
+
+def save_over(img, path):
+    """Replace path with img in the original's format, never truncating it."""
+    fmt = bytes(QImageReader.imageFormat(path)).decode() or None
+    part = path + '.part'
+    if not img.save(part, fmt, JPEG_QUALITY if fmt == 'jpeg' else -1):
+        if os.path.exists(part):
+            os.remove(part)
+        return False
+    os.replace(part, path)
+    return True