feat: Real-ESRGAN texture upscaling (ONNX) (#405)

qml/TexturePropertiesPanel.qml

@@ -316,13 +316,52 @@ GroupBox {
                 }
             }
 
+            // #405: AI super-resolution of the current texture (Real-ESRGAN).
+            // Shown on an ONNX build; writes <stem>_upscaled_xN.png next to the source.
+            // `upscaling` disables the scale buttons + shows a Cancel while a run
+            // is in flight (it can take minutes on CPU for large textures).
+            property bool upscaling: false
+            ThemedButton {
+                text: "Upscale 2×"
+                visible: MaterialEditorQML.aiPbrAvailable() && !parent.upscaling
+                enabled: MaterialEditorQML.textureName !== ""
+                    && MaterialEditorQML.textureName !== "*Select a texture*"
+                onClicked: { parent.upscaling = true; pbrStatus.text = "Upscaling 2×…"
+                             MaterialEditorQML.upscaleCurrentTexture(2) }
+            }
+            ThemedButton {
+                text: "Upscale 4×"
+                visible: MaterialEditorQML.aiPbrAvailable() && !parent.upscaling
+                enabled: MaterialEditorQML.textureName !== ""
+                    && MaterialEditorQML.textureName !== "*Select a texture*"
+                onClicked: { parent.upscaling = true; pbrStatus.text = "Upscaling 4×…"
+                             MaterialEditorQML.upscaleCurrentTexture(4) }
+            }
+            ThemedButton {
+                text: "Cancel"
+                visible: parent.upscaling
+                onClicked: { pbrStatus.text = "Cancelling…"; MaterialEditorQML.cancelUpscale() }
+            }
+
             Connections {
                 target: MaterialEditorQML
                 function onPbrSynthCompleted(result) {
                     pbrStatus.text = result.fromCache ? "PBR maps ready (cached)."
                                                       : "PBR maps generated."
                 }
                 function onPbrSynthError(err) { pbrStatus.text = "PBR: " + err }
+                function onUpscaleDownloading() { pbrStatus.text = "Downloading upscale model…" }
+                function onUpscaleProgress(done, total) {
+                    pbrStatus.text = "Upscaling… tile " + done + "/" + total
+                }
+                function onUpscaleCompleted(path) {
+                    pbrStatus.parent.upscaling = false
+                    pbrStatus.text = "Upscaled → " + path.split('/').pop()
+                }
+                function onUpscaleError(err) {
+                    pbrStatus.parent.upscaling = false
+                    pbrStatus.text = "Upscale: " + err
+                }
             }
         }
 

scripts/export-pbrify-onnx.py

