Skip to content

Commit 3de731a

Browse files
feat: Real-ESRGAN texture upscaling (ONNX) (#405) (#749)
* feat(onnx): Real-ESRGAN texture upscaling — core + facade + CLI (#405) ONNX-backed 2x/4x super-resolution, reusing the #404 ONNX infra (ENABLE_ONNX, ModelDownloader, NCHW conversion). - TextureUpscaler (Ogre-free, like PbrMapSynth): scale-aware overlapping-tile upscale that composites results in OUTPUT space with a feathered seam blend. Detects the scale factor from the model's output/input ratio at runtime (validates the output tensor element count before copying). Reuses PbrMapSynth::toNCHW / nchwToRgb. - AIAssistManager: Map enum extended with UpscaleX2/UpscaleX4 (BSD-3 Real-ESRGAN models, downloaded on first use from the same HF repo); upscaleTexture(srcPath, scale, overwrite) ensures+runs+caches, writes <stem>_upscaled.png, emits upscaleStarted/Completed/Error. Sentry breadcrumb ai.assist.upscale. - CLI: `qtmesh material --texture low.png --upscale {2|4} [-o high.png]` (cmdMaterialUpscale, delegates to the facade). ENABLE_ONNX-guarded. Model: Real-ESRGAN x4plus / x2plus (BSD-3-Clause, xinntao), exported to ONNX and hosted at fernandotonon/QtMeshEditor-models. Verified end-to-end: 256x256 → 1024x1024 (4x) and 128 → 256 (2x), model auto-downloaded from HF; bad factor / missing texture rejected with usage errors. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(#405): Real-ESRGAN upscaling — MCP + GUI + tests + docs - MCP: upscale_texture tool (texture_path, scale 2/4, overwrite) → AIAssistManager. - GUI: "Upscale 2× / 4×" buttons in the Material Editor Texture Properties panel (ONNX-only), via MaterialEditorQML::upscaleCurrentTexture which relays the facade's upscale signals. - Tests: CLIPipeline upscale coverage (missing texture / bad factor / non-numeric / no-model-fails-clean, offline-guarded) + TextureUpscaler error-contract tests. - scripts/export-realesrgan-onnx.py: one-time offline .pth→ONNX exporter (BSD-3 Real-ESRGAN x4plus/x2plus, pinned release assets). NOT shipped. - docs: CLAUDE.md CLI line + architecture entry; refreshed the #404 hosting note (models are now hosted on HF) + the QTMESH_PBR_NO_DOWNLOAD offline guard. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(onnx): bundle ONNX Runtime SONAME symlinks so the packaged binary loads (#405) scan-assets-qtmesh failed: the Linux .deb/Docker binary aborts with "libonnxruntime.so.1: cannot open shared object file". Root cause (latent since #404 turned ENABLE_ONNX on for the Linux release): the POST_BUILD copied only the single RESOLVED versioned lib (libonnxruntime.so.1.20.1) next to the binary, but the loader requests the SONAME (libonnxruntime.so.1), which wasn't shipped — so `./bin/*.so*` packaging never included a name the binary could load. - OnnxRuntime.cmake exposes QTMESH_ONNX_LIB_DIR. - The app + UnitTests POST_BUILD now glob-copy every libonnxruntime.so* / .dylib / .dll (versioned file + SONAME symlinks) next to the binary, so the packaged .deb/Docker image resolves libonnxruntime.so.1 at runtime. Verified on macOS: both libonnxruntime.1.20.1.dylib and libonnxruntime.dylib are now copied next to the binary (previously only one). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(#405): address PR review — upscale robustness, threading, supply chain CodeRabbit findings on the Real-ESRGAN work: Major: - Scale-aware cache: output is now <stem>_upscaled_x{2,4}.png so a cached 2× result is never returned for a 4× request (and vice versa). - GUI no longer freezes: upscaleCurrentTexture ensures the model on the GUI thread (download needs an event loop), then runs the pure-CPU tiled inference on a std::thread and marshals the result back via a queued invocation. Added AIAssistManager::ensureUpscaleModel for the main-thread ensure step. (CLI/MCP keep the synchronous path.) - Preserve source alpha: cutout/opacity textures keep their alpha (upscaled nearest-to-match and reapplied) instead of coming back fully opaque. - Guard the output-canvas allocation: compute size in 64-bit, cap at 256 Mpx, and catch std::bad_alloc so a huge input / bad scale fails cleanly instead of overflowing int or terminating outside the Ort handler. - CLI honors -o strictly: rename→copy fallback (cross-device), and return exit 1 if the requested output file can't be produced (was silently exit 0). - Export scripts verify SHA-256 of each .pth before deserializing (.pth is a code-execution boundary) — both realesrgan + pbrify scripts. Minor: - Validate the detected scale is a uniform integer factor on BOTH axes, and that each tile returns exactly tw*scale × th*scale (else fail, not silent crop). - CLAUDE.md: correct the stale "empty default" model-base-URL wording (now the hosted HF repo). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(#405): upscale progress + cancel + download status + multithreading The GUI gave no feedback during an upscale (a 2h run looked hung), and CPU inference was pinned to a single core. - TextureUpscaler: optional ProgressFn(done,total) reported per tile; returning false cancels (error="cancelled"). SetIntraOpNumThreads now uses hardware_concurrency-1 instead of 1 — a 256→1024 4× drops from ~2 min to ~7.5s (~7 cores). (Scoped to upscaling; PbrMapSynth stays single-threaded — its maps are small/fast.) - MaterialEditorQML: upscaleCurrentTexture runs on a worker, emits upscaleDownloading (first-run model fetch), upscaleProgress (per tile), and upscaleCompleted/Error; cancelUpscale() sets a shared atomic the worker's progress callback checks. Model-ensure stays on the GUI thread (download needs an event loop). - QML: status shows "Downloading upscale model…" / "Upscaling… tile X/Y", the scale buttons disable mid-run, and a Cancel button is shown. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent d3d7c7c commit 3de731a

18 files changed

Lines changed: 919 additions & 7 deletions

CLAUDE.md

Lines changed: 3 additions & 1 deletion
Large diffs are not rendered by default.

qml/TexturePropertiesPanel.qml

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -316,13 +316,52 @@ GroupBox {
316316
}
317317
}
318318

319+
// #405: AI super-resolution of the current texture (Real-ESRGAN).
320+
// Shown on an ONNX build; writes <stem>_upscaled_xN.png next to the source.
321+
// `upscaling` disables the scale buttons + shows a Cancel while a run
322+
// is in flight (it can take minutes on CPU for large textures).
323+
property bool upscaling: false
324+
ThemedButton {
325+
text: "Upscale 2×"
326+
visible: MaterialEditorQML.aiPbrAvailable() && !parent.upscaling
327+
enabled: MaterialEditorQML.textureName !== ""
328+
&& MaterialEditorQML.textureName !== "*Select a texture*"
329+
onClicked: { parent.upscaling = true; pbrStatus.text = "Upscaling 2×…"
330+
MaterialEditorQML.upscaleCurrentTexture(2) }
331+
}
332+
ThemedButton {
333+
text: "Upscale 4×"
334+
visible: MaterialEditorQML.aiPbrAvailable() && !parent.upscaling
335+
enabled: MaterialEditorQML.textureName !== ""
336+
&& MaterialEditorQML.textureName !== "*Select a texture*"
337+
onClicked: { parent.upscaling = true; pbrStatus.text = "Upscaling 4×…"
338+
MaterialEditorQML.upscaleCurrentTexture(4) }
339+
}
340+
ThemedButton {
341+
text: "Cancel"
342+
visible: parent.upscaling
343+
onClicked: { pbrStatus.text = "Cancelling…"; MaterialEditorQML.cancelUpscale() }
344+
}
345+
319346
Connections {
320347
target: MaterialEditorQML
321348
function onPbrSynthCompleted(result) {
322349
pbrStatus.text = result.fromCache ? "PBR maps ready (cached)."
323350
: "PBR maps generated."
324351
}
325352
function onPbrSynthError(err) { pbrStatus.text = "PBR: " + err }
353+
function onUpscaleDownloading() { pbrStatus.text = "Downloading upscale model…" }
354+
function onUpscaleProgress(done, total) {
355+
pbrStatus.text = "Upscaling… tile " + done + "/" + total
356+
}
357+
function onUpscaleCompleted(path) {
358+
pbrStatus.parent.upscaling = false
359+
pbrStatus.text = "Upscaled → " + path.split('/').pop()
360+
}
361+
function onUpscaleError(err) {
362+
pbrStatus.parent.upscaling = false
363+
pbrStatus.text = "Upscale: " + err
364+
}
326365
}
327366
}
328367

