From 0dce0df1ee63b2f158805c1a1f1a3743cc4a104b Mon Sep 17 00:00:00 2001
From: d8ahazard <d8ahazard@gmail.com>
Date: Thu, 29 Sep 2022 17:46:23 -0500
Subject: 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.
---
 modules/upscaler.py | 121 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 121 insertions(+)
 create mode 100644 modules/upscaler.py

(limited to 'modules/upscaler.py')

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)]
+
-- 
cgit v1.2.3


From 8f1b3153180da3b626529a7d077eac63d4b2e5e4 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Fri, 30 Sep 2022 09:46:52 +0300
Subject: fix bugs in the PR

---
 modules/sd_models.py | 2 +-
 modules/upscaler.py  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

(limited to 'modules/upscaler.py')

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 3f3f6b7c..4b9000a4 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -96,7 +96,7 @@ def list_models():
         print(f"Checkpoint in --ckpt argument not found (Possible it was moved to {model_path}: {cmd_ckpt}", file=sys.stderr)
     for filename in model_list:
         h = model_hash(filename)
-        title = modeltitle(filename, h)
+        title, model_name = modeltitle(filename, h)
         checkpoints_list[title] = CheckpointInfo(filename, title, h, model_name)
 
 def get_closet_checkpoint_match(searchString):
diff --git a/modules/upscaler.py b/modules/upscaler.py
index d698282f..7b51844f 100644
--- a/modules/upscaler.py
+++ b/modules/upscaler.py
@@ -60,7 +60,7 @@ class Upscaler:
                 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)
+            img = img.resize((dest_w, dest_h), resample=LANCZOS)
 
         return img
 
-- 
cgit v1.2.3


From a5e7b371d61bc063062b636236f1f37c264f4115 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Fri, 30 Sep 2022 10:38:48 +0300
Subject: fix the bug with broken rescaling in PR

---
 modules/images.py   | 14 +++++++++++---
 modules/upscaler.py |  2 +-
 2 files changed, 12 insertions(+), 4 deletions(-)

(limited to 'modules/upscaler.py')

diff --git a/modules/images.py b/modules/images.py
index a6538dbe..6430cfec 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -219,9 +219,17 @@ def resize_image(resize_mode, im, width, height):
         if opts.upscaler_for_img2img is None or opts.upscaler_for_img2img == "None" or im.mode == 'L':
             return im.resize((w, h), resample=LANCZOS)
 
-        upscaler = [x for x in shared.sd_upscalers if x.name == opts.upscaler_for_img2img][0]
-        scale = w / im.width
-        return upscaler.scaler.upscale(im, scale)
+        upscalers = [x for x in shared.sd_upscalers if x.name == opts.upscaler_for_img2img]
+        assert len(upscalers) > 0, f"could not find upscaler named {opts.upscaler_for_img2img}"
+
+        upscaler = upscalers[0]
+        scale = max(w / im.width, h / im.height)
+        upscaled = upscaler.scaler.upscale(im, scale)
+
+        if upscaled.width != w or upscaled.height != h:
+            upscaled = im.resize((w, h), resample=LANCZOS)
+
+        return upscaled
 
     if resize_mode == 0:
         res = resize(im, width, height)
diff --git a/modules/upscaler.py b/modules/upscaler.py
index 7b51844f..d9d7c5e2 100644
--- a/modules/upscaler.py
+++ b/modules/upscaler.py
@@ -60,7 +60,7 @@ class Upscaler:
                 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)
+            img = img.resize((int(dest_w), int(dest_h)), resample=LANCZOS)
 
         return img
 
-- 
cgit v1.2.3