1580158
/
kohya_ss
mirror of https://github.com/bmaltais/kohya_ss
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update CUDA and PyTorch configuration in Dockerfile (#3195)

pull/3229/head
JoeYang1412 2025-05-03 08:21:51 +08:00 committed by GitHub
parent 950ee5e7b5
commit e69834cd6f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 11 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
Dockerfile
View File

@ -26,11 +26,13 @@ RUN --mount=type=cache,id=apt-$TARGETARCH$TARGETVARIANT,sharing=locked,target=/v
# The versions must align and be in sync with the requirements_linux_docker.txt
# hadolint ignore=SC2102
RUN --mount=type=cache,id=pip-$TARGETARCH$TARGETVARIANT,sharing=locked,target=/root/.cache/pip \
pip install -U --extra-index-url https://download.pytorch.org/whl/cu121 --extra-index-url https://pypi.nvidia.com \
torch==2.1.2 torchvision==0.16.2 \
xformers==0.0.23.post1 \
ninja \
pip setuptools wheel
pip install --upgrade pip setuptools wheel ninja &&\
pip install -U \
--index-url https://download.pytorch.org/whl/cu124 \
--extra-index-url https://pypi.nvidia.com \
torch==2.5.0+cu124 \
torchvision==0.20.0+cu124 &&\
pip install -U xformers --index-url https://download.pytorch.org/whl/cu124
# Install requirements
RUN --mount=type=cache,id=pip-$TARGETARCH$TARGETVARIANT,sharing=locked,target=/root/.cache/pip \
@ -60,25 +62,17 @@ ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
WORKDIR /tmp
ENV CUDA_VERSION=12.1.1
ENV NV_CUDA_CUDART_VERSION=12.1.105-1
ENV NVIDIA_REQUIRE_CUDA=cuda>=12.1
ENV NV_CUDA_COMPAT_PACKAGE=cuda-compat-12-1
# Install CUDA partially
ADD https://developer.download.nvidia.com/compute/cuda/repos/debian11/x86_64/cuda-keyring_1.0-1_all.deb .
ADD https://developer.download.nvidia.com/compute/cuda/repos/debian12/x86_64/cuda-keyring_1.1-1_all.deb .
RUN --mount=type=cache,id=apt-$TARGETARCH$TARGETVARIANT,sharing=locked,target=/var/cache/apt \
--mount=type=cache,id=aptlists-$TARGETARCH$TARGETVARIANT,sharing=locked,target=/var/lib/apt/lists \
dpkg -i cuda-keyring_1.0-1_all.deb && \
rm cuda-keyring_1.0-1_all.deb && \
dpkg -i cuda-keyring_1.1-1_all.deb && \
rm cuda-keyring_1.1-1_all.deb && \
sed -i 's/^Components: main$/& contrib/' /etc/apt/sources.list.d/debian.sources && \
apt-get update && \
apt-get install -y --no-install-recommends \
# Installing the whole CUDA typically increases the image size by approximately **8GB**.
# To decrease the image size, we opt to install only the necessary libraries.
# Here is the package list for your reference: https://developer.download.nvidia.com/compute/cuda/repos/debian11/x86_64
# !If you experience any related issues, replace the following line with `cuda-12-1` to obtain the complete CUDA package.
cuda-cudart-12-1=${NV_CUDA_CUDART_VERSION} ${NV_CUDA_COMPAT_PACKAGE} libcusparse-12-1 libnvjitlink-12-1
cuda-toolkit-12-4
# Install runtime dependencies
RUN --mount=type=cache,id=apt-$TARGETARCH$TARGETVARIANT,sharing=locked,target=/var/cache/apt \