feat: Update Dockerfile to use PyTorch base image and improve depende…

…ncy installation with caching

Dockerfile

@@ -1,26 +1,23 @@
-FROM nvidia/cuda:12.1.0-base-ubuntu22.04
+FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel
 
 # Set working directory
 WORKDIR /app
 
-# Install system dependencies
+# Copy dependency files first
+COPY requirements.txt pyproject.toml ./
+
+# Install system dependencies and Python dependencies
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
+    --mount=type=cache,target=/root/.cache/pip \
     apt-get update && apt-get install -y \
-    python3.10 \
-    python3-pip \
     ffmpeg \
     libsndfile1 \
     git \
     espeak-ng \
-    && rm -rf /var/lib/apt/lists/*
-
-# Copy dependency files first
-COPY requirements.txt pyproject.toml ./
-
-# Install Python dependencies with caching
-RUN --mount=type=cache,target=/root/.cache/pip \
-    pip3 install --no-cache-dir -r requirements.txt
+    && rm -rf /var/lib/apt/lists/* \
+    && pip install uv \
+    && uv pip install -r requirements.txt
 
 # Clone the zonos repository and update submodules
 RUN git clone https://github.com/Zyphra/Zonos.git /app/zonos \
@@ -29,17 +26,17 @@ RUN git clone https://github.com/Zyphra/Zonos.git /app/zonos \
     && cd ..
 
 # Install Zonos and its dependencies
-RUN pip3 install /app/zonos \
-    && pip3 install kanjize>=1.5.0 \
+RUN uv pip install --system -e /app/zonos \
+    && uv pip install kanjize>=1.5.0 \
     inflect>=7.5.0 \
     phonemizer>=3.3.0 \
     sudachidict-full>=20241021 \
     sudachipy>=0.6.10
 
 # Install GPU dependencies with caching
 RUN --mount=type=cache,target=/root/.cache/pip \
-    FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip3 install flash-attn --no-build-isolation \
-    && pip3 install --no-cache-dir \
+    FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE uv pip install flash-attn --no-build-isolation \
+    && uv pip install --no-cache-dir \
        https://github.com/state-spaces/mamba/releases/download/v2.2.4/mamba_ssm-2.2.4+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl \
        https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.5.0.post8/causal_conv1d-1.5.0.post8+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl