mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
fix: load RealESRGAN pretrained weights for actual AI upscaling
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:
@@ -1,6 +1,7 @@
|
||||
"""Image upscaling with Real-ESRGAN fallback to Lanczos."""
|
||||
import sys
|
||||
import json
|
||||
import os
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
REALESRGAN_MODEL_PATH = os.environ.get(
|
||||
"REALESRGAN_MODEL_PATH",
|
||||
"/opt/models/realesrgan/RealESRGAN_x4plus.pth",
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
input_path = sys.argv[1]
|
||||
output_path = sys.argv[2]
|
||||
@@ -30,20 +37,24 @@ def main():
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
if not os.path.exists(REALESRGAN_MODEL_PATH):
|
||||
raise FileNotFoundError(f"RealESRGAN model not found: {REALESRGAN_MODEL_PATH}")
|
||||
|
||||
use_gpu = gpu_available()
|
||||
device = torch.device("cuda" if use_gpu else "cpu")
|
||||
|
||||
# RealESRGAN_x4plus is a 4x model internally
|
||||
model = RRDBNet(
|
||||
num_in_ch=3,
|
||||
num_out_ch=3,
|
||||
num_feat=64,
|
||||
num_block=23,
|
||||
num_grow_ch=32,
|
||||
scale=scale,
|
||||
scale=4,
|
||||
)
|
||||
upsampler = RealESRGANer(
|
||||
scale=scale,
|
||||
model_path=None,
|
||||
scale=4,
|
||||
model_path=REALESRGAN_MODEL_PATH,
|
||||
model=model,
|
||||
half=use_gpu,
|
||||
device=device,
|
||||
@@ -57,8 +68,8 @@ def main():
|
||||
emit_progress(95, "Saving result")
|
||||
result.save(output_path)
|
||||
method = "realesrgan"
|
||||
except (ImportError, Exception):
|
||||
# Fallback to Lanczos upscaling
|
||||
except ImportError:
|
||||
# RealESRGAN not installed - fall back to Lanczos
|
||||
emit_progress(50, "Upscaling with Lanczos")
|
||||
img_upscaled = img.resize(new_size, Image.LANCZOS)
|
||||
emit_progress(95, "Saving result")
|
||||
|
||||
Reference in New Issue
Block a user