From ad5afcaae0b47e9e68b49aacf04cc3ad59d41a8e Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Wed, 29 Mar 2023 16:46:03 -0500
Subject: Save/restore working webui/extension configs

---
 webui.py | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

(limited to 'webui.py')

diff --git a/webui.py b/webui.py
index b570895f..b8f9a2c1 100644
--- a/webui.py
+++ b/webui.py
@@ -5,6 +5,7 @@ import importlib
 import signal
 import re
 import warnings
+import json
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
@@ -37,7 +38,7 @@ if ".dev" in torch.__version__ or "+git" in torch.__version__:
     torch.__long_version__ = torch.__version__
     torch.__version__ = re.search(r'[\d.]+[\d]', torch.__version__).group(0)
 
-from modules import shared, devices, sd_samplers, upscaler, extensions, localization, ui_tempdir, ui_extra_networks
+from modules import shared, devices, sd_samplers, upscaler, extensions, localization, ui_tempdir, ui_extra_networks, config_states
 import modules.codeformer_model as codeformer
 import modules.face_restoration
 import modules.gfpgan_model as gfpgan
@@ -105,6 +106,17 @@ def initialize():
     localization.list_localizations(cmd_opts.localizations_dir)
     startup_timer.record("list extensions")
 
+    config_state_file = shared.opts.restore_config_state_file
+    shared.opts.restore_config_state_file = ""
+    shared.opts.save(shared.config_filename)
+
+    if os.path.isfile(config_state_file):
+        print(f"*** About to restore extension state from file: {config_state_file}")
+        with open(config_state_file, "r", encoding="utf-8") as f:
+            config_state = json.load(f)
+            config_states.restore_extension_state(config_state)
+        startup_timer.record("restore extension config")
+
     if cmd_opts.ui_debug_mode:
         shared.sd_upscalers = upscaler.UpscalerLanczos().scalers
         modules.scripts.load_scripts()
@@ -301,6 +313,17 @@ def webui():
         extensions.list_extensions()
         startup_timer.record("list extensions")
 
+        config_state_file = shared.opts.restore_config_state_file
+        shared.opts.restore_config_state_file = ""
+        shared.opts.save(shared.config_filename)
+
+        if os.path.isfile(config_state_file):
+            print(f"*** About to restore extension state from file: {config_state_file}")
+            with open(config_state_file, "r", encoding="utf-8") as f:
+                config_state = json.load(f)
+                config_states.restore_extension_state(config_state)
+            startup_timer.record("restore extension config")
+
         localization.list_localizations(cmd_opts.localizations_dir)
 
         modelloader.forbid_loaded_nonbuiltin_upscalers()
-- 
cgit v1.2.3