Fixing docker build for DeepSpeed, introducing build for conda environment

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%.*}"
+if [ "$CLEAN" = true ]; then
+  rm -rf ${SCRIPT_DIR=}/docker/conda/build
+  rm -rf ${SCRIPT_DIR=}/docker/deepspeed/build
+fi
 
-SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
-cd $SCRIPT_DIR
-
-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}
+
+$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} with PyTorch ${PYTORCH_VERSION}"
-
+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."