Update Dockerfile

Dockerfile

@@ -1,9 +1,23 @@
-# Use the NVIDIA CUDA runtime as a base image
-FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04
+FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-devel
 
 # Set environment variables
 ENV DEBIAN_FRONTEND=noninteractive
 
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    git \
+    cmake \
+    build-essential \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    ffmpeg \
+    python3.8 \
+    python3-pip \
+    python3.8-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create a symlink for python
+RUN ln -s /usr/bin/python3 /usr/bin/python
+
 RUN useradd -m -u 1000 user
 
 USER user
@@ -19,23 +33,21 @@ ENV HOME=/home/user \
     GRADIO_SHARE=False \
 	SYSTEM=spaces
 
+# Set CUDA_HOME environment variable
+ENV CUDA_HOME=/usr/local/cuda-11.8
+ENV TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6+PTX;8.9;9.0"
+ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
+
+# Set the environment variable to specify the GPU device
+ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
+ENV CUDA_VISIBLE_DEVICES=0
+
+
 # Set the working directory to the user's home directory
 WORKDIR $HOME/app
 
 # Install system dependencies as root
 USER root
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    git \
-    cmake \
-    build-essential \
-    libgl1-mesa-glx \
-    libglib2.0-0 \
-    ffmpeg \
-    python3.8 \
-    python3-pip \
-    python3.8-dev \
-    && rm -rf /var/lib/apt/lists/*
-
 
 # Set Python 3.8 as the default python and pip versions
 RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.8 1
@@ -43,7 +55,7 @@ RUN update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1
 
 USER user
 # Clone the repository (adjust the URL if needed)
-RUN git clone --recursive https://github.com/jnjaby/KEEP.git .
+RUN git clone --recursive https://github.com/jnjaby/KEEP.git $HOME/app
 
 # Copy the app.py script into the container
 COPY app.py .
@@ -64,10 +76,5 @@ USER user
 RUN pip install ffmpeg-python
 RUN pip install dlib
 
-# Set the environment variable to specify the GPU device
-ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
-ENV CUDA_VISIBLE_DEVICES=0
-
-
 # Command to run your application
 CMD ["python", "app.py"]