Use model loader with stable-diffusion too.

Hook the model loader into the SD_models file.
Add default url/download if checkpoint is not found.
Add matching stablediffusion-models-path argument.
Add message that --ckpt-dir will be removed in the future, but have it pipe to stablediffusion-models-path for now.
Update help strings for models-path args so they're more or less uniform.
Move sd_model "setup" call to webUI with the others.
Ensure "cleanup_models" method moves existing models to the new locations, including SD, and that we aren't deleting folders that still have stuff in them.

1 files changed, 18 insertions, 7 deletions

diff --git a/modules/modelloader.py b/modules/modelloader.py
index 9520a681..2ee364f0 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -45,7 +45,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
                     if file not in existing:
                         path = os.path.join(place, file)
                         existing.append(path)
-        if model_url is not None:
+        if model_url is not None and len(existing) == 0:
             if dl_name is not None:
                 model_file = load_file_from_url(url=model_url, model_dir=model_path, file_name=dl_name, progress=True)
             else:
@@ -69,7 +69,13 @@ def friendly_name(file: str):
 
 
 def cleanup_models():
+    # This code could probably be more efficient if we used a tuple list or something to store the src/destinations
+    # and then enumerate that, but this works for now. In the future, it'd be nice to just have every "model" scaler
+    # somehow auto-register and just do these things...
     root_path = script_path
+    src_path = models_path
+    dest_path = os.path.join(models_path, "Stable-diffusion")
+    move_files(src_path, dest_path, ".ckpt")
     src_path = os.path.join(root_path, "ESRGAN")
     dest_path = os.path.join(models_path, "ESRGAN")
     move_files(src_path, dest_path)
@@ -84,20 +90,25 @@ def cleanup_models():
     move_files(src_path, dest_path)
 
 
-def move_files(src_path: str, dest_path: str):
+def move_files(src_path: str, dest_path: str, ext_filter: str = None):
     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))
+                fullpath = os.path.join(src_path, file)
+                if os.path.isfile(fullpath):
+                    print(f"Checking file {file} in {src_path}")
+                    if ext_filter is not None:
+                        if ext_filter not in file:
+                            continue
+                    print(f"Moving {file} from {src_path} to {dest_path}.")
                     try:
                         shutil.move(fullpath, dest_path)
                     except:
                         pass
-            print("Removing folder: %s" % src_path)
-            shutil.rmtree(src_path, True)
+            if len(os.listdir(src_path)) == 0:
+                print(f"Removing empty folder: {src_path}")
+                shutil.rmtree(src_path, True)
     except:
         pass
\ No newline at end of file