Save/restore working webui/extension configs

modules/ui_extensions.py

@@ -2,6 +2,7 @@ import json
 import os.path
 import sys
 import time
+from datetime import datetime
 import traceback
 
 import git
@@ -11,7 +12,8 @@ import html
 import shutil
 import errno
 
-from modules import extensions, shared, paths
+from modules import extensions, shared, paths, config_states
+from modules.paths_internal import config_states_dir
 from modules.call_queue import wrap_gradio_gpu_call
 
 available_extensions = {"extensions": []}
@@ -30,6 +32,9 @@ def apply_and_restart(disable_list, update_list, disable_all):
     update = json.loads(update_list)
     assert type(update) == list, f"wrong update_list data for apply_and_restart: {update_list}"
 
+    if update:
+        save_config_state("Backup (pre-update)")
+
     update = set(update)
 
     for ext in extensions.extensions:
@@ -50,6 +55,48 @@ def apply_and_restart(disable_list, update_list, disable_all):
     shared.state.need_restart = True
 
 
+def save_config_state(name):
+    current_config_state = config_states.get_config()
+    if not name:
+        name = "Config"
+    current_config_state["name"] = name
+    filename = os.path.join(config_states_dir, datetime.now().strftime("%Y_%m_%d-%H_%M_%S") + ".json")
+    print(f"Saving backup of webui/extension state to {filename}.")
+    with open(filename, "w", encoding="utf-8") as f:
+        json.dump(current_config_state, f)
+    config_states.list_config_states()
+    new_value = next(iter(config_states.all_config_states.keys()), "Current")
+    new_choices = ["Current"] + list(config_states.all_config_states.keys())
+    return gr.Dropdown.update(value=new_value, choices=new_choices), f"<span>Saved current webui/extension state to '{filename}'</span>"
+
+
+def restore_config_state(confirmed, config_state_name, restore_type):
+    if config_state_name == "Current":
+        return "<span>Select a config to restore from.</span>"
+    if not confirmed:
+        return "<span>Cancelled.</span>"
+
+    check_access()
+
+    save_config_state("Backup (pre-restore)")
+
+    config_state = config_states.all_config_states[config_state_name]
+
+    print(f"Restoring webui state from backup: {restore_type}")
+
+    if restore_type == "extensions" or restore_type == "both":
+        shared.opts.restore_config_state_file = config_state["filename"]
+        shared.opts.save(shared.config_filename)
+
+    if restore_type == "webui" or restore_type == "both":
+        config_states.restore_webui_config(config_state)
+
+    shared.state.interrupt()
+    shared.state.need_restart = True
+
+    return ""
+
+
 def check_updates(id_task, disable_list):
     check_access()
 
@@ -121,6 +168,117 @@ def extension_table():
     return code
 
 
