CHANGED METHOD: Replace multi-stage Docker build with single-stage runtime installation approach

This commit significantly simplifies the Docker image architecture by removing the complex multi-stage build process that was causing build failures and compatibility issues across different GPU generations.

Key changes:
- Replace multi-stage builder pattern with runtime-based Sage Attention installation via enhanced entrypoint.sh
- Downgrade from CUDA 12.9 to CUDA 12.8 for broader GPU compatibility (RTX 30+ series)
- Remove pre-built wheel installation in favor of dynamic source compilation during container startup
- Add comprehensive multi-GPU detection and mixed-generation support in entrypoint script
- Integrate intelligent build caching with rebuild detection when GPU configuration changes
- Remove --use-sage-attention from default CMD to allow flexible runtime configuration

Architecture improvements:
- Single FROM nvidia/cuda:12.8.0-devel-ubuntu24.04 (was multi-stage with runtime + devel)
- Simplified package installation without build/runtime separation
- Enhanced Python 3.12 setup with proper symlinks
- Removed complex git SHA resolution and cache-busting mechanisms

Performance optimizations:
- Dynamic CUDA architecture targeting (TORCH_CUDA_ARCH_LIST) based on detected GPUs
- Intelligent Triton version selection (3.2 for RTX 20, latest for RTX 30+)
- Parallel compilation settings moved to environment variables
- Reduced Docker layer count for faster builds and smaller image size

The previous multi-stage approach was abandoned due to:
- Frequent build failures across different CUDA environments
- Complex dependency management between builder and runtime stages
- Inability to handle mixed GPU generations at build time
- Excessive build times and debugging complexity

This runtime-based approach provides better flexibility, reliability, and user experience while maintaining optimal performance through intelligent GPU detection and version selection.

Dockerfile

@@ -1,63 +1,38 @@
-# ------------------------------------------------------------------------------
+# Use NVIDIA CUDA 12.8 devel image for compilation support
-# Builder: CUDA devel image to compile SageAttention from source (Ubuntu 24.04)
+FROM nvidia/cuda:12.8.0-devel-ubuntu24.04
-# ------------------------------------------------------------------------------
-FROM nvidia/cuda:12.9.0-devel-ubuntu24.04 AS builder
 
