From ac5ed0ad61a541d73974fd11bf362fdc3b8f3986 Mon Sep 17 00:00:00 2001
From: XXXXRT666 <157766680+XXXXRT666@users.noreply.github.com>
Date: Wed, 30 Apr 2025 23:37:22 +0100
Subject: [PATCH] .

---
 Dockerfile      | 3 ++-
 docker_build.sh | 7 +++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index ed53883d..ddefd9b0 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,6 +1,7 @@
 ARG CUDA_VERSION=12.4
+ARG CUDA_BASE=runtime
 
-FROM nvidia/cuda:${CUDA_VERSION}.1-cudnn-devel-ubuntu22.04
+FROM nvidia/cuda:${CUDA_VERSION}.1-cudnn-${CUDA_BASE}-ubuntu22.04
 
 LABEL maintainer="XXXXRT"
 LABEL version="V4-0501"
diff --git a/docker_build.sh b/docker_build.sh
index 88bf1a3c..8aa9a063 100644
--- a/docker_build.sh
+++ b/docker_build.sh
@@ -67,9 +67,16 @@ done
 
 TARGETPLATFORM=$(uname -m | grep -q 'x86_64' && echo "linux/amd64" || echo "linux/arm64")
 
+if [ $LITE = true ]; then
+    CUDA_BASE="runtime"
+else
+    CUDA_BASE="devel"
+fi
+
 docker build \
     --build-arg CUDA_VERSION=$CUDA_VERSION \
     --build-arg LITE=$LITE \
     --build-arg TARGETPLATFORM="$TARGETPLATFORM" \
+    --build-arg CUDA_BASE=$CUDA_BASE \
     -t "${USER}/gpt-sovits:local" \
     .