From f4defde7526aa7767e550adf5b59cb1b7b09c782 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Sun, 11 Jun 2023 17:11:06 -0300
Subject: [PATCH] Add a menu for installing extensions

---
 modules/github.py | 28 ++++++++++++++++++++++++++++
 server.py         | 11 +++++++++++
 2 files changed, 39 insertions(+)
 create mode 100644 modules/github.py

diff --git a/modules/github.py b/modules/github.py
new file mode 100644
index 0000000..7960c8f
--- /dev/null
+++ b/modules/github.py
@@ -0,0 +1,28 @@
+import os
+import subprocess
+
+def clone_or_pull_repository(github_url):
+    repository_folder = "extensions"
+    repo_name = github_url.split("/")[-1].split(".")[0]
+
+    # Check if the repository folder exists
+    if not os.path.exists(repository_folder):
+        os.makedirs(repository_folder)
+
+    repo_path = os.path.join(repository_folder, repo_name)
+
+    # Check if the repository is already cloned
+    if os.path.exists(repo_path):
+        # Perform a 'git pull' to update the repository
+        try:
+            pull_output = subprocess.check_output(["git", "-C", repo_path, "pull"], stderr=subprocess.STDOUT)
+            return pull_output.decode()
+        except subprocess.CalledProcessError as e:
+            return str(e)
+
+    # Clone the repository
+    try:
+        clone_output = subprocess.check_output(["git", "clone", github_url, repo_path], stderr=subprocess.STDOUT)
+        return clone_output.decode()
+    except subprocess.CalledProcessError as e:
+        return str(e)
diff --git a/server.py b/server.py
index fec59cd..166d211 100644
--- a/server.py
+++ b/server.py
@@ -46,6 +46,7 @@ from PIL import Image
 import modules.extensions as extensions_module
 from modules import chat, shared, training, ui, utils
 from modules.extensions import apply_extensions
+from modules.github import clone_or_pull_repository
 from modules.html_generator import chat_html_wrapper
 from modules.LoRA import add_lora_to_model
 from modules.models import load_model, unload_model
@@ -833,10 +834,20 @@ def create_interface():
             with gr.Row():
                 with gr.Column():
                     shared.gradio['extensions_menu'] = gr.CheckboxGroup(choices=utils.get_available_extensions(), value=shared.args.extensions, label="Available extensions", info='Note that some of these extensions may require manually installing Python requirements through the command: pip install -r extensions/extension_name/requirements.txt', elem_classes='checkboxgroup-table')
+
                 with gr.Column():
                     shared.gradio['bool_menu'] = gr.CheckboxGroup(choices=bool_list, value=bool_active, label="Boolean command-line flags", elem_classes='checkboxgroup-table')
 
             shared.gradio['reset_interface'] = gr.Button("Apply and restart the interface")
+            with gr.Row():
+                extension_name = gr.Textbox(lines=1, label='Install or update an extension', info='Enter the GitHub URL below. For a list of extensions, see: https://github.com/oobabooga/text-generation-webui-extensions ⚠️  WARNING ⚠️ : extensions can execute arbitrary code. Make sure to inspect their source code before activating them.')
+                extension_install = gr.Button('Install or update', elem_classes="small-button")
+
+            extension_status = gr.Markdown()
+
+            extension_install.click(
+                clone_or_pull_repository, extension_name, extension_status, show_progress=False).then(
+                lambda: gr.update(choices=utils.get_available_extensions(), value=shared.args.extensions), outputs=shared.gradio['extensions_menu'])
 
             # Reset interface event
             shared.gradio['reset_interface'].click(