Fix and document AUTO_UPDATE var

ai-dock · Aug 16, 2024 · a86c554 · a86c554
1 parent aa3b4c7
commit a86c554
Show file tree

Hide file tree

Showing 10 changed files with 96 additions and 60 deletions.
diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
@@ -17,7 +17,7 @@ jobs:
       matrix:
         build:
           # Undeclared release tag finds latest from GitHub tags
-          - {latest: "true", tag: "v24.1.4", python: "3.10", pytorch: "2.3.0", cuda: "12.1.1-base"}
+          - {latest: "false", kohya: "v24.1.6", python: "3.10", pytorch: "2.4.0", cuda: "12.1.1-base"}
 
     steps:
       -
@@ -34,43 +34,54 @@ jobs:
       -
         name: Env Setter
         run: |
-          echo "PACKAGE_NAME=${GITHUB_REPOSITORY,,}" >> ${GITHUB_ENV}
+          REPO=${GITHUB_REPOSITORY,,}
+          echo "REPO_NAMESPACE=${REPO%%/*}" >> ${GITHUB_ENV}
+          echo "REPO_NAME=${REPO#*/}" >> ${GITHUB_ENV}
       -
         name: Checkout
         uses: actions/checkout@v3
       -
         name: Permissions fixes
         run: |
-          reponame="$(basename ${GITHUB_REPOSITORY})"
-          target="${HOME}/work/${reponame}/${reponame}/build/COPY*"
+          target="${HOME}/work/${{ env.REPO_NAME }}/${{ env.REPO_NAME }}/build/COPY*"
           chmod -R ug+rwX ${target}
+      -
+        name: Login to DockerHub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ vars.DOCKERHUB_USER }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
       -
         name: Login to GitHub Container Registry
-        uses: docker/login-action@v2
+        uses: docker/login-action@v3
         with:
           registry: ghcr.io
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
       -
         name: Set tags
         run: |
-          img_path="ghcr.io/${{ env.PACKAGE_NAME }}"
+          img_path_ghcr="ghcr.io/${{ env.REPO_NAMESPACE }}/${{ env.REPO_NAME }}"
+          img_path_dhub="${{ vars.DOCKERHUB_USER }}/${{ env.REPO_NAME }}-cuda"
           
-          if [[ -z '${{ matrix.build.tag }}' ]]; then
-                KOHYA_TAG="$(curl -s https://api.github.com/repos/bmaltais/kohya_ss/tags | jq -r '.[0].name')"
+          if [[ -z '${{ matrix.build.kohya }}' ]]; then
+                KOHYA_BUILD_REF="$(curl -s https://api.github.com/repos/bmaltais/kohya_ss/tags | jq -r '.[0].name')"
             else
-                KOHYA_TAG="${{ matrix.build.tag }}"
+                KOHYA_BUILD_REF="${{ matrix.build.kohya }}"
             fi
-            [ -z "$KOHYA_TAG" ] && { echo "Error: KOHYA_TAG is empty. Exiting script." >&2; exit 1; }
-            echo "KOHYA_TAG=${KOHYA_TAG}" >> ${GITHUB_ENV}
+            [ -z "$KOHYA_BUILD_REF" ] && { echo "Error: KOHYA_BUILD_REF is empty. Exiting script." >&2; exit 1; }
+            echo "KOHYA_BUILD_REF=${KOHYA_BUILD_REF}" >> ${GITHUB_ENV}
 
             base_tag="v2-cuda-${{ matrix.build.cuda }}-${{ env.UBUNTU_VERSION }}"
 
             if [[ ${{ matrix.build.latest }} == "true" ]]; then
                 echo "Marking latest"
-                TAGS="${img_path}:${base_tag}-${KOHYA_TAG}, ${img_path}:${base_tag}, ${img_path}:latest, ${img_path}:latest-cuda"
+                # GHCR.io
+                TAGS="${img_path}:${base_tag}-${KOHYA_BUILD_REF}, ${img_path}:${base_tag}, ${img_path}:latest, ${img_path}:latest-cuda"
+                # Docker.io Tags
+                TAGS="${TAGS}, ${img_path_dhub}:${KOHYA_BUILD_REF}, ${img_path_dhub}:latest"
             else
