From 7b448364d15ec8a871243179ebfa0e8b001deee6 Mon Sep 17 00:00:00 2001
From: clsferguson <48876201+clsferguson@users.noreply.github.com>
Date: Sun, 21 Sep 2025 22:15:28 -0600
Subject: [PATCH] fix(build): use CUDA devel builder + venv to build and bundle
 SageAttention 2.2 wheel; make launch flag effective
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Switch the builder stage to nvidia/cuda:12.9.0-devel-ubuntu24.04 and create a Python 3.12 venv to avoid PEP 668 “externally managed” errors, install Torch 2.8.0+cu129 in that venv, and build a cp312 SageAttention 2.2 wheel from upstream; copy and install the wheel in the slim runtime so --use-sage-attention works at startup.
This resolves prior build failures on Debian Trixie slim where CUDA toolkits were unavailable and fixes runtime ModuleNotFoundError by ensuring the module is present in the exact interpreter ComfyUI uses.
---
 Dockerfile | 36 +++++++++++++-----------------------
 1 file changed, 13 insertions(+), 23 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 919c77c98..16a7cc22f 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -5,24 +5,26 @@ FROM nvidia/cuda:12.9.0-devel-ubuntu24.04 AS sage-builder
 
 ENV DEBIAN_FRONTEND=noninteractive \
     PIP_DISABLE_PIP_VERSION_CHECK=1 \
-    PIP_NO_CACHE_DIR=1
+    PIP_NO_CACHE_DIR=1 \
+    VENV=/opt/venv
 
-# Python 3.12 and build tools (Ubuntu 24.04 ships Python 3.12)
+# 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/*
 
-# Make 'python' point to Python 3
-RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 1
+# Create a venv to avoid PEP 668 'externally-managed-environment'
+RUN python3 -m venv "$VENV"
+ENV PATH="$VENV/bin:$PATH"
 
 WORKDIR /tmp/sage
 
-# Match runtime Torch (cu129) before building the extension so ABIs align
+# 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
 
-# Shallow clone latest SageAttention and build a cp312 wheel
+# 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
 
@@ -31,7 +33,6 @@ RUN git clone --depth 1 https://github.com/thu-ml/SageAttention.git . \
 # --------------------------
 FROM python:3.12.11-slim-trixie
 
-# Environment
 ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONUNBUFFERED=1 \
     COMFY_AUTO_INSTALL=1 \
@@ -40,34 +41,25 @@ 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 \
-    fonts-dejavu-core \
-    fontconfig \
-    util-linux \
+    git build-essential cmake \
+    libgl1 libglx-mesa0 libglib2.0-0 \
+    fonts-dejavu-core fontconfig util-linux \
  && rm -rf /var/lib/apt/lists/*
 
 # Create runtime user/group
 RUN groupadd --gid 1000 appuser \
  && useradd --uid 1000 --gid 1000 --create-home --shell /bin/bash appuser
 
-# Workdir
 WORKDIR /app/ComfyUI
 
-# Leverage layer caching: install deps before copying full tree
+# Install core deps
 COPY requirements.txt ./
-
-# Core Python deps (Torch CUDA 12.9, ComfyUI reqs), media/NVML libs
 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
 
-# Bring in the SageAttention 2.2 wheel compiled in the builder stage and install it
+# Install the SageAttention wheel built in the builder stage
 COPY --from=sage-builder /dist/sageattention-*.whl /tmp/
 RUN python -m pip install /tmp/sageattention-*.whl
 
@@ -80,8 +72,6 @@ RUN chmod +x /entrypoint.sh \
  && chown -R appuser:appuser /app /home/appuser /entrypoint.sh
 
 EXPOSE 8188
-
-# Start as root so entrypoint can adjust ownership and drop privileges
 USER root
 ENTRYPOINT ["/entrypoint.sh"]
 CMD ["python", "main.py", "--listen", "0.0.0.0", "--use-sage-attention"]