Merge pull request #9518 from yike5460/master

add branch support for extension installation

1 files changed, 14 insertions, 6 deletions

diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py
index efd6cda2..b402bc8b 100644
--- a/modules/ui_extensions.py
+++ b/modules/ui_extensions.py
@@ -129,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'

@@ -150,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:

@@ -376,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],

                 )