From 105b96e575c051984636057e2a71029b8ad6e25a Mon Sep 17 00:00:00 2001
From: Kishor Prins <me@68x.ca>
Date: Mon, 14 Apr 2025 18:27:08 -0400
Subject: [PATCH] Update to ROCm 6.4 and use uv to install dependencies

---
 docker/rocm/Dockerfile         | 38 +++++++++++++++++++++++++---------
 docker/rocm/docker-compose.yml |  6 ++++++
 pyproject.toml                 | 18 +++++-----------
 3 files changed, 39 insertions(+), 23 deletions(-)

diff --git a/docker/rocm/Dockerfile b/docker/rocm/Dockerfile
index 6da250e..a24f336 100644
--- a/docker/rocm/Dockerfile
+++ b/docker/rocm/Dockerfile
@@ -23,15 +23,30 @@ RUN apt-get update && apt upgrade -y && apt-get install -y --no-install-recommen
 RUN  mkdir -p /app/api/src/models/v1_0
 WORKDIR /app
 
+# Install UV using the installer script
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    mv /root/.local/bin/uv /usr/local/bin/ && \
+    mv /root/.local/bin/uvx /usr/local/bin/
+
+# Create non-root user and set up directories and permissions
+RUN useradd -m -u 1001 appuser && \
+    mkdir -p /app/api/src/models/v1_0 && \
+    chown -R appuser:appuser /app
+
+USER appuser
+WORKDIR /app
+
 # Copy dependency files
-COPY pyproject.toml ./pyproject.toml
+COPY --chown=appuser:appuser pyproject.toml ./pyproject.toml
 
-# Install dependencies
-ENV CAUSAL_CONV1D_FORCE_BUILD=TRUE \
-    HIP_ARCHITECTURES=gfx1100
+ENV PHONEMIZER_ESPEAK_PATH=/usr/bin \
+    PHONEMIZER_ESPEAK_DATA=/usr/share/espeak-ng-data \
+    ESPEAK_DATA_PATH=/usr/share/espeak-ng-data
 
-RUN pip3 install --upgrade pip && \
-    pip3 install -e .
+# Install dependencies with GPU extras (using cache mounts)
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv venv --python 3.10 && \
+    uv sync --extra rocm
 
 # Copy project files including models
 COPY --chown=appuser:appuser api ./api
@@ -39,16 +54,19 @@ COPY --chown=appuser:appuser web ./web
 COPY --chown=appuser:appuser docker/scripts/ ./
 RUN chmod +x ./entrypoint.sh
 
-# Set environment variables
+# Set all environment variables in one go
 ENV PYTHONUNBUFFERED=1 \
     PYTHONPATH=/app:/app/api \
-    USE_GPU=true \
-    DOWNLOAD_MODEL=true
+    PATH="/app/.venv/bin:$PATH" \
+    UV_LINK_MODE=copy \
+    USE_GPU=true
 
+ENV DOWNLOAD_MODEL=true
 # Download model if enabled
 RUN if [ "$DOWNLOAD_MODEL" = "true" ]; then \
     python download_model.py --output api/src/models/v1_0; \
     fi
 
+ENV DEVICE="gpu"
 # Run FastAPI server through entrypoint.sh
-CMD ["python3", "-m", "uvicorn", "api.src.main:app", "--host", "0.0.0.0", "--port", "8880", "--log-level", "debug"]
+CMD ["./entrypoint.sh"]
diff --git a/docker/rocm/docker-compose.yml b/docker/rocm/docker-compose.yml
index fbe0e37..2ed0f39 100644
--- a/docker/rocm/docker-compose.yml
+++ b/docker/rocm/docker-compose.yml
@@ -8,6 +8,12 @@ services:
         - seccomp:unconfined
       cap_add:
         - SYS_PTRACE
+      group_add:
+      # NOTE: These groups are the group ids for: video, input, and render
+      # Numbers can be found via running: getent group $GROUP_NAME | cut -d: -f3
+        - 44
+        - 993
+        - 996
       restart: 'always'
       volumes:
         - ./kokoro-tts/config:/root/.config/miopen
diff --git a/pyproject.toml b/pyproject.toml
index ec10b3b..39c0ec8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,7 +14,7 @@ dependencies = [
     "python-dotenv==1.0.1",
     "sqlalchemy==2.0.27",
     # ML/DL Base
-    "numpy>=1.26.0",
+    "numpy==1.26.4",
     "scipy==1.14.1",
     # Audio processing
     "soundfile==0.13.0",
@@ -47,8 +47,10 @@ gpu = [
     "torch==2.6.0+cu124",
 ]
 rocm = [
-    "torch==2.5.1+rocm6.2",
-    "pytorch-triton-rocm==3.1.0",
+    "torch @ https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4/torch-2.6.0%2Brocm6.4.0.git2fb0ac2b-cp310-cp310-linux_x86_64.whl",
+    "pytorch-triton-rocm @ https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4/pytorch_triton_rocm-3.2.0%2Brocm6.4.0.git6da9e660-cp310-cp310-linux_x86_64.whl",
+    "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4/torchvision-0.21.0%2Brocm6.4.0.git4040d51f-cp310-cp310-linux_x86_64.whl",
+    "torchaudio @ https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4/torchaudio-2.6.0%2Brocm6.4.0.gitd8831425-cp310-cp310-linux_x86_64.whl",
 ]
 cpu = [
     "torch==2.6.0",
@@ -75,11 +77,6 @@ conflicts = [
 torch = [
     { index = "pytorch-cpu", extra = "cpu" },
     { index = "pytorch-cuda", extra = "gpu" },
-    { index = "pytorch-rocm", extra = "rocm" },
-]
-# uv would not install torch for rocm6.2 without this explicit transient dependency
-pytorch-triton-rocm = [
-    { index = "pytorch-rocm", extra = "rocm" },
 ]
 
 [[tool.uv.index]]
@@ -92,11 +89,6 @@ name = "pytorch-cuda"
 url = "https://download.pytorch.org/whl/cu124"
 explicit = true
 
-[[tool.uv.index]]
-name = "pytorch-rocm"
-url = "https://download.pytorch.org/whl/rocm6.2"
-explicit = true
-
 [build-system]
 requires = ["setuptools>=61.0"]
 build-backend = "setuptools.build_meta"