Merge branch 'master' into focal-point-cropping

1 files changed, 27 insertions, 7 deletions

diff --git a/modules/sd_models.py b/modules/sd_models.py
index eae22e87..e697bb72 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -7,8 +7,9 @@ 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

 

 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:

@@ -154,6 +155,9 @@ def get_state_dict_from_checkpoint(pl_sd):
     return pl_sd

 

 

+vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"}

+

+

 def load_model_weights(model, checkpoint_info):

     checkpoint_file = checkpoint_info.filename

     sd_model_hash = checkpoint_info.hash

@@ -185,7 +189,7 @@ def load_model_weights(model, checkpoint_info):
         if os.path.exists(vae_file):

             print(f"Loading VAE weights from: {vae_file}")

             vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location)

-            vae_dict = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss"}

+            vae_dict = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss" and k not in vae_ignore_keys}

             model.first_stage_model.load_state_dict(vae_dict)

 

         model.first_stage_model.to(devices.dtype_vae)

@@ -203,14 +207,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)

 

@@ -222,6 +238,9 @@ def load_model():
     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

@@ -234,9 +253,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()

+        load_model(checkpoint_info)

         return shared.sd_model

 

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

@@ -249,6 +268,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)