From c0ca6348e8489651df861a101142805c213c66a0 Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Sun, 31 Dec 2023 00:04:47 +0200
Subject: load_spandrel_model: always return a model descriptor

---
 modules/modelloader.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

(limited to 'modules')

diff --git a/modules/modelloader.py b/modules/modelloader.py
index 0b89d682..8bcee08c 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -1,8 +1,9 @@
 from __future__ import annotations
 
+import importlib
 import logging
 import os
-import importlib
+from typing import TYPE_CHECKING
 from urllib.parse import urlparse
 
 import torch
@@ -10,6 +11,8 @@ import torch
 from modules import shared
 from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, UpscalerNone
 
+if TYPE_CHECKING:
+    import spandrel
 
 logger = logging.getLogger(__name__)
 
@@ -142,17 +145,17 @@ def load_spandrel_model(
     half: bool = False,
     dtype: str | None = None,
     expected_architecture: str | None = None,
-):
+) -> spandrel.ModelDescriptor:
     import spandrel
-    model = spandrel.ModelLoader(device=device).load_from_file(path)
-    if expected_architecture and model.architecture != expected_architecture:
+    model_descriptor = spandrel.ModelLoader(device=device).load_from_file(path)
+    if expected_architecture and model_descriptor.architecture != expected_architecture:
         logger.warning(
-            f"Model {path!r} is not a {expected_architecture!r} model (got {model.architecture!r})",
+            f"Model {path!r} is not a {expected_architecture!r} model (got {model_descriptor.architecture!r})",
         )
     if half:
-        model = model.model.half()
+        model_descriptor.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
+        model_descriptor.model.to(dtype=dtype)
+    model_descriptor.model.eval()
+    logger.debug("Loaded %s from %s (device=%s, half=%s, dtype=%s)", model_descriptor, path, device, half, dtype)
+    return model_descriptor
-- 
cgit v1.2.1


From 777af661a21821994993df3ef566b01df2bb61a0 Mon Sep 17 00:00:00 2001
From: Aarni Koskela <akx@iki.fi>
Date: Sun, 31 Dec 2023 00:09:51 +0200
Subject: Be more clear about Spandrel model nomenclature

---
 modules/gfpgan_model.py     | 10 ++++++----
 modules/modelloader.py      |  2 +-
 modules/realesrgan_model.py |  4 ++--
 modules/upscaler_utils.py   |  2 +-
 4 files changed, 10 insertions(+), 8 deletions(-)

(limited to 'modules')

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__)
 
-- 
cgit v1.2.1