Refactoring docker images, updating CUDA to 12.6 and pytorch 2.6

.gitignore

@@ -173,5 +173,4 @@ cython_debug/
 .idea/*
 
 # Docker build
-docker/conda/build
 docker/deepspeed/build

DOCKER_README.md

@@ -36,7 +36,7 @@ The Docker image currently works on Windows and Linux, optionally supporting NVI
 4. Download CUDA toolkit keyring: `wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-keyring_1.1-1_all.deb`
 5. Install keyring: `sudo dpkg -i cuda-keyring_1.1-1_all.deb`
 6. Update package list: `sudo apt-get update`
-7. Install CUDA toolkit: `sudo apt-get -y install cuda-toolkit-12-4`
+7. Install CUDA toolkit: `sudo apt-get -y install cuda-toolkit-12-6`
 8. Install Docker Desktop using WSL2 as the backend
 9. Restart
 10. If you wish to monitor the terminal remotely via SSH, follow [this guide](https://www.hanselman.com/blog/how-to-ssh-into-wsl2-on-windows-10-from-an-external-machine).

Dockerfile

@@ -1,54 +1,19 @@
-FROM continuumio/miniconda3:24.7.1-0
+ARG GITHUB_REPOSITORY
+FROM ${GITHUB_REPOSITORY}alltalk_environment:latest
 
 # Argument to choose the model: piper, vits, xtts
 ARG TTS_MODEL="xtts"
 ENV TTS_MODEL=$TTS_MODEL
 
-ARG ALLTALK_DIR=/opt/alltalk
+ARG DEEPSPEED_VERSION=0.16.2
+ENV DEEPSPEED_VERSION=$DEEPSPEED_VERSION
 
-SHELL ["/bin/bash", "-l", "-c"]
+ARG ALLTALK_DIR=/opt/alltalk
-ENV SHELL=/bin/bash
-ENV HOST=0.0.0.0
-ENV DEBIAN_FRONTEND=noninteractive
-ENV CUDA_DOCKER_ARCH=all
-ENV NVIDIA_VISIBLE_DEVICES=all
-ENV CONDA_AUTO_UPDATE_CONDA="false"
 
 ENV GRADIO_SERVER_NAME="0.0.0.0"
 
-##############################################################################
-# Installation/Basic Utilities
-##############################################################################
-RUN <<EOR
-    apt-get update
-    apt-get upgrade -y
-    apt-get install --no-install-recommends -y \
-      espeak-ng \
-      curl \
-      wget \
-      jq \
-      vim
-
-    apt-get clean && rm -rf /var/lib/apt/lists/*
-EOR
-
 WORKDIR ${ALLTALK_DIR}
 
-##############################################################################
-# Create a conda environment and install dependencies:
-##############################################################################
-COPY docker/conda/build/environment-*.yml environment.yml
-RUN <<EOR
-    RESULT=$( { conda env create -f environment.yml ; } 2>&1 )
-
-    if [ $? -ne 0 ] ; then
-      echo "Failed to install conda dependencies: $RESULT"
-      exit 1
-    fi
-
-    conda clean -a && pip cache purge
-EOR
-
 ##############################################################################
 # Install python dependencies (cannot use --no-deps because requirements are not complete)
 ##############################################################################
@@ -59,7 +24,7 @@ ENV PIP_CACHE_DIR=${ALLTALK_DIR}/pip_cache
 RUN <<EOR
     conda activate alltalk
 
-    mkdir ${ALLTALK_DIR}k/pip_cache
+    mkdir ${ALLTALK_DIR}/pip_cache
     pip install --no-cache-dir --cache-dir=${ALLTALK_DIR}/pip_cache -r system/requirements/requirements_standalone.txt
     pip install --no-cache-dir --cache-dir=${ALLTALK_DIR}/pip_cache --upgrade gradio==4.32.2
     # Parler:
@@ -77,10 +42,11 @@ RUN <<EOR
     DEEPSPEED_WHEEL=$(realpath /tmp/deepspeed/*.whl)
     conda activate alltalk
 
-    RESULT=$( { CFLAGS="-I$CONDA_PREFIX/include/" LDFLAGS="-L$CONDA_PREFIX/lib/" \
+    echo "Using precompiled DeepSpeed wheel at ${DEEPSPEED_WHEEL}"
-      pip install --no-cache-dir ${DEEPSPEED_WHEEL} ; } 2>&1 )
+    CFLAGS="-I$CONDA_PREFIX/include/" LDFLAGS="-L$CONDA_PREFIX/lib/" \
+      pip install --no-cache-dir ${DEEPSPEED_WHEEL}
 
-    if echo $RESULT | grep -izq error ; then
+    if [ $? -ne 0 ] ; then
       echo "Failed to install pip dependencies: $RESULT"
       exit 1
     fi
@@ -118,7 +84,6 @@ conda activate alltalk
 python diagnostics.py
 EOF
     chmod +x start_alltalk.sh
-    chmod +x start_environment.sh
     chmod +x start_finetune.sh
     chmod +x start_diagnostics.sh
 EOR

docker-build.sh

@@ -10,7 +10,6 @@ DOCKER_TAG=latest
 CLEAN=false
 
 # Create required build directories if they don't exist
-mkdir -p ${SCRIPT_DIR=}/docker/conda/build
 mkdir -p ${SCRIPT_DIR=}/docker/deepspeed/build
 echo "Created build directories"
 
@@ -29,10 +28,20 @@ while [ "$#" -gt 0 ]; do
       TTS_MODEL="$2"
       shift
       ;;
+    --deepspeed-version)
+      DEEPSPEED_VERSION="$2"
+      shift
+      ;;
     --tag)
       DOCKER_TAG="$2"
       shift
       ;;
+    --github-repository)
+      if [ -n "${GITHUB_REPOSITORY}" ] && ! [[ $GITHUB_REPOSITORY =~ ^--.* ]]; then
+        GITHUB_REPOSITORY="$2"
+        shift
+      fi
+      ;;
     --clean)
       CLEAN=true
       ;;
@@ -45,20 +54,34 @@ while [ "$#" -gt 0 ]; do
 done
 
 if [ "$CLEAN" = true ]; then
-  rm -rf ${SCRIPT_DIR=}/docker/conda/build
   rm -rf ${SCRIPT_DIR=}/docker/deepspeed/build
   # Recreate directories after clean
-  mkdir -p ${SCRIPT_DIR=}/docker/conda/build
   mkdir -p ${SCRIPT_DIR=}/docker/deepspeed/build
   echo "Cleaned and recreated build directories"  
 fi
 
-$SCRIPT_DIR/docker/conda/build-conda-env.sh \
+echo "Building base environment image"
+$SCRIPT_DIR/docker/base/build-base-env.sh \
   --cuda-version ${CUDA_VERSION} \
-  --python-version ${PYTHON_VERSION}
+  --python-version ${PYTHON_VERSION} \
+  --github-repository ${GITHUB_REPOSITORY} \
+  --tag ${DOCKER_TAG}
 
+if [ $? -ne 0 ]; then
+  echo "Failed to build base environment image"
+  exit 1
+fi
+
+echo "Building DeepSpeed"
 $SCRIPT_DIR/docker/deepspeed/build-deepspeed.sh \
-  --python-version ${PYTHON_VERSION}
+  --python-version ${PYTHON_VERSION} \
+  --github-repository ${GITHUB_REPOSITORY} \
+  --tag ${DOCKER_TAG}
+
+if [ $? -ne 0 ]; then
+  echo "Failed to build DeepSpeed"
+  exit 1
+fi
 
 echo "Starting docker build process using TTS model '${TTS_MODEL}' and docker tag '${DOCKER_TAG}'"
 echo "Building for CUDA $CUDA_VERSION using python ${PYTHON_VERSION}"
@@ -68,7 +91,8 @@ docker buildx \
   --progress=plain \
   --build-arg TTS_MODEL=$TTS_MODEL \
   --build-arg ALLTALK_DIR=$ALLTALK_DIR \
-  -t alltalk_beta:${DOCKER_TAG} \
+  --build-arg DEEPSPEED_VERSION=$DEEPSPEED_VERSION \
+  -t ${GITHUB_REPOSITORY}alltalk_beta:${DOCKER_TAG} \
   .
 
 echo "Docker build process finished. Use docker-start.sh to start the container."

docker-start.sh

@@ -29,6 +29,12 @@ while [ "$#" -gt 0 ]; do
       DOCKER_TAG="$2"
       shift
       ;;
+    --github-repository)
+      if [ -n "${GITHUB_REPOSITORY}" ] && ! [[ $GITHUB_REPOSITORY =~ ^--.* ]]; then
+        GITHUB_REPOSITORY="$2"
+        shift
+      fi
+      ;;
     *)
       # Allow to pass arbitrary arguments to docker as well to be flexible:
       ADDITIONAL_ARGS+=( $1 )
@@ -65,4 +71,4 @@ docker run \
   --name alltalk \
  "${DOCKER_ARGS[@]}" \
  "${ADDITIONAL_ARGS[@]}" \
-  alltalk_beta:${DOCKER_TAG} &> /dev/stdout
+  ${GITHUB_REPOSITORY}alltalk_beta:${DOCKER_TAG} &> /dev/stdout

docker/base/Dockerfile

@@ -1,6 +1,6 @@
 FROM continuumio/miniconda3:24.7.1-0 AS build
 
-ARG CUDA_VERSION=12.4.1
+ARG CUDA_VERSION=12.6.0
 ENV CUDA_VERSION=$CUDA_VERSION
 
 ARG PYTHON_VERSION=3.11.11
@@ -14,6 +14,8 @@ ENV CUDA_DOCKER_ARCH=all
 ENV NVIDIA_VISIBLE_DEVICES=all
 ENV CONDA_AUTO_UPDATE_CONDA="false"
 
+ENV ALLTALK_DIR=/opt/alltalk
+
 ##############################################################################
 # Installation/Basic Utilities
 ##############################################################################
@@ -35,24 +37,24 @@ EOR
 ##############################################################################
 RUN <<EOR
     CUDA_SHORT_VERSION=${CUDA_VERSION%.*}
-
     conda create -y -n "alltalk" -c conda-forge python=${PYTHON_VERSION}
     conda activate alltalk
-    RESULT=$( { conda install -y \
+    conda install -y \
       gcc_linux-64 \
       gxx_linux-64 \
-      pytorch \
-      pytorch-cuda=${CUDA_SHORT_VERSION} \
-      torchvision \
-      torchaudio \
       libaio \
-      nvidia/label/cuda-${CUDA_SHORT_VERSION}.0::cuda-toolkit \
       faiss-gpu=1.9.0 \
+      cuda=${CUDA_VERSION} cuda-tools=${CUDA_VERSION} cuda-toolkit=${CUDA_VERSION} cuda-version=${CUDA_SHORT_VERSION} \
+      cuda-command-line-tools=${CUDA_VERSION} cuda-compiler=${CUDA_VERSION} cuda-runtime=${CUDA_VERSION} \
+      cuda-libraries=${CUDA_VERSION} cuda-libraries-dev=${CUDA_VERSION} cuda-libraries-static=${CUDA_VERSION} \
+      cuda-cudart=${CUDA_SHORT_VERSION} cuda-cudart_linux-64=${CUDA_SHORT_VERSION} \
+      cudnn=9.3 \
+      py-cpuinfo=9 \
       conda-forge::ffmpeg=7.1.0 \
       conda-forge::portaudio=19.7.0 \
       -c pytorch \
-      -c anaconda \
+      -c nvidia \
-      -c nvidia ; } 2>&1 )
+      -c anaconda
 
     if [ $? -ne 0 ] ; then
       echo "Failed to install conda dependencies: $RESULT"
@@ -62,16 +64,17 @@ RUN <<EOR
     conda clean -a && pip cache purge
 EOR
 
-##############################################################################
-# Export conda environment:
-##############################################################################
 RUN <<EOR
-  conda activate alltalk
+  CUDA_SHORT_VERSION=${CUDA_VERSION%.*}
-  conda env export > /environment-cu-${CUDA_VERSION}-cp-${PYTHON_VERSION}.yml
+  CUDA_SHORT_VERSION_NO_DOT=${CUDA_SHORT_VERSION//./}
-EOR
 
-##############################################################################
+  conda activate alltalk
-# Make environment file available to docker build output:
+  mkdir -p ${ALLTALK_DIR}/pip_cache
-##############################################################################
+  pip install --no-cache-dir --cache-dir=${ALLTALK_DIR}/pip_cache \
-FROM scratch
+      torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu${CUDA_SHORT_VERSION_NO_DOT}
-COPY --from=build /environment-*.yml /
+
+  if [ $? -ne 0 ] ; then
+    echo "Failed to install pip dependencies: $RESULT"
+    exit 1
+  fi
+EOR

docker/base/build-base-env.sh

@@ -5,10 +5,7 @@ cd $SCRIPT_DIR
 
 . ${SCRIPT_DIR=}/../variables.sh
 
-if [ -f $SCRIPT_DIR/build/environment-cu-${CUDA_VERSION}-cp-${PYTHON_VERSION}.yml ]; then
+DOCKER_TAG=latest
-   echo "Environment file already exists - skipping..."
-   exit 0
-fi
 
 # Parse arguments
 while [ "$#" -gt 0 ]; do
@@ -21,8 +18,17 @@ while [ "$#" -gt 0 ]; do
       PYTHON_VERSION="$2"
       shift
       ;;
+    --tag)
+      DOCKER_TAG="$2"
+      shift
+      ;;
+    --github-repository)
+      if [ -n "${GITHUB_REPOSITORY}" ] && ! [[ $GITHUB_REPOSITORY =~ ^--.* ]]; then
+        GITHUB_REPOSITORY="$2"
+        shift
+      fi
+      ;;
     *)
-      # Allow to pass arbitrary arguments to docker as well to be flexible:
       echo "Unknown argument '$1'"
       exit 1
       ;;
@@ -32,12 +38,10 @@ done
 
 echo "Building conda environment using python ${PYTHON_VERSION} with CUDA ${CUDA_VERSION}"
 
-rm -rf build # make sure to properly clean up
-mkdir -p build
-
 docker buildx \
   build \
+  --progress=plain \
   --build-arg CUDA_VERSION=$CUDA_VERSION \
   --build-arg PYTHON_VERSION=$PYTHON_VERSION \
-  --output=$SCRIPT_DIR/build \
+  -t ${GITHUB_REPOSITORY}alltalk_environment:${DOCKER_TAG} \
   .

docker/deepspeed/Dockerfile

@@ -1,53 +1,15 @@
-FROM continuumio/miniconda3:24.7.1-0
+ARG GITHUB_REPOSITORY
+FROM ${GITHUB_REPOSITORY}alltalk_environment:latest
 
-ARG DEEPSPEED_VERSION=0.16.2
+ARG DEEPSPEED_VERSION=0.16.3
 ENV DEEPSPEED_VERSION=$DEEPSPEED_VERSION
 
-SHELL ["/bin/bash", "-l", "-c"]
-ENV SHELL=/bin/bash
-ENV HOST=0.0.0.0
-ENV DEBIAN_FRONTEND=noninteractive
-ENV CUDA_DOCKER_ARCH=all
-ENV NVIDIA_VISIBLE_DEVICES=all
-ENV CONDA_AUTO_UPDATE_CONDA="false"
-
 ##############################################################################
 # Directories:
 ##############################################################################
 ENV STAGE_DIR=/tmp
 RUN mkdir -p ${STAGE_DIR}
 
-##############################################################################
-# Installation/Basic Utilities
-##############################################################################
-RUN <<EOR
-  apt-get update
-  apt-get upgrade -y
-  apt-get install --no-install-recommends -y \
-      espeak-ng \
-      curl \
-      wget \
-      jq \
-      vim
-
-   apt-get clean && rm -rf /var/lib/apt/lists/*
-EOR
-
-##############################################################################
-# Create a conda environment and install dependencies:
-##############################################################################
-COPY build/environment-*.yml environment.yml
-RUN <<EOR
-    RESULT=$( { conda env create -f environment.yml ; } 2>&1 )
-
-    if [ $? -ne 0 ] ; then
-      echo "Failed to install conda dependencies: $RESULT"
-      exit 1
-    fi
-
-    conda clean -a && pip cache purge
-EOR
-
 ##############################################################################
 # Install python dependencies
 ##############################################################################
@@ -56,8 +18,7 @@ RUN <<EOR
   conda activate alltalk
 
   pip install \
-      deepspeed-kernels \
+      deepspeed-kernels
-      scikit-learn
 EOR
 
 
@@ -69,6 +30,15 @@ RUN <<EOR
   cd ${STAGE_DIR}/DeepSpeed
   git checkout .
   git checkout "tags/v${DEEPSPEED_VERSION}" -b "v${DEEPSPEED_VERSION}"
+
+  # Patching some files due to bug https://github.com/microsoft/DeepSpeed/issues/6709
+  START_LINE=`awk '/#ifndef BF16_AVAILABLE/{ print NR; exit }' csrc/transformer/inference/csrc/gelu.cu`
+  TO_LINE=$((START_LINE + 2))
+  sed -i "${START_LINE},${TO_LINE}d" ./csrc/transformer/inference/csrc/gelu.cu
+
+  START_LINE=`awk '/#ifndef BF16_AVAILABLE/{ print NR; exit }' csrc/transformer/inference/csrc/transform.cu`
+  TO_LINE=$((START_LINE + 2))
+  sed -i "${START_LINE},${TO_LINE}d" ./csrc/transformer/inference/csrc/transform.cu
 EOR
 
 ##############################################################################

docker/deepspeed/build-deepspeed.sh

@@ -5,6 +5,8 @@ cd $SCRIPT_DIR
 
 . ${SCRIPT_DIR=}/../variables.sh
 
+DOCKER_TAG=latest
+
 # Parse arguments
 while [ "$#" -gt 0 ]; do
   case "$1" in
@@ -16,8 +18,17 @@ while [ "$#" -gt 0 ]; do
       DEEPSPEED_VERSION="$2"
       shift
       ;;
+    --tag)
+      DOCKER_TAG="$2"
+      shift
+      ;;
+    --github-repository)
+      if [ -n "${GITHUB_REPOSITORY}" ] && ! [[ $GITHUB_REPOSITORY =~ ^--.* ]]; then
+        GITHUB_REPOSITORY="$2"
+        shift
+      fi
+      ;;
     *)
-      # Allow to pass arbitrary arguments to docker as well to be flexible:
       echo "Unknown argument '$1'"
       exit 1
       ;;
@@ -33,22 +44,16 @@ then
    exit 0
 fi
 
-CONDA_ENV=${SCRIPT_DIR=}/../conda/build/environment-cu-${CUDA_VERSION}-cp-${PYTHON_VERSION}.yml
-if [ ! -f ${CONDA_ENV} ]; then
-    echo "No conda environment found. Please run 'build-conda-env.sh' first."
-    exit 1
-fi
-
 echo "Building DeepSpeed $DEEPSPEED_VERSION using python ${PYTHON_VERSION}"
 
 rm -rf build # make sure to properly clean up
 mkdir -p build
-cp ${CONDA_ENV} ${SCRIPT_DIR=}/build
 
 docker buildx \
   build \
+  --build-arg GITHUB_REPOSITORY=$GITHUB_REPOSITORY \
   --build-arg DEEPSPEED_VERSION=$DEEPSPEED_VERSION \
-  -t deepspeed:$DEEPSPEED_VERSION \
+  -t ${GITHUB_REPOSITORY}alltalk_deepspeed:${DOCKER_TAG} \
   .
 
 docker run \
@@ -57,4 +62,4 @@ docker run \
   --gpus=all \
   --name deepspeed \
   -v $SCRIPT_DIR/build:/deepspeed \
-  deepspeed:$DEEPSPEED_VERSION
+  ${GITHUB_REPOSITORY}alltalk_deepspeed:${DOCKER_TAG}

docker/variables.sh

@@ -1,6 +1,7 @@
 #!/usr/bin/env bash
 
-export CUDA_VERSION=12.4.1
+export CUDA_VERSION=12.6.0
 export PYTHON_VERSION=3.11.11
 export DEEPSPEED_VERSION=0.16.2
 export ALLTALK_DIR=/opt/alltalk
+export GITHUB_REPOSITORY=

system/tts_engines/rvc/lib/utils.py

@@ -2,6 +2,8 @@ import numpy as np
 import re
 import unicodedata
 from fairseq import checkpoint_utils
+from fairseq.data import Dictionary
+import torch
 
 import logging
 
@@ -66,6 +68,7 @@ def load_embedding(embedder_model):
         #print("MODEL PATH IS", model_path)
         
         # Load model ensemble and task
+        torch.serialization.add_safe_globals([Dictionary])
         models = checkpoint_utils.load_model_ensemble_and_task(
             [f"{model_path}"],
             suffix="",