Various UI fixes in config state tab

2 files changed, 226 insertions, 19 deletions

diff --git a/modules/config_states.py b/modules/config_states.py
new file mode 100644
index 00000000..2ea00929
--- /dev/null
+++ b/modules/config_states.py
@@ -0,0 +1,200 @@
+"""
+Supports saving and restoring webui and extensions from a known working set of commits
+"""
+
+import os
+import sys
+import traceback
+import json
+import time
+import tqdm
+
+from datetime import datetime
+from collections import OrderedDict
+import git
+
+from modules import shared, extensions
+from modules.paths_internal import extensions_dir, extensions_builtin_dir, script_path, config_states_dir
+
+
+all_config_states = OrderedDict()
+
+
+def list_config_states():
+    global all_config_states
+
+    all_config_states.clear()
+    os.makedirs(config_states_dir, exist_ok=True)
+
+    config_states = []
+    for filename in os.listdir(config_states_dir):
+        if filename.endswith(".json"):
+            path = os.path.join(config_states_dir, filename)
+            with open(path, "r", encoding="utf-8") as f:
+                j = json.load(f)
+                j["filepath"] = path
+                config_states.append(j)
+
+    config_states = list(sorted(config_states, key=lambda cs: cs["created_at"], reverse=True))
+
+    for cs in config_states:
+        timestamp = time.asctime(time.gmtime(cs["created_at"]))
+        name = cs.get("name", "Config")
+        full_name = f"{name}: {timestamp}"
+        all_config_states[full_name] = cs
+
+    return all_config_states
+
+
+def get_webui_config():
+    webui_repo = None
+
+    try:
+        if os.path.exists(os.path.join(script_path, ".git")):
+            webui_repo = git.Repo(script_path)
+    except Exception:
+        print(f"Error reading webui git info from {script_path}:", file=sys.stderr)
+        print(traceback.format_exc(), file=sys.stderr)
+
+    webui_remote = None
+    webui_commit_hash = None
+    webui_commit_date = None
+    webui_branch = None
+    if webui_repo and not webui_repo.bare:
+        try:
+            webui_remote = next(webui_repo.remote().urls, None)
+            head = webui_repo.head.commit
+            webui_commit_date = webui_repo.head.commit.committed_date
+            webui_commit_hash = head.hexsha
+            webui_branch = webui_repo.active_branch.name
+
+        except Exception:
+            webui_remote = None
+
+    return {
+        "remote": webui_remote,
+        "commit_hash": webui_commit_hash,
+        "commit_date": webui_commit_date,
+        "branch": webui_branch,
+    }
+
+
+def get_extension_config():
+    ext_config = {}
+
+    for ext in extensions.extensions:
+        entry = {
+            "name": ext.name,
+            "path": ext.path,
+            "enabled": ext.enabled,
+            "is_builtin": ext.is_builtin,
+            "remote": ext.remote,
+            "commit_hash": ext.commit_hash,
+            "commit_date": ext.commit_date,
+            "branch": ext.branch,
+            "have_info_from_repo": ext.have_info_from_repo
+        }
+
+        ext_config[ext.name] = entry
+
+    return ext_config
+
+
+def get_config():
+    creation_time = datetime.now().timestamp()
+    webui_config = get_webui_config()
+    ext_config = get_extension_config()
+
+    return {
+        "created_at": creation_time,
+        "webui": webui_config,
+        "extensions": ext_config
+    }
+
+
+def restore_webui_config(config):
+    print("* Restoring webui state...")
+
+    if "webui" not in config:
+        print("Error: No webui data saved to config")
+        return
+
+    webui_config = config["webui"]
+
+    if "commit_hash" not in webui_config:
+        print("Error: No commit saved to webui config")
+        return
+
+    webui_commit_hash = webui_config.get("commit_hash", None)
+    webui_repo = None
+
+    try:
+        if os.path.exists(os.path.join(script_path, ".git")):
+            webui_repo = git.Repo(script_path)
+    except Exception:
+        print(f"Error reading webui git info from {script_path}:", file=sys.stderr)
+        print(traceback.format_exc(), file=sys.stderr)
+        return
+
+    try:
+        webui_repo.git.fetch(all=True)
+        webui_repo.git.reset(webui_commit_hash, hard=True)
+        print(f"* Restored webui to commit {webui_commit_hash}.")
+    except Exception:
+        print(f"Error restoring webui to commit {webui_commit_hash}:", file=sys.stderr)
+        print(traceback.format_exc(), file=sys.stderr)
+
+
+def restore_extension_config(config):
+    print("* Restoring extension state...")
+
+    if "extensions" not in config:
+        print("Error: No extension data saved to config")
+        return
+
+    ext_config = config["extensions"]
+
+    results = []
+    disabled = []
+
+    for ext in tqdm.tqdm(extensions.extensions):
+        if ext.is_builtin:
+            continue
+
+        ext.read_info_from_repo()
+        current_commit = ext.commit_hash
+
+        if ext.name not in ext_config:
+            ext.disabled = True
+            disabled.append(ext.name)
+            results.append((ext, current_commit[:8], False, "Saved extension state not found in config, marking as disabled"))
+            continue
+
+        entry = ext_config[ext.name]
+
+        if "commit_hash" in entry and entry["commit_hash"]:
+            try:
+                ext.fetch_and_reset_hard(entry["commit_hash"])
+                ext.read_info_from_repo()
+                if current_commit != entry["commit_hash"]:
+                    results.append((ext, current_commit[:8], True, entry["commit_hash"][:8]))
+            except Exception as ex:
+                results.append((ext, current_commit[:8], False, ex))
+        else:
+            results.append((ext, current_commit[:8], False, "No commit hash found in config"))
+
+        if not entry.get("enabled", False):
+            ext.disabled = True
+            disabled.append(ext.name)
+        else:
+            ext.disabled = False
+
+    shared.opts.disabled_extensions = disabled
+    shared.opts.save(shared.config_filename)
+
+    print("* Finished restoring extensions. Results:")
+    for ext, prev_commit, success, result in results:
+        if success:
+            print(f"  + {ext.name}: {prev_commit} -> {result}")
+        else:
+            print(f"  ! {ext.name}: FAILURE ({result})")
diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py
index ed677b3e..b94b3a3a 100644
--- a/modules/ui_extensions.py
+++ b/modules/ui_extensions.py
@@ -17,6 +17,7 @@ from modules.paths_internal import config_states_dir
 from modules.call_queue import wrap_gradio_gpu_call

 

 available_extensions = {"extensions": []}