-                TAGS="${img_path}:${base_tag}-${KOHYA_TAG}"
+                TAGS="${img_path}:${base_tag}-${KOHYA_BUILD_REF}, ${img_path_dhub}:${KOHYA_BUILD_REF}"
             fi
           echo "TAGS=${TAGS}" >> ${GITHUB_ENV}
       -
@@ -82,7 +93,7 @@ jobs:
             IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.build.python }}-v2-cuda-${{ matrix.build.cuda }}-${{ env.UBUNTU_VERSION }}
             PYTHON_VERSION=${{ matrix.build.python }}
             PYTORCH_VERSION=${{ matrix.build.pytorch }}
-            KOHYA_TAG=${{ env.KOHYA_TAG }}
+            KOHYA_BUILD_REF=${{ env.KOHYA_BUILD_REF }}
           push: true
           provenance: false
           tags: ${{ env.TAGS }}
@@ -93,7 +104,7 @@ jobs:
       fail-fast: false
       matrix:
         build:
-          - {latest: "true", tag: "v24.1.4", python: "3.10", pytorch: "2.3.0", rocm: "6.0-core"}
+          - {latest: "false", kohya: "v24.1.6", python: "3.10", pytorch: "2.3.1", rocm: "6.0-core"}
     steps:
       -
         name: Free Space
@@ -131,21 +142,24 @@ jobs:
         run: |
           img_path="ghcr.io/${{ env.PACKAGE_NAME }}"
           
-          if [[ -z '${{ matrix.build.tag }}' ]]; then
-                KOHYA_TAG="$(curl -s https://api.github.com/repos/bmaltais/kohya_ss/tags | jq -r '.[0].name')"
+          if [[ -z '${{ matrix.build.kohya }}' ]]; then
+                KOHYA_BUILD_REF="$(curl -s https://api.github.com/repos/bmaltais/kohya_ss/tags | jq -r '.[0].name')"
             else
-                KOHYA_TAG="${{ matrix.build.tag }}"
+                KOHYA_BUILD_REF="${{ matrix.build.kohya }}"
             fi
-            [ -z "$KOHYA_TAG" ] && { echo "Error: KOHYA_TAG is empty. Exiting script." >&2; exit 1; }
-            echo "KOHYA_TAG=${KOHYA_TAG}" >> ${GITHUB_ENV}
+            [ -z "$KOHYA_BUILD_REF" ] && { echo "Error: KOHYA_BUILD_REF is empty. Exiting script." >&2; exit 1; }
+            echo "KOHYA_BUILD_REF=${KOHYA_BUILD_REF}" >> ${GITHUB_ENV}
 
             base_tag="v2-rocm-${{ matrix.build.rocm }}-${{ env.UBUNTU_VERSION }}"
 
             if [[ ${{ matrix.build.latest }} == "true" ]]; then
                 echo "Marking latest"
-                TAGS="${img_path}:${base_tag}-${KOHYA_TAG}, ${img_path}:${base_tag}, ${img_path}:latest-rocm"
+                # GHCR.io Tags
+                TAGS="${img_path}:${base_tag}-${KOHYA_BUILD_REF}, ${img_path}:${base_tag}, ${img_path}:latest-rocm"
+                # Docker.io Tags
+                TAGS="${TAGS}, ${img_path_dhub}:${KOHYA_BUILD_REF}, ${img_path_dhub}:latest"
             else
-                TAGS="${img_path}:${base_tag}-${KOHYA_TAG}"
+                TAGS="${img_path}:${base_tag}-${KOHYA_BUILD_REF}, ${img_path_dhub}:${KOHYA_BUILD_REF}"
             fi
           echo "TAGS=${TAGS}" >> ${GITHUB_ENV}
       -
