diff --git a/app/services/realesrgan_bridge.py b/app/services/realesrgan_bridge.py index c02679e..1494c7b 100644 --- a/app/services/realesrgan_bridge.py +++ b/app/services/realesrgan_bridge.py @@ -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,