Update Dockerfile

Dockerfile

@@ -1,23 +1,9 @@
-FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-devel
+# Use the NVIDIA CUDA runtime as a base image
+FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04
 
 # 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
@@ -33,58 +19,53 @@ 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 working directory to the user's home directory
+WORKDIR $HOME/app
 
-# Set the environment variable to specify the GPU device
-ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
-ENV CUDA_VISIBLE_DEVICES=0
+# 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 the working directory to the user's home directory
-WORKDIR $HOME
+# Set Python 3.8 as the default python and pip versions
+RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.8 1
+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 app
+RUN git clone --recursive https://github.com/jnjaby/KEEP.git .
 
-# Set the working directory to the app directory
-WORKDIR $HOME/app
+# Copy the app.py script into the container
+COPY app.py .
+COPY requirements_HF.txt .
 
-# Copy the app.py script and requirements file into the container
-COPY --chown=user:user app.py .
-COPY --chown=user:user requirements_HF.txt .
+# Install Python dependencies from requirements.txt
+RUN pip install --upgrade pip
+RUN pip install -r requirements_HF.txt
+RUN pip install gradio
+RUN pip install cupy-wheel
 
-# Debug: List contents of the app directory and its subdirectories
-RUN echo "Contents of $HOME/app:" && \
-    ls -R $HOME/app
+USER root
+# Install basicsr (assuming setup.py is in basicsr directory)
+RUN python basicsr/setup.py develop
 
-# Install Python dependencies
-RUN pip install --upgrade pip && \
-    pip install -r requirements_HF.txt && \
-    pip install gradio && \
-    pip install "numpy<1.25,>=1.18" && \
-    pip install cupy-cuda11x && \
-    pip install ffmpeg-python dlib
+USER user
+# Install additional Python packages
+RUN pip install ffmpeg-python
+RUN pip install dlib
 
-# Install basicsr
-RUN cd basicsr && \
-    if [ -f "VERSION" ]; then \
-        echo "VERSION file found at basicsr/VERSION"; \
-        cat VERSION; \
-        # Create the VERSION file in the location expected by setup.py
-        mkdir -p basicsr && cp VERSION basicsr/VERSION; \
-    else \
-        echo "VERSION file not found, this is unexpected"; \
-        exit 1; \
-    fi && \
-    echo "Current working directory: $(pwd)" && \
-    echo "Contents of current directory:" && \
-    ls -la && \
-    echo "Contents of basicsr subdirectory:" && \
-    ls -la basicsr && \
-    echo "Running setup.py with verbose output" && \
-    python setup.py develop -v
+# 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"]