Adding Github actions for building the docker images

.github/actions/docker-setup/action.yml

@@ -0,0 +1,70 @@
+name: 'Setting up docker'
+
+inputs:
+  DOCKERHUB_USERNAME:
+    required: true
+  DOCKERHUB_REPO_NAME:
+    required: true
+  DOCKERHUB_TOKEN:
+    required: true
+
+runs:
+  using: "composite"
+  steps:
+    - name: Free Disk Space (Ubuntu)
+      uses: jlumbroso/free-disk-space@v1.3.1
+      with:
+        tool-cache: true
+        swap-storage: false
+        large-packages: false
+    - name: Read variables
+      shell: bash
+      run: |
+        source ./docker/variables.sh
+        for key in "${!ALLTALK_VARS[@]}"
+        do
+          echo "${key}=${ALLTALK_VARS[${key}]}"
+          echo "${key}=${ALLTALK_VARS[${key}]}" >> $GITHUB_ENV
+        done
+
+    - name: Extract version from Git Tag
+      shell: bash {0}
+      run: |
+        # Use the git tag if existing, otherwise 'latest':
+        DOCKER_TAG=$( echo "${{ github.ref }}" | grep -q tags | cut -d'/' -f 3 )
+        if [ -z "${DOCKER_TAG}" ]; then
+          DOCKER_TAG="latest"
+        fi
+        echo "Using DOCKER_TAG=${DOCKER_TAG}"
+        echo "DOCKER_TAG=${DOCKER_TAG}" >> $GITHUB_ENV
+
+    - name: Extract Docker Hub User and Repo Name
+      shell: bash
+      env:
+        DOCKERHUB_USERNAME: ${{ inputs.DOCKERHUB_USERNAME }}
+        DOCKERHUB_REPO_NAME: ${{ inputs.DOCKERHUB_REPO_NAME }}
+      run: |
+        if [ -n "${DOCKERHUB_USERNAME}" ]; then
+          echo "Using Github repository variable 'DOCKERHUB_USERNAME' as Docker hub username with value: '${DOCKERHUB_USERNAME}'"
+          echo "DOCKERHUB_USERNAME=${DOCKERHUB_USERNAME}" >> $GITHUB_ENV
+        else
+          echo "Using Github repository name as Docker hub username"
+          echo "DOCKERHUB_USERNAME=$(echo ${{ github.repository }} | cut -d'/' -f 1)" >> $GITHUB_ENV
+        fi
+
+        if [ -n "${DOCKERHUB_REPO_NAME}" ]; then
+          echo "Using Github repository variable 'DOCKERHUB_REPO_NAME' as Docker hub repo name with value: '${DOCKERHUB_REPO_NAME}'"
+          echo "DOCKERHUB_REPO_NAME=${DOCKERHUB_REPO_NAME}" >> $GITHUB_ENV
+        else
+          echo "Using Github repository name as Docker hub repo name"
+          echo "DOCKERHUB_REPO_NAME=$(echo ${{ github.repository }} | cut -d'/' -f 2)" >> $GITHUB_ENV
+        fi
+
+    - name: Set up Docker Buildx
+      uses: docker/setup-buildx-action@v3
+
+    - name: Login to DockerHub
+      uses: docker/login-action@v3
+      with:
+        username: ${{ inputs.DOCKERHUB_USERNAME }}
+        password: ${{ inputs.DOCKERHUB_TOKEN }}

.github/workflows/publish-docker-v2.yml

