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
|
import numpy as np
|
||||||
|
|
||||||
from app.config import settings
|
from app.config import settings
|
||||||
|
from app.services import model_manager
|
||||||
|
import urllib.request
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -45,6 +47,24 @@ class RealESRGANBridge:
|
|||||||
model_path = os.path.join(settings.models_dir, f'{model_name}.pth')
|
model_path = os.path.join(settings.models_dir, f'{model_name}.pth')
|
||||||
if not os.path.exists(model_path):
|
if not os.path.exists(model_path):
|
||||||
logger.warning(f'Model not found at {model_path}, will attempt to auto-download')
|
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
|
# Load model
|
||||||
model = RRDBNet(
|
model = RRDBNet(
|
||||||
@@ -58,7 +78,7 @@ class RealESRGANBridge:
|
|||||||
|
|
||||||
self.upsampler = RealESRGANer(
|
self.upsampler = RealESRGANer(
|
||||||
scale=scale,
|
scale=scale,
|
||||||
model_path=model_path if os.path.exists(model_path) else None,
|
model_path=model_path,
|
||||||
model=model,
|
model=model,
|
||||||
tile=settings.tile_size,
|
tile=settings.tile_size,
|
||||||
tile_pad=settings.tile_pad,
|
tile_pad=settings.tile_pad,
|
||||||
|
|||||||
Reference in New Issue
Block a user