From 44438c60e55344a93322d2665e51cdaedb0d52e4 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Mon, 25 Sep 2023 13:12:35 -0700
Subject: [PATCH] Add INSTALL_EXTENSIONS environment variable

---
 README.md    |  2 +-
 one_click.py | 24 +++++++++++++++---------
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index deba55af..70f2446c 100644
--- a/README.md
+++ b/README.md
@@ -58,7 +58,7 @@ To define persistent command-line flags like `--listen` or `--api`, edit the `CM
 * There is no need to run any of those scripts as admin/root.
 * For additional instructions about AMD setup, WSL setup, and nvcc installation, consult [this page](https://github.com/oobabooga/text-generation-webui/blob/main/docs/One-Click-Installers.md).
 * The installer has been tested mostly on NVIDIA GPUs. If you can find a way to improve it for your AMD/Intel Arc/Mac Metal GPU, you are highly encouraged to submit a PR to this repository. The main file to be edited is `one_click.py`.
-* For automated installation, you can use the `GPU_CHOICE` and `LAUNCH_AFTER_INSTALL` environment variables. For instance: `GPU_CHOICE=A LAUNCH_AFTER_INSTALL=False ./start_linux.sh`.
+* For automated installation, you can use the `GPU_CHOICE`, `LAUNCH_AFTER_INSTALL`, and `INSTALL_EXTENSIONS` environment variables. For instance: `GPU_CHOICE=A LAUNCH_AFTER_INSTALL=False INSTALL_EXTENSIONS=False ./start_linux.sh`.
 
 ### Manual installation using Conda
 
diff --git a/one_click.py b/one_click.py
index d2516c5e..ab3876f7 100644
--- a/one_click.py
+++ b/one_click.py
@@ -183,16 +183,22 @@ def update_requirements(initial_installation=False):
 
     run_cmd("git pull --autostash", assert_success=True, environment=True)
 
-    # Initial installation only: install the extensions requirements
-    if initial_installation:
-        extensions = next(os.walk("extensions"))[1]
-        for extension in extensions:
-            if extension in ['superbooga']:  # No wheels available for requirements
-                continue
+    # Extensions requirements are installed only during the initial install by default.
+    # That can be changed with the INSTALL_EXTENSIONS environment variable.
+    install_extensions = os.environ.get("INSTALL_EXTENSIONS", "false").lower() in ("yes", "y", "true", "1", "t", "on")
+    if initial_installation or install_extensions:
+        if not install_extensions:
+            print_big_message("Will not install extensions due to INSTALL_EXTENSIONS environment variable.")
+        else:
+            print("Installing extensions requirements.")
+            extensions = next(os.walk("extensions"))[1]
+            for extension in extensions:
+                if extension in ['superbooga']:  # No wheels available for requirements
+                    continue
 
-            extension_req_path = os.path.join("extensions", extension, "requirements.txt")
-            if os.path.exists(extension_req_path):
-                run_cmd("python -m pip install -r " + extension_req_path + " --upgrade", assert_success=True, environment=True)
+                extension_req_path = os.path.join("extensions", extension, "requirements.txt")
+                if os.path.exists(extension_req_path):
+                    run_cmd("python -m pip install -r " + extension_req_path + " --upgrade", assert_success=True, environment=True)
 
     # Detect the PyTorch version
     torver = torch_version()