From 3a2b0d4b2d818b80915c946b9e451b8938d9302c Mon Sep 17 00:00:00 2001
From: Yeung Yiu Hung <446404+darkcl@users.noreply.github.com>
Date: Tue, 28 Oct 2025 14:48:55 +0800
Subject: [PATCH] fix docker image to run on CUDA 12.8

---
 .dockerignore       |  3 +++
 Dockerfile          | 24 +++++++-----------------
 docker-compose.yaml |  2 +-
 requirements.txt    |  1 +
 4 files changed, 12 insertions(+), 18 deletions(-)
 create mode 100644 .dockerignore

diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 00000000..3d6bb92b
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,3 @@
+checkpoints/
+outputs/
+logs/
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
index c22ca468..39f850a7 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.6.0-runtime-ubuntu22.04 AS base
+FROM nvidia/cuda:12.8.0-runtime-ubuntu22.04 AS base
 
 # Set environment variables
 ENV PYTHONDONTWRITEBYTECODE=1 \
@@ -14,6 +14,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     python3-venv \
     python3-dev \
     build-essential \
+    ffmpeg \
     git \
     curl \
     wget \
@@ -25,30 +26,19 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 RUN python -m venv /opt/venv
 ENV PATH="/opt/venv/bin:$PATH"
 
-# Create a non-root user to run the application
-RUN useradd -m -u 1001 appuser
-
 # Set working directory
 WORKDIR /app
 
-# Clone the repository
-RUN git clone https://github.com/ace-step/ACE-Step.git .
+COPY . .
 
 # Install specific PyTorch version compatible with CUDA 12.6
 RUN pip3 install --no-cache-dir --upgrade pip && \
     pip3 install --no-cache-dir hf_transfer peft && \
-    pip3 install --no-cache-dir -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu126
+    pip3 install --no-cache-dir -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu128
 RUN pip3 install --no-cache-dir .
 
-# Ensure target directories for volumes exist and have correct initial ownership
-RUN mkdir -p /app/outputs /app/checkpoints /app/logs && \
-    chown -R appuser:appuser /app/outputs /app/checkpoints /app/logs
-
-# Change ownership of app files to appuser
-RUN chown -R appuser:appuser /app
+RUN mkdir -p /app/outputs /app/checkpoints /app/logs
 
-# Switch to non-root user
-USER appuser
 
 # Expose the port the app runs on
 EXPOSE 7865
@@ -57,7 +47,7 @@ VOLUME [ "/app/checkpoints", "/app/outputs", "/app/logs" ]
 
 # Set healthcheck
 HEALTHCHECK --interval=60s --timeout=10s --start-period=5s --retries=5 \
-  CMD curl -f http://localhost:7865/ || exit 1
+    CMD curl -f http://localhost:7865/ || exit 1
 
 # Command to run the application with GPU support
-CMD ["python3", "acestep/gui.py", "--server_name", "0.0.0.0", "--bf16", "true"]
+CMD ["python3", "acestep/gui.py", "--server_name", "0.0.0.0", "--bf16", "true"]
\ No newline at end of file
diff --git a/docker-compose.yaml b/docker-compose.yaml
index abb28f43..31d4d06f 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -2,7 +2,7 @@
 services:
   &name ace-step:
     build:
-      context: https://github.com/ace-step/ACE-Step.git
+      context: . 
       dockerfile: Dockerfile
     container_name: *name
     hostname: *name
diff --git a/requirements.txt b/requirements.txt
index 15d61cd8..fb780476 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,6 +10,7 @@ pytorch_lightning==2.5.1
 soundfile==0.13.1
 torch
 torchaudio
+torchcodec
 torchvision
 tqdm
 transformers==4.50.0