Merge pull request #1371 from d8ahazard/master

Fix LDSR, BSRGAN, recursive SD checkpoint Loading

1 files changed, 13 insertions, 2 deletions

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)