Add LDSR and "GoLatent?" Upscaling (#763)

* Add LDSR Upscaling

4 files changed, 75 insertions, 0 deletions

diff --git a/modules/ldsr_model.py b/modules/ldsr_model.py
new file mode 100644
index 00000000..95e84659
--- /dev/null
+++ b/modules/ldsr_model.py
@@ -0,0 +1,67 @@
+import os
+import sys
+import traceback
+from collections import namedtuple
+
+from basicsr.utils.download_util import load_file_from_url
+
+import modules.images
+from modules import shared
+from modules.paths import script_path
+
+LDSRModelInfo = namedtuple("LDSRModelInfo", ["name", "location", "model", "netscale"])
+
+ldsr_models = []
+have_ldsr = False
+LDSR_obj = None
+
+
+class UpscalerLDSR(modules.images.Upscaler):
+    def __init__(self, steps):
+        self.steps = steps
+        self.name = "LDSR"
+
+    def do_upscale(self, img):
+        return upscale_with_ldsr(img)
+
+
+def add_lsdr():
+    modules.shared.sd_upscalers.append(UpscalerLDSR(100))
+
+
+def setup_ldsr():
+    path = modules.paths.paths.get("LDSR", None)
+    if path is None:
+        return
+    global have_ldsr
+    global LDSR_obj
+    try:
+        from LDSR import LDSR
+        model_url = "https://heibox.uni-heidelberg.de/f/578df07c8fc04ffbadf3/?dl=1"
+        yaml_url = "https://heibox.uni-heidelberg.de/f/31a76b13ea27482981b4/?dl=1"
+        repo_path = 'latent-diffusion/experiments/pretrained_models/'
+        model_path = load_file_from_url(url=model_url, model_dir=os.path.join("repositories", repo_path),
+                                        progress=True, file_name="model.chkpt")
+        yaml_path = load_file_from_url(url=yaml_url, model_dir=os.path.join("repositories", repo_path),
+                                       progress=True, file_name="project.yaml")
+        have_ldsr = True
+        LDSR_obj = LDSR(model_path, yaml_path)
+
+
+    except Exception:
+        print("Error importing LDSR:", file=sys.stderr)
+        print(traceback.format_exc(), file=sys.stderr)
+        have_ldsr = False
+
+
+def upscale_with_ldsr(image):
+    setup_ldsr()
+    if not have_ldsr or LDSR_obj is None:
+        return image
+
+    ddim_steps = shared.opts.ldsr_steps
+    pre_scale = shared.opts.ldsr_pre_down
+    post_scale = shared.opts.ldsr_post_down
+
+    image = LDSR_obj.super_resolution(image, ddim_steps, pre_scale, post_scale)
+    return image
diff --git a/modules/paths.py b/modules/paths.py
index 97c17a98..3a19f9e5 100644
--- a/modules/paths.py
+++ b/modules/paths.py
@@ -19,6 +19,7 @@ path_dirs = [
     (os.path.join(sd_path, '../taming-transformers'), 'taming', 'Taming Transformers'),

     (os.path.join(sd_path, '../CodeFormer'), 'inference_codeformer.py', 'CodeFormer'),

     (os.path.join(sd_path, '../BLIP'), 'models/blip.py', 'BLIP'),

+    (os.path.join(sd_path, '../latent-diffusion'), 'LDSR.py', 'LDSR'),

 ]

 

 paths = {}

diff --git a/modules/shared.py b/modules/shared.py
index 6ce831b7..280d6c4c 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -144,6 +144,12 @@ class Options:
         "ESRGAN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for ESRGAN upscalers. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),

         "SWIN_tile": OptionInfo(192, "Tile size for all SwinIR.", gr.Slider, {"minimum": 16, "maximum": 512, "step": 16}),

         "SWIN_tile_overlap": OptionInfo(8, "Tile overlap, in pixels for SwinIR. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}),

+        "ldsr_steps": OptionInfo(100, "LDSR processing steps. Lower = faster",

+                                          gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}),

+        "ldsr_pre_down":OptionInfo(1, "LDSR Pre-process downssample scale. 1 = no down-sampling, 4 = 1/4 scale.",

+                                          gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),

+        "ldsr_post_down":OptionInfo(1, "LDSR Post-process down-sample scale. 1 = no down-sampling, 4 = 1/4 scale.",

+                                          gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),

         "random_artist_categories": OptionInfo([], "Allowed categories for random artists selection when using the Roll button", gr.CheckboxGroup, {"choices": artist_db.categories()}),

         "upscale_at_full_resolution_padding": OptionInfo(16, "Inpainting at full resolution: padding, in pixels, for the masked region.", gr.Slider, {"minimum": 0, "maximum": 128, "step": 4}),

         "upscaler_for_hires_fix": OptionInfo(None, "Upscaler for highres. fix", gr.Radio, lambda: {"choices": [x.name for x in sd_upscalers]}),

diff --git a/modules/ui.py b/modules/ui.py
index 0fbb081d..e3cd7f3a 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -23,6 +23,7 @@ from modules.shared import opts, cmd_opts
 import modules.shared as shared

 from modules.sd_samplers import samplers, samplers_for_img2img

 import modules.realesrgan_model as realesrgan

+import modules.ldsr_model

 import modules.scripts

 import modules.gfpgan_model

 import modules.codeformer_model