Dockerfile.gpu

# GPU overlay on top of existing openmontage image
FROM localhost/openmontage:latest

# Install CUDA toolkit runtime (for PyTorch CUDA support in WSL2)
RUN apt-get update && apt-get install -y --no-install-recommends \
    wget ca-certificates \
    && wget -qO /tmp/cuda-keyring.deb https://developer.download.nvidia.com/compute/cuda/repos/debian12/x86_64/cuda-keyring_1.1-1_all.deb \
    && dpkg -i /tmp/cuda-keyring.deb \
    && apt-get update \
    && apt-get install -y --no-install-recommends \
    cuda-libraries-12-4 \
    && apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/cuda-keyring.deb

# CUDA visibility
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility

# PyTorch with CUDA 12.4 + GPU deps
RUN pip install --no-cache-dir \
    torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 \
    && pip install --no-cache-dir diffusers transformers accelerate

ENTRYPOINT ["/bin/bash"]