Merge branch 'master' into gallery-styling

1 files changed, 14 insertions, 13 deletions

diff --git a/modules/sd_models.py b/modules/sd_models.py
index ab014efb..5f992064 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -8,14 +8,11 @@ from omegaconf import OmegaConf
 

 from ldm.util import instantiate_from_config

 

-from modules import shared, modelloader

+from modules import shared, modelloader, devices

 from modules.paths import models_path

 

 model_dir = "Stable-diffusion"

 model_path = os.path.abspath(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 = {}

@@ -30,12 +27,10 @@ except Exception:
     pass

 

 

-def setup_model(dirname):

-    global user_dir

-    user_dir = dirname

+def setup_model():

     if not os.path.exists(model_path):

         os.makedirs(model_path)

-    checkpoints_list.clear()

+

     list_models()

 

 

@@ -45,13 +40,13 @@ def checkpoint_tiles():
 

 def list_models():

     checkpoints_list.clear()

-    model_list = modelloader.load_models(model_path=model_path, model_url=model_url, command_path=user_dir, ext_filter=[".ckpt"], download_name=model_name)

+    model_list = modelloader.load_models(model_path=model_path, command_path=shared.cmd_opts.ckpt_dir, ext_filter=[".ckpt"])

 

     def modeltitle(path, shorthash):

         abspath = os.path.abspath(path)

 

-        if user_dir is not None and abspath.startswith(user_dir):

-            name = abspath.replace(user_dir, '')

+        if shared.cmd_opts.ckpt_dir is not None and abspath.startswith(shared.cmd_opts.ckpt_dir):

+            name = abspath.replace(shared.cmd_opts.ckpt_dir, '')

         elif abspath.startswith(model_path):

             name = abspath.replace(model_path, '')

         else:

@@ -69,6 +64,7 @@ def list_models():
         h = model_hash(cmd_ckpt)

         title, short_model_name = modeltitle(cmd_ckpt, h)

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

+        shared.opts.data['sd_model_checkpoint'] = title

     elif cmd_ckpt is not None and cmd_ckpt != shared.default_sd_model_file:

         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:

@@ -105,8 +101,11 @@ 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)

+        if shared.cmd_opts.ckpt is not None:

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

+        print(f" - directory {model_path}", file=sys.stderr)

+        if shared.cmd_opts.ckpt_dir is not None:

+            print(f" - directory {os.path.abspath(shared.cmd_opts.ckpt_dir)}", 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)

 

@@ -133,6 +132,8 @@ def load_model_weights(model, checkpoint_file, sd_model_hash):
     if not shared.cmd_opts.no_half:

         model.half()

 

+    devices.dtype = torch.float32 if shared.cmd_opts.no_half else torch.float16

+

     model.sd_model_hash = sd_model_hash

     model.sd_model_checkpint = checkpoint_file