diff --git a/Dockerfile.blackwell b/Dockerfile.blackwell new file mode 100644 index 000000000..fbb40c0d1 --- /dev/null +++ b/Dockerfile.blackwell @@ -0,0 +1,75 @@ +FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04 + +# Blackwell (RTX 50xx) requires sm_120 architecture +ARG CUDA_ARCHITECTURES="12.0" + +ENV DEBIAN_FRONTEND=noninteractive + +# Install system dependencies +RUN apt update && \ + apt install -y \ + python3 python3-pip git wget curl cmake ninja-build \ + libgl1 libglib2.0-0 ffmpeg && \ + apt clean + +WORKDIR /workspace + +COPY requirements.txt . + +# Upgrade pip first +RUN pip install --upgrade pip setuptools wheel + +# Install PyTorch 2.7.1 with CUDA 12.8 for Blackwell +RUN pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu128 + +# Install requirements +RUN pip install -r requirements.txt + +# Install SageAttention from git (patch GPU detection for Blackwell) +ENV TORCH_CUDA_ARCH_LIST="${CUDA_ARCHITECTURES}" +ENV FORCE_CUDA="1" +ENV MAX_JOBS="1" + +COPY <