Update Dockerfile

Dockerfile

@@ -7,18 +7,30 @@ FROM python:3.9
 RUN useradd -m -u 1000 user
 WORKDIR /app
 
+# Set the CUDA_HOME environment variable
+ENV CUDA_HOME=/usr/local/cuda
+
+# Install CUDA dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    cuda-command-line-tools-11-8 \
+    libcublas-11-8 \
+    libcublas-dev-11-8 \
+    libcudnn8=8.4.1.50-1+cuda11.8 \
+    libcudnn8-dev=8.4.1.50-1+cuda11.8 \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
+
 # Copy and install Python dependencies
 COPY --chown=user ./requirements.txt requirements.txt
 RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
-# Clone and install the modified gaussian splatting repository
-RUN git clone --recursive https://github.com/ashawkey/diff-gaussian-rasterization \
-    && pip install ./diff-gaussian-rasterization
-
 # Install specific versions of PyTorch and other dependencies
 RUN pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118 \
     && pip install -U xformers --index-url https://download.pytorch.org/whl/cu118
 
+# Clone and install the modified gaussian splatting repository
+RUN git clone --recursive https://github.com/ashawkey/diff-gaussian-rasterization \
+    && pip install ./diff-gaussian-rasterization
+
 # Install nvdiffrast for mesh extraction
 RUN pip install git+https://github.com/NVlabs/nvdiffrast