From 247f58a5e740a7bd3980815961425b778d77ec28 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 17 Sep 2022 12:05:04 +0300
Subject: add support for switching model checkpoints at runtime

---
 webui.py | 61 ++++++++++---------------------------------------------------
 1 file changed, 10 insertions(+), 51 deletions(-)

(limited to 'webui.py')

diff --git a/webui.py b/webui.py
index add72123..ff8997db 100644
--- a/webui.py
+++ b/webui.py
@@ -3,13 +3,8 @@ import threading
 
 from modules.paths import script_path
 
-import torch
-from omegaconf import OmegaConf
-
 import signal
 
-from ldm.util import instantiate_from_config
-
 from modules.shared import opts, cmd_opts, state
 import modules.shared as shared
 import modules.ui
@@ -24,6 +19,7 @@ import modules.extras
 import modules.lowvram
 import modules.txt2img
 import modules.img2img
+import modules.sd_models
 
 
 modules.codeformer_model.setup_codeformer()
@@ -33,29 +29,17 @@ shared.face_restorers.append(modules.face_restoration.FaceRestoration())
 esrgan.load_models(cmd_opts.esrgan_models_path)
 realesrgan.setup_realesrgan()
 
+queue_lock = threading.Lock()
 
-def load_model_from_config(config, ckpt, verbose=False):
-    print(f"Loading model [{shared.sd_model_hash}] from {ckpt}")
-    pl_sd = torch.load(ckpt, map_location="cpu")
-    if "global_step" in pl_sd:
-        print(f"Global Step: {pl_sd['global_step']}")
-    sd = pl_sd["state_dict"]
 
-    model = instantiate_from_config(config.model)
-    m, u = model.load_state_dict(sd, strict=False)
-    if len(m) > 0 and verbose:
-        print("missing keys:")
-        print(m)
-    if len(u) > 0 and verbose:
-        print("unexpected keys:")
-        print(u)
-    if cmd_opts.opt_channelslast:
-        model = model.to(memory_format=torch.channels_last)
-    model.eval()
-    return model
+def wrap_queued_call(func):
+    def f(*args, **kwargs):
+        with queue_lock:
+            res = func(*args, **kwargs)
 
+        return res
 
-queue_lock = threading.Lock()
+    return f
 
 
 def wrap_gradio_gpu_call(func):
@@ -80,33 +64,8 @@ def wrap_gradio_gpu_call(func):
 
 modules.scripts.load_scripts(os.path.join(script_path, "scripts"))
 
-try:
-    # this silences the annoying "Some weights of the model checkpoint were not used when initializing..." message at start.
-
-    from transformers import logging
-
-    logging.set_verbosity_error()
-except Exception:
-    pass
-
-with open(cmd_opts.ckpt, "rb") as file:
-    import hashlib
-    m = hashlib.sha256()
-
-    file.seek(0x100000)
-    m.update(file.read(0x10000))
-    shared.sd_model_hash = m.hexdigest()[0:8]
-
-sd_config = OmegaConf.load(cmd_opts.config)
-shared.sd_model = load_model_from_config(sd_config, cmd_opts.ckpt)
-shared.sd_model = (shared.sd_model if cmd_opts.no_half else shared.sd_model.half())
-
-if cmd_opts.lowvram or cmd_opts.medvram:
-    modules.lowvram.setup_for_low_vram(shared.sd_model, cmd_opts.medvram)
-else:
-    shared.sd_model = shared.sd_model.to(shared.device)
-
-modules.sd_hijack.model_hijack.hijack(shared.sd_model)
+shared.sd_model = modules.sd_models.load_model()
+shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(shared.sd_model)))
 
 
 def webui():
-- 
cgit v1.2.3