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.

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)
-                    fullpath = os.path.join(src_path, file)
+                if os.path.isfile(fullpath):
-                    print("Moving file: %s to %s" % (fullpath, dest_path))
+                    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)
+            if len(os.listdir(src_path)) == 0:
-            shutil.rmtree(src_path, True)
+                print(f"Removing empty folder: {src_path}")
+                shutil.rmtree(src_path, True)
     except:
         pass

modules/sd_models.py

@@ -8,7 +8,13 @@ from omegaconf import OmegaConf
 
 from ldm.util import instantiate_from_config
 
-from modules import shared
+from modules import shared, modelloader
+from modules.paths import models_path
+
+model_dir = "Stable-diffusion"
+model_path = os.path.join(models_path, model_dir)
+model_name = "sd-v1-4.ckpt"
+model_url = "https://drive.yerf.org/wl/?id=EBfTrmcCCUAGaQBXVIj5lJmEhjoP1tgl&mode=grid&download=1"
 
 CheckpointInfo = namedtuple("CheckpointInfo", ['filename', 'title', 'hash'])
 checkpoints_list = {}
@@ -23,23 +29,28 @@ except Exception:
     pass
 
 
-def list_models():
+def modeltitle(path, h):
+    abspath = os.path.abspath(path)
+
+    if abspath.startswith(model_dir):
+        name = abspath.replace(model_dir, '')
+    else:
+        name = os.path.basename(path)
+
+    if name.startswith("\\") or name.startswith("/"):
+        name = name[1:]
+
+    return f'{name} [{h}]'
+
+
+def setup_model(dirname):
+    global model_path
+    global model_name
+    global model_url
+    if not os.path.exists(model_path):
+        os.makedirs(model_path)
     checkpoints_list.clear()
-
+    model_list = modelloader.load_models(model_path, model_url, dirname, model_name, ext_filter=".ckpt")
-    model_dir = os.path.abspath(shared.cmd_opts.ckpt_dir)
-
-    def modeltitle(path, h):
-        abspath = os.path.abspath(path)
-
-        if abspath.startswith(model_dir):
-            name = abspath.replace(model_dir, '')
-        else:
-            name = os.path.basename(path)
-
-        if name.startswith("\\") or name.startswith("/"):
-            name = name[1:]
-
-        return f'{name} [{h}]'
 
     cmd_ckpt = shared.cmd_opts.ckpt
     if os.path.exists(cmd_ckpt):
@@ -47,13 +58,12 @@ def list_models():
         title = modeltitle(cmd_ckpt, h)
         checkpoints_list[title] = CheckpointInfo(cmd_ckpt, title, h)
     elif cmd_ckpt is not None and cmd_ckpt != shared.default_sd_model_file:
-        print(f"Checkpoint in --ckpt argument not found: {cmd_ckpt}", file=sys.stderr)
+        print(f"Checkpoint in --ckpt argument not found (Possible it was moved to {model_path}: {cmd_ckpt}", file=sys.stderr)
 
-    if os.path.exists(model_dir):
+    for filename in model_list:
-        for filename in glob.glob(model_dir + '/**/*.ckpt', recursive=True):
+        h = model_hash(filename)
-            h = model_hash(filename)
+        title = modeltitle(filename, h)
-            title = modeltitle(filename, h)
+        checkpoints_list[title] = CheckpointInfo(filename, title, h)
-            checkpoints_list[title] = CheckpointInfo(filename, title, h)
 
 
 def model_hash(filename):

modules/shared.py

@@ -20,7 +20,8 @@ model_path = os.path.join(script_path, 'models')
 parser = argparse.ArgumentParser()
 parser.add_argument("--config", type=str, default=os.path.join(sd_path, "configs/stable-diffusion/v1-inference.yaml"), help="path to config which constructs model",)
 parser.add_argument("--ckpt", type=str, default=sd_model_file, help="path to checkpoint of stable diffusion model; this checkpoint will be added to the list of checkpoints and loaded by default if you don't have a checkpoint selected in settings",)
-parser.add_argument("--ckpt-dir", type=str, default=model_path, help="path to directory with stable diffusion checkpoints",)
+# This should be deprecated, but we'll leave it for a few iterations
+parser.add_argument("--ckpt-dir", type=str, default=None, help="Path to directory with stable diffusion checkpoints (Deprecated, use '--stablediffusion-models-path'", )
 parser.add_argument("--gfpgan-dir", type=str, help="GFPGAN directory", default=('./src/gfpgan' if os.path.exists('./src/gfpgan') else './GFPGAN'))
 parser.add_argument("--gfpgan-model", type=str, help="GFPGAN model file name", default=None)
 parser.add_argument("--no-half", action='store_true', help="do not switch the model to 16-bit floats")
