External/SnapOtter
1
0
Fork 0
mirror of https://github.com/snapotter-hq/SnapOtter.git synced 2026-08-03 07:46:42 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix: load RealESRGAN pretrained weights for actual AI upscaling

Browse Source 
model_path was None, so the model had random weights and always fell back
to Lanczos. Now loads RealESRGAN_x4plus.pth from /opt/models/realesrgan/
(configurable via REALESRGAN_MODEL_PATH env var). Only falls back to
Lanczos on ImportError, not blanket Exception.
This commit is contained in:
Siddharth Kumar Sah
2026-04-10 00:18:13 +08:00
parent b575243e9a
commit fa9569c920
1 changed files with 16 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ai/python/upscale.py
+16 -5
View File
@@ -1,6 +1,7 @@
"""Image upscaling with Real-ESRGAN fallback to Lanczos.""" """Image upscaling with Real-ESRGAN fallback to Lanczos."""
import sys import sys
import json import json
import os
def emit_progress(percent, stage): def emit_progress(percent, stage):
@@ -8,6 +9,12 @@ def emit_progress(percent, stage):
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True) print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
REALESRGAN_MODEL_PATH = os.environ.get(
"REALESRGAN_MODEL_PATH",
"/opt/models/realesrgan/RealESRGAN_x4plus.pth",
)
def main(): def main():
input_path = sys.argv[1] input_path = sys.argv[1]
output_path = sys.argv[2] output_path = sys.argv[2]
@@ -30,20 +37,24 @@ def main():
import numpy as np import numpy as np
import torch import torch
if not os.path.exists(REALESRGAN_MODEL_PATH):
raise FileNotFoundError(f"RealESRGAN model not found: {REALESRGAN_MODEL_PATH}")
use_gpu = gpu_available() use_gpu = gpu_available()
device = torch.device("cuda" if use_gpu else "cpu") device = torch.device("cuda" if use_gpu else "cpu")
# RealESRGAN_x4plus is a 4x model internally
model = RRDBNet( model = RRDBNet(
num_in_ch=3, num_in_ch=3,
num_out_ch=3, num_out_ch=3,
num_feat=64, num_feat=64,
num_block=23, num_block=23,
num_grow_ch=32, num_grow_ch=32,
scale=scale, scale=4,
) )
upsampler = RealESRGANer( upsampler = RealESRGANer(
scale=scale, scale=4,
model_path=None, model_path=REALESRGAN_MODEL_PATH,
model=model, model=model,
half=use_gpu, half=use_gpu,
device=device, device=device,
@@ -57,8 +68,8 @@ def main():
emit_progress(95, "Saving result") emit_progress(95, "Saving result")
result.save(output_path) result.save(output_path)
method = "realesrgan" method = "realesrgan"
except (ImportError, Exception): except ImportError:
# Fallback to Lanczos upscaling # RealESRGAN not installed - fall back to Lanczos
emit_progress(50, "Upscaling with Lanczos") emit_progress(50, "Upscaling with Lanczos")
img_upscaled = img.resize(new_size, Image.LANCZOS) img_upscaled = img.resize(new_size, Image.LANCZOS)
emit_progress(95, "Saving result") emit_progress(95, "Saving result")