Merge branch 'master' into features-to-readme

1 files changed, 45 insertions, 7 deletions

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 3aa21ec1..b1c91b0d 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -9,6 +9,7 @@ from ldm.util import instantiate_from_config
 

 from modules import shared, modelloader, devices

 from modules.paths import models_path

+from modules.sd_hijack_inpainting import do_inpainting_hijack, should_hijack_inpainting

 

 model_dir = "Stable-diffusion"

 model_path = os.path.abspath(os.path.join(models_path, model_dir))

@@ -20,7 +21,7 @@ checkpoints_loaded = collections.OrderedDict()
 try:

     # this silences the annoying "Some weights of the model checkpoint were not used when initializing..." message at start.

 

-    from transformers import logging

+    from transformers import logging, CLIPModel

 

     logging.set_verbosity_error()

 except Exception:

@@ -122,9 +123,34 @@ def select_checkpoint():
     return checkpoint_info

 

 

+chckpoint_dict_replacements = {

+    'cond_stage_model.transformer.embeddings.': 'cond_stage_model.transformer.text_model.embeddings.',

+    'cond_stage_model.transformer.encoder.': 'cond_stage_model.transformer.text_model.encoder.',

+    'cond_stage_model.transformer.final_layer_norm.': 'cond_stage_model.transformer.text_model.final_layer_norm.',

+}

+

+

+def transform_checkpoint_dict_key(k):

+    for text, replacement in chckpoint_dict_replacements.items():

+        if k.startswith(text):

+            k = replacement + k[len(text):]

+

+    return k

+

+

 def get_state_dict_from_checkpoint(pl_sd):

     if "state_dict" in pl_sd:

-        return pl_sd["state_dict"]

+        pl_sd = pl_sd["state_dict"]

+

+    sd = {}

+    for k, v in pl_sd.items():

+        new_key = transform_checkpoint_dict_key(k)

+

+        if new_key is not None:

+            sd[new_key] = v

+

+    pl_sd.clear()

+    pl_sd.update(sd)

 

     return pl_sd

 

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

 

         sd = get_state_dict_from_checkpoint(pl_sd)

-        model.load_state_dict(sd, strict=False)

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

 

         if shared.cmd_opts.opt_channelslast:

             model.to(memory_format=torch.channels_last)

@@ -178,14 +204,26 @@ def load_model_weights(model, checkpoint_info):
     model.sd_checkpoint_info = checkpoint_info

 

 

-def load_model():

+def load_model(checkpoint_info=None):

     from modules import lowvram, sd_hijack

-    checkpoint_info = select_checkpoint()

+    checkpoint_info = checkpoint_info or select_checkpoint()

 

     if checkpoint_info.config != shared.cmd_opts.config:

         print(f"Loading config from: {checkpoint_info.config}")

 

     sd_config = OmegaConf.load(checkpoint_info.config)

+    

+    if should_hijack_inpainting(checkpoint_info):

+        # Hardcoded config for now...

+        sd_config.model.target = "ldm.models.diffusion.ddpm.LatentInpaintDiffusion"

+        sd_config.model.params.use_ema = False

+        sd_config.model.params.conditioning_key = "hybrid"

+        sd_config.model.params.unet_config.params.in_channels = 9

+

+        # Create a "fake" config with a different name so that we know to unload it when switching models.

+        checkpoint_info = checkpoint_info._replace(config=checkpoint_info.config.replace(".yaml", "-inpainting.yaml"))

+

+    do_inpainting_hijack()

     sd_model = instantiate_from_config(sd_config.model)

     load_model_weights(sd_model, checkpoint_info)

 

@@ -209,9 +247,9 @@ def reload_model_weights(sd_model, info=None):
     if sd_model.sd_model_checkpoint == checkpoint_info.filename:

         return

 

-    if sd_model.sd_checkpoint_info.config != checkpoint_info.config:

+    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()

+        shared.sd_model = load_model(checkpoint_info)

         return shared.sd_model

 

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