@@ -20,27 +20,49 @@
 Then host the .onnx files and point AIAssistManager's model URLs at them.
 """
 import argparse
+import hashlib
 import os
 import sys
 import urllib.request
 
+# (stem, sha256) — content-verified before deserializing (.pth is a
+# code-execution boundary). Hashes verified against the pinned PBRIFY_REF commit.
 MODELS = {
-    "normal":    "1x-PBRify_NormalV3",
-    "roughness": "1x-PBRify_RoughnessV2",
-    "height":    "1x-PBRify_Height",
+    "normal":    ("1x-PBRify_NormalV3",
+                  "b0a18270da765f02eaae3c228203bee0677fc28b2854a562515e4aae9c61223b"),
+    "roughness": ("1x-PBRify_RoughnessV2",
+                  "7003c39041af64cdb77d5120bd560a2878538fbb14a07657d8fd12aac5773679"),
+    "height":    ("1x-PBRify_Height",
+                  "5b973ecb8bae9d96d14d77b8a8f1d88fb6a8580bcc1bb55d2872811acdc4277d"),
 }
 # Pin to a specific commit (not the mutable `main`) so exports are reproducible
 # and the source can't change under us. Bump deliberately when re-exporting.
 PBRIFY_REF = "190db5378909749bdbad0f951b5724ba066ea32d"
 BASE_URL = "https://github.com/Kim2091/PBRify_Remix/raw/" + PBRIFY_REF + "/Models/{name}.pth"
 
 
+def sha256(path: str) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
 def download(name: str, dest: str) -> None:
     url = BASE_URL.format(name=name)
     print(f"  downloading {url}")
     urllib.request.urlretrieve(url, dest)
 
 
+def verify(path: str, expected: str) -> None:
+    got = sha256(path)
+    if got != expected:
+        raise SystemExit(
+            f"SHA-256 mismatch for {path}\n  expected {expected}\n  got      {got}\n"
+            "Refusing to deserialize a .pth that doesn't match the pinned hash.")
+
+
 def export_one(pth_path: str, onnx_path: str) -> None:
     import torch
     from spandrel import ModelLoader
@@ -82,11 +104,12 @@ def main() -> int:
     os.makedirs(args.pth_dir, exist_ok=True)
     os.makedirs(args.out_dir, exist_ok=True)
 
-    for slot, name in MODELS.items():
+    for slot, (name, digest) in MODELS.items():
         print(f"=== {slot}: {name} ===")
         pth = os.path.join(args.pth_dir, name + ".pth")
         if args.download or not os.path.exists(pth):
             download(name, pth)
+        verify(pth, digest)  # before deserializing (.pth = code-exec boundary)
         export_one(pth, os.path.join(args.out_dir, name + ".onnx"))
     print("ALL EXPORTS OK")
     return 0

scripts/export-realesrgan-onnx.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""Convert the BSD-3 Real-ESRGAN weights (.pth) to ONNX for #405.
+
+ONE-TIME, OFFLINE dev tool — NOT shipped; the app runs the resulting .onnx in
+C++ via ONNX Runtime (src/TextureUpscaler.cpp). Mirrors export-pbrify-onnx.py.
+
+Models (BSD-3-Clause, Xintao Wang — https://github.com/xinntao/Real-ESRGAN,
+LICENSE has no code/weights carve-out):
+    RealESRGAN_x4plus.pth   4x super-resolution (RRDBNet)
+    RealESRGAN_x2plus.pth   2x super-resolution (RRDBNet)
+Both are 3-channel in -> 3-channel out, output H*scale x W*scale, values [0,1].
+
+Usage:
+    python3 -m venv venv
+    ./venv/bin/pip install torch spandrel onnx onnxruntime onnxscript
+    ./venv/bin/python scripts/export-realesrgan-onnx.py --download --out-dir dist/esrgan_onnx
+Then host the .onnx files and point AIAssistManager's model base URL at them.
+"""
+import argparse
+import os
+import sys
+import hashlib
+import urllib.request
+
+# (filename-stem, release tag, sha256) — pinned release assets (immutable URLs)
+# AND content-verified: .pth deserialization is a code-execution boundary, so a
+# compromised release must not be loaded. Hashes verified against the upstream
+# BSD-3 xinntao/Real-ESRGAN release assets.
+MODELS = {
+    "x4": ("RealESRGAN_x4plus", "v0.1.0",
+           "4fa0d38905f75ac06eb49a7951b426670021be3018265fd191d2125df9d682f1"),
+    "x2": ("RealESRGAN_x2plus", "v0.2.1",
+           "49fafd45f8fd7aa8d31ab2a22d14d91b536c34494a5cfe31eb5d89c2fa266abb"),
+}
+BASE_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/{tag}/{name}.pth"
+
+
+def sha256(path: str) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def download(name: str, tag: str, dest: str) -> None:
+    url = BASE_URL.format(tag=tag, name=name)
+    print(f"  downloading {url}")
+    urllib.request.urlretrieve(url, dest)
+
+
+def verify(path: str, expected: str) -> None:
+    got = sha256(path)
+    if got != expected:
+        raise SystemExit(
+            f"SHA-256 mismatch for {path}\n  expected {expected}\n  got      {got}\n"
+            "Refusing to deserialize a .pth that doesn't match the pinned hash.")
+
+
+def export_one(pth_path: str, onnx_path: str) -> None:
+    import torch
+    from spandrel import ModelLoader
+    import onnxruntime as ort
+
+    desc = ModelLoader().load_from_file(pth_path)
+    net = desc.model.eval()
+    print(f"  arch={getattr(getattr(desc,'architecture',None),'name','?')} "
+          f"scale={getattr(desc,'scale',None)}")
+
+    dummy = torch.rand(1, 3, 64, 64)
+    torch.onnx.export(
+        net, dummy, onnx_path, opset_version=18, dynamo=False,
+        input_names=["input"], output_names=["output"],
+        dynamic_axes={"input": {0: "b", 2: "h", 3: "w"},
+                      "output": {0: "b", 2: "h", 3: "w"}})
+
+    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+    r = sess.run(None, {sess.get_inputs()[0].name: dummy.numpy()})[0]
+    print(f"  -> {onnx_path} ({os.path.getsize(onnx_path)} bytes); "
+          f"64x64 -> {r.shape[2]}x{r.shape[3]} ({r.shape[2] // 64}x)")
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--download", action="store_true")
+    ap.add_argument("--pth-dir", default=".")
+    ap.add_argument("--out-dir", default="esrgan_onnx")
+    args = ap.parse_args()
+
+    os.makedirs(args.pth_dir, exist_ok=True)
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    for key, (name, tag, digest) in MODELS.items():
+        print(f"=== {key}: {name} ===")
+        pth = os.path.join(args.pth_dir, name + ".pth")
+        if args.download or not os.path.exists(pth):
+            download(name, tag, pth)
+        verify(pth, digest)  # before deserializing (.pth = code-exec boundary)
+        export_one(pth, os.path.join(args.out_dir, name + ".onnx"))
+    print("ALL EXPORTS OK")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/AIAssistManager.cpp

