Merge pull request #1371 from d8ahazard/master

Fix LDSR, BSRGAN, recursive SD checkpoint Loading

4 files changed, 17 insertions, 7 deletions

diff --git a/modules/bsrgan_model.py b/modules/bsrgan_model.py
index 47346f31..e62c6657 100644
--- a/modules/bsrgan_model.py
+++ b/modules/bsrgan_model.py
@@ -69,7 +69,7 @@ class UpscalerBSRGAN(modules.upscaler.Upscaler):
         if not os.path.exists(filename) or filename is None:
             print(f"BSRGAN: Unable to load model from {filename}", file=sys.stderr)
             return None
-        model = RRDBNet(in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=2)  # define network
+        model = RRDBNet(in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=4)  # define network
         model.load_state_dict(torch.load(filename), strict=True)
         model.eval()
         for k, v in model.named_parameters():
diff --git a/modules/ldsr_model.py b/modules/ldsr_model.py
index 877e7e73..1c1070fc 100644
--- a/modules/ldsr_model.py
+++ b/modules/ldsr_model.py
@@ -22,8 +22,20 @@ class UpscalerLDSR(Upscaler):
         self.scalers = [scaler_data]
 
     def load_model(self, path: str):
+        # Remove incorrect project.yaml file if too big
+        yaml_path = os.path.join(self.model_path, "project.yaml")
+        old_model_path = os.path.join(self.model_path, "model.pth")
+        new_model_path = os.path.join(self.model_path, "model.ckpt")
+        if os.path.exists(yaml_path):
+            statinfo = os.stat(yaml_path)
+            if statinfo.st_size >= 10485760:
+                print("Removing invalid LDSR YAML file.")
+                os.remove(yaml_path)
+        if os.path.exists(old_model_path):
+            print("Renaming model from model.pth to model.ckpt")
+            os.rename(old_model_path, new_model_path)
         model = load_file_from_url(url=self.model_url, model_dir=self.model_path,
-                                   file_name="model.pth", progress=True)
+                                   file_name="model.ckpt", progress=True)
         yaml = load_file_from_url(url=self.yaml_url, model_dir=self.model_path,
                                   file_name="project.yaml", progress=True)
 
@@ -41,5 +53,4 @@ class UpscalerLDSR(Upscaler):
             print("NO LDSR!")
             return img
         ddim_steps = shared.opts.ldsr_steps
-        pre_scale = shared.opts.ldsr_pre_down
         return ldsr.super_resolution(img, ddim_steps, self.scale)
diff --git a/modules/ldsr_model_arch.py b/modules/ldsr_model_arch.py
index 7faac6e1..14db5076 100644
--- a/modules/ldsr_model_arch.py
+++ b/modules/ldsr_model_arch.py
@@ -98,9 +98,7 @@ class LDSR:
         im_og = image
         width_og, height_og = im_og.size
         # If we can adjust the max upscale size, then the 4 below should be our variable
-        print("Foo")
         down_sample_rate = target_scale / 4
-        print(f"Downsample rate is {down_sample_rate}")
         wd = width_og * down_sample_rate
         hd = height_og * down_sample_rate
         width_downsampled_pre = int(wd)
@@ -111,7 +109,7 @@ class LDSR:
                 f'Downsampling from [{width_og}, {height_og}] to [{width_downsampled_pre}, {height_downsampled_pre}]')
             im_og = im_og.resize((width_downsampled_pre, height_downsampled_pre), Image.LANCZOS)
         else:
-            print(f"Down sample rate is 1 from {target_scale} / 4")
+            print(f"Down sample rate is 1 from {target_scale} / 4 (Not downsampling)")
         logs = self.run(model["model"], im_og, diffusion_steps, eta)
 
         sample = logs["sample"]
diff --git a/modules/modelloader.py b/modules/modelloader.py
index 1106aeb7..8c862b42 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -1,3 +1,4 @@
+import glob
 import os
 import shutil
 import importlib
@@ -41,7 +42,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
 
         for place in places:
             if os.path.exists(place):
-                for file in os.listdir(place):
+                for file in glob.iglob(place + '**/**', recursive=True):
                     full_path = os.path.join(place, file)
                     if os.path.isdir(full_path):
                         continue