Merge pull request #1109 from d8ahazard/ModelLoader

Model Loader, Fixes

1 files changed, 50 insertions, 11 deletions

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 0e7ed905..3f3f6b7c 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -8,7 +8,14 @@ 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"

+user_dir = None

 

 CheckpointInfo = namedtuple("CheckpointInfo", ['filename', 'title', 'hash', 'model_name'])

 checkpoints_list = {}

@@ -23,14 +30,47 @@ except Exception:
     pass

 

 

+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

+    global user_dir

+    global model_list

+    user_dir = dirname

+    if not os.path.exists(model_path):

+        os.makedirs(model_path)

+    checkpoints_list.clear()

+    list_models()

+

+

 def checkpoint_tiles():

     return sorted([x.title for x in checkpoints_list.values()])

 

 

 def list_models():

+    global model_path

+    global model_url

+    global model_name

+    global user_dir

     checkpoints_list.clear()

-

-    model_dir = os.path.abspath(shared.cmd_opts.ckpt_dir)

+    model_list = modelloader.load_models(model_path=model_path,model_url=model_url,command_path= user_dir,

+                                         ext_filter=[".ckpt"], download_name=model_name)

+    print(f"Model list: {model_list}")

+    model_dir = os.path.abspath(model_path)

 

     def modeltitle(path, h):

         abspath = os.path.abspath(path)

@@ -53,13 +93,11 @@ def list_models():
         title, model_name = modeltitle(cmd_ckpt, h)

         checkpoints_list[title] = CheckpointInfo(cmd_ckpt, title, h, model_name)

     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)

-

-    if os.path.exists(model_dir):

-        for filename in glob.glob(model_dir + '/**/*.ckpt', recursive=True):

-            h = model_hash(filename)

-            title, model_name = modeltitle(filename, h)

-            checkpoints_list[title] = CheckpointInfo(filename, title, h, model_name)

+        print(f"Checkpoint in --ckpt argument not found (Possible it was moved to {model_path}: {cmd_ckpt}", file=sys.stderr)

+    for filename in model_list:

+        h = model_hash(filename)

+        title = modeltitle(filename, h)

+        checkpoints_list[title] = CheckpointInfo(filename, title, h, model_name)

 

 def get_closet_checkpoint_match(searchString):

     applicable = sorted([info for info in checkpoints_list.values() if searchString in info.title], key = lambda x:len(x.title))

@@ -69,6 +107,7 @@ def get_closet_checkpoint_match(searchString):
 

 def model_hash(filename):

     try:

+        print(f"Opening: {filename}")

         with open(filename, "rb") as file:

             import hashlib

             m = hashlib.sha256()

@@ -89,7 +128,7 @@ def select_checkpoint():
     if len(checkpoints_list) == 0:

         print(f"No checkpoints found. When searching for checkpoints, looked at:", file=sys.stderr)

         print(f" - file {os.path.abspath(shared.cmd_opts.ckpt)}", file=sys.stderr)

-        print(f" - directory {os.path.abspath(shared.cmd_opts.ckpt_dir)}", file=sys.stderr)

+        print(f" - directory {os.path.abspath(shared.cmd_opts.stablediffusion_models_path)}", file=sys.stderr)

         print(f"Can't run without a checkpoint. Find and place a .ckpt file into any of those locations. The program will exit.", file=sys.stderr)

         exit(1)