@@ -0,0 +1,123 @@
+name: Publish Docker v2
+
+on:
+  workflow_dispatch:
+  release:
+    types: [ published ]
+  push:
+    branches:
+      - 'alltalkbeta'
+    tags:
+      - 'v*'
+
+# Limit the concurrency of entire workflow runs for a specific branch:
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  build_base_environment_docker_image:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Setup docker
+        uses: ./.github/actions/docker-setup
+        with:
+          DOCKERHUB_USERNAME: ${{ vars.DOCKERHUB_USERNAME }}
+          DOCKERHUB_REPO_NAME: ${{ vars.DOCKERHUB_REPO_NAME }}
+          DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}
+
+      - name: Docker meta
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: |
+            ${{ env.DOCKERHUB_USERNAME }}/${{ env.DOCKERHUB_REPO_NAME }}_environment
+          tags: |
+            # Use sematic versioning on push tag event:
+            type=semver,pattern={{version}}
+            # set latest tag for main and alltalkbeta branch:
+            type=raw,value=latest,enable=${{ github.ref == format('refs/heads/{0}', 'main') || github.ref == format('refs/heads/{0}', 'alltalkbeta') }}
+
+      - uses: dorny/paths-filter@v3
+        id: changes
+        with:
+          filters: |
+            docker-base:
+              - 'docker/variables.sh'
+              - 'docker/base/Dockerfile'
+            docker-workflow:
+              - '.github/workflows/publish-docker*.yml'
+
+      - name: Build and push base Docker image
+        if: |
+          steps.changes.outputs.docker-base == 'true' ||
+          steps.changes.outputs.docker-workflow == 'true' ||
+          github.event_name == 'release' ||
+          (github.event_name == 'push' && contains(github.ref, 'refs/tags/'))
+        uses: docker/build-push-action@v6
+        with:
+          context: ./docker/base
+          file: ./docker/base/Dockerfile
+          platforms: "linux/amd64"
+          push: true
+          build-args: |
+            CUDA_VERSION=${{ env.CUDA_VERSION }}
+            PYTHON_VERSION=${{ env.PYTHON_VERSION }}
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+          cache-from: type=registry,ref=${{ env.DOCKERHUB_USERNAME }}/${{ env.DOCKERHUB_REPO_NAME }}_environment:buildcache
+          cache-to: type=registry,ref=${{ env.DOCKERHUB_USERNAME }}/${{ env.DOCKERHUB_REPO_NAME }}_environment:buildcache,mode=max
+
+  build_main_docker_image:
+    needs: build_base_environment_docker_image
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        tts: [ xtts, vits, piper ]
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Setup docker
+        uses: ./.github/actions/docker-setup
+        with:
+          DOCKERHUB_USERNAME: ${{ vars.DOCKERHUB_USERNAME }}
+          DOCKERHUB_REPO_NAME: ${{ vars.DOCKERHUB_REPO_NAME }}
+          DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}
+
+      - name: Docker meta
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: |
+            ${{ env.DOCKERHUB_USERNAME }}/${{ env.DOCKERHUB_REPO_NAME }}
+          flavor: |
+            suffix=-${{ matrix.tts }},onlatest=true
+          tags: |
+            # Use sematic versioning on push tag event:
+            type=semver,pattern={{version}}
+            # set latest tag for main and alltalkbeta branch:
+            type=raw,value=latest,enable=${{ github.ref == format('refs/heads/{0}', 'main') || github.ref == format('refs/heads/{0}', 'alltalkbeta') }}
+
+      - name: Build and push base Docker image
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: ./Dockerfile
+          platforms: "linux/amd64"
+          push: true
+          build-args: |
+            TTS_MODEL=${{ matrix.tts }}
+            ALLTALK_DIR=${{ env.ALLTALK_DIR }}
+            DEEPSPEED_VERSION=${{ env.DEEPSPEED_VERSION }}
+            DOCKER_TAG=${{ env.DOCKER_TAG }}
+            GITHUB_REPOSITORY=${{ env.DOCKERHUB_USERNAME }}/
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+          cache-from: type=registry,ref=${{ env.DOCKERHUB_USERNAME }}/${{ env.DOCKERHUB_REPO_NAME }}:buildcache-${{ matrix.tts }}
+          cache-to: type=registry,ref=${{ env.DOCKERHUB_USERNAME }}/${{ env.DOCKERHUB_REPO_NAME }}:buildcache-${{ matrix.tts }},mode=max

Dockerfile

@@ -1,5 +1,6 @@
 ARG GITHUB_REPOSITORY
-FROM ${GITHUB_REPOSITORY}alltalk_environment:latest
+ARG DOCKER_TAG=latest
+FROM ${GITHUB_REPOSITORY}alltalk_tts_environment:${DOCKER_TAG}
 
 # Argument to choose the model: piper, vits, xtts
 ARG TTS_MODEL="xtts"
