From 9a99126f688e212664b00d06c628bae0666be31c Mon Sep 17 00:00:00 2001
From: ekoc <erol.koc@ergon.ch>
Date: Fri, 13 Dec 2024 22:23:37 +0100
Subject: [PATCH 1/4] Adding Dockerfile and scripts for building and starting

---
 DOCKER_README.md           |  72 +++++++++++++++++++
 Dockerfile                 | 138 +++++++++++++++++++++++++++++++++++++
 docker-build.sh            |  33 +++++++++
 docker-start.sh            |  61 ++++++++++++++++
 docker_confignew.json      |   1 +
 docker_default_config.json |   7 ++
 6 files changed, 312 insertions(+)
 create mode 100644 DOCKER_README.md
 create mode 100644 Dockerfile
 create mode 100755 docker-build.sh
 create mode 100755 docker-start.sh
 create mode 100644 docker_confignew.json
 create mode 100644 docker_default_config.json

diff --git a/DOCKER_README.md b/DOCKER_README.md
new file mode 100644
index 0000000..62a32b6
--- /dev/null
+++ b/DOCKER_README.md
@@ -0,0 +1,72 @@
+# Docker
+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:
+  - Downloads XTTS as default TTS engine
+  - Enables RVC by default
+  - Downloads all supported RVC models
+  - Enables deepspeed by default
+- Starting the Docker image should only a few seconds due to all the steps that were already executed during build.
+
+## Docker for Linux
+
+### Ubuntu Specific Setup for GPUs
+1. Make sure the latest nvidia drivers are installed: `sudo ubuntu-drivers install`
+1. Install Docker your preferred way. One way to do it is to follow the official documentation [here](https://docs.docker.com/engine/install/ubuntu/#uninstall-old-versions).
+    - Start by uninstalling the old versions
+    - Follow the "apt" repository installation method
+    - Check that everything is working with the "hello-world" container
+1. If, when launching the docker contain, you have an error message saying that the GPU cannot be used, you might have to install [Nvidia Docker Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html).
+    - Install with the "apt" method
+    - Run the docker configuration command
+      ```sudo nvidia-ctk runtime configure --runtime=docker```
+    - Restart docker
+
+## Docker for Windows (WSL2)
+### Windows Specific Setup for GPUs
+> Make sure your Nvidia drivers are up to date: https://www.nvidia.com/download/index.aspx
+1. Install WSL2 in PowerShell with `wsl --install` and restart
+2. Open PowerShell, type and enter ```ubuntu```.  It should now load you into wsl2
+3. Remove the original nvidia cache key: `sudo apt-key del 7fa2af80`
+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`
+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).
+11. Open PowerShell, type ```ubuntu```, [then follow below](#building-and-running-in-docker)
+
+## Building and Running in Docker
+
+1. Open a terminal (or Ubuntu WSL) and go where you cloned the repo
+3. Build the image with `./docker-build.sh`
+4. Start the container with `./docker-start.sh`
+5. Visit `http://localhost:7851/` or remotely with `http://<ip>:7851`
+
+## Arguments for building and starting docker
+There are various arguments to customize the build and start of the docker image.
+
+### Arguments for `docker-build.sh`
+- `--tts_model` allows to choose the TTS model that is used by default. Valid values are `piper`, `vits`, `xtts`. Defaults to `xtts`.
+  - Example: `docker-build.sh --tts_model piper`
+- `--tag` allows to choose the docker tag. Defaults to `latest`.
+  - Example: `docker-build.sh --tag mytag`
+
+### 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
+    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`
+- `--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`
+- `--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`. 
+- Since the above commands only address the most important options, you might pass additional arbitrary docker commands
+    to the `docker-start.sh`.
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 0000000..69ac137
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,138 @@
+FROM continuumio/miniconda3:24.7.1-0
+
+# Argument to choose the model: piper, vits, xtts
+ARG TTS_MODEL="xtts"
+ENV TTS_MODEL=$TTS_MODEL
+
+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"
+
+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
+
+# Create a conda environment and install dependencies:
+ARG INSTALL_ENV_DIR=/alltalk/alltalk_environment/env
+ENV CONDA_AUTO_UPDATE_CONDA="false"
+RUN <<EOR
+    conda create -y -n "alltalk" -c conda-forge python=3.11.9
+    conda activate alltalk
+    conda install -y \
+      gcc_linux-64 \
+      gxx_linux-64 \
+      pytorch=2.2.1 \
+      torchvision \
+      torchaudio \
+      pytorch-cuda=12.1 \
+      libaio \
+      nvidia/label/cuda-12.1.0::cuda-toolkit=12.1 \
+      faiss-gpu=1.9.0 \
+      conda-forge::ffmpeg=7.1.0 \
+      conda-forge::portaudio=19.7.0 \
+      -c pytorch \
+      -c anaconda \
+      -c nvidia | grep -zq PackagesNotFoundError && exit 1
+    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
+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
+
+    # Deepspeed:
+    curl -LO https://github.com/erew123/alltalk_tts/releases/download/DeepSpeed-14.0/deepspeed-0.14.2+cu121torch2.2-cp311-cp311-manylinux_2_24_x86_64.whl
+    CFLAGS="-I$CONDA_PREFIX/include/" LDFLAGS="-L$CONDA_PREFIX/lib/" \
+      pip install --no-cache-dir --cache-dir=/alltalk/pip_cache deepspeed-0.14.2+cu121torch2.2-cp311-cp311-manylinux_2_24_x86_64.whl
+    rm -f deepspeed-0.14.2+cu121torch2.2-cp311-cp311-manylinux_2_24_x86_64.whl
+
+    # Parler:
+    pip install --no-cache-dir --no-deps --cache-dir=/alltalk/pip_cache -r system/requirements/requirements_parler.txt
+
+    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
+source ~/.bashrc
+
+# Merging config from docker_confignew.json into confignew.json:
+jq -s '.[0] * .[1] * .[2]' confignew.json docker_default_config.json docker_confignew.json  > confignew.json.tmp
+mv confignew.json.tmp confignew.json
+
+conda activate alltalk
+python script.py
+EOF
+    cat << EOF > start_finetune.sh
+#!/usr/bin/env bash
+source ~/.bashrc
+export TRAINER_TELEMETRY=0
+conda activate alltalk
+python finetune.py
+EOF
+    cat << EOF > start_diagnostics.sh
+#!/usr/bin/env bash
+source ~/.bashrc
+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
+
+COPY . .
+
+# Create script to execute firstrun.py:
+RUN echo $'#!/usr/bin/env bash \n\
+source ~/.bashrc \n\
+conda activate alltalk \n\
+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 /root/.triton/autotune
+
+# Enabling 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:
+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:
+ENTRYPOINT ["sh", "-c", "./start_alltalk.sh"]
diff --git a/docker-build.sh b/docker-build.sh
new file mode 100755
index 0000000..bb60d10
--- /dev/null
+++ b/docker-build.sh
@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+
+TTS_MODEL=xtts
+DOCKER_TAG=latest
+
+# Parse arguments
+while [ "$#" -gt 0 ]; do
+  case "$1" in
+    --tts_model)
+      TTS_MODEL="$2"
+      shift
+      ;;
+    --tag)
+      DOCKER_TAG="$2"
+      shift
+      ;;
+    *)
+      printf '%s\n' "Invalid argument ($1)"
+      exit 1
+      ;;
+  esac
+  shift
+done
+
+echo "Starting docker build process using TTS model '${TTS_MODEL}' and docker tag '${DOCKER_TAG}'"
+
+docker buildx \
+  build \
+  --build-arg TTS_MODEL=$TTS_MODEL \
+  -t alltalk_beta:${DOCKER_TAG} \
+  .
+
+echo "Docker build process finished"
\ No newline at end of file
diff --git a/docker-start.sh b/docker-start.sh
new file mode 100755
index 0000000..d76486a
--- /dev/null
+++ b/docker-start.sh
@@ -0,0 +1,61 @@
+#!/usr/bin/env bash
+
+ALLTALK_DIR="/alltalk"
+WITH_UI=true
+declare -a ADDITIONAL_ARGS=()
+
+# Parse arguments
+while [ "$#" -gt 0 ]; do
+  case "$1" in
+    --config)
+      CONFIG="$2"
+      shift
+      ;;
+    --voices)
+      VOICES="$2"
+      shift
+      ;;
+    --rvc_voices)
+      RVC_VOICES="$2"
+      shift
+      ;;
+    --no_ui)
+      WITH_UI=false
+      ;;
+    *)
+      # Allow to pass arbitrary arguments to docker as well to be flexible:
+      ADDITIONAL_ARGS+=( $1 )
+      ;;
+  esac
+  shift
+done
+
+# Compose docker arguments based on user input to the script:
+declare -a DOCKER_ARGS=()
+
+if [[ -n $CONFIG ]]; then
+  # Mount the config file to docker_confignew.json:
+  DOCKER_ARGS+=( -v ${CONFIG}:${ALLTALK_DIR}/docker_confignew.json )
+fi
+
+if [[ -n $VOICES ]]; then
+  DOCKER_ARGS+=( -v ${VOICES}:${ALLTALK_DIR}/voices )
+fi
+
+if [[ -n $RVC_VOICES ]]; then
+  DOCKER_ARGS+=( -v ${RVC_VOICES}:${ALLTALK_DIR}/models/rvc_voices )
+fi
+
+if [ "$WITH_UI" = true ] ; then
+    DOCKER_ARGS+=( -p 7852:7852 )
+fi
+
+docker run \
+  --rm \
+  -it \
+  -p 7851:7851 \
+  --gpus=all \
+  --name alltalk \
+ "${DOCKER_ARGS[@]}" \
+ "${ADDITIONAL_ARGS[@]}" \
+  alltalk_beta:latest &> /dev/stdout
\ No newline at end of file
diff --git a/docker_confignew.json b/docker_confignew.json
new file mode 100644
index 0000000..9e26dfe
--- /dev/null
+++ b/docker_confignew.json
@@ -0,0 +1 @@
+{}
\ No newline at end of file
diff --git a/docker_default_config.json b/docker_default_config.json
new file mode 100644
index 0000000..e14efdf
--- /dev/null
+++ b/docker_default_config.json
@@ -0,0 +1,7 @@
+{
+  "delete_output_wavs": "1",
+  "rvc_settings": {
+    "rvc_enabled": true,
+    "f0method": "rmvpe"
+  }
+}
\ No newline at end of file

From 7522fa87d6cb0a6ed20ce4d92dd44a308f98ab9f Mon Sep 17 00:00:00 2001
From: ekoc <erol.koc@ergon.ch>
Date: Thu, 19 Dec 2024 11:09:03 +0100
Subject: [PATCH 2/4] Adding docker build for DeepSpeed

---
 .gitignore                   |   5 +-
 Dockerfile                   |  61 +++++++++++++++------
 deepspeed/Dockerfile         | 103 +++++++++++++++++++++++++++++++++++
 deepspeed/build-deepspeed.sh |  65 ++++++++++++++++++++++
 docker-build.sh              |  32 +++++++++++
 5 files changed, 249 insertions(+), 17 deletions(-)
 create mode 100644 deepspeed/Dockerfile
 create mode 100755 deepspeed/build-deepspeed.sh

diff --git a/.gitignore b/.gitignore
index 563c227..4428008 100644
--- a/.gitignore
+++ b/.gitignore
@@ -170,4 +170,7 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 
-.idea/*
\ No newline at end of file
+.idea/*
+
+# Deepspeed build
+deepspeed/build
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
index 69ac137..d865510 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -4,12 +4,22 @@ 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
+
 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
 
 RUN <<EOR
     apt-get update
@@ -30,23 +40,30 @@ WORKDIR /alltalk
 ARG INSTALL_ENV_DIR=/alltalk/alltalk_environment/env
 ENV CONDA_AUTO_UPDATE_CONDA="false"
 RUN <<EOR
-    conda create -y -n "alltalk" -c conda-forge python=3.11.9
+    CUDA_SHORT_VERSION=${CUDA_VERSION%.*}
+
+    conda create -y -n "alltalk" -c conda-forge python=${PYTHON_VERSION}
     conda activate alltalk
-    conda install -y \
+    RESULT=$( { conda install -y \
       gcc_linux-64 \
       gxx_linux-64 \
-      pytorch=2.2.1 \
+      pytorch=${PYTORCH_VERSION} \
+      pytorch-cuda=${CUDA_SHORT_VERSION} \
       torchvision \
       torchaudio \
-      pytorch-cuda=12.1 \
       libaio \
-      nvidia/label/cuda-12.1.0::cuda-toolkit=12.1 \
+      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 | grep -zq PackagesNotFoundError && exit 1
+      -c nvidia ; } 2>&1 )
+
+    if echo $RESULT | grep -izq error ; then
+      echo "Failed to install conda dependencies 2: $RESULT"
+      exit 1
+    fi
     conda clean -a && pip cache purge
 EOR
 
@@ -61,22 +78,34 @@ RUN <<EOR
     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
+    # Parler:
+    pip install --no-cache-dir --cache-dir=/alltalk/pip_cache -r system/requirements/requirements_parler.txt
 
-    # Deepspeed:
-    curl -LO https://github.com/erew123/alltalk_tts/releases/download/DeepSpeed-14.0/deepspeed-0.14.2+cu121torch2.2-cp311-cp311-manylinux_2_24_x86_64.whl
-    CFLAGS="-I$CONDA_PREFIX/include/" LDFLAGS="-L$CONDA_PREFIX/lib/" \
-      pip install --no-cache-dir --cache-dir=/alltalk/pip_cache deepspeed-0.14.2+cu121torch2.2-cp311-cp311-manylinux_2_24_x86_64.whl
-    rm -f deepspeed-0.14.2+cu121torch2.2-cp311-cp311-manylinux_2_24_x86_64.whl
+    conda clean --all --force-pkgs-dirs -y && pip cache purge
+EOR
 
-    # Parler:
-    pip install --no-cache-dir --no-deps --cache-dir=/alltalk/pip_cache -r system/requirements/requirements_parler.txt
+# Deepspeed:
+RUN mkdir -p /tmp/deepseped
+COPY deepspeed/build/*.whl /tmp/deepspeed/
+RUN <<EOR
+    DEEPSPEED_WHEEL=$(realpath /tmp/deepspeed/*.whl)
+    conda activate alltalk
+
+    RESULT=$( { CFLAGS="-I$CONDA_PREFIX/include/" LDFLAGS="-L$CONDA_PREFIX/lib/" \
+      pip install --no-cache-dir ${DEEPSPEED_WHEEL} ; } 2>&1 )
+
+    if echo $RESULT | grep -izq error ; then
+      echo "Failed to install pip dependencies: $RESULT"
+      exit 1
+    fi
 
+    rm ${DEEPSPEED_WHEEL}
     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
+### 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
diff --git a/deepspeed/Dockerfile b/deepspeed/Dockerfile
new file mode 100644
index 0000000..26df530
--- /dev/null
+++ b/deepspeed/Dockerfile
@@ -0,0 +1,103 @@
+ARG CUDA_VERSION=12.1.1
+FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
+
+ARG PYTHON_VERSION=3.11
+ENV PYTHON_VERSION=$PYTHON_VERSION
+
+ARG PYTORCH_VERSION=2.2.1
+ENV PYTORCH_VERSION=$PYTORCH_VERSION
+
+ARG DEEPSPEED_VERSION=0.16.1
+ENV DEEPSPEED_VERSION=$DEEPSPEED_VERSION
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+##############################################################################
+# Directories:
+##############################################################################
+ENV STAGE_DIR=/tmp
+RUN mkdir -p ${STAGE_DIR}
+
+##############################################################################
+# Installation/Basic Utilities
+##############################################################################
+SHELL ["/bin/bash", "-l", "-c"]
+ENV SHELL=/bin/bash
+
+RUN <<EOR
+  apt-get update
+  apt-get install --no-install-recommends -y \
+    software-properties-common \
+    ca-certificates
+
+  add-apt-repository ppa:git-core/ppa -y # for latest git
+  add-apt-repository ppa:deadsnakes/ppa -y # for python
+
+  apt-get update
+  apt-get upgrade -y
+  apt-get install --no-install-recommends -y \
+    build-essential \
+    autotools-dev \
+    pdsh \
+    cmake \
+    g++ \
+    gcc \
+    curl \
+    wget \
+    vim \
+    unzip \
+    llvm-dev \
+    git \
+    python${PYTHON_VERSION}-dev \
+    libcupti-dev \
+    libaio-dev
+EOR
+
+##############################################################################
+# Python & pip
+##############################################################################
+RUN <<EOR
+  # Correct symlinks to use the proper python version:
+  PYTHON_MAJOR_VERSION=${PYTHON_VERSION%%.*}
+  rm -f /usr/bin/python${PYTHON_MAJOR_VERSION}
+  ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python${PYTHON_MAJOR_VERSION}
+  ln -s /usr/bin/python3 /usr/bin/python
+
+  curl -O https://bootstrap.pypa.io/pip/3.7/get-pip.py
+  python get-pip.py
+  rm get-pip.py
+  pip install --upgrade pip
+EOR
+
+# Minimal dependencies needed to build deepspeed:
+RUN pip install \
+    deepspeed-kernels \
+    scikit-learn \
+    torch==${PYTORCH_VERSION}
+
+##############################################################################
+# DeepSpeed
+##############################################################################
+RUN <<EOR
+  git clone https://github.com/microsoft/DeepSpeed.git ${STAGE_DIR}/DeepSpeed
+  cd ${STAGE_DIR}/DeepSpeed
+  git checkout .
+  git checkout "tags/v${DEEPSPEED_VERSION}" -b "v${DEEPSPEED_VERSION}"
+EOR
+
+##############################################################################
+# DeepSpeed build file
+##############################################################################
+RUN <<EOR
+  cat << EOF > build_deepspeed.sh
+#!/usr/bin/env bash
+mkdir -p /deepspeed
+cd ${STAGE_DIR}/DeepSpeed
+DS_BUILD_OPS=1 python setup.py build_ext -j8 bdist_wheel
+mv ${STAGE_DIR}/DeepSpeed/dist/*.whl /deepspeed/
+EOF
+EOR
+
+RUN chmod +x /build_deepspeed.sh
+
+ENTRYPOINT ["/build_deepspeed.sh"]
\ No newline at end of file
diff --git a/deepspeed/build-deepspeed.sh b/deepspeed/build-deepspeed.sh
new file mode 100755
index 0000000..35977bb
--- /dev/null
+++ b/deepspeed/build-deepspeed.sh
@@ -0,0 +1,65 @@
+#!/usr/bin/env bash
+
+CUDA_VERSION=12.1.1
+PYTHON_VERSION=3.11
+PYTORCH_VERSION=2.2.1
+DEEPSPEED_VERSION=0.16.1
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+cd $SCRIPT_DIR
+
+# Parse arguments
+while [ "$#" -gt 0 ]; do
+  case "$1" in
+    --cuda-version)
+      CUDA_VERSION="$2"
+      shift
+      ;;
+    --python-version)
+      PYTHON_VERSION="$2"
+      shift
+      ;;
+    --pytorch-version)
+      PYTORCH_VERSION="$2"
+      shift
+      ;;
+    --deepspeed-version)
+      DEEPSPEED_VERSION="$2"
+      shift
+      ;;
+    *)
+      # Allow to pass arbitrary arguments to docker as well to be flexible:
+      echo "Unknown argument '$1'"
+      exit 1
+      ;;
+  esac
+  shift
+done
+
+PYTHON_VERSION_NO_DOT=${PYTHON_VERSION//./}
+if [[ -n $(find build -name "deepspeed-${DEEPSPEED_VERSION}*-cp${PYTHON_VERSION_NO_DOT}-cp${PYTHON_VERSION_NO_DOT}-*.whl") ]]
+then
+   echo "DeepSpeed was already built - skipping..."
+   exit 0
+fi
+
+echo "Building DeepSpeed $DEEPSPEED_VERSION for CUDA $CUDA_VERSION using python ${PYTHON_VERSION} with PyTorch ${PYTORCH_VERSION}"
+
+rm -rf build # make sure to properly clean up - we only want 1 wheel at the time
+mkdir -p build
+docker buildx \
+  build \
+  --build-arg CUDA_VERSION=$CUDA_VERSION \
+  --build-arg PYTHON_VERSION=$PYTHON_VERSION \
+  --build-arg PYTORCH_VERSION=$PYTORCH_VERSION \
+  --build-arg DEEPSPEED_VERSION=$DEEPSPEED_VERSION \
+  -t deepspeed:cu-$CUDA_VERSION-ds-$DEEPSPEED_VERSION \
+  .
+
+docker run \
+  --rm \
+  -it \
+  --gpus=all \
+  --name deepspeed \
+  -v $SCRIPT_DIR/build:/deepspeed \
+  deepspeed:cu-$CUDA_VERSION-ds-$DEEPSPEED_VERSION
\ No newline at end of file
diff --git a/docker-build.sh b/docker-build.sh
index bb60d10..3c0e1aa 100755
--- a/docker-build.sh
+++ b/docker-build.sh
@@ -1,11 +1,26 @@
 #!/usr/bin/env bash
 
 TTS_MODEL=xtts
+CUDA_VERSION=12.1.1
+PYTHON_VERSION=3.11.9
+PYTORCH_VERSION=2.2.1
 DOCKER_TAG=latest
 
 # Parse arguments
 while [ "$#" -gt 0 ]; do
   case "$1" in
+    --cuda-version)
+      CUDA_VERSION="$2"
+      shift
+      ;;
+    --python-version)
+      PYTHON_VERSION="$2"
+      shift
+      ;;
+    --pytorch-version)
+      PYTORCH_VERSION="$2"
+      shift
+      ;;
     --tts_model)
       TTS_MODEL="$2"
       shift
@@ -22,11 +37,28 @@ while [ "$#" -gt 0 ]; do
   shift
 done
 
+echo "$PYTHON_VERSION -> ${PYTHON_VERSION%.*}"
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+cd $SCRIPT_DIR
+
+PYTHON_MAJOR_MINOR=${PYTHON_VERSION%.*}
+$SCRIPT_DIR/deepspeed/build-deepspeed.sh \
+  --cuda-version ${CUDA_VERSION} \
+  --python-version ${PYTHON_MAJOR_MINOR} \
+  --pytorch-version ${PYTORCH_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}"
+
 
 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 \
   -t alltalk_beta:${DOCKER_TAG} \
   .
 

From c7f9b04826d49369f8590ea5a11473f6de381c2a Mon Sep 17 00:00:00 2001
From: ekoc <erol.koc@ergon.ch>
Date: Fri, 20 Dec 2024 21:18:20 +0100
Subject: [PATCH 3/4] Fixing docker build for DeepSpeed, introducing build for
 conda environment

---
 .gitignore                                    |  5 +-
 DOCKER_README.md                              | 20 ++--
 Dockerfile                                    | 88 ++++++++---------
 docker-build.sh                               | 42 ++++----
 docker-start.sh                               | 11 ++-
 docker/conda/Dockerfile                       | 77 +++++++++++++++
 docker/conda/build-conda-env.sh               | 43 +++++++++
 {deepspeed => docker/deepspeed}/Dockerfile    | 96 +++++++++----------
 .../deepspeed}/build-deepspeed.sh             | 37 ++++---
 docker/variables.sh                           |  6 ++
 10 files changed, 272 insertions(+), 153 deletions(-)
 create mode 100644 docker/conda/Dockerfile
 create mode 100755 docker/conda/build-conda-env.sh
 rename {deepspeed => docker/deepspeed}/Dockerfile (56%)
 rename {deepspeed => docker/deepspeed}/build-deepspeed.sh (54%)
 create mode 100644 docker/variables.sh

diff --git a/.gitignore b/.gitignore
index 4428008..99a590e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -172,5 +172,6 @@ cython_debug/
 
 .idea/*
 
-# Deepspeed build
-deepspeed/build
\ No newline at end of file
+# Docker build
+docker/conda/build
+docker/deepspeed/build
\ No newline at end of file
diff --git a/DOCKER_README.md b/DOCKER_README.md
index 62a32b6..e2343fa 100644
--- a/DOCKER_README.md
+++ b/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`.
diff --git a/Dockerfile b/Dockerfile
index d865510..71b83d9 100644
--- a/Dockerfile
+++ b/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"]
diff --git a/docker-build.sh b/docker-build.sh
index 3c0e1aa..0d52e0a 100755
--- a/docker-build.sh
+++ b/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"
\ No newline at end of file
+echo "Docker build process finished. Use docker-start.sh to start the container."
\ No newline at end of file
diff --git a/docker-start.sh b/docker-start.sh
index d76486a..0069d07 100755
--- a/docker-start.sh
+++ b/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
\ No newline at end of file
+  alltalk_beta:${DOCKER_TAG} &> /dev/stdout
\ No newline at end of file
diff --git a/docker/conda/Dockerfile b/docker/conda/Dockerfile
new file mode 100644
index 0000000..57ff742
--- /dev/null
+++ b/docker/conda/Dockerfile
@@ -0,0 +1,77 @@
+FROM continuumio/miniconda3:24.7.1-0 AS build
+
+ARG CUDA_VERSION=12.4.1
+ENV CUDA_VERSION=$CUDA_VERSION
+
+ARG PYTHON_VERSION=3.11.11
+ENV PYTHON_VERSION=$PYTHON_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"
+
+##############################################################################
+# 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:
+##############################################################################
+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-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 )
+
+    if echo $RESULT | grep -izq error ; then
+      echo "Failed to install conda dependencies: $RESULT"
+      exit 1
+    fi
+
+    conda clean -a && pip cache purge
+EOR
+
+##############################################################################
+# Export conda environment:
+##############################################################################
+RUN <<EOR
+  conda activate alltalk
+  conda env export > /environment-cu-${CUDA_VERSION}-cp-${PYTHON_VERSION}.yml
+EOR
+
+##############################################################################
+# Make environment file available to docker build output:
+##############################################################################
+FROM scratch
+COPY --from=build /environment-*.yml /
\ No newline at end of file
diff --git a/docker/conda/build-conda-env.sh b/docker/conda/build-conda-env.sh
new file mode 100755
index 0000000..6ff6df5
--- /dev/null
+++ b/docker/conda/build-conda-env.sh
@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+cd $SCRIPT_DIR
+
+. ${SCRIPT_DIR=}/../variables.sh
+
+if [ -f $SCRIPT_DIR/build/environment-cu-${CUDA_VERSION}-cp-${PYTHON_VERSION}.yml ]; then
+   echo "Environment file already exists - skipping..."
+   exit 0
+fi
+
+# Parse arguments
+while [ "$#" -gt 0 ]; do
+  case "$1" in
+    --cuda-version)
+      CUDA_VERSION="$2"
+      shift
+      ;;
+    --python-version)
+      PYTHON_VERSION="$2"
+      shift
+      ;;
+    *)
+      # Allow to pass arbitrary arguments to docker as well to be flexible:
+      echo "Unknown argument '$1'"
+      exit 1
+      ;;
+  esac
+  shift
+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 \
+  --build-arg CUDA_VERSION=$CUDA_VERSION \
+  --build-arg PYTHON_VERSION=$PYTHON_VERSION \
+  --output=$SCRIPT_DIR/build \
+  .
\ No newline at end of file
diff --git a/deepspeed/Dockerfile b/docker/deepspeed/Dockerfile
similarity index 56%
rename from deepspeed/Dockerfile
rename to docker/deepspeed/Dockerfile
index 26df530..05c256c 100644
--- a/deepspeed/Dockerfile
+++ b/docker/deepspeed/Dockerfile
@@ -1,16 +1,15 @@
-ARG CUDA_VERSION=12.1.1
-FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
+FROM continuumio/miniconda3:24.7.1-0
 
