| FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04 | |
| WORKDIR /code | |
| # Install system dependencies | |
| RUN apt-get update && apt-get install -y \ | |
| python3.10 \ | |
| python3-pip \ | |
| git | |
| # Clone their repo | |
| COPY . /code/ | |
| # Install PyTorch with CUDA support - updated to match requirements.txt | |
| RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 | |
| # Install their requirements | |
| RUN pip3 install -r requirements.txt | |
| # Set environment variables | |
| ENV NVIDIA_VISIBLE_DEVICES=all | |
| ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility | |
| ENV CUDA_VISIBLE_DEVICES=0 | |
| ENV ATTN_BACKEND=xformers # Use xformers instead of flash-attn | |
| ENV SPCONV_ALGO=native | |
| CMD ["python3", "app.py"] |