From 19f9c1b1acf16f5ceee7f9ceb853ca7d968be902 Mon Sep 17 00:00:00 2001
From: Atinoda <61033436+Atinoda@users.noreply.github.com>
Date: Mon, 12 Jun 2023 23:28:04 +0100
Subject: [PATCH] Implement persistent extensions with optional runtime build

- Extensions are now explicitly copied from source to named volume
- Optional rebuild at runtime allows for extensions not present during build
- Closes #3
---
 Dockerfile                            |  5 +++--
 README.md                             | 11 ++++++++-
 config/extensions/.gitkeep            |  0
 docker-compose.yml                    |  2 ++
 docker-compose.yml.build              |  2 ++
 scripts/docker-entrypoint.sh          |  8 ++++++-
 scripts/extensions_runtime_rebuild.sh | 32 +++++++++++++++++++++++++++
 7 files changed, 56 insertions(+), 4 deletions(-)
 create mode 100644 config/extensions/.gitkeep
 create mode 100644 scripts/extensions_runtime_rebuild.sh

diff --git a/Dockerfile b/Dockerfile
index 4df003c..cb7c3ab 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -57,9 +57,10 @@ ENV PYTHONUNBUFFERED=1
 ARG BUILD_DATE
 ENV BUILD_DATE=$BUILD_DATE
 RUN echo "$BUILD_DATE" > /build_date.txt
