Use Spandrel for upscaling and face restoration architectures (aside from GFPGAN and LDSR)

This commit is contained in:
Aarni Koskela
2023-12-25 14:43:51 +02:00
parent e472383acb
commit b0f5934234
19 changed files with 277 additions and 3768 deletions

View File

@@ -1,5 +1,6 @@
from __future__ import annotations
import logging
import os
import shutil
import importlib
@@ -10,6 +11,9 @@ from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, Upscale
from modules.paths import script_path, models_path
logger = logging.getLogger(__name__)
def load_file_from_url(
url: str,
*,
@@ -177,3 +181,15 @@ def load_upscalers():
# Special case for UpscalerNone keeps it at the beginning of the list.
key=lambda x: x.name.lower() if not isinstance(x.scaler, (UpscalerNone, UpscalerLanczos, UpscalerNearest)) else ""
)
def load_spandrel_model(path, *, device, half: bool = False, dtype=None):
import spandrel
model = spandrel.ModelLoader(device=device).load_from_file(path)
if half:
model = model.model.half()
if dtype:
model = model.model.to(dtype=dtype)
model.eval()
logger.debug("Loaded %s from %s (device=%s, half=%s, dtype=%s)", model, path, device, half, dtype)
return model