@@ -1,6 +1,7 @@
 #include "AIAssistManager.h"
 
 #include "NormalMapGenerator.h"
+#include "TextureUpscaler.h"
 #include "ModelDownloader.h"
 #include "SentryReporter.h"
 
@@ -32,6 +33,8 @@ const char* mapDownloadLabel(AIAssistManager::Map m) {
         case AIAssistManager::Map::Normal:    return "PBR Normal";
         case AIAssistManager::Map::Roughness: return "PBR Roughness";
         case AIAssistManager::Map::Height:    return "PBR Height";
+        case AIAssistManager::Map::UpscaleX2: return "Upscale 2x";
+        case AIAssistManager::Map::UpscaleX4: return "Upscale 4x";
     }
     return "PBR";
 }
@@ -43,6 +46,8 @@ QString AIAssistManager::mapModelFile(Map map)
         case Map::Normal:    return QStringLiteral("1x-PBRify_NormalV3.onnx");
         case Map::Roughness: return QStringLiteral("1x-PBRify_RoughnessV2.onnx");
         case Map::Height:    return QStringLiteral("1x-PBRify_Height.onnx");
+        case Map::UpscaleX2: return QStringLiteral("RealESRGAN_x2plus.onnx");
+        case Map::UpscaleX4: return QStringLiteral("RealESRGAN_x4plus.onnx");
     }
     return {};
 }
@@ -299,3 +304,57 @@ QVariantMap AIAssistManager::synthesizePbrMapsQml(const QString& albedoPath,
     if (o.contains("overwriteCache"))    opts.overwriteCache    = o["overwriteCache"].toBool();
     return synthesizePbrMaps(albedoPath, opts).toVariantMap();
 }
