Holy $hit.

Yep.

Fix gfpgan_model_arch requirement(s).
Add Upscaler base class, move from images.
Add a lot of methods to Upscaler.
Re-work all the child upscalers to be proper classes.
Add BSRGAN scaler.
Add ldsr_model_arch class, removing the dependency for another repo that just uses regular latent-diffusion stuff.
Add one universal method that will always find and load new upscaler models without having to add new "setup_model" calls. Still need to add command line params, but that could probably be automated.
Add a "self.scale" property to all Upscalers so the scalers themselves can do "things" in response to the requested upscaling size.
Ensure LDSR doesn't get stuck in a longer loop of "upscale/downscale/upscale" as we try to reach the target upscale size.
Add typehints for IDE sanity.
PEP-8 improvements.
Moar.

1 files changed, 121 insertions, 0 deletions

diff --git a/modules/upscaler.py b/modules/upscaler.py
new file mode 100644
index 00000000..d698282f
--- /dev/null
+++ b/modules/upscaler.py
@@ -0,0 +1,121 @@
+import os
+from abc import abstractmethod
+
+import PIL
+import numpy as np
+import torch
+from PIL import Image
+
+import modules.shared
+from modules import modelloader, shared
+
+LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
+from modules.paths import models_path
+
+
+class Upscaler:
+    name = None
+    model_path = None
+    model_name = None
+    model_url = None
+    enable = True
+    filter = None
+    model = None
+    user_path = None
+    scalers: []
+    tile = True
+
+    def __init__(self, create_dirs=False):
+        self.mod_pad_h = None
+        self.tile_size = modules.shared.opts.ESRGAN_tile
+        self.tile_pad = modules.shared.opts.ESRGAN_tile_overlap
+        self.device = modules.shared.device
+        self.img = None
+        self.output = None
+        self.scale = 1
+        self.half = not modules.shared.cmd_opts.no_half
+        self.pre_pad = 0
+        self.mod_scale = None
+        if self.name is not None and create_dirs:
+            self.model_path = os.path.join(models_path, self.name)
+            if not os.path.exists(self.model_path):
+                os.makedirs(self.model_path)
+
+        try:
+            import cv2
+            self.can_tile = True
+        except:
+            pass
+
+    @abstractmethod
+    def do_upscale(self, img: PIL.Image, selected_model: str):
+        return img
+
+    def upscale(self, img: PIL.Image, scale: int, selected_model: str = None):
+        self.scale = scale
+        dest_w = img.width * scale
+        dest_h = img.height * scale
+        for i in range(3):
+            if img.width >= dest_w and img.height >= dest_h:
+                break
+            img = self.do_upscale(img, selected_model)
+        if img.width != dest_w or img.height != dest_h:
+            img = img.resize(dest_w, dest_h, resample=LANCZOS)
+
+        return img
+
+    @abstractmethod
+    def load_model(self, path: str):
+        pass
+
+    def find_models(self, ext_filter=None) -> list:
+        return modelloader.load_models(model_path=self.model_path, model_url=self.model_url, command_path=self.user_path)
+
+    def update_status(self, prompt):
+        print(f"\nextras: {prompt}", file=shared.progress_print_out)
+
+
+class UpscalerData:
+    name = None
+    data_path = None
+    scale: int = 4
+    scaler: Upscaler = None
+    model: None
+
+    def __init__(self, name: str, path: str, upscaler: Upscaler = None, scale: int = 4, model=None):
+        self.name = name
+        self.data_path = path
+        self.scaler = upscaler
+        self.scale = scale
+        self.model = model
+
+
+class UpscalerNone(Upscaler):
+    name = "None"
+    scalers = []
+
+    def load_model(self, path):
+        pass
+
+    def do_upscale(self, img, selected_model=None):
+        return img
+
+    def __init__(self, dirname=None):
+        super().__init__(False)
+        self.scalers = [UpscalerData("None", None, self)]
+
+
+class UpscalerLanczos(Upscaler):
+    scalers = []
+
+    def do_upscale(self, img, selected_model=None):
+        return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=LANCZOS)
+
+    def load_model(self, _):
+        pass
+
+    def __init__(self, dirname=None):
+        super().__init__(False)
+        self.name = "Lanczos"
+        self.scalers = [UpscalerData("Lanczos", None, self)]
+