1 files changed, 36 insertions, 10 deletions
diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py
index da7e79f0..b402bc8b 100644
--- a/modules/ui_extensions.py
+++ b/modules/ui_extensions.py
@@ -21,7 +21,7 @@ def check_access():
     assert not shared.cmd_opts.disable_extension_access, "extension access disabled because of command line flags"
 
 
-def apply_and_restart(disable_list, update_list):
+def apply_and_restart(disable_list, update_list, disable_all):
     check_access()
 
     disabled = json.loads(disable_list)
@@ -43,6 +43,7 @@ def apply_and_restart(disable_list, update_list):
             print(traceback.format_exc(), file=sys.stderr)
 
     shared.opts.disabled_extensions = disabled
+    shared.opts.disable_all_extensions = disable_all
     shared.opts.save(shared.config_filename)
 
     shared.state.interrupt()
@@ -63,6 +64,9 @@ def check_updates(id_task, disable_list):
 
         try:
             ext.check_updates()
+        except FileNotFoundError as e:
+            if 'FETCH_HEAD' not in str(e):
+                raise
         except Exception:
             print(f"Error checking updates for {ext.name}:", file=sys.stderr)
             print(traceback.format_exc(), file=sys.stderr)
@@ -87,6 +91,8 @@ def extension_table():
     """
 
     for ext in extensions.extensions:
+        ext.read_info_from_repo()
+
         remote = f"""<a href="{html.escape(ext.remote or '')}" target="_blank">{html.escape("built-in" if ext.is_builtin else ext.remote or '')}</a>"""
 
         if ext.can_update:
@@ -94,9 +100,13 @@ def extension_table():
         else:
             ext_status = ext.status
 
+        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)"'
+
         code += f"""
             <tr>
-                <td><label><input class="gr-check-radio gr-checkbox" name="enable_{html.escape(ext.name)}" type="checkbox" {'checked="checked"' if ext.enabled else ''}>{html.escape(ext.name)}</label></td>
+                <td><label{style}><input class="gr-check-radio gr-checkbox" name="enable_{html.escape(ext.name)}" type="checkbox" {'checked="checked"' if ext.enabled else ''}>{html.escape(ext.name)}</label></td>
                 <td>{remote}</td>
                 <td>{ext.version}</td>
                 <td{' class="extension_status"' if ext.remote is not None else ''}>{ext_status}</td>
@@ -119,7 +129,7 @@ def normalize_git_url(url):
     return url
 
 
-def install_extension_from_url(dirname, url):
+def install_extension_from_url(dirname, branch_name, url):
     check_access()
 
     assert url, 'No URL specified'
@@ -140,10 +150,17 @@ def install_extension_from_url(dirname, url):
 
     try:
         shutil.rmtree(tmpdir, True)
-        with git.Repo.clone_from(url, tmpdir) as repo:
-            repo.remote().fetch()
-            for submodule in repo.submodules:
-                submodule.update()
+        if branch_name == '':
+            # if no branch is specified, use the default branch
+            with git.Repo.clone_from(url, tmpdir) as repo:
+                repo.remote().fetch()
+                for submodule in repo.submodules:
+                    submodule.update()
+        else:
+            with git.Repo.clone_from(url, tmpdir, branch=branch_name) as repo:
+                repo.remote().fetch()
+                for submodule in repo.submodules:
+                    submodule.update()
         try:
             os.rename(tmpdir, target_dir)
         except OSError as err:
@@ -289,16 +306,24 @@ def create_ui():
                 with gr.Row(elem_id="extensions_installed_top"):
                     apply = gr.Button(value="Apply and restart UI", variant="primary")
                     check = gr.Button(value="Check for updates")
+                    extensions_disable_all = gr.Radio(label="Disable all extensions", choices=["none", "extra", "all"], value=shared.opts.disable_all_extensions, elem_id="extensions_disable_all")
                     extensions_disabled_list = gr.Text(elem_id="extensions_disabled_list", visible=False).style(container=False)
                     extensions_update_list = gr.Text(elem_id="extensions_update_list", visible=False).style(container=False)
 
-                info = gr.HTML()
+                html = ""
+                if shared.opts.disable_all_extensions != "none":
+                    html = """
+<span style="color: var(--primary-400);">
+    "Disable all extensions" was set, change it to "none" to load all extensions again
+</span>
+                    """
+                info = gr.HTML(html)
                 extensions_table = gr.HTML(lambda: extension_table())
 
                 apply.click(
                     fn=apply_and_restart,
                     _js="extensions_apply",
-                    inputs=[extensions_disabled_list, extensions_update_list],
+                    inputs=[extensions_disabled_list, extensions_update_list, extensions_disable_all],
                     outputs=[],
                 )
 
@@ -358,13 +383,14 @@ def create_ui():
 
             with gr.TabItem("Install from URL"):
                 install_url = gr.Text(label="URL for extension's git repository")
+                install_branch = gr.Text(label="Specific branch name", placeholder="Leave empty for default main branch")
                 install_dirname = gr.Text(label="Local directory name", placeholder="Leave empty for auto")
                 install_button = gr.Button(value="Install", variant="primary")
                 install_result = gr.HTML(elem_id="extension_install_result")
 
                 install_button.click(
                     fn=modules.ui.wrap_gradio_call(install_extension_from_url, extra_outputs=[gr.update()]),
-                    inputs=[install_dirname, install_url],
+                    inputs=[install_dirname, install_branch, install_url],
                     outputs=[extensions_table, install_result],
                 )