fix: auto-download default model and avoid passing None to RealESRGANer
All checks were successful
Build and Push Docker Image / build-gpu (push) Successful in 1m25s
Build and Push Docker Image / build-cpu (push) Successful in 1m7s

This commit is contained in:
Developer
2026-02-16 21:54:23 +01:00
parent 7c4a3c5abf
commit dd861a2346

View File

@@ -7,6 +7,8 @@ import cv2
import numpy as np
from app.config import settings
from app.services import model_manager
import urllib.request
logger = logging.getLogger(__name__)
@@ -45,6 +47,24 @@ class RealESRGANBridge:
model_path = os.path.join(settings.models_dir, f'{model_name}.pth')
if not os.path.exists(model_path):
logger.warning(f'Model not found at {model_path}, will attempt to auto-download')
if settings.auto_model_download:
# Try to locate known model URL and download it
try:
meta = model_manager.KNOWN_MODELS.get(model_name)
if meta and meta.get('url'):
os.makedirs(settings.models_dir, exist_ok=True)
url = meta['url']
logger.info(f'Downloading model {model_name} from {url}')
urllib.request.urlretrieve(url, model_path)
logger.info(f'Model downloaded to {model_path}')
else:
logger.error(f'No download URL known for model: {model_name}')
except Exception as e:
logger.error(f'Automatic model download failed: {e}', exc_info=True)
if not os.path.exists(model_path):
logger.error(f'Model file still not found: {model_path} - aborting initialization')
return False
# Load model
model = RRDBNet(
@@ -58,7 +78,7 @@ class RealESRGANBridge:
self.upsampler = RealESRGANer(
scale=scale,
model_path=model_path if os.path.exists(model_path) else None,
model_path=model_path,
model=model,
tile=settings.tile_size,
tile_pad=settings.tile_pad,