+# Copy and enable all scripts
+COPY ./scripts /scripts
+RUN chmod +x /scripts/*
 # Run
-COPY ./scripts/docker-entrypoint.sh /scripts/docker-entrypoint.sh
-RUN chmod +x /scripts/docker-entrypoint.sh
 ENTRYPOINT ["/scripts/docker-entrypoint.sh"]
 
 
diff --git a/README.md b/README.md
index 626be51..8cdf0a9 100644
--- a/README.md
+++ b/README.md
@@ -52,7 +52,9 @@ Three commonly used ports are exposed:
 *Extensions may use additional ports - check the application documentation for more details.*
 
 ### Volumes
-The provided example docker compose maps several volumes from the local `config` directory into the container: `loras, models, presets, prompts, training`. If these folders are empty, they will be initialised when the container is run.
+The provided example docker compose maps several volumes from the local `config` directory into the container: `loras, models, presets, prompts, training, extensions`. If these folders are empty, they will be initialised when the container is run.
+
+Extensions will persist their state between container launches if you use the mapped folder - **but they will not automatically update when a new image is released, so this feature is disabled by default.** You must uncomment the mapping in `docker-compose.yml` to enable this extension persistence.
 
 *If you are getting an error about missing files, try clearing these folders and letting the service re-populate them.*
 
@@ -61,6 +63,13 @@ Extra launch arguments can be defined in the environment variable `EXTRA_LAUNCH_
 
 *Launch arguments should be defined as a space-separated list, just like writing them on the command line. These arguments are passed to the `server.py` module.*
 
+### Runtime extension build
+Extra launch arguments can be defined in the environment variable `BUILD_EXTENSIONS_LIVE` (e.g., `"silero_tts whisper_stt"`, will rebuild those extensions at launch). This feature may be useful if you are developing a third-party extension and need its dependencies refresh at launch.
+
+**Startup times will be much slower** if you use this feature, because it will rebuild the named extensions every time the container is started (i.e., don't use this feature unless you are certain that you need it.)
+
+*Extension names for runtime build should be defined as a space-separated list.*
+
 ### Updates
 These projects are moving quickly! To update to the most recent version on Docker hub, pull the latest image:
 
diff --git a/config/extensions/.gitkeep b/config/extensions/.gitkeep
new file mode 100644
index 0000000..e69de29
diff --git a/docker-compose.yml b/docker-compose.yml
index 2907e86..de6e4c1 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -5,6 +5,7 @@ services:
     container_name: text-generation-webui
     environment:
       - EXTRA_LAUNCH_ARGS="--listen --verbose" # Custom launch args (e.g., --model MODEL_NAME)
+#      - BUILD_EXTENSIONS_LIVE="silero_tts whisper_stt" # Install named extensions during every container launch. THIS WILL SIGNIFICANLTLY SLOW LAUNCH TIME.
     ports:
       - 7860:7860  # Default web port
 #      - 5000:5000  # Default API port
@@ -17,6 +18,7 @@ services:
       - ./config/prompts:/app/prompts
       - ./config/softprompts:/app/softprompts
       - ./config/training:/app/training
+      - ./config/extensions:/app/extensions
     logging:
       driver:  json-file
       options:
diff --git a/docker-compose.yml.build b/docker-compose.yml.build
index 5570fbb..cf8a236 100644
--- a/docker-compose.yml.build
+++ b/docker-compose.yml.build
@@ -9,6 +9,7 @@ services:
     container_name: text-generation-webui
     environment:
       - EXTRA_LAUNCH_ARGS="--listen --verbose" # Custom launch args (e.g., --model MODEL_NAME)
+#      - BUILD_EXTENSIONS_LIVE="silero_tts whisper_stt" # Install named extensions during every container launch. THIS WILL SIGNIFICANLTLY SLOW LAUNCH TIME.
     ports:
       - 7860:7860  # Default web port
 #      - 5000:5000  # Default API port
@@ -21,6 +22,7 @@ services:
       - ./config/prompts:/app/prompts
       - ./config/softprompts:/app/softprompts
       - ./config/training:/app/training
+      - ./config/extensions:/app/extensions
     logging:
       driver:  json-file
       options:
diff --git a/scripts/docker-entrypoint.sh b/scripts/docker-entrypoint.sh
index bb0049f..201a4da 100644
--- a/scripts/docker-entrypoint.sh
+++ b/scripts/docker-entrypoint.sh
@@ -10,7 +10,7 @@ function ctrl_c {
 trap ctrl_c SIGTERM SIGINT SIGQUIT SIGHUP
 
 # Generate default configs if empty
-CONFIG_DIRECTORIES=("loras" "models" "presets" "prompts" "training/datasets" "training/formats")
+CONFIG_DIRECTORIES=("loras" "models" "presets" "prompts" "training/datasets" "training/formats" "extensions")
 for config_dir in "${CONFIG_DIRECTORIES[@]}"; do
   if [ -z "$(ls /app/"$config_dir")" ]; then
     echo "*** Initialising config for: '$config_dir' ***"
@@ -19,6 +19,12 @@ for config_dir in "${CONFIG_DIRECTORIES[@]}"; do
   fi
 done
 
+# Runtime extension build
+if [[ -n "$BUILD_EXTENSIONS_LIVE" ]]; then
+  eval "live_extensions=($BUILD_EXTENSIONS_LIVE)"
+  . /scripts/extensions_runtime_rebuild.sh $live_extensions
+fi
+
 # Print variant
 VARIANT=$(cat /variant.txt)
 echo "=== Running text-generation-webui variant: '$VARIANT' ===" 
diff --git a/scripts/extensions_runtime_rebuild.sh b/scripts/extensions_runtime_rebuild.sh
new file mode 100644
index 0000000..9d8a052
--- /dev/null
+++ b/scripts/extensions_runtime_rebuild.sh
@@ -0,0 +1,32 @@
+#!/bin/bash
+
+# Note starting directory
+cur_dir=$(pwd)
+
+# Specify the directory containing the top-level folders
+directory="/app/extensions"
+
+# Iterate over the extensions passed in args
+for extension in "$@"; do
+    echo $extension
+    folder="$directory/$extension"
+    if [ -d "$folder" ]; then
+        # Change directory to the current folder
+        cd "$folder"
+
+        # Check if requirements.txt file exists
+        if [ -f "requirements.txt" ]; then
+            echo "Live installing requirements for $folder..."
+            pip3 install -r requirements.txt
+            echo "Requirements installed in $folder"
+        else
+            echo "Skipping live install of $folder: requirements.txt not found"
+        fi
+
+        # Change back to the original directory
+        cd "$directory"
+    fi
+done
+
+# Return to starting directory
+cd $cur_dir