@@ -157,7 +171,7 @@ jobs:
             IMAGE_BASE=ghcr.io/ai-dock/python:${{ matrix.build.python }}-v2-rocm-${{ matrix.build.rocm }}-${{ env.UBUNTU_VERSION }}
             PYTHON_VERSION=${{ matrix.build.python }}
             PYTORCH_VERSION=${{ matrix.build.pytorch }}
-            KOHYA_TAG=${{ env.KOHYA_TAG }}
+            KOHYA_BUILD_REF=${{ env.KOHYA_BUILD_REF }}
           push: true
           provenance: false
           tags: ${{ env.TAGS }}
diff --git a/README.md b/README.md
@@ -23,12 +23,12 @@ Tags follow these patterns:
 ##### _CUDA_
 - `:v2-cuda-[x.x.x]-base-[ubuntu-version]`
 
-- `:latest-cuda` &rarr; `:v2-cuda-12.1.1-base-22.04-v24.1.4`
+- `:latest-cuda` &rarr; `:v2-cuda-12.1.1-base-22.04`
 
 ##### _ROCm_
 - `:v2-rocm-[x.x.x]-core-[ubuntu-version]`
 
-- `:latest-rocm` &rarr; `:v2-rocm-6.0-core-22.04-v24.1.4`
+- `:latest-rocm` &rarr; `:v2-rocm-6.0-core-22.04`
 
 
 Browse [here](https://github.com/ai-dock/kohya_ss/pkgs/container/kohya_ss) for an image suitable for your target environment.
@@ -42,12 +42,12 @@ Supported Platforms: `NVIDIA CUDA`, `AMD ROCm`
 | Variable                 | Description |
 | ------------------------ | ----------- |
 | `AUTO_UPDATE`            | Update Kohya_ss on startup (default `false`) |
-| `KOHYA_BRANCH`           | Kohya_ss branch/commit hash for auto update. (default `master`) |
-| `KOHYA_FLAGS`            | Startup flags |
+| `KOHYA_ARGS`             | Startup arguments |
 | `KOHYA_PORT_HOST`        | Kohya's GUI port (default `7860`) |
+| `KOHYA_REF`              | Git reference for auto update. Accepts branch, tag or commit hash. Default: latest release |
 | `KOHYA_URL`              | Override `$DIRECT_ADDRESS:port` with URL for Kohya's GUI |
+| `TENSORBOARD_ARGS`       | Startup arguments (default `--logdir /opt/kohya_ss/logs`) |
 | `TENSORBOARD_PORT_HOST`  | Tensorboard port (default `6006`) |
-| `TENSORBOARD_FLAGS`      | Startup flags (default `--logdir /opt/kohya_ss/logs`) |
 | `TENSORBOARD_URL`        | Override `$DIRECT_ADDRESS:port` with URL for Tensorboard |
 
 See the base environment variables [here](https://github.com/ai-dock/base-image/wiki/2.0-Environment-Variables) for more configuration options.
@@ -71,7 +71,7 @@ The following services will be launched alongside the [default services](https:/
 
 The service will launch on port `7860` unless you have specified an override with `KOHYA_PORT`.
 
-You can set startup flags by using variable `KOHYA_FLAGS`.
+You can set startup arguments by using variable `KOHYA_ARGS`.
 
 To manage this service you can use `supervisorctl [start|stop|restart] kohya_ss`.
 

diff --git a/build/COPY_ROOT_1/opt/ai-dock/bin/build/layer1/common.sh b/build/COPY_ROOT_1/opt/ai-dock/bin/build/layer1/common.sh
@@ -9,16 +9,16 @@ build_common_main() {
 
 build_common_install_kohya_ss() {
     # Get latest tag from GitHub if not provided
-    if [[ -z $KOHYA_TAG ]]; then
-        export KOHYA_TAG="$(curl -s https://api.github.com/repos/bmaltais/kohya_ss/tags | \
+    if [[ -z $KOHYA_BUILD_REF ]]; then
+        export KOHYA_BUILD_REF="$(curl -s https://api.github.com/repos/bmaltais/kohya_ss/tags | \
             jq -r '.[0].name')"
-        env-store KOHYA_TAG
+        env-store KOHYA_BUILD_REF
     fi
 
     cd /opt
     git clone --recursive https://github.com/bmaltais/kohya_ss
     cd /opt/kohya_ss
-    git checkout "$KOHYA_TAG"
+    git checkout "$KOHYA_BUILD_REF"
     printf "\n%s\n" '#myTensorButton, #myTensorButtonStop {display:none!important;}' >> assets/style.css
     "$KOHYA_VENV_PIP" install --no-cache-dir \
         tensorboard \

diff --git a/build/COPY_ROOT_1/opt/ai-dock/bin/preflight.d/10-default.sh b/build/COPY_ROOT_1/opt/ai-dock/bin/preflight.d/10-default.sh
@@ -3,12 +3,16 @@
 
 function preflight_main() {
     source /opt/ai-dock/bin/venv-set.sh kohya
+
+    # egg-links break with symlinks - Address that here
+    [[ -d ${WORKSPACE}/kohya_ss ]] && sed -i "s|/opt/|${WORKSPACE}|g" $KOHYA_VENV/lib/python3.10/site-packages/library.egg-link
+
     preflight_configure_accelerate
     preflight_update_kohya_ss
-    printf "%s" "${KOHYA_FLAGS}" > /etc/kohya_ss_flags.conf
-    export TENSORBOARD_FLAGS=${TENSORBOARD_FLAGS:-"--logdir /opt/kohya_ss/logs"}
-    env-store TENSORBOARD_FLAGS
-    printf "%s" "${TENSORBOARD_FLAGS}" > /etc/tensorboard_flags.conf
+    printf "%s" "${KOHYA_ARGS}" > /etc/kohya_ss_args.conf
+    export TENSORBOARD_ARGS=${TENSORBOARD_ARGS:-"--logdir /opt/kohya_ss/logs"}
+    env-store TENSORBOARD_ARGS
+    printf "%s" "${TENSORBOARD_ARGS}" > /etc/tensorboard_args.conf
 }
 
 function preflight_configure_accelerate() {

diff --git a/...ROOT_1/opt/ai-dock/bin/set-kohya-flags.sh → ..._ROOT_1/opt/ai-dock/bin/set-kohya-args.sh b/...ROOT_1/opt/ai-dock/bin/set-kohya-flags.sh → ..._ROOT_1/opt/ai-dock/bin/set-kohya-args.sh
@@ -1,4 +1,4 @@
 #!/bin/bash
 
-echo "$@" > /etc/kohya_ss_flags.conf
+echo "$@" > /etc/kohya_ss_args.conf
 supervisorctl restart kohya_ss
diff --git a/build/COPY_ROOT_1/opt/ai-dock/bin/supervisor-kohya_ss.sh b/build/COPY_ROOT_1/opt/ai-dock/bin/supervisor-kohya_ss.sh
@@ -41,12 +41,12 @@ function start() {
 
     printf "Starting $SERVICE_NAME...\n"
 
-    PLATFORM_FLAGS=
+    PLATFORM_ARGS=
     if [[ $XPU_TARGET = "AMD_GPU" ]]; then
-        PLATFORM_FLAGS="--use-rocm"
+        PLATFORM_ARGS="--use-rocm"
     fi
 
-    BASE_FLAGS="--headless"
+    BASE_ARGS="--headless"
 
     # Delay launch until micromamba is ready
     if [[ -f /run/workspace_sync || -f /run/container_config ]]; then
@@ -70,13 +70,13 @@ function start() {
     fuser -k -SIGKILL ${LISTEN_PORT}/tcp > /dev/null 2>&1 &
     wait -n
 
-    FLAGS_COMBINED="${PLATFORM_FLAGS} ${BASE_FLAGS} $(cat /etc/kohya_ss_flags.conf)"
+    ARGS_COMBINED="${PLATFORM_ARGS} ${BASE_ARGS} $(cat /etc/kohya_ss_args.conf)"
     printf "Starting %s...\n" "${SERVICE_NAME}"
 
     cd /opt/kohya_ss
     source "$KOHYA_VENV/bin/activate"
     LD_PRELOAD=libtcmalloc.so python kohya_gui.py \
-        ${FLAGS_COMBINED} --server_port ${LISTEN_PORT}
+        ${ARGS_COMBINED} --server_port ${LISTEN_PORT}
 }
 
 start 2>&1
diff --git a/build/COPY_ROOT_1/opt/ai-dock/bin/supervisor-tensorboard.sh b/build/COPY_ROOT_1/opt/ai-dock/bin/supervisor-tensorboard.sh
@@ -66,9 +66,9 @@ function start() {
     printf "Starting %s...\n" "${SERVICE_NAME}"
 
     source "$KOHYA_VENV/bin/activate"
-    FLAGS_COMBINED="$(cat /etc/tensorboard_flags.conf)"
+    ARGS_COMBINED="$(cat /etc/tensorboard_args.conf)"
     LD_PRELOAD=libtcmalloc.so tensorboard \
-        $FLAGS_COMBINED \
+        $ARGS_COMBINED \
         --port ${LISTEN_PORT}
 }
 

diff --git a/build/COPY_ROOT_1/opt/ai-dock/bin/update-kohya_ss.sh b/build/COPY_ROOT_1/opt/ai-dock/bin/update-kohya_ss.sh
@@ -1,29 +1,37 @@
 #!/bin/bash
 umask 002
-branch=master
 
 source /opt/ai-dock/etc/environment.sh
 source /opt/ai-dock/bin/venv-set.sh kohya
 
-if [[ -n "${KOHYA_BRANCH}" ]]; then
-    branch="${KOHYA_BRANCH}"
+if [[ -n "${KOHYA_REF}" ]]; then
+    ref="${KOHYA_REF}"
+else
+    # The latest tagged release
+    ref="$(curl -s https://api.github.com/repos/bmaltais/kohya_ss/tags | \
+            jq -r '.[0].name')"
 fi
 
-# -b flag has priority
-while getopts b: flag
+# -r argument has priority
+while getopts r: flag
 do
     case "${flag}" in
-        b) branch="$OPTARG";;
+        r) ref="$OPTARG";;
     esac
 done
 
-printf "Updating Kohya's GUI (${branch})...\n"
+[[ -n $ref ]] || { echo "Failed to get update target"; exit 1; }
+
+printf "Updating Kohya's GUI (${ref})...\n"
 
 cd /opt/kohya_ss
+git stash
 git fetch --tags
-git checkout ${branch}
+git checkout ${ref}
 git pull
 git submodule update --recursive
 
+printf "\n%s\n" '#myTensorButton, #myTensorButtonStop {display:none!important;}' >> assets/style.css
+
 "$KOHYA_VENV_PIP" install --no-cache-dir \
     -r requirements.txt
diff --git a/config/provisioning/default.sh b/config/provisioning/default.sh
@@ -13,8 +13,7 @@ PIP_PACKAGES=(
   )
 
 CHECKPOINT_MODELS=(
-    "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt"
-    #"https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.ckpt"
+    "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned.ckpt"
     #"https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors"
     #"https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/resolve/main/sd_xl_refiner_1.0.safetensors"
 )
@@ -78,7 +77,17 @@ function provisioning_print_end() {
 
 # Download from $1 URL to $2 file path
 function provisioning_download() {
-    wget -qnc --content-disposition --show-progress -e dotbytes="${3:-4M}" -P "$2" "$1"
+    if [[ -n $HF_TOKEN && $1 =~ ^https://([a-zA-Z0-9_-]+\.)?huggingface\.co(/|$|\?) ]]; then
+        auth_token="$HF_TOKEN"
+    elif 
+        [[ -n $CIVITAI_TOKEN && $1 =~ ^https://([a-zA-Z0-9_-]+\.)?civitai\.com(/|$|\?) ]]; then
+        auth_token="$CIVITAI_TOKEN"
+    fi
+    if [[ -n $auth_token ]];then
+        wget --header="Authorization: Bearer $auth_token" -qnc --content-disposition --show-progress -e dotbytes="${3:-4M}" -P "$2" "$1"
+    else
+        wget -qnc --content-disposition --show-progress -e dotbytes="${3:-4M}" -P "$2" "$1"
+    fi
 }
 
 provisioning_start
diff --git a/docker-compose.yaml b/docker-compose.yaml
@@ -7,8 +7,8 @@ services:
       context: ./build
       args:
         PYTHON_VERSION: ${PYTHON_VERSION:-3.10}
-        PYTORCH_VERSION: ${PYTORCH_VERSION:-2.3.0}
-        KOHYA_TAG: ${KOHYA_TAG:-}
+        PYTORCH_VERSION: ${PYTORCH_VERSION:-2.4.0}
+        KOHYA_BUILD_REF: ${KOHYA_BUILD_REF:-}
         IMAGE_BASE: ${IMAGE_BASE:-ghcr.io/ai-dock/python:${PYTHON_VERSION:-3.10}-v2-cuda-12.1.1-base-22.04}
       tags:
         - "ghcr.io/ai-dock/kohya_ss:${IMAGE_TAG:-cuda-12.1.1-base-22.04}"
@@ -54,6 +54,7 @@ services:
         - CF_TUNNEL_TOKEN=${CF_TUNNEL_TOKEN:-}
         - CF_QUICK_TUNNELS=${CF_QUICK_TUNNELS:-true}
         - WEB_ENABLE_AUTH=${WEB_ENABLE_AUTH:-true}
+        - WEB_ENABLE_HTTPS=${WEB_ENABLE_HTTPS:-false}
         - WEB_USER=${WEB_USER:-user}
         - WEB_PASSWORD=${WEB_PASSWORD:-password}
         - SSH_PORT_HOST=${SSH_PORT_HOST:-2222}
@@ -62,7 +63,7 @@ services:
         - SERVICEPORTAL_METRICS_PORT=${SERVICEPORTAL_METRICS_PORT:-21111}
         - SERVICEPORTAL_URL=${SERVICEPORTAL_URL:-}
         - KOHYA_BRANCH=${KOHYA_BRANCH:-}
-        - KOHYA_FLAGS=${KOHYA_FLAGS:-}
+        - KOHYA_ARGS=${KOHYA_ARGS:-}
         - KOHYA_PORT_HOST=${KOHYA_PORT_HOST:-7860}
         - KOHYA_PORT_LOCAL=${KOHYA_PORT_LOCAL:-17860}
         - KOHYA_METRICS_PORT=${KOHYA_METRICS_PORT:-27860}
@@ -74,7 +75,7 @@ services:
         - SYNCTHING_UI_PORT_HOST=${SYNCTHING_UI_PORT_HOST:-8384}
         - SYNCTHING_TRANSPORT_PORT_HOST=${SYNCTHING_TRANSPORT_PORT_HOST:-22999}
         - SYNCTHING_URL=${SYNCTHING_URL:-}
-        - TENSORBOARD_FLAGS=${TENSORBOARD_FLAGS:-"--logdir /opt/kohya_ss/logs"}
+        - TENSORBOARD_ARGS=${TENSORBOARD_ARGS:-"--logdir /opt/kohya_ss/logs"}
         - TENSORBOARD_PORT_HOST=${TENSORBOARD_PORT_HOST:-6006}
         - TENSORBOARD_PORT_LOCAL=${TENSORBOARD_PORT_LOCAL:-16006}
         - TENSORBOARD_METRICS_PORT=${TENSORBOARD_METRICS_PORT:-26006}