Fixing docker build for DeepSpeed, introducing build for conda enviro…

…nment

.gitignore

@@ -172,5 +172,6 @@ cython_debug/
 
 .idea/*
 
-# Deepspeed build
-deepspeed/build
+# Docker build
+docker/conda/build
+docker/deepspeed/build

DOCKER_README.md

@@ -2,9 +2,11 @@
 The Docker image currently works on Windows and Linux, optionally supporting NVIDIA GPUs.
 
 ## General Remarks
-- The resulting Docker image is 22 GB in size. Building might require even more disk space temporarily.
-- Build time depends on your hardware and internet connection. Expect at least 10min to be normal.
-- The Docker build:
+- The resulting Docker image is 21 GB in size. Building might require even more disk space temporarily.
+    - Another 15 GB is required for building DeepSpeed
+- Build time depends on your hardware and internet connection. Expect at least 20-30min to be normal for a full build.
+  - This includes building the conda environment as well as DeepSpeed, which is the basis for the alltalk Docker image.
+- The Docker build for alltalk:
   - Downloads XTTS as default TTS engine
   - Enables RVC by default
   - Downloads all supported RVC models
@@ -55,18 +57,22 @@ There are various arguments to customize the build and start of the docker image
   - Example: `docker-build.sh --tts_model piper`
 - `--tag` allows to choose the docker tag. Defaults to `latest`.
   - Example: `docker-build.sh --tag mytag`
+- `--clean` allows remove existing dependency build like conda environment or DeepSpeed.
+    - Example: `docker-build.sh --clean`
 
 ### Arguments for `docker-start.sh`
 - `--config` lets you choose a config JSON file which can subset of `confignew.json`. This allows you to change only 
   few values and leave the rest as defined in the default `confignew.json` file.
-  - Example: `docker-build.sh --config /my/config/file.json` with content `{"branding": "My Brand "}` will just change
+  - Example: `docker-start.sh --config /my/config/file.json` with content `{"branding": "My Brand "}` will just change
     the branding in `confignew.json`.
 - `--voices` lets you add voices for the TTS engine in WAV format. You have to specify the folder containing all
   voice files.
-  - Example: `docker-build.sh --voices /my/voices/dir`
+  - Example: `docker-start.sh --voices /my/voices/dir`
 - `--rvc_voices` similar to voices, this option lets you pick the folder containing the RVC models.
-  - Example: `docker-build.sh --rvc_vices /my/rvc/voices/dir`
+  - Example: `docker-start.sh --rvc_vices /my/rvc/voices/dir`
 - `--no_ui` allows you to not expose port 7852 for the gradio interface. Note that you still have to set `launch_gradio`
-  to `false` via JSON file passed to `--config`. 
+  to `false` via JSON file passed to `--config`.
+- `--tag` allows to choose the docker tag of the image to run. Defaults to `latest`.
+    - Example: `docker-start.sh --tag mytag`
 - Since the above commands only address the most important options, you might pass additional arbitrary docker commands
     to the `docker-start.sh`.

Dockerfile

@@ -4,23 +4,21 @@ FROM continuumio/miniconda3:24.7.1-0
 ARG TTS_MODEL="xtts"
 ENV TTS_MODEL=$TTS_MODEL
 
-ARG CUDA_VERSION="12.1.1"
-ENV CUDA_VERSION=$CUDA_VERSION
-
-ARG PYTHON_VERSION=3.11.9
-ENV PYTHON_VERSION=$PYTHON_VERSION
-
-ARG PYTORCH_VERSION=2.2.1
-ENV PYTORCH_VERSION=$PYTORCH_VERSION
+ARG ALLTALK_DIR=/opt/alltalk
 
 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 GRADIO_SERVER_NAME="0.0.0.0"
 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
@@ -34,59 +32,47 @@ RUN <<EOR
     apt-get clean && rm -rf /var/lib/apt/lists/*
 EOR
 
-WORKDIR /alltalk
+WORKDIR ${ALLTALK_DIR}
 
+##############################################################################
 # Create a conda environment and install dependencies:
-ARG INSTALL_ENV_DIR=/alltalk/alltalk_environment/env
-ENV CONDA_AUTO_UPDATE_CONDA="false"
+##############################################################################
+COPY docker/conda/build/environment-*.yml environment.yml
 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 \
-      gcc_linux-64 \
-      gxx_linux-64 \
-      pytorch=${PYTORCH_VERSION} \
-      pytorch-cuda=${CUDA_SHORT_VERSION} \
-      torchvision \
-      torchaudio \
-      libaio \
-      nvidia/label/cuda-${CUDA_SHORT_VERSION}.0::cuda-toolkit \
-      faiss-gpu=1.9.0 \
-      conda-forge::ffmpeg=7.1.0 \
-      conda-forge::portaudio=19.7.0 \
-      -c pytorch \
-      -c anaconda \
-      -c nvidia ; } 2>&1 )
+    RESULT=$( { conda env create -f environment.yml ; } 2>&1 )
 
     if echo $RESULT | grep -izq error ; then
-      echo "Failed to install conda dependencies 2: $RESULT"
+      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)
+##############################################################################
 COPY system/config system/config
 COPY system/requirements/requirements_standalone.txt system/requirements/requirements_standalone.txt
 COPY system/requirements/requirements_parler.txt system/requirements/requirements_parler.txt
-ENV PIP_CACHE_DIR=/alltalk/pip_cache
+ENV PIP_CACHE_DIR=${ALLTALK_DIR}/pip_cache
 RUN <<EOR
     conda activate alltalk
 
-    mkdir /alltalk/pip_cache
-    pip install --no-cache-dir --cache-dir=/alltalk/pip_cache -r system/requirements/requirements_standalone.txt
-    pip install --no-cache-dir --cache-dir=/alltalk/pip_cache --upgrade gradio==4.32.2
+    mkdir ${ALLTALK_DIR}k/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:
-    pip install --no-cache-dir --cache-dir=/alltalk/pip_cache -r system/requirements/requirements_parler.txt
+    pip install --no-cache-dir --cache-dir=${ALLTALK_DIR}/pip_cache -r system/requirements/requirements_parler.txt
 
     conda clean --all --force-pkgs-dirs -y && pip cache purge
 EOR
 
-# Deepspeed:
+##############################################################################
+# Install DeepSpeed
+##############################################################################
 RUN mkdir -p /tmp/deepseped
-COPY deepspeed/build/*.whl /tmp/deepspeed/
+COPY docker/deepspeed/build/*.whl /tmp/deepspeed/
 RUN <<EOR
     DEEPSPEED_WHEEL=$(realpath /tmp/deepspeed/*.whl)
     conda activate alltalk
@@ -103,11 +89,9 @@ RUN <<EOR
     conda clean --all --force-pkgs-dirs -y && pip cache purge
 EOR
 
-### Deepspeed requires cutlass:
-###RUN git clone --depth 1 --branch "v3.5.1" https://github.com/NVIDIA/cutlass /alltalk/cutlass
-###ENV CUTLASS_PATH=/alltalk/cutlass
-
+##############################################################################
 # Writing scripts to start alltalk:
+##############################################################################
 RUN <<EOR
     cat << EOF > start_alltalk.sh
 #!/usr/bin/env bash
@@ -141,7 +125,9 @@ EOR
 
 COPY . .
 
-# Create script to execute firstrun.py:
+##############################################################################
+# Create script to execute firstrun.py and run it:
+##############################################################################
 RUN echo $'#!/usr/bin/env bash \n\
 source ~/.bashrc \n\
 conda activate alltalk \n\
@@ -150,18 +136,24 @@ python ./system/config/firstrun.py $@' > ./start_firstrun.sh
 RUN chmod +x start_firstrun.sh
 RUN ./start_firstrun.sh --tts_model $TTS_MODEL
 
-RUN mkdir -p /alltalk/outputs
+RUN mkdir -p ${ALLTALK_DIR}/outputs
 RUN mkdir -p /root/.triton/autotune
 
-# Enabling deepspeed for all models:
+##############################################################################
+# Enable deepspeed for all models:
+##############################################################################
 RUN find . -name model_settings.json -exec sed -i -e 's/"deepspeed_enabled": false/"deepspeed_enabled": true/g' {} \;
 
-# Downloading all RVC models:
+##############################################################################
+# Download all RVC models:
+##############################################################################
 RUN <<EOR
   jq -r '.[]' system/tts_engines/rvc_files.json > /tmp/rvc_files.txt
   xargs -n 1 curl --create-dirs --output-dir models/rvc_base -LO < /tmp/rvc_files.txt
   rm -f /tmp/rvc_files.txt
 EOR
 
-## Start alltalk:
+##############################################################################
+# Start alltalk:
+##############################################################################
 ENTRYPOINT ["sh", "-c", "./start_alltalk.sh"]

docker-build.sh

@@ -1,10 +1,13 @@
 #!/usr/bin/env bash
 
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+cd $SCRIPT_DIR
+
+. ${SCRIPT_DIR=}/docker/variables.sh
+
 TTS_MODEL=xtts
-CUDA_VERSION=12.1.1
-PYTHON_VERSION=3.11.9
-PYTORCH_VERSION=2.2.1
 DOCKER_TAG=latest
+CLEAN=false
 
 # Parse arguments
 while [ "$#" -gt 0 ]; do
@@ -17,10 +20,6 @@ while [ "$#" -gt 0 ]; do
       PYTHON_VERSION="$2"
       shift
       ;;
-    --pytorch-version)
-      PYTORCH_VERSION="$2"
-      shift
-      ;;
     --tts_model)
       TTS_MODEL="$2"
       shift
@@ -29,6 +28,9 @@ while [ "$#" -gt 0 ]; do
       DOCKER_TAG="$2"
       shift
       ;;
+    --clean)
+      CLEAN=true
+      ;;
     *)
       printf '%s\n' "Invalid argument ($1)"
       exit 1
@@ -37,29 +39,27 @@ while [ "$#" -gt 0 ]; do
   shift
 done
 
-echo "$PYTHON_VERSION -> ${PYTHON_VERSION%.*}"
-
-SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
-cd $SCRIPT_DIR
+if [ "$CLEAN" = true ]; then
+  rm -rf ${SCRIPT_DIR=}/docker/conda/build
+  rm -rf ${SCRIPT_DIR=}/docker/deepspeed/build
+fi
 
-PYTHON_MAJOR_MINOR=${PYTHON_VERSION%.*}
-$SCRIPT_DIR/deepspeed/build-deepspeed.sh \
+$SCRIPT_DIR/docker/conda/build-conda-env.sh \
   --cuda-version ${CUDA_VERSION} \
-  --python-version ${PYTHON_MAJOR_MINOR} \
-  --pytorch-version ${PYTORCH_VERSION}
+  --python-version ${PYTHON_VERSION}
 
-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} with PyTorch ${PYTORCH_VERSION}"
+$SCRIPT_DIR/docker/deepspeed/build-deepspeed.sh \
+  --python-version ${PYTHON_VERSION}
 
+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}"
 
 docker buildx \
   build \
   --progress=plain \
   --build-arg TTS_MODEL=$TTS_MODEL \
-  --build-arg CUDA_VERSION=$CUDA_VERSION \
-  --build-arg PYTHON_VERSION=$PYTHON_VERSION \
-  --build-arg PYTORCH_VERSION=$PYTORCH_VERSION \
+  --build-arg ALLTALK_DIR=$ALLTALK_DIR \
   -t alltalk_beta:${DOCKER_TAG} \
   .
 
-echo "Docker build process finished"
+echo "Docker build process finished. Use docker-start.sh to start the container."

docker-start.sh

@@ -1,7 +1,10 @@
 #!/usr/bin/env bash
 
-ALLTALK_DIR="/alltalk"
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+. ${SCRIPT_DIR=}/docker/variables.sh
+
 WITH_UI=true
+DOCKER_TAG=latest
 declare -a ADDITIONAL_ARGS=()
 
 # Parse arguments
@@ -22,6 +25,10 @@ while [ "$#" -gt 0 ]; do
     --no_ui)
       WITH_UI=false
       ;;
+    --tag)
+      DOCKER_TAG="$2"
+      shift
+      ;;
     *)
       # Allow to pass arbitrary arguments to docker as well to be flexible:
       ADDITIONAL_ARGS+=( $1 )
@@ -58,4 +65,4 @@ docker run \
   --name alltalk \
  "${DOCKER_ARGS[@]}" \
  "${ADDITIONAL_ARGS[@]}" \
-  alltalk_beta:latest &> /dev/stdout
+  alltalk_beta:${DOCKER_TAG} &> /dev/stdout