+
+QString AIAssistManager::ensureUpscaleModel(int scale)
+{
+    const Map m = (scale == 2) ? Map::UpscaleX2 : Map::UpscaleX4;
+    ensureModelBlocking(m);   // event-loop driven; call on the GUI thread
+    const QString p = modelPath(m);
+    return QFileInfo::exists(p) ? p : QString();
+}
+
+QString AIAssistManager::upscaleTexture(const QString& srcPath, int scale, bool overwrite)
+{
+    SentryReporter::addBreadcrumb(QStringLiteral("ai.assist.upscale"),
+        QStringLiteral("upscale %1 x%2").arg(QFileInfo(srcPath).fileName()).arg(scale));
+    emit upscaleStarted();
+
+    auto failUp = [&](const QString& msg) -> QString {
+        emit upscaleError(msg);
+        return {};
+    };
+    const QFileInfo fi(srcPath);
+    if (!fi.exists())
+        return failUp(tr("Texture not found: %1").arg(srcPath));
+    if (scale != 2 && scale != 4)
+        return failUp(tr("Upscale factor must be 2 or 4."));
+
+    // Scale-specific cache name so a cached 2× result is never returned for a
+    // 4× request (and vice versa) on the overwrite=false path.
+    const QString outPath = QDir(fi.absolutePath())
+        .filePath(fi.completeBaseName()
+                  + QStringLiteral("_upscaled_x%1.png").arg(scale));
+    if (!overwrite && QFileInfo::exists(outPath)) {  // cache: skip re-upscale
+        emit upscaleCompleted(outPath);
+        return outPath;
+    }
+
+#ifndef ENABLE_ONNX
+    return failUp(tr("Texture upscaling is not enabled. Rebuild with -DENABLE_ONNX=ON."));
+#else
+    const QImage src(srcPath);
+    if (src.isNull())
+        return failUp(tr("Could not load image: %1").arg(srcPath));
+
+    const Map m = (scale == 2) ? Map::UpscaleX2 : Map::UpscaleX4;
+    ensureModelBlocking(m);
+    const TextureUpscaler::Result res =
+        TextureUpscaler::upscale(src, modelPath(m), {});
+    if (!res.ok)
+        return failUp(res.error.isEmpty() ? tr("Upscale failed.") : res.error);
+    if (!res.image.save(outPath, "PNG"))
+        return failUp(tr("Could not write the upscaled image."));
+    emit upscaleCompleted(outPath);
+    return outPath;
+#endif
+}

src/AIAssistManager.h

@@ -53,8 +53,9 @@ class AIAssistManager : public QObject
     static AIAssistManager* instance();
     static AIAssistManager* qmlInstance(QQmlEngine* engine, QJSEngine* scriptEngine);
 
-    /// The three per-map PBRify models (each a separate CC0 SPAN .onnx).
-    enum class Map { Normal, Roughness, Height };
+    /// Per-model ONNX files (each a separate download). PBRify maps (#404, CC0)
+    /// + Real-ESRGAN upscalers (#405, BSD-3).
+    enum class Map { Normal, Roughness, Height, UpscaleX2, UpscaleX4 };
 
     /// True only when the binary was compiled with ENABLE_ONNX.
     Q_INVOKABLE bool isAvailable() const;
@@ -77,12 +78,30 @@ class AIAssistManager : public QObject
     Q_INVOKABLE QVariantMap synthesizePbrMapsQml(const QString& albedoPath,
                                                  const QVariantMap& opts = {});
 
+    // ── #405: Real-ESRGAN texture upscaling ─────────────────────────────────
+    /// Upscale the texture at `srcPath` by `scale` (2 or 4) via the BSD-3
+    /// Real-ESRGAN ONNX model (downloaded on first use). Writes
+    /// `<stem>_upscaled.png` next to the source and returns its path.
+    /// Synchronous; emits upscaleStarted/Completed/Error. Cached: an existing
+    /// output is reused unless `overwrite`.
+    Q_INVOKABLE QString upscaleTexture(const QString& srcPath, int scale = 4,
+                                       bool overwrite = false);
+
+    /// Ensure the 2×/4× upscale model is present (download + block) — MUST be
+    /// called on a thread with an event loop (the GUI thread). Lets the
+    /// GUI fetch the model first, then run the pure-CPU inference on a worker.
+    /// Returns the model path, or empty if it couldn't be made available.
+    QString ensureUpscaleModel(int scale);
+
 signals:
     void modelReadyChanged();
     void modelDownloadProgress(qint64 received, qint64 total);
     void synthesisStarted();
     void synthesisCompleted(QVariantMap result);
     void synthesisError(const QString& error);
+    void upscaleStarted();
+    void upscaleCompleted(const QString& outputPath);
+    void upscaleError(const QString& error);
 
 private:
     explicit AIAssistManager(QObject* parent = nullptr);