From f299645aeeb65fcddde2d136fd550b6b01ffebb3 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 4 Sep 2022 18:54:12 +0300
Subject: ESRGAN support

---
 webui.py | 43 +++++++++++++++++++++++++++++++------------
 1 file changed, 31 insertions(+), 12 deletions(-)

(limited to 'webui.py')

diff --git a/webui.py b/webui.py
index d79b5966..dbc9dd54 100644
--- a/webui.py
+++ b/webui.py
@@ -21,17 +21,14 @@ import modules.processing as processing
 import modules.sd_hijack
 import modules.gfpgan_model as gfpgan
 import modules.realesrgan_model as realesrgan
+import modules.esrgan_model as esrgan
 import modules.images as images
 import modules.lowvram
 import modules.txt2img
 import modules.img2img
 
 
-shared.sd_upscalers = {
-    "RealESRGAN": lambda img: realesrgan.upscale_with_realesrgan(img, 2, 0),
-    "Lanczos": lambda img: img.resize((img.width*2, img.height*2), resample=images.LANCZOS),
-    "None": lambda img: img
-}
+esrgan.load_models(cmd_opts.esrgan_models_path)
 realesrgan.setup_realesrgan()
 gfpgan.setup_gfpgan()
 
@@ -54,26 +51,48 @@ def load_model_from_config(config, ckpt, verbose=False):
     model.eval()
     return model
 
+cached_images = {}
 
-def run_extras(image, GFPGAN_strength, RealESRGAN_upscaling, RealESRGAN_model_index):
+def run_extras(image, gfpgan_strength, upscaling_resize, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):
     processing.torch_gc()
 
     image = image.convert("RGB")
 
     outpath = opts.outdir_samples or opts.outdir_extras_samples
 
-    if gfpgan.have_gfpgan is not None and GFPGAN_strength > 0:
-
+    if gfpgan.have_gfpgan is not None and gfpgan_strength > 0:
         restored_img = gfpgan.gfpgan_fix_faces(np.array(image, dtype=np.uint8))
         res = Image.fromarray(restored_img)
 
-        if GFPGAN_strength < 1.0:
-            res = Image.blend(image, res, GFPGAN_strength)
+        if gfpgan_strength < 1.0:
+            res = Image.blend(image, res, gfpgan_strength)
+
+        image = res
+
+    if upscaling_resize != 1.0:
+        def upscale(image, scaler_index, resize):
+            small = image.crop((image.width // 2, image.height // 2, image.width // 2 + 10, image.height // 2 + 10))
+            pixels = tuple(np.array(small).flatten().tolist())
+            key = (resize, scaler_index, image.width, image.height) + pixels
+
+            c = cached_images.get(key)
+            if c is None:
+                upscaler = shared.sd_upscalers[scaler_index]
+                c = upscaler.upscale(image, image.width * resize, image.height * resize)
+                cached_images[key] = c
+
+            return c
+
+        res = upscale(image, extras_upscaler_1, upscaling_resize)
+
+        if extras_upscaler_2 != 0 and extras_upscaler_2_visibility>0:
+            res2 = upscale(image, extras_upscaler_2, upscaling_resize)
+            res = Image.blend(res, res2, extras_upscaler_2_visibility)
 
         image = res
 
-    if realesrgan.have_realesrgan and RealESRGAN_upscaling != 1.0:
-        image = realesrgan.upscale_with_realesrgan(image, RealESRGAN_upscaling, RealESRGAN_model_index)
+    while len(cached_images) > 2:
+        del cached_images[next(iter(cached_images.keys()))]
 
     images.save_image(image, outpath, "", None, '', opts.samples_format, short_filename=True, no_prompt=True)
 
-- 
cgit v1.2.1