Merge branch 'AUTOMATIC1111:master' into master

1 files changed, 18 insertions, 8 deletions

diff --git a/modules/sd_models.py b/modules/sd_models.py
index d99dbce8..f86dc3ed 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -3,11 +3,12 @@ import os.path
 import sys

 from collections import namedtuple

 import torch

+import re

 from omegaconf import OmegaConf

 

 from ldm.util import instantiate_from_config

 

-from modules import shared, modelloader, devices

+from modules import shared, modelloader, devices, script_callbacks

 from modules.paths import models_path

 from modules.sd_hijack_inpainting import do_inpainting_hijack, should_hijack_inpainting

 

@@ -35,8 +36,10 @@ def setup_model():
     list_models()

 

 

-def checkpoint_tiles():

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

+def checkpoint_tiles(): 

+    convert = lambda name: int(name) if name.isdigit() else name.lower() 

+    alphanumeric_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)] 

+    return sorted([x.title for x in checkpoints_list.values()], key = alphanumeric_key)

 

 

 def list_models():

@@ -170,7 +173,9 @@ def load_model_weights(model, checkpoint_info):
             print(f"Global Step: {pl_sd['global_step']}")

 

         sd = get_state_dict_from_checkpoint(pl_sd)

-        missing, extra = model.load_state_dict(sd, strict=False)

+        del pl_sd

+        model.load_state_dict(sd, strict=False)

+        del sd

 

         if shared.cmd_opts.opt_channelslast:

             model.to(memory_format=torch.channels_last)

@@ -194,9 +199,10 @@ def load_model_weights(model, checkpoint_info):
 

         model.first_stage_model.to(devices.dtype_vae)

 

-        checkpoints_loaded[checkpoint_info] = model.state_dict().copy()

-        while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache:

-            checkpoints_loaded.popitem(last=False)  # LRU

+        if shared.opts.sd_checkpoint_cache > 0:

+            checkpoints_loaded[checkpoint_info] = model.state_dict().copy()

+            while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache:

+                checkpoints_loaded.popitem(last=False)  # LRU

     else:

         print(f"Loading weights [{sd_model_hash}] from cache")

         checkpoints_loaded.move_to_end(checkpoint_info)

@@ -238,6 +244,9 @@ def load_model(checkpoint_info=None):
     sd_hijack.model_hijack.hijack(sd_model)

 

     sd_model.eval()

+    shared.sd_model = sd_model

+

+    script_callbacks.model_loaded_callback(sd_model)

 

     print(f"Model loaded.")

     return sd_model

@@ -252,7 +261,7 @@ def reload_model_weights(sd_model, info=None):
 

     if sd_model.sd_checkpoint_info.config != checkpoint_info.config or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info):

         checkpoints_loaded.clear()

-        shared.sd_model = load_model(checkpoint_info)

+        load_model(checkpoint_info)

         return shared.sd_model

 

     if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:

@@ -265,6 +274,7 @@ def reload_model_weights(sd_model, info=None):
     load_model_weights(sd_model, checkpoint_info)

 

     sd_hijack.model_hijack.hijack(sd_model)

+    script_callbacks.model_loaded_callback(sd_model)

 

     if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:

         sd_model.to(devices.device)