fix: auto-download default model and avoid passing None to RealESRGANer
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user