Cleanup existing directories, fixes

1 files changed, 38 insertions, 0 deletions

diff --git a/modules/modelloader.py b/modules/modelloader.py
index d59fbe05..9520a681 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -1,8 +1,11 @@
 import os
+import shutil
 from urllib.parse import urlparse
 
 from basicsr.utils.download_util import load_file_from_url
 
+from modules.paths import script_path, models_path
+
 
 def load_models(model_path: str, model_url: str = None, command_path: str = None, dl_name: str = None, existing=None,
                 ext_filter=None) -> list:
@@ -63,3 +66,38 @@ def friendly_name(file: str):
     model_name, extension = os.path.splitext(file)
     model_name = model_name.replace("_", " ").title()
     return model_name
+
+
+def cleanup_models():
+    root_path = script_path
+    src_path = os.path.join(root_path, "ESRGAN")
+    dest_path = os.path.join(models_path, "ESRGAN")
+    move_files(src_path, dest_path)
+    src_path = os.path.join(root_path, "gfpgan")
+    dest_path = os.path.join(models_path, "GFPGAN")
+    move_files(src_path, dest_path)
+    src_path = os.path.join(root_path, "SwinIR")
+    dest_path = os.path.join(models_path, "SwinIR")
+    move_files(src_path, dest_path)
+    src_path = os.path.join(root_path, "repositories/latent-diffusion/experiments/pretrained_models/")
+    dest_path = os.path.join(models_path, "LDSR")
+    move_files(src_path, dest_path)
+
+
+def move_files(src_path: str, dest_path: str):
+    try:
+        if not os.path.exists(dest_path):
+            os.makedirs(dest_path)
+        if os.path.exists(src_path):
+            for file in os.listdir(src_path):
+                if os.path.isfile(file):
+                    fullpath = os.path.join(src_path, file)
+                    print("Moving file: %s to %s" % (fullpath, dest_path))
+                    try:
+                        shutil.move(fullpath, dest_path)
+                    except:
+                        pass
+            print("Removing folder: %s" % src_path)
+            shutil.rmtree(src_path, True)
+    except:
+        pass
\ No newline at end of file