@@ -36,12 +37,23 @@ EOR
 ##############################################################################
 # Install DeepSpeed
 ##############################################################################
-RUN mkdir -p /tmp/deepseped
+RUN mkdir -p /tmp/deepspeed
-COPY docker/deepspeed/build/*.whl /tmp/deepspeed/
+COPY docker/deepspeed/build*/*.whl /tmp/deepspeed/
 RUN <<EOR
-    DEEPSPEED_WHEEL=$(realpath /tmp/deepspeed/*.whl)
+    DEEPSPEED_WHEEL=$(realpath -q /tmp/deepspeed/*.whl)
     conda activate alltalk
 
+    # Download DeepSpeed wheel if it was not built locally:
+    if [ -z "${DEEPSPEED_WHEEL}" ] || [ ! -f $DEEPSPEED_WHEEL ] ; then
+      echo "Downloading pre-built DeepSpeed wheel"
+      CURL_ERROR=$( { curl --output-dir /tmp/deepspeed -fLO "https://github.com/erew123/alltalk_tts/releases/download/DeepSpeed-for-docker/deepspeed-0.16.2+b344c04d-cp311-cp311-linux_x86_64.whl" ; } 2>&1 )
+      if [ $? -ne 0 ] ; then
+        echo "Failed to download DeepSpeed: $CURL_ERROR"
+        exit 1
+      fi
+      DEEPSPEED_WHEEL=$(realpath -q /tmp/deepspeed/*.whl)
+    fi
+
     echo "Using precompiled DeepSpeed wheel at ${DEEPSPEED_WHEEL}"
     CFLAGS="-I$CONDA_PREFIX/include/" LDFLAGS="-L$CONDA_PREFIX/lib/" \
       pip install --no-cache-dir ${DEEPSPEED_WHEEL}
@@ -75,7 +87,7 @@ EOF
 source ~/.bashrc
 export TRAINER_TELEMETRY=0
 conda activate alltalk
-python finetune.py
+python -m trainer.distribute --script finetune.py
 EOF
     cat << EOF > start_diagnostics.sh
 #!/usr/bin/env bash

docker-build.sh

@@ -8,6 +8,7 @@ cd $SCRIPT_DIR
 TTS_MODEL=xtts
 DOCKER_TAG=latest
 CLEAN=false
+LOCAL_DEEPSPEED_BUILD=false
 
 # Create required build directories if they don't exist
 mkdir -p ${SCRIPT_DIR=}/docker/deepspeed/build
@@ -37,11 +38,14 @@ while [ "$#" -gt 0 ]; do
       shift
       ;;
     --github-repository)
-      if [ -n "${GITHUB_REPOSITORY}" ] && ! [[ $GITHUB_REPOSITORY =~ ^--.* ]]; then
+      if [ -n "$2" ] && ! [[ $2 =~ ^--.* ]]; then
         GITHUB_REPOSITORY="$2"
         shift
       fi
       ;;
+    --local-deepspeed-build)
+      LOCAL_DEEPSPEED_BUILD=true
+      ;;
     --clean)
       CLEAN=true
       ;;
@@ -72,15 +76,17 @@ if [ $? -ne 0 ]; then
   exit 1
 fi
 
-echo "Building DeepSpeed"
+if [ "$LOCAL_DEEPSPEED_BUILD" = true ]; then
-$SCRIPT_DIR/docker/deepspeed/build-deepspeed.sh \
+  echo "Building DeepSpeed"
-  --python-version ${PYTHON_VERSION} \
+  $SCRIPT_DIR/docker/deepspeed/build-deepspeed.sh \
-  --github-repository ${GITHUB_REPOSITORY} \
+    --python-version ${PYTHON_VERSION} \
-  --tag ${DOCKER_TAG}
+    --github-repository ${GITHUB_REPOSITORY} \
+    --tag ${DOCKER_TAG}
 
-if [ $? -ne 0 ]; then
+  if [ $? -ne 0 ]; then
-  echo "Failed to build DeepSpeed"
+    echo "Failed to build DeepSpeed"
-  exit 1
+    exit 1
+  fi
 fi
 
 echo "Starting docker build process using TTS model '${TTS_MODEL}' and docker tag '${DOCKER_TAG}'"
@@ -92,7 +98,8 @@ docker buildx \
   --build-arg TTS_MODEL=$TTS_MODEL \
   --build-arg ALLTALK_DIR=$ALLTALK_DIR \
   --build-arg DEEPSPEED_VERSION=$DEEPSPEED_VERSION \
-  -t ${GITHUB_REPOSITORY}alltalk_beta:${DOCKER_TAG} \
+  --build-arg DOCKER_TAG=$DOCKER_TAG \
+  -t ${GITHUB_REPOSITORY}alltalk_tts:${DOCKER_TAG} \
   .
 
 echo "Docker build process finished. Use docker-start.sh to start the container."

docker-start.sh

@@ -30,7 +30,7 @@ while [ "$#" -gt 0 ]; do
       shift
       ;;
     --github-repository)
-      if [ -n "${GITHUB_REPOSITORY}" ] && ! [[ $GITHUB_REPOSITORY =~ ^--.* ]]; then
+      if [ -n "$2" ] && ! [[ $2 =~ ^--.* ]]; then
         GITHUB_REPOSITORY="$2"
         shift
       fi
@@ -71,4 +71,4 @@ docker run \
   --name alltalk \
  "${DOCKER_ARGS[@]}" \
  "${ADDITIONAL_ARGS[@]}" \
-  ${GITHUB_REPOSITORY}alltalk_beta:${DOCKER_TAG} &> /dev/stdout
+  ${GITHUB_REPOSITORY}alltalk_tts:${DOCKER_TAG} &> /dev/stdout

docker/base/build-base-env.sh

@@ -43,5 +43,5 @@ docker buildx \
   --progress=plain \
   --build-arg CUDA_VERSION=$CUDA_VERSION \
   --build-arg PYTHON_VERSION=$PYTHON_VERSION \
-  -t ${GITHUB_REPOSITORY}alltalk_environment:${DOCKER_TAG} \
+  -t ${GITHUB_REPOSITORY}alltalk_tts_environment:${DOCKER_TAG} \
   .

docker/deepspeed/Dockerfile

@@ -1,5 +1,6 @@
 ARG GITHUB_REPOSITORY
-FROM ${GITHUB_REPOSITORY}alltalk_environment:latest
+ARG DOCKER_TAG=latest
+FROM ${GITHUB_REPOSITORY}alltalk_tts_environment:${DOCKER_TAG}
 
 ARG DEEPSPEED_VERSION=0.16.3
 ENV DEEPSPEED_VERSION=$DEEPSPEED_VERSION

docker/deepspeed/build-deepspeed.sh

@@ -61,5 +61,6 @@ docker run \
   -it \
   --gpus=all \
   --name deepspeed \
+  --build-arg DOCKER_TAG=$DOCKER_TAG \
   -v $SCRIPT_DIR/build:/deepspeed \
   ${GITHUB_REPOSITORY}alltalk_deepspeed:${DOCKER_TAG}

docker/variables.sh

@@ -1,7 +1,17 @@
 #!/usr/bin/env bash
 
-export CUDA_VERSION=12.6.0
+declare -A ALLTALK_VARS
-export PYTHON_VERSION=3.11.11
+ALLTALK_VARS["CUDA_VERSION"]=12.6.0
-export DEEPSPEED_VERSION=0.16.2
+ALLTALK_VARS["PYTHON_VERSION"]=3.11.11
-export ALLTALK_DIR=/opt/alltalk
+ALLTALK_VARS["DEEPSPEED_VERSION"]=0.16.2
-export GITHUB_REPOSITORY=
+ALLTALK_VARS["ALLTALK_DIR"]=/opt/alltalk
+ALLTALK_VARS["GITHUB_REPOSITORY"]=
+
+# Export single variables (needed by Docker locally)
+for key in "${!ALLTALK_VARS[@]}"
+do
+  export "${key}=${ALLTALK_VARS[${key}]}"
+done
+
+# Export the entire associative array (needed by Github action)
+export ALLTALK_VARS