scripts/export-pbrify-onnx.py

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,27 +20,49 @@
2020
Then host the .onnx files and point AIAssistManager's model URLs at them.
2121
"""
2222
import argparse
23+
import hashlib
2324
import os
2425
import sys
2526
import urllib.request
2627

28+
# (stem, sha256) — content-verified before deserializing (.pth is a
29+
# code-execution boundary). Hashes verified against the pinned PBRIFY_REF commit.
2730
MODELS = {
28-
"normal": "1x-PBRify_NormalV3",
29-
"roughness": "1x-PBRify_RoughnessV2",
30-
"height": "1x-PBRify_Height",
31+
"normal": ("1x-PBRify_NormalV3",
32+
"b0a18270da765f02eaae3c228203bee0677fc28b2854a562515e4aae9c61223b"),
33+
"roughness": ("1x-PBRify_RoughnessV2",
34+
"7003c39041af64cdb77d5120bd560a2878538fbb14a07657d8fd12aac5773679"),
35+
"height": ("1x-PBRify_Height",
36+
"5b973ecb8bae9d96d14d77b8a8f1d88fb6a8580bcc1bb55d2872811acdc4277d"),
3137
}
3238
# Pin to a specific commit (not the mutable `main`) so exports are reproducible
3339
# and the source can't change under us. Bump deliberately when re-exporting.
3440
PBRIFY_REF = "190db5378909749bdbad0f951b5724ba066ea32d"
3541
BASE_URL = "https://github.com/Kim2091/PBRify_Remix/raw/" + PBRIFY_REF + "/Models/{name}.pth"
3642

3743

44+
def sha256(path: str) -> str:
45+
h = hashlib.sha256()
46+
with open(path, "rb") as f:
47+
for chunk in iter(lambda: f.read(1 << 20), b""):
48+
h.update(chunk)
49+
return h.hexdigest()
50+
51+
3852
def download(name: str, dest: str) -> None:
3953
url = BASE_URL.format(name=name)
4054
print(f" downloading {url}")
4155
urllib.request.urlretrieve(url, dest)
4256

4357

58+
def verify(path: str, expected: str) -> None:
59+
got = sha256(path)
60+
if got != expected:
61+
raise SystemExit(
62+
f"SHA-256 mismatch for {path}\n expected {expected}\n got {got}\n"
63+
"Refusing to deserialize a .pth that doesn't match the pinned hash.")
64+
65+
4466
def export_one(pth_path: str, onnx_path: str) -> None:
4567
import torch
4668
from spandrel import ModelLoader
@@ -82,11 +104,12 @@ def main() -> int:
82104
os.makedirs(args.pth_dir, exist_ok=True)
83105
os.makedirs(args.out_dir, exist_ok=True)
84106

85-
for slot, name in MODELS.items():
107+
for slot, (name, digest) in MODELS.items():
86108
print(f"=== {slot}: {name} ===")
87109
pth = os.path.join(args.pth_dir, name + ".pth")
88110
if args.download or not os.path.exists(pth):
89111
download(name, pth)
112+
verify(pth, digest) # before deserializing (.pth = code-exec boundary)
90113
export_one(pth, os.path.join(args.out_dir, name + ".onnx"))
91114
print("ALL EXPORTS OK")
92115
return 0

scripts/export-realesrgan-onnx.py

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
#!/usr/bin/env python3
2+
"""Convert the BSD-3 Real-ESRGAN weights (.pth) to ONNX for #405.
3+
4+
ONE-TIME, OFFLINE dev tool — NOT shipped; the app runs the resulting .onnx in
5+
C++ via ONNX Runtime (src/TextureUpscaler.cpp). Mirrors export-pbrify-onnx.py.
6+
7+
Models (BSD-3-Clause, Xintao Wang — https://github.com/xinntao/Real-ESRGAN,
8+
LICENSE has no code/weights carve-out):
9+
RealESRGAN_x4plus.pth 4x super-resolution (RRDBNet)
10+
RealESRGAN_x2plus.pth 2x super-resolution (RRDBNet)
11+
Both are 3-channel in -> 3-channel out, output H*scale x W*scale, values [0,1].
12+
13+
Usage:
14+
python3 -m venv venv
15+
./venv/bin/pip install torch spandrel onnx onnxruntime onnxscript
16+
./venv/bin/python scripts/export-realesrgan-onnx.py --download --out-dir dist/esrgan_onnx
17+
Then host the .onnx files and point AIAssistManager's model base URL at them.
18+
"""
19+
import argparse
20+
import os
21+
import sys
22+
import hashlib
23+
import urllib.request
24+
25+
# (filename-stem, release tag, sha256) — pinned release assets (immutable URLs)
26+
# AND content-verified: .pth deserialization is a code-execution boundary, so a
27+
# compromised release must not be loaded. Hashes verified against the upstream
28+
# BSD-3 xinntao/Real-ESRGAN release assets.
29+
MODELS = {
30+
"x4": ("RealESRGAN_x4plus", "v0.1.0",
31+
"4fa0d38905f75ac06eb49a7951b426670021be3018265fd191d2125df9d682f1"),
32+
"x2": ("RealESRGAN_x2plus", "v0.2.1",
33+
"49fafd45f8fd7aa8d31ab2a22d14d91b536c34494a5cfe31eb5d89c2fa266abb"),
34+
}
35+
BASE_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/{tag}/{name}.pth"
36+
37+
38+
def sha256(path: str) -> str:
39+
h = hashlib.sha256()
40+
with open(path, "rb") as f:
41+
for chunk in iter(lambda: f.read(1 << 20), b""):
42+
h.update(chunk)
43+
return h.hexdigest()
44+
45+
46+
def download(name: str, tag: str, dest: str) -> None:
47+
url = BASE_URL.format(tag=tag, name=name)
48+
print(f" downloading {url}")
49+
urllib.request.urlretrieve(url, dest)
50+
51+
52+
def verify(path: str, expected: str) -> None:
53+
got = sha256(path)
54+
if got != expected:
55+
raise SystemExit(
56+
f"SHA-256 mismatch for {path}\n expected {expected}\n got {got}\n"
57+
"Refusing to deserialize a .pth that doesn't match the pinned hash.")
58+
59+
60+
def export_one(pth_path: str, onnx_path: str) -> None:
61+
import torch
62+
from spandrel import ModelLoader
63+
import onnxruntime as ort
64+
65+
desc = ModelLoader().load_from_file(pth_path)
66+
net = desc.model.eval()
67+
print(f" arch={getattr(getattr(desc,'architecture',None),'name','?')} "
68+
f"scale={getattr(desc,'scale',None)}")
69+
70+
dummy = torch.rand(1, 3, 64, 64)
71+
torch.onnx.export(
72+
net, dummy, onnx_path, opset_version=18, dynamo=False,
73+
input_names=["input"], output_names=["output"],
74+
dynamic_axes={"input": {0: "b", 2: "h", 3: "w"},
75+
"output": {0: "b", 2: "h", 3: "w"}})
76+
77+
sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
78+
r = sess.run(None, {sess.get_inputs()[0].name: dummy.numpy()})[0]
79+
print(f" -> {onnx_path} ({os.path.getsize(onnx_path)} bytes); "
80+
f"64x64 -> {r.shape[2]}x{r.shape[3]} ({r.shape[2] // 64}x)")
81+
82+
83+
def main() -> int:
84+
ap = argparse.ArgumentParser(description=__doc__)
85+
ap.add_argument("--download", action="store_true")
86+
ap.add_argument("--pth-dir", default=".")
87+
ap.add_argument("--out-dir", default="esrgan_onnx")
88+
args = ap.parse_args()
89+
90+
os.makedirs(args.pth_dir, exist_ok=True)
91+
os.makedirs(args.out_dir, exist_ok=True)
92+
93+
for key, (name, tag, digest) in MODELS.items():
94+
print(f"=== {key}: {name} ===")
95+
pth = os.path.join(args.pth_dir, name + ".pth")
96+
if args.download or not os.path.exists(pth):
97+
download(name, tag, pth)
98+
verify(pth, digest) # before deserializing (.pth = code-exec boundary)
99+
export_one(pth, os.path.join(args.out_dir, name + ".onnx"))
100+
print("ALL EXPORTS OK")
101+
return 0
102+
103+
104+
if __name__ == "__main__":
105+
sys.exit(main())

src/AIAssistManager.cpp

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#include "AIAssistManager.h"
22

33
#include "NormalMapGenerator.h"
4+
#include "TextureUpscaler.h"
45
#include "ModelDownloader.h"
56
#include "SentryReporter.h"
67

@@ -32,6 +33,8 @@ const char* mapDownloadLabel(AIAssistManager::Map m) {
3233
case AIAssistManager::Map::Normal: return "PBR Normal";
3334
case AIAssistManager::Map::Roughness: return "PBR Roughness";
3435
case AIAssistManager::Map::Height: return "PBR Height";
36+
case AIAssistManager::Map::UpscaleX2: return "Upscale 2x";
37+
case AIAssistManager::Map::UpscaleX4: return "Upscale 4x";
3538
}
3639
return "PBR";
3740
}
@@ -43,6 +46,8 @@ QString AIAssistManager::mapModelFile(Map map)
4346
case Map::Normal: return QStringLiteral("1x-PBRify_NormalV3.onnx");
4447
case Map::Roughness: return QStringLiteral("1x-PBRify_RoughnessV2.onnx");
4548
case Map::Height: return QStringLiteral("1x-PBRify_Height.onnx");
49+
case Map::UpscaleX2: return QStringLiteral("RealESRGAN_x2plus.onnx");
50+
case Map::UpscaleX4: return QStringLiteral("RealESRGAN_x4plus.onnx");
4651
}
4752
return {};
4853
}
@@ -299,3 +304,57 @@ QVariantMap AIAssistManager::synthesizePbrMapsQml(const QString& albedoPath,
299304
if (o.contains("overwriteCache")) opts.overwriteCache = o["overwriteCache"].toBool();
300305
return synthesizePbrMaps(albedoPath, opts).toVariantMap();
301306
}
307+
308+
QString AIAssistManager::ensureUpscaleModel(int scale)
309+
{
310+
const Map m = (scale == 2) ? Map::UpscaleX2 : Map::UpscaleX4;
311+
ensureModelBlocking(m); // event-loop driven; call on the GUI thread
312+
const QString p = modelPath(m);
313+
return QFileInfo::exists(p) ? p : QString();
314+
}
315+
316+
QString AIAssistManager::upscaleTexture(const QString& srcPath, int scale, bool overwrite)
317+
{
318+
SentryReporter::addBreadcrumb(QStringLiteral("ai.assist.upscale"),
319+
QStringLiteral("upscale %1 x%2").arg(QFileInfo(srcPath).fileName()).arg(scale));
320+
emit upscaleStarted();
321+
322+
auto failUp = [&](const QString& msg) -> QString {
323+
emit upscaleError(msg);
324+
return {};
325+
};
326+
const QFileInfo fi(srcPath);
327+
if (!fi.exists())
328+
return failUp(tr("Texture not found: %1").arg(srcPath));
329+
if (scale != 2 && scale != 4)
330+
return failUp(tr("Upscale factor must be 2 or 4."));
331+
332+
// Scale-specific cache name so a cached 2× result is never returned for a
333+
// 4× request (and vice versa) on the overwrite=false path.
334+
const QString outPath = QDir(fi.absolutePath())
335+
.filePath(fi.completeBaseName()
336+
+ QStringLiteral("_upscaled_x%1.png").arg(scale));
337+
if (!overwrite && QFileInfo::exists(outPath)) { // cache: skip re-upscale
338+
emit upscaleCompleted(outPath);
339+
return outPath;
340+
}
341+
342+
#ifndef ENABLE_ONNX
343+
return failUp(tr("Texture upscaling is not enabled. Rebuild with -DENABLE_ONNX=ON."));
344+
#else
345+
const QImage src(srcPath);
346+
if (src.isNull())
347+
return failUp(tr("Could not load image: %1").arg(srcPath));
348+
349+
const Map m = (scale == 2) ? Map::UpscaleX2 : Map::UpscaleX4;
350+
ensureModelBlocking(m);
351+
const TextureUpscaler::Result res =
352+
TextureUpscaler::upscale(src, modelPath(m), {});
353+
if (!res.ok)
354+
return failUp(res.error.isEmpty() ? tr("Upscale failed.") : res.error);
355+
if (!res.image.save(outPath, "PNG"))
356+
return failUp(tr("Could not write the upscaled image."));
357+
emit upscaleCompleted(outPath);
358+
return outPath;
359+
#endif
360+
}

src/AIAssistManager.h

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,8 +53,9 @@ class AIAssistManager : public QObject
5353
static AIAssistManager* instance();
5454
static AIAssistManager* qmlInstance(QQmlEngine* engine, QJSEngine* scriptEngine);
5555

56-
/// The three per-map PBRify models (each a separate CC0 SPAN .onnx).
57-
enum class Map { Normal, Roughness, Height };
56+
/// Per-model ONNX files (each a separate download). PBRify maps (#404, CC0)
57+
/// + Real-ESRGAN upscalers (#405, BSD-3).
58+
enum class Map { Normal, Roughness, Height, UpscaleX2, UpscaleX4 };
5859

5960
/// True only when the binary was compiled with ENABLE_ONNX.
6061
Q_INVOKABLE bool isAvailable() const;
@@ -77,12 +78,30 @@ class AIAssistManager : public QObject
7778
Q_INVOKABLE QVariantMap synthesizePbrMapsQml(const QString& albedoPath,
7879
const QVariantMap& opts = {});
7980

81+
// ── #405: Real-ESRGAN texture upscaling ─────────────────────────────────
82+
/// Upscale the texture at `srcPath` by `scale` (2 or 4) via the BSD-3
83+
/// Real-ESRGAN ONNX model (downloaded on first use). Writes
84+
/// `<stem>_upscaled.png` next to the source and returns its path.
85+
/// Synchronous; emits upscaleStarted/Completed/Error. Cached: an existing
86+
/// output is reused unless `overwrite`.
87+
Q_INVOKABLE QString upscaleTexture(const QString& srcPath, int scale = 4,
88+
bool overwrite = false);
89+
90+
/// Ensure the 2×/4× upscale model is present (download + block) — MUST be
91+
/// called on a thread with an event loop (the GUI thread). Lets the
92+
/// GUI fetch the model first, then run the pure-CPU inference on a worker.
93+
/// Returns the model path, or empty if it couldn't be made available.
94+
QString ensureUpscaleModel(int scale);
95+
8096
signals:
8197
void modelReadyChanged();
8298
void modelDownloadProgress(qint64 received, qint64 total);
8399
void synthesisStarted();
84100
void synthesisCompleted(QVariantMap result);
85101
void synthesisError(const QString& error);
102+
void upscaleStarted();
103+
void upscaleCompleted(const QString& outputPath);
104+
void upscaleError(const QString& error);
86105

87106
private:
88107
explicit AIAssistManager(QObject* parent = nullptr);

0 commit comments

Comments
 (0)