-ARG DEBIAN_FRONTEND=noninteractive
+# Environment
-# Configurable ref; use "main" or a tag/branch; set this at build time if needed
+ENV DEBIAN_FRONTEND=noninteractive \
-ARG SAGE_REF=main
+    PYTHONUNBUFFERED=1 \
-# Cache-buster to force re-resolving the latest commit
-ARG SAGE_FORCE_REFRESH=0
-
-# System deps for build
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    ca-certificates git curl \
-    build-essential cmake ninja-build pkg-config \
-    python3 python3-pip python3-venv python3-dev \
- && rm -rf /var/lib/apt/lists/*
-
-# Ensure recent pip toolchain
-ENV PIP_DISABLE_PIP_VERSION_CHECK=1 PIP_NO_CACHE_DIR=1
-RUN python3 -m pip install --upgrade pip setuptools wheel
-
-# Python deps: Torch (CUDA 12.9), Triton
-RUN python3 -m pip install --upgrade \
-    torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu129 \
- && python3 -m pip install "triton>=3.0.0"
-
-# Resolve latest commit for the desired ref and build a wheel for SageAttention 2.x
-RUN set -eux; \
-    echo "CACHE_BUSTER=${SAGE_FORCE_REFRESH}"; \
-    SAGE_SHA="$(git ls-remote https://github.com/thu-ml/SageAttention.git "${SAGE_REF}" | awk '{print $1}' | head -n1)"; \
-    echo "Resolved SageAttention ${SAGE_REF} -> ${SAGE_SHA}"; \
-    git clone https://github.com/thu-ml/SageAttention.git /build/SageAttention; \
-    cd /build/SageAttention; \
-    git checkout "${SAGE_SHA}"; \
-    export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32; \
-    python3 -m pip wheel --no-build-isolation --no-deps -w /wheels .
-
-# ------------------------------------------------------------------------------
-# Final image: CUDA runtime (Ubuntu 24.04) + Python 3 + ComfyUI + SageAttention
-# ------------------------------------------------------------------------------
-FROM nvidia/cuda:12.9.0-runtime-ubuntu24.04
-
-ARG DEBIAN_FRONTEND=noninteractive
-ENV PYTHONUNBUFFERED=1 \
     COMFY_AUTO_INSTALL=1 \
     PIP_DISABLE_PIP_VERSION_CHECK=1 \
-    PIP_NO_CACHE_DIR=1
+    PIP_NO_CACHE_DIR=1 \
+    EXT_PARALLEL=4 \
+    NVCC_APPEND_FLAGS="--threads 8" \
+    MAX_JOBS=32 \
+    SAGE_ATTENTION_AVAILABLE=0
 
-# System deps
+# System deps including Python 3.12
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    ca-certificates git curl \
+    python3.12 \
-    build-essential cmake \
+    python3.12-dev \
-    libgl1 libglx-mesa0 libglib2.0-0 fonts-dejavu-core fontconfig util-linux \
+    python3.12-venv \
-    python3 python3-pip python3-venv python3-dev \
+    python3-pip \
+    git \
+    build-essential \
+    cmake \
+    libgl1 \
+    libglx-mesa0 \
+    libglib2.0-0 \
+    fonts-dejavu-core \
+    fontconfig \
+    util-linux \
+    wget \
+    curl \
+ && ln -sf /usr/bin/python3.12 /usr/bin/python \
+ && ln -sf /usr/bin/python3.12 /usr/bin/python3 \
  && rm -rf /var/lib/apt/lists/*
 
-# Convenience: `python`/`pip` aliases
-RUN ln -s /usr/bin/python3 /usr/local/bin/python && ln -s /usr/bin/pip3 /usr/local/bin/pip || true
-
 # Create runtime user/group
 RUN groupadd --gid 1000 appuser \
  && useradd --uid 1000 --gid 1000 --create-home --shell /bin/bash appuser
@@ -65,26 +40,23 @@ RUN groupadd --gid 1000 appuser \
 # Workdir
 WORKDIR /app/ComfyUI
 
-# Copy requirements early to leverage caching
+# Leverage layer caching: install deps before copying full tree
 COPY requirements.txt ./
 
-# Core Python deps (Torch CUDA 12.9, app reqs), media/NVML libs, Triton
+# Core Python deps (torch CUDA 12.8, 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 torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu128 \
+ && python -m pip install triton \
  && python -m pip install -r requirements.txt \
- && python -m pip install imageio-ffmpeg "av>=14.2" nvidia-ml-py \
+ && python -m pip install imageio-ffmpeg "av>=14.2" nvidia-ml-py
- && python -m pip install "triton>=3.0.0"
-
-# Install SageAttention wheel built in the builder
-COPY --from=builder /wheels /tmp/wheels
-RUN python -m pip install --no-cache-dir /tmp/wheels/*.whl && rm -rf /tmp/wheels
 
 # Copy the application
 COPY . .
 
-# Entrypoint script
+# Entrypoint
 COPY entrypoint.sh /entrypoint.sh
-RUN chmod +x /entrypoint.sh && chown -R appuser:appuser /app /home/appuser /entrypoint.sh
+RUN chmod +x /entrypoint.sh \
+ && chown -R appuser:appuser /app /home/appuser /entrypoint.sh
 
 EXPOSE 8188
 
@@ -92,5 +64,5 @@ EXPOSE 8188
 USER root
 ENTRYPOINT ["/entrypoint.sh"]
 
-# Enable SageAttention by default; override with USE_SAGE_ATTENTION=0 if desired
+# Default command - entrypoint will add --use-sage-attention if available
-CMD ["python", "main.py", "--listen", "0.0.0.0", "--use-sage-attention"]
+CMD ["python", "main.py", "--listen", "0.0.0.0"]