diff options
| -rw-r--r-- | extensions-builtin/SwinIR/scripts/swinir_model.py | 6 | ||||
| -rw-r--r-- | modules/gfpgan_model.py | 10 | ||||
| -rw-r--r-- | modules/modelloader.py | 2 | ||||
| -rw-r--r-- | modules/realesrgan_model.py | 4 | ||||
| -rw-r--r-- | modules/upscaler_utils.py | 2 |
5 files changed, 13 insertions, 11 deletions
diff --git a/extensions-builtin/SwinIR/scripts/swinir_model.py b/extensions-builtin/SwinIR/scripts/swinir_model.py index aae159af..95c7ec64 100644 --- a/extensions-builtin/SwinIR/scripts/swinir_model.py +++ b/extensions-builtin/SwinIR/scripts/swinir_model.py @@ -71,7 +71,7 @@ class UpscalerSwinIR(Upscaler): else: filename = path - model = modelloader.load_spandrel_model( + model_descriptor = modelloader.load_spandrel_model( filename, device=self._get_device(), dtype=devices.dtype, @@ -79,10 +79,10 @@ class UpscalerSwinIR(Upscaler): ) if getattr(opts, 'SWIN_torch_compile', False): try: - model = torch.compile(model) + model_descriptor.model.compile() except Exception: logger.warning("Failed to compile SwinIR model, fallback to JIT", exc_info=True) - return model + return model_descriptor def _get_device(self): return devices.get_device_for('swinir') diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py index 48f8ad5e..445b0409 100644 --- a/modules/gfpgan_model.py +++ b/modules/gfpgan_model.py @@ -3,6 +3,8 @@ from __future__ import annotations import logging
import os
+import torch
+
from modules import (
devices,
errors,
@@ -25,7 +27,7 @@ class FaceRestorerGFPGAN(face_restoration_utils.CommonFaceRestoration): def get_device(self):
return devices.device_gfpgan
- def load_net(self) -> None:
+ def load_net(self) -> torch.Module:
for model_path in modelloader.load_models(
model_path=self.model_path,
model_url=model_url,
@@ -34,13 +36,13 @@ class FaceRestorerGFPGAN(face_restoration_utils.CommonFaceRestoration): ext_filter=['.pth'],
):
if 'GFPGAN' in os.path.basename(model_path):
- net = modelloader.load_spandrel_model(
+ model = modelloader.load_spandrel_model(
model_path,
device=self.get_device(),
expected_architecture='GFPGAN',
).model
- net.different_w = True # see https://github.com/chaiNNer-org/spandrel/pull/81
- return net
+ model.different_w = True # see https://github.com/chaiNNer-org/spandrel/pull/81
+ return model
raise ValueError("No GFPGAN model found")
def restore(self, np_image):
diff --git a/modules/modelloader.py b/modules/modelloader.py index 8bcee08c..a7194137 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -143,7 +143,7 @@ def load_spandrel_model( *, device: str | torch.device | None, half: bool = False, - dtype: str | None = None, + dtype: str | torch.dtype | None = None, expected_architecture: str | None = None, ) -> spandrel.ModelDescriptor: import spandrel diff --git a/modules/realesrgan_model.py b/modules/realesrgan_model.py index 65f2e880..4d35b695 100644 --- a/modules/realesrgan_model.py +++ b/modules/realesrgan_model.py @@ -36,14 +36,14 @@ class UpscalerRealESRGAN(Upscaler): errors.report(f"Unable to load RealESRGAN model {path}", exc_info=True)
return img
- mod = modelloader.load_spandrel_model(
+ model_descriptor = modelloader.load_spandrel_model(
info.local_data_path,
device=self.device,
half=(not cmd_opts.no_half and not cmd_opts.upcast_sampling),
expected_architecture="ESRGAN", # "RealESRGAN" isn't a specific thing for Spandrel
)
return upscale_with_model(
- mod,
+ model_descriptor,
img,
tile_size=opts.ESRGAN_tile,
tile_overlap=opts.ESRGAN_tile_overlap,
diff --git a/modules/upscaler_utils.py b/modules/upscaler_utils.py index dde5d7ad..174c9bc3 100644 --- a/modules/upscaler_utils.py +++ b/modules/upscaler_utils.py @@ -6,7 +6,7 @@ import torch import tqdm from PIL import Image -from modules import devices, images +from modules import images logger = logging.getLogger(__name__) |