@@ -34,12 +35,13 @@ parser.add_argument("--always-batch-cond-uncond", action='store_true', help="dis
 parser.add_argument("--unload-gfpgan", action='store_true', help="does not do anything.")
 parser.add_argument("--precision", type=str, help="evaluate at this precision", choices=["full", "autocast"], default="autocast")
 parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site (doesn't work for me but you might have better luck)")
-parser.add_argument("--codeformer-models-path", type=str, help="Path to directory with codeformer model(s)", default=os.path.join(model_path, 'Codeformer'))
+parser.add_argument("--codeformer-models-path", type=str, help="Path to directory with codeformer model file(s).", default=os.path.join(model_path, 'Codeformer'))
-parser.add_argument("--gfpgan-models-path", type=str, help="Path to directory with GFPGAN model(s)", default=os.path.join(model_path, 'GFPGAN'))
+parser.add_argument("--gfpgan-models-path", type=str, help="Path to directory with GFPGAN model file(s).", default=os.path.join(model_path, 'GFPGAN'))
-parser.add_argument("--esrgan-models-path", type=str, help="Path to directory with ESRGAN models", default=os.path.join(model_path, 'ESRGAN'))
+parser.add_argument("--esrgan-models-path", type=str, help="Path to directory with ESRGAN model file(s).", default=os.path.join(model_path, 'ESRGAN'))
-parser.add_argument("--realesrgan-models-path", type=str, help="Path to directory with RealESRGAN models", default=os.path.join(model_path, 'RealESRGAN'))
+parser.add_argument("--realesrgan-models-path", type=str, help="Path to directory with RealESRGAN model file(s).", default=os.path.join(model_path, 'RealESRGAN'))
-parser.add_argument("--swinir-models-path", type=str, help="Path to directory with SwinIR models", default=os.path.join(model_path, 'SwinIR'))
+parser.add_argument("--stablediffusion-models-path", type=str, help="Path to directory with Stable-diffusion checkpoints.", default=os.path.join(model_path, 'SwinIR'))
-parser.add_argument("--ldsr-models-path", type=str, help="Path to directory with LDSR models", default=os.path.join(model_path, 'LDSR'))
+parser.add_argument("--swinir-models-path", type=str, help="Path to directory with SwinIR model file(s).", default=os.path.join(model_path, 'SwinIR'))
+parser.add_argument("--ldsr-models-path", type=str, help="Path to directory with LDSR model file(s).", default=os.path.join(model_path, 'LDSR'))
 parser.add_argument("--opt-split-attention", action='store_true', help="force-enables cross-attention layer optimization. By default, it's on for torch.cuda and off for other torch devices.")
 parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer optimization")
 parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of split attention optimization that does not consume all the VRAM it can find")
@@ -57,7 +59,10 @@ parser.add_argument("--autolaunch", action='store_true', help="open the webui UR
 parser.add_argument("--use-textbox-seed", action='store_true', help="use textbox for seeds in UI (no up/down, but possible to input long seeds)", default=False)
 
 cmd_opts = parser.parse_args()
-
+if cmd_opts.ckpt_dir is not None:
+    print("The 'ckpt-dir' arg is deprecated in favor of the 'stablediffusion-models-path' argument and will be "
+          "removed in a future release. Please use the new option if you wish to use a custom checkpoint directory.")
+    cmd_opts.__setattr__("stablediffusion-models-path", cmd_opts.ckpt_dir)
 device = get_optimal_device()
 
 batch_cond_uncond = cmd_opts.always_batch_cond_uncond or not (cmd_opts.lowvram or cmd_opts.medvram)
@@ -65,6 +70,7 @@ parallel_processing_allowed = not cmd_opts.lowvram and not cmd_opts.medvram
 
 config_filename = cmd_opts.ui_settings_file
 
+
 class State:
     interrupted = False
     job = ""
@@ -98,8 +104,8 @@ prompt_styles = modules.styles.StyleDatabase(styles_filename)
 interrogator = modules.interrogate.InterrogateModels("interrogate")
 
 face_restorers = []
-
+# This was moved to webui.py with the other model "setup" calls.
-modules.sd_models.list_models()
+# modules.sd_models.list_models()
 
 
 def realesrgan_models_names():
@@ -195,7 +201,7 @@ options_templates.update(options_section(('system', "System"), {
 options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "sd_model_checkpoint": OptionInfo(None, "Stable Diffusion checkpoint", gr.Radio, lambda: {"choices": [x.title for x in modules.sd_models.checkpoints_list.values()]}),
     "img2img_color_correction": OptionInfo(False, "Apply color correction to img2img results to match original colors."),
-    "save_images_before_color_correction": OptionInfo(False, "Save a copy of image before applying color correction to img2img results"),    
+    "save_images_before_color_correction": OptionInfo(False, "Save a copy of image before applying color correction to img2img results"),
     "img2img_fix_steps": OptionInfo(False, "With img2img, do exactly the amount of steps the slider specifies (normally you'd do less with less denoising)."),
     "enable_quantization": OptionInfo(False, "Enable quantization in K samplers for sharper and cleaner results. This may change existing seeds. Requires restart to apply."),
     "enable_emphasis": OptionInfo(True, "Use (text) to make model pay more attention to text and [text] to make it pay less attention"),

webui.py

@@ -23,6 +23,7 @@ from modules.paths import script_path
 from modules.shared import cmd_opts
 
 modelloader.cleanup_models()
+modules.sd_models.setup_model(cmd_opts.stablediffusion_models_path)
 codeformer.setup_model(cmd_opts.codeformer_models_path)
 gfpgan.setup_model(cmd_opts.gfpgan_models_path)
 shared.face_restorers.append(modules.face_restoration.FaceRestoration())