wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-04-15 13:02:35 +08:00
Code Issues Actions 17 Packages Projects Releases Wiki Activity

build(docker): compile SageAttention 2.2 on slim trixie with Debian CUDA toolkit; install wheel into runtime

Browse Source 
Switch to a two-stage build that uses python:3.12-slim-trixie as both builder and runtime, enabling contrib/non-free/non-free-firmware in APT to install Debian’s nvidia-cuda-toolkit (nvcc) for compiling SageAttention 2.2 from source. Install Torch cu129 in the builder and build a cp312 wheel, then copy and install that wheel into the slim runtime so --use-sage-attention works at startup. This removes the heavy CUDA devel base, avoids a venv by permitting pip system installs during build, and keeps the final image minimal while ensuring ABI alignment with Torch cu129.
This commit is contained in:
clsferguson 2025-09-21 22:42:46 -06:00 committed by GitHub
parent 7b448364d1
commit bc0e12819d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 28 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
Dockerfile
View File

@ -1,32 +1,33 @@
# --------------------------
# Stage 1: build SageAttention 2.2 wheel from source with nvcc available
# Stage 1: build SageAttention 2.2 wheel (Debian trixie + nvcc)
# --------------------------
FROM nvidia/cuda:12.9.0-devel-ubuntu24.04 AS sage-builder
FROM python:3.12.11-slim-trixie AS sage-builder
ENV DEBIAN_FRONTEND=noninteractive \
PIP_DISABLE_PIP_VERSION_CHECK=1 \
PIP_NO_CACHE_DIR=1 \
VENV=/opt/venv
PIP_NO_CACHE_DIR=1
# Python 3.12 and build tools
RUN apt-get update && apt-get install -y --no-install-recommends \
python3 python3-pip python3-venv \
git build-essential cmake \
&& rm -rf /var/lib/apt/lists/*
# Create a venv to avoid PEP 668 'externally-managed-environment'
RUN python3 -m venv "$VENV"
ENV PATH="$VENV/bin:$PATH"
# Enable contrib/non-free/non-free-firmware, install CUDA toolkit + build deps
RUN set -eux; \
sed -i 's/ main$/ main contrib non-free non-free-firmware/g' /etc/apt/sources.list; \
apt-get update; \
apt-get install -y --no-install-recommends \
ca-certificates curl git build-essential cmake \
nvidia-cuda-toolkit \
; \
rm -rf /var/lib/apt/lists/*
WORKDIR /tmp/sage
# Install Torch (cu129) in the venv before building the extension so ABIs align
RUN python -m pip install --upgrade pip setuptools wheel \
&& python -m pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu129
# Install Torch cu129 in system site (PEP 668: allow system installs)
RUN python -m pip install --upgrade pip setuptools wheel --break-system-packages && \
python -m pip install torch torchvision torchaudio \
--extra-index-url https://download.pytorch.org/whl/cu129 \
--break-system-packages
# Shallow clone latest SageAttention and build a cp312 wheel from source
RUN git clone --depth 1 https://github.com/thu-ml/SageAttention.git . \
&& python -m pip wheel . --no-deps --no-build-isolation -w /dist
# Shallow clone SageAttention and build a cp312 wheel
RUN git clone --depth 1 https://github.com/thu-ml/SageAttention.git . && \
python -m pip wheel . --no-deps --no-build-isolation -w /dist --break-system-packages
# --------------------------
# Stage 2: runtime image (slim)
@ -41,8 +42,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
# System deps
RUN apt-get update && apt-get install -y --no-install-recommends \
git build-essential cmake \
libgl1 libglx-mesa0 libglib2.0-0 \
git build-essential cmake libgl1 libglx-mesa0 libglib2.0-0 \
fonts-dejavu-core fontconfig util-linux \
&& rm -rf /var/lib/apt/lists/*
@ -52,12 +52,13 @@ RUN groupadd --gid 1000 appuser \
WORKDIR /app/ComfyUI
# Install core deps
# Install core deps (Torch cu129 must match builder)
COPY requirements.txt ./
RUN python -m pip install --upgrade pip setuptools wheel \
&& python -m pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu129 \
&& python -m pip install -r requirements.txt \
&& python -m pip install imageio-ffmpeg "av>=14.2" nvidia-ml-py
RUN python -m pip install --upgrade pip setuptools wheel && \
python -m pip install torch torchvision torchaudio \
--extra-index-url https://download.pytorch.org/whl/cu129 && \
python -m pip install -r requirements.txt && \
python -m pip install imageio-ffmpeg "av>=14.2" nvidia-ml-py
# Install the SageAttention wheel built in the builder stage
COPY --from=sage-builder /dist/sageattention-*.whl /tmp/
@ -66,7 +67,7 @@ RUN python -m pip install /tmp/sageattention-*.whl
# Copy the application
COPY . .
# Entrypoint
# Entrypoint and launch
COPY entrypoint.sh /entrypoint.sh
RUN chmod +x /entrypoint.sh \
&& chown -R appuser:appuser /app /home/appuser /entrypoint.sh