From 0dce0df1ee63b2f158805c1a1f1a3743cc4a104b Mon Sep 17 00:00:00 2001 From: d8ahazard 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.1 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/upscaler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'modules/upscaler.py') 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.1 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/upscaler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'modules/upscaler.py') 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.1