Remove LDSR project.yaml if too large.

If we accidentally downloaded the wrong file and saved it as project.yaml, this will delete it so it can be re-downloaded.

1 files changed, 7 insertions, 0 deletions

diff --git a/modules/ldsr_model.py b/modules/ldsr_model.py
index 877e7e73..4d8687c2 100644
--- a/modules/ldsr_model.py
+++ b/modules/ldsr_model.py
@@ -22,6 +22,13 @@ 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")
+        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)
         model = load_file_from_url(url=self.model_url, model_dir=self.model_path,
                                    file_name="model.pth", progress=True)
         yaml = load_file_from_url(url=self.yaml_url, model_dir=self.model_path,