+def update_config_states_table(state_name):
+    if state_name == "Current":
+        config_state = config_states.get_config()
+    else:
+        config_state = config_states.all_config_states[state_name]
+
+    config_name = config_state.get("name", "Config")
+    created_date = time.asctime(time.gmtime(config_state["created_at"]))
+
+    code = f"""<!-- {time.time()} -->"""
+
+    webui_remote = config_state["webui"]["remote"] or ""
+    webui_branch = config_state["webui"]["branch"]
+    webui_commit_hash = config_state["webui"]["commit_hash"]
+    if webui_commit_hash:
+        webui_commit_hash = webui_commit_hash[:8]
+    else:
+        webui_commit_hash = "<unknown>"
+    webui_commit_date = config_state["webui"]["commit_date"]
+    if webui_commit_date:
+        webui_commit_date = time.asctime(time.gmtime(webui_commit_date))
+    else:
+        webui_commit_date = "<unknown>"
+
+    code += f"""<h2>Config Backup: {config_name}</h2>
+      <span>Created at: {created_date}</span>"""
+
+    code += f"""<h2>WebUI State</h2>
+      <table id="config_state_webui">
+        <thead>
+            <tr>
+                <th>URL</th>
+                <th>Branch</th>
+                <th>Commit</th>
+                <th>Date</th>
+            </tr>
+        </thead>
+        <tbody>
+            <tr>
+                <td>{webui_remote}</td>
+                <td>{webui_branch}</td>
+                <td>{webui_commit_hash}</td>
+                <td>{webui_commit_date}</td>
+            </tr>
+        </tbody>
+      </table>
+    """
+
+    code += """<h2>Extension State</h2>
+      <table id="config_state_extensions">
+        <thead>
+            <tr>
+                <th>Extension</th>
+                <th>URL</th>
+                <th>Branch</th>
+                <th>Commit</th>
+                <th>Date</th>
+            </tr>
+        </thead>
+        <tbody>
+    """
+
+    ext_map = {ext.name: ext for ext in extensions.extensions}
+
+    for ext_name, ext_conf in config_state["extensions"].items():
+        ext_remote = ext_conf["remote"] or ""
+        ext_branch = ext_conf["branch"] or "<unknown>"
+        ext_enabled = ext_conf["enabled"]
+        ext_commit_hash = ext_conf["commit_hash"] or "<unknown>"
+        ext_commit_date = ext_conf["commit_date"]
+        if ext_commit_date:
+            ext_commit_date = time.asctime(time.gmtime(ext_commit_date))
+        else:
+            ext_commit_date = "<unknown>"
+
+        remote = f"""<a href="{html.escape(ext_remote)}" target="_blank">{html.escape(ext_remote or '')}</a>"""
+
+        style_enabled = ""
+        style_remote = ""
+        style_branch = ""
+        style_commit = ""
+        if ext_name in ext_map:
+            current_ext = ext_map[ext_name]
+            current_ext.read_info_from_repo()
+            if current_ext.enabled != ext_enabled:
+                style_enabled = ' style="color: var(--primary-400)"'
+            if current_ext.remote != ext_remote:
+                style_remote = ' style="color: var(--primary-400)"'
+            if current_ext.branch != ext_branch:
+                style_branch = ' style="color: var(--primary-400)"'
+            if current_ext.commit_hash != ext_commit_hash:
+                style_commit = ' style="color: var(--primary-400)"'
+
+        code += f"""
+            <tr>
+                <td><label{style_enabled}><input class="gr-check-radio gr-checkbox" type="checkbox" disabled="true" {'checked="checked"' if ext_enabled else ''}>{html.escape(ext_name)}</label></td>
+                <td><label{style_remote}>{remote}</label></td>
+                <td><label{style_branch}>{ext_branch}</label></td>
+                <td><label{style_commit}>{ext_commit_hash[:8]}</label></td>
+                <td><label{style_commit}>{ext_commit_date}</label></td>
+            </tr>
+    """
+
+    code += """
+        </tbody>
+    </table>
+    """
+
+    return code
+
+
 def normalize_git_url(url):
     if url is None:
         return ""
@@ -292,6 +450,8 @@ def refresh_available_extensions_from_data(hide_tags, sort_column, filter_text="
 def create_ui():
     import modules.ui
 
+    config_states.list_config_states()
+
     with gr.Blocks(analytics_enabled=False) as ui:
         with gr.Tabs(elem_id="tabs_extensions") as tabs:
             with gr.TabItem("Installed"):
@@ -386,4 +546,28 @@ def create_ui():
                     outputs=[extensions_table, install_result],
                 )
 
+            with gr.TabItem("Backup/Restore"):
+                with gr.Row(elem_id="extensions_backup_top_row"):
+                    config_states_list = gr.Dropdown(label="Saved Configs", elem_id="extension_backup_saved_configs", value="Current", choices=["Current"] + list(config_states.all_config_states.keys()))
+                    modules.ui.create_refresh_button(config_states_list, config_states.list_config_states, lambda: {"choices": ["Current"] + list(config_states.all_config_states.keys())}, "refresh_config_states")
+                    config_restore_type = gr.Radio(label="State to restore", choices=["extensions", "webui", "both"], value="extensions", elem_id="extension_backup_restore_type")
+                    config_restore_button = gr.Button(value="Restore Selected Config", variant="primary", elem_id="extension_backup_restore")
+                with gr.Row(elem_id="extensions_backup_top_row2"):
+                    config_save_name = gr.Textbox("", placeholder="Config Name", show_label=False)
+                    config_save_button = gr.Button(value="Save Current Config")
+
+                config_states_info = gr.HTML("")
+                config_states_table = gr.HTML(lambda: update_config_states_table("Current"))
+
+                config_save_button.click(fn=save_config_state, inputs=[config_save_name], outputs=[config_states_list, config_states_info])
+
+                dummy_component = gr.Label(visible=False)
+                config_restore_button.click(fn=restore_config_state, _js="config_state_confirm_restore", inputs=[dummy_component, config_states_list, config_restore_type], outputs=[config_states_info])
+
+                config_states_list.change(
+                    fn=update_config_states_table,
+                    inputs=[config_states_list],
+                    outputs=[config_states_table],
+                )
+
     return ui