-ARG PYTHON_VERSION=3.11
-ENV PYTHON_VERSION=$PYTHON_VERSION
-
-ARG PYTORCH_VERSION=2.2.1
-ENV PYTORCH_VERSION=$PYTORCH_VERSION
-
-ARG DEEPSPEED_VERSION=0.16.1
+ARG DEEPSPEED_VERSION=0.16.2
 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:
@@ -21,62 +20,49 @@ RUN mkdir -p ${STAGE_DIR}
 ##############################################################################
 # Installation/Basic Utilities
 ##############################################################################
-SHELL ["/bin/bash", "-l", "-c"]
-ENV SHELL=/bin/bash
-
 RUN <<EOR
   apt-get update
+  apt-get upgrade -y
   apt-get install --no-install-recommends -y \
-    software-properties-common \
-    ca-certificates
+      espeak-ng \
+      curl \
+      wget \
+      jq \
+      vim
 
-  add-apt-repository ppa:git-core/ppa -y # for latest git
-  add-apt-repository ppa:deadsnakes/ppa -y # for python
+   apt-get clean && rm -rf /var/lib/apt/lists/*
+EOR
 
-  apt-get update
-  apt-get upgrade -y
-  apt-get install --no-install-recommends -y \
-    build-essential \
-    autotools-dev \
-    pdsh \
-    cmake \
-    g++ \
-    gcc \
-    curl \
-    wget \
-    vim \
-    unzip \
-    llvm-dev \
-    git \
-    python${PYTHON_VERSION}-dev \
-    libcupti-dev \
-    libaio-dev
+##############################################################################
+# 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 echo $RESULT | grep -izq error ; then
+      echo "Failed to install conda dependencies: $RESULT"
+      exit 1
+    fi
+
+    conda clean -a && pip cache purge
 EOR
 
 ##############################################################################
-# Python & pip
+# Install python dependencies
 ##############################################################################
+ENV PIP_CACHE_DIR=/tmp/pip_cache
 RUN <<EOR
-  # Correct symlinks to use the proper python version:
-  PYTHON_MAJOR_VERSION=${PYTHON_VERSION%%.*}
-  rm -f /usr/bin/python${PYTHON_MAJOR_VERSION}
-  ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python${PYTHON_MAJOR_VERSION}
-  ln -s /usr/bin/python3 /usr/bin/python
-
-  curl -O https://bootstrap.pypa.io/pip/3.7/get-pip.py
-  python get-pip.py
-  rm get-pip.py
-  pip install --upgrade pip
+  conda activate alltalk
+
+  pip install \
+      deepspeed-kernels \
+      scikit-learn
 EOR
 
-# Minimal dependencies needed to build deepspeed:
-RUN pip install \
-    deepspeed-kernels \
-    scikit-learn \
-    torch==${PYTORCH_VERSION}
 
 ##############################################################################
-# DeepSpeed
+# Clone DeepSpeed sources
 ##############################################################################
 RUN <<EOR
   git clone https://github.com/microsoft/DeepSpeed.git ${STAGE_DIR}/DeepSpeed
@@ -86,12 +72,15 @@ RUN <<EOR
 EOR
 
 ##############################################################################
-# DeepSpeed build file
+# Create DeepSpeed build file
 ##############################################################################
 RUN <<EOR
   cat << EOF > build_deepspeed.sh
 #!/usr/bin/env bash
+source ~/.bashrc
 mkdir -p /deepspeed
+
+conda activate alltalk
 cd ${STAGE_DIR}/DeepSpeed
 DS_BUILD_OPS=1 python setup.py build_ext -j8 bdist_wheel
 mv ${STAGE_DIR}/DeepSpeed/dist/*.whl /deepspeed/
@@ -100,4 +89,7 @@ EOR
 
 RUN chmod +x /build_deepspeed.sh
 
+##############################################################################
+# Build DeepSpeed:
+##############################################################################
 ENTRYPOINT ["/build_deepspeed.sh"]
\ No newline at end of file
diff --git a/deepspeed/build-deepspeed.sh b/docker/deepspeed/build-deepspeed.sh
similarity index 54%
rename from deepspeed/build-deepspeed.sh
rename to docker/deepspeed/build-deepspeed.sh
index 35977bb..8def559 100755
--- a/deepspeed/build-deepspeed.sh
+++ b/docker/deepspeed/build-deepspeed.sh
@@ -1,28 +1,17 @@
 #!/usr/bin/env bash
 
-CUDA_VERSION=12.1.1
-PYTHON_VERSION=3.11
-PYTORCH_VERSION=2.2.1
-DEEPSPEED_VERSION=0.16.1
-
 SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 cd $SCRIPT_DIR
 
+. ${SCRIPT_DIR=}/../variables.sh
+
 # Parse arguments
 while [ "$#" -gt 0 ]; do
   case "$1" in
-    --cuda-version)
-      CUDA_VERSION="$2"
-      shift
-      ;;
     --python-version)
       PYTHON_VERSION="$2"
       shift
       ;;
-    --pytorch-version)
-      PYTORCH_VERSION="$2"
-      shift
-      ;;
     --deepspeed-version)
       DEEPSPEED_VERSION="$2"
       shift
@@ -36,24 +25,30 @@ while [ "$#" -gt 0 ]; do
   shift
 done
 
-PYTHON_VERSION_NO_DOT=${PYTHON_VERSION//./}
+PYTHON_SHORT_VERSION=${PYTHON_VERSION%.*}
+PYTHON_VERSION_NO_DOT=${PYTHON_SHORT_VERSION//./}
 if [[ -n $(find build -name "deepspeed-${DEEPSPEED_VERSION}*-cp${PYTHON_VERSION_NO_DOT}-cp${PYTHON_VERSION_NO_DOT}-*.whl") ]]
 then
    echo "DeepSpeed was already built - skipping..."
    exit 0
 fi
 
-echo "Building DeepSpeed $DEEPSPEED_VERSION for CUDA $CUDA_VERSION using python ${PYTHON_VERSION} with PyTorch ${PYTORCH_VERSION}"
+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 - we only want 1 wheel at the time
+rm -rf build # make sure to properly clean up
 mkdir -p build
+cp ${CONDA_ENV} ${SCRIPT_DIR=}/build
+
 docker buildx \
   build \
-  --build-arg CUDA_VERSION=$CUDA_VERSION \
-  --build-arg PYTHON_VERSION=$PYTHON_VERSION \
-  --build-arg PYTORCH_VERSION=$PYTORCH_VERSION \
   --build-arg DEEPSPEED_VERSION=$DEEPSPEED_VERSION \
-  -t deepspeed:cu-$CUDA_VERSION-ds-$DEEPSPEED_VERSION \
+  -t deepspeed:$DEEPSPEED_VERSION \
   .
 
 docker run \
@@ -62,4 +57,4 @@ docker run \
   --gpus=all \
   --name deepspeed \
   -v $SCRIPT_DIR/build:/deepspeed \
-  deepspeed:cu-$CUDA_VERSION-ds-$DEEPSPEED_VERSION
\ No newline at end of file
+  deepspeed:$DEEPSPEED_VERSION
\ No newline at end of file
diff --git a/docker/variables.sh b/docker/variables.sh
new file mode 100644
index 0000000..9cc0d0f
--- /dev/null
+++ b/docker/variables.sh
@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+
+export CUDA_VERSION=12.4.1
+export PYTHON_VERSION=3.11.11
+export DEEPSPEED_VERSION=0.16.2
+export ALLTALK_DIR=/opt/alltalk
\ No newline at end of file

From 78b0a51c4e0f9a3b1ce6bbf6b9a19ff3413f9bac Mon Sep 17 00:00:00 2001
From: ekoc <erol.koc@ergon.ch>
Date: Fri, 20 Dec 2024 21:41:06 +0100
Subject: [PATCH 4/4] Fixing a typo in the readme

---
 DOCKER_README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/DOCKER_README.md b/DOCKER_README.md
index e2343fa..02f7a99 100644
--- a/DOCKER_README.md
+++ b/DOCKER_README.md
@@ -69,7 +69,7 @@ There are various arguments to customize the build and start of the docker image
   voice files.
   - 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-start.sh --rvc_vices /my/rvc/voices/dir`
+  - Example: `docker-start.sh --rvc_voices /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`.
 - `--tag` allows to choose the docker tag of the image to run. Defaults to `latest`.