+STYLE_PRIMARY = ' style="color: var(--primary-400)"'

 

 

 def check_access():

@@ -67,7 +68,7 @@ def save_config_state(name):
     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>"

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

@@ -78,14 +79,12 @@ def restore_config_state(confirmed, config_state_name, restore_type):
 

     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}")

+    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.restore_config_state_file = config_state["filepath"]

         shared.opts.save(shared.config_filename)

 

     if restore_type == "webui" or restore_type == "both":

@@ -149,7 +148,7 @@ def extension_table():
 

         style = ""

         if shared.opts.disable_all_extensions == "extra" and not ext.is_builtin or shared.opts.disable_all_extensions == "all":

-            style = ' style="color: var(--primary-400)"'

+            style = STYLE_PRIMARY

 

         code += f"""

             <tr>

@@ -181,17 +180,25 @@ def update_config_states_table(state_name):
 

     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_hash = config_state["webui"]["commit_hash"] or "<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>"

 

+    current_webui = config_states.get_webui_config()

+

+    style_remote = ""

+    style_branch = ""

+    style_commit = ""

+    if current_webui["remote"] != webui_remote:

+        style_remote = STYLE_PRIMARY

+    if current_webui["branch"] != webui_branch:

+        style_branch = STYLE_PRIMARY

+    if current_webui["commit_hash"] != webui_commit_hash:

+        style_commit = STYLE_PRIMARY

+

     code += f"""<h2>Config Backup: {config_name}</h2>

       <span>Created at: {created_date}</span>"""

 

@@ -207,10 +214,10 @@ def update_config_states_table(state_name):
         </thead>

         <tbody>

             <tr>

-                <td>{webui_remote}</td>

-                <td>{webui_branch}</td>

-                <td>{webui_commit_hash}</td>

-                <td>{webui_commit_date}</td>

+                <td><label{style_remote}>{webui_remote}</label></td>

+                <td><label{style_branch}>{webui_branch}</label></td>

+                <td><label{style_commit}>{webui_commit_hash[:8]}</label></td>

+                <td><label{style_commit}>{webui_commit_date}</label></td>

             </tr>

         </tbody>

       </table>

@@ -253,13 +260,13 @@ def update_config_states_table(state_name):
             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)"'

+                style_enabled = STYLE_PRIMARY

             if current_ext.remote != ext_remote:

-                style_remote = ' style="color: var(--primary-400)"'

+                style_remote = STYLE_PRIMARY

             if current_ext.branch != ext_branch:

-                style_branch = ' style="color: var(--primary-400)"'

+                style_branch = STYLE_PRIMARY

             if current_ext.commit_hash != ext_commit_hash:

-                style_commit = ' style="color: var(--primary-400)"'

+                style_commit = STYLE_PRIMARY

 

         code += f"""

             <tr>