Updates:

Upgraded base allows provisioning script to alter init.sh vars

No longer deleting most of the horde shell scripts. UX improved

add vars to prevent running as a service and for installing dev dependencies

.env

@@ -6,10 +6,12 @@ IMAGE_TAG="latest"
 ### Horde Vars ###
 
 BRIDGE_API_KEY=0000000000
-BRIDGE_WORKER_NAME=""
+BRIDGE_WORKER_NAME="Docker Dreamer"
 BRIDGE_MAX_THREADS=1
 BRIDGE_MAX_POWER=8
 BRIDGE_MODELS_TO_LOAD="Top 2"
+HORDE_DEV=false
+HORDE_SKIP_SERVICE=false
 HORDE_WEBUI_PORT_HOST=7860
 HORDE_WORKER_BRANCH=main
 WEB_USER=admin
@@ -20,6 +22,7 @@ WEB_PASSWORD=password
 WORKSPACE=/workspace/
 CF_TUNNEL_TOKEN=
 CF_QUICK_TUNNELS=true
+GPU_COUNT=
 JUPYTER_PORT=8888
 JUPYTER_PORT_HOST=8888
 JUPYTER_TOKEN=password

.github/workflows/docker-build.yml

@@ -24,7 +24,6 @@ jobs:
         pytorch:
           - "2.0.1"
         cuda:
-          - "11.7.1"
           - "11.8.0"
         level:
           - "base"

README.md

@@ -114,6 +114,8 @@ You can use the included `cloudflared` service to make secure connections withou
 | `CF_TUNNEL_TOKEN`     | Cloudflare zero trust tunnel token - See [documentation](https://developers.cloudflare.com/cloudflare-one/connections/connect-networks/). |
 | `CF_QUICK_TUNNELS`    | Create ephemeral Cloudflare tunnels for web services (default `false`) |
 | `GPU_COUNT`           | Limit the number of available GPUs |
+| `HORDE_DEV`           | Set `true` to install development dependencies. Defaults to `false` |
+| `HORDE_SKIP_SERVICE`  | Set `true` to prevent the worker launching on container start. Defaults to `false` |
 | `HORDE_WORKER_BRANCH` | AI-Horde-Worker branch/commit hash. Defaults to `main` |
 | `JUPYTER_MODE`        | `lab` (default), `notebook` |
 | `JUPYTER_PORT`        | Set an alternative port (default `8888`) |
@@ -139,7 +141,7 @@ It can be useful to perform certain actions when starting a container, such as c
 
 You can use the environment variable `PROVISIONING_SCRIPT` to specify the URL of a script you'd like to run.
 
-If you are running locally you may instead opt to mount an executable script at `/opt/ai-dock/bin/provisioning.sh`.
+If you are running locally you may instead opt to mount a script at `/opt/ai-dock/bin/provisioning.sh`.
 
 >[!NOTE]  
 >`supervisord` will not spawn any processes until the provisioning script has completed.
@@ -217,7 +219,7 @@ A dreamer (stable diffusion worker) will be created for every NVIDIA GPU found i
 
 The first worker (GPU 0) will be named according to your specification in `BRIDGE_WORKER_NAME`. Extra workers will have this name appended with '#2, #3...'
 
-The terminal UI is disabled. If you would like to use it you'll need to SSH (or docker exec) into your container and stop the dreamer service(s) with `supervisorctl`. You can then edit your `bridgeData.yaml` and start the worker manually.
+The terminal UI is disabled when running as a service. If you'd like to use it you should set `HORDE_SKIP_SERVICE=true` and access the container by SSH (or docker exec) and launch the worker manually.
 
 ### Web UI
 

build/COPY_ROOT/etc/supervisor/supervisord/conf.d/dreamer.conf

@@ -5,7 +5,7 @@ numprocs=%(ENV_GPU_COUNT)s
 directory=/opt/AI-Horde-Worker
 priority=2000
 autostart=true
-startsecs=10
+startsecs=0
 startretries=3
 autorestart=unexpected
 stopsignal=TERM

build/COPY_ROOT/opt/ai-dock/bin/build/layer0/nvidia.sh

@@ -11,12 +11,10 @@ install_horde() {
     micromamba run -n horde ${PIP_INSTALL} \
         torch=="${PYTORCH_VERSION}" \
         xformers \
-        hordelib \
         nvidia-ml-py3 \
         ruamel.yaml
-
-    micromamba run -n horde ${PIP_INSTALL} \
-        -r /opt/AI-Horde-Worker/requirements.txt
+
+    /opt/horde/bin/update-horde-worker.sh
 }
 
 main "$@"; exit

build/COPY_ROOT/opt/ai-dock/bin/preflight.sh

@@ -1,52 +1,61 @@
-#!/bin/bash
-
-function main() {
-    move_to_workspace
-    copy_notebook
-    update_horde
-    set_safe_envs
-    write_config
+#!/bin/false
+
+# This file will be sourced in init.sh
+
+function preflight_main() {
+    preflight_move_to_workspace
+    preflight_copy_notebook
+    preflight_update_horde
+    preflight_set_safe_envs
+    preflight_write_config
 }
 
-function move_to_workspace() {
+function preflight_move_to_workspace() {
     if [[ $WORKSPACE_MOUNTED = "true" ]]; then
-        if [[ ! -d "${WORKSPACE}AI-Horde-Worker" ]]; then
+        if [[ ! -d ${WORKSPACE}AI-Horde-Worker ]]; then
             mv /opt/AI-Horde-Worker ${WORKSPACE}
         fi
             rm -rf /opt/AI-Horde-Worker
-            ln -s "${WORKSPACE}/AI-Horde-Worker" /opt/AI-Horde-Worker
+            ln -s ${WORKSPACE}AI-Horde-Worker /opt/AI-Horde-Worker
     fi
 }
 
-function copy_notebook() {
+function preflight_copy_notebook() {
     if micromamba env list | grep 'jupyter' > /dev/null 2>&1;  then
-        if [[ ! -f "${WORKSPACE}/dreamer.ipynb" ]]; then
+        if [[ ! -f "${WORKSPACE}dreamer.ipynb" ]]; then
             cp /usr/local/share/ai-dock/dreamer.ipynb ${WORKSPACE}
         fi
     fi
 }
 
-function update_horde() {
-    /opt/horde/bin/update-hordelib.sh
+function preflight_update_horde() {
     /opt/horde/bin/update-horde-worker.sh
 }
 
-function set_safe_envs() {
+function preflight_set_safe_envs() {
+    DEFAULT_WORKER_NAME="Docker Dreamer"
+
     # Ensure worker_name for bridge_stable_diffusion.py -n
-    if [[ -z $BRIDGE_WORKER_NAME || "$BRIDGE_WORKER_NAME" = $DEFAULT_WORKER_NAME ]]; then
+    if [[ -z $BRIDGE_WORKER_NAME || $BRIDGE_WORKER_NAME = "$DEFAULT_WORKER_NAME" ]]; then
         export BRIDGE_WORKER_NAME="$DEFAULT_WORKER_NAME $(uuidgen -r)"
     fi
 
-    if [[ -n $BRIDGE_DREAMER_NAME ]]; then
+    if [[ -n $BRIDGE_DREAMER_NAME && $BRIDGE_DREAMER_NAME != $DEFAULT_WORKER_NAME ]]; then
         export BRIDGE_WORKER_NAME="${BRIDGE_DREAMER_NAME}"
     fi
 
+    if [[ -z $BRIDGE_DREAMER_NAME ]]; then
+        export BRIDGE_DREAMER_NAME="$BRIDGE_WORKER_NAME"
+    fi
+
     # Terminal UI will not work when running as a service
-    export BRIDGE_DISABLE_TERMINAL_UI=true
+    if [[ $HORDE_SKIP_SERVICE != "true" ]]; then
+        export BRIDGE_DISABLE_TERMINAL_UI=true
+    fi
 }
 
-function write_config() {
+function preflight_write_config() {
     micromamba run -n horde python /opt/horde/scripts/write-config.py
 }
 
-main "$@"; exit
+preflight_main "$@"

build/COPY_ROOT/opt/ai-dock/bin/supervisor-dreamer.sh

@@ -8,6 +8,15 @@ if [[ -z $PROC_NUM ]]; then
     PROC_NUM=0
 fi
 
+if [[ -z $PROC_NAME ]]; then
+    PROC_NAME="dreamer"
+fi
+
+if [[ -n $HORDE_SKIP_SERVICE && $HORDE_SKIP_SERVICE != "false" ]]; then
+    printf "Skipping startup for %s %s\n" "${PROC_NAME}" ${PROC_NUM}
+    exit 0
+fi
+
 if [[ $PROC_NUM -eq 0 ]]; then
     worker_name="$BRIDGE_WORKER_NAME"
 else

build/COPY_ROOT/opt/horde/bin/update-horde-worker.sh

@@ -22,18 +22,35 @@ git checkout ${branch}
 git pull
 
 # Use default gradio css
-rm webui.css
+if [[ -f webui.css ]]; then
+    git update-index --assume-unchanged webui.css
+    rm webui.css
+fi
+# Ensure the horde shell scripts work without alteration
+if [[ -f update-runtime.sh ]]; then
+    git update-index --assume-unchanged update-runtime.sh
+    rm update-runtime.sh
+fi
 
-# Overide shell scripts in case of interactive use. Defaults will break things - See docs
-rm update-runtime.sh
-rm horde-bridge.sh
-rm runtime.sh
-ln -s /opt/horde/bin/worker-update-runtime.sh update-runtime.sh
-ln -s /opt/horde/bin/worker-horde-bridge.sh horde-bridge.sh
-ln -s /opt/horde/bin/worker-runtime.sh runtime.sh
+rm -rf /opt/AI-Horde-Worker/conda > /dev/null 2>&1
+rm -rf /opt/AI-Horde-Worker/bin > /dev/null 2>&1
+cp -rf /opt/horde/override/bin .
+cp -rf /opt/horde/override/conda .
 
 # We pre-installed pytorch and want a universal version across all micromamba environments (if more than one)
-cp requirements.txt reqs-mod-dreamer.txt
-sed -i '/^torch.*[\W|=|>|<]*$/d' reqs-mod-dreamer.txt
 
-micromamba run -n horde pip --no-cache-dir install -r reqs-mod-dreamer.txt
+cp requirements.txt reqs-ai-horde-dreamer.txt
+sed -i '/^--extra-index-url.*[\W|=|>|<]*$/d' reqs-ai-horde-dreamer.txt
+sed -i '/^torch.*[\W|=|>|<]*$/d' reqs-ai-horde-dreamer.txt
+sed -i '/^xformers.*[\W|=|>|<]*$/d' reqs-ai-horde-dreamer.txt
+micromamba run -n horde pip --no-cache-dir install -r reqs-ai-horde-dreamer.txt
+
+if [[ -n $HORDE_DEV && $HORDE_DEV != "false" ]]; then
+    cp requirements.dev.txt reqs-ai-horde-dreamer-dev.txt
+    sed -i '/^--extra-index-url.*[\W|=|>|<]*$/d' reqs-ai-horde-dreamer-dev.txt
+    sed -i '/^torch.*[\W|=|>|<]*$/d' reqs-ai-horde-dreamer-dev.txt
+    sed -i '/^xformers.*[\W|=|>|<]*$/d' reqs-ai-horde-dreamer-dev.txt
+    micromamba run -n horde pip --no-cache-dir install -r reqs-ai-horde-dreamer-dev.txt
+fi
+
+rm reqs-ai-horde-dreamer*.txt

build/COPY_ROOT/opt/horde/bin/update-runtime.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+update-hordelib.sh
+update-horde-worker.sh

build/COPY_ROOT/opt/horde/override/bin/micromamba

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+if [[ $8 = "hordelib.benchmark" ]]; then
+    micromamba run -n horde python -m hordelib.benchmark
+elif [[ $8 = *".py" ]]; then
+    micromamba run -n horde python -s "${@:8}"
+else
+    printf "The docker override does not understand your use case - Please raise an issue at https://github.com/ai-dock/ai-horde-dreamer\n"
+fi

build/COPY_ROOT/opt/horde/override/conda/envs/linux/bin/python

@@ -0,0 +1 @@
+Just a regular file - Here to make runtime.sh proceed without updating

build/COPY_ROOT/opt/horde/override/readme.txt

@@ -0,0 +1,6 @@
+This directory contains files to be copied into /opt/AI-Horde-Worker
+
+The purpose is to make the horde shell scripts work without having to
+modify them and cause problems with git pull
+
+update-runtime.sh will still be deleted because it wants to clobber bin/micromamba

build/Dockerfile

@@ -3,16 +3,13 @@
 ARG IMAGE_BASE="ghcr.io/ai-dock/pytorch:2.0.1-py3.10-cuda-11.8.0-base-22.04"
 FROM ${IMAGE_BASE}
 
-LABEL org.opencontainers.image.source https://github.com/ai-dock/pytorch
+LABEL org.opencontainers.image.source https://github.com/ai-dock/ai-horde-dreamer
 
-LABEL org.opencontainers.image.description "Jupyter with Python kernels and pytorch pre-installed"
+LABEL org.opencontainers.image.description "AI Horde image generation worker"
 
 LABEL maintainer="Rob Ballantyne <[email protected]>"
 
-# Note: the default mamba environment is set by ai-dock/python
-
-# Add new paths at front 
-#ENV PATH=/opt/another/bin:$PATH
+ENV PATH=/opt/horde/bin:$PATH
 
 # Copy early so we can use scripts in the build - Changes to these files will invalidate the cache and cause a rebuild.
 COPY ./COPY_ROOT/ /
@@ -23,7 +20,6 @@ COPY ./COPY_ROOT/ /
 ARG IMAGE_BASE
 RUN bash -c /opt/ai-dock/bin/build/layer0/init.sh
 
-ENV PATH=/opt/horde/bin:$PATH
 ENV MAMBA_DEFAULT_ENV=horde
 
 # Keep init.sh as-is and place additional logic in /opt/ai-dock/bin/preflight.sh

config/provisioning.sh

@@ -1,6 +1,7 @@
-#!/bin/bash
+#!/bin/false
 
-# You can edit this file and make it do something useful
+# This file will be sourced in init.sh
+# You can edit below here and make it do something useful
 
 printf "Hello world!\n"
 

docker-compose.yaml

@@ -48,6 +48,9 @@ services:
     environment:
         # AI-Horde-Worker branch or commit hash
         - HORDE_WORKER_BRANCH=${HORDE_WORKER_BRANCH}
+        - HORDE_DEV=${HORDE_DEV}
+        - HORDE_SKIP_SERVICE=${HORDE_SKIP_SERVICE}
+        - GPU_COUNT=${GPU_COUNT}
 
         # Web UI auth credentials - Change the values.
         - WEB_USER=${WEB_USER}