Update Dockerfile

Dockerfile

@@ -1,34 +1,13 @@
 # Use the NVIDIA CUDA runtime as a base image
 FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04
-#FROM nvidia/cuda:11.8.0-devel-ubuntu22.04
 
 # Set environment variables
 ENV DEBIAN_FRONTEND=noninteractive
 
+# Create a non-root user
 RUN useradd -m -u 1000 user
 
-USER user
-
-# Set environment variables including PATH for user-installed binaries
-ENV HOME=/home/user \
-    CUDA_HOME=/usr/local/cuda \
-    PATH=/home/user/.local/bin:${PATH} \
-    LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} \
-    LIBRARY_PATH=${CUDA_HOME}/lib64/stubs:${LIBRARY_PATH} \
-    PYTHONPATH=/home/user/.local/lib/python3.9/site-packages:$HOME/app \
-    PYTHONUNBUFFERED=1 \
-    GRADIO_ALLOW_FLAGGING=never \
-    GRADIO_NUM_PORTS=1 \
-    GRADIO_SERVER_NAME=0.0.0.0 \
-    GRADIO_THEME=huggingface \
-    GRADIO_SHARE=False \
-    SYSTEM=spaces
-
-# 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 \
@@ -46,41 +25,57 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.9 1
 RUN update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1
 
-# Configure git to trust the directory
-RUN git config --global --add safe.directory /home/user/app
-RUN chown -R user:user /home/user/app
-
+# Switch to non-root user
 USER user
-# Clone the repository (adjust the URL if needed)
+
+# Create virtual environment
+RUN python -m venv /home/user/venv
+ENV VIRTUAL_ENV=/home/user/venv
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+
+# Set environment variables
+ENV HOME=/home/user \
+    CUDA_HOME=/usr/local/cuda \
+    LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH} \
+    LIBRARY_PATH=${CUDA_HOME}/lib64/stubs:${LIBRARY_PATH} \
+    PYTHONUNBUFFERED=1 \
+    GRADIO_ALLOW_FLAGGING=never \
+    GRADIO_NUM_PORTS=1 \
+    GRADIO_SERVER_NAME=0.0.0.0 \
+    GRADIO_THEME=huggingface \
+    GRADIO_SHARE=False \
+    SYSTEM=spaces
+
+# Set working directory
+WORKDIR $HOME/app
+
+# Install PyTorch first
+RUN pip install --no-cache-dir torch==1.7.1 torchvision --extra-index-url https://download.pytorch.org/whl/cu113 && \
+    python -c "import torch; print('PyTorch version:', torch.__version__)"
+
+# Clone the repository
 RUN git clone --recursive https://github.com/jnjaby/KEEP.git .
 
-# Copy the app.py script into the container
+# Copy application files
 COPY app.py .
 COPY 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
-
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements_HF.txt
+RUN pip install --no-cache-dir gradio
 
-# Install torch explicitly first
-RUN pip install torch>=1.7.1 torchvision --extra-index-url https://download.pytorch.org/whl/cu113
-# Verify torch installation
-RUN python -c "import torch; print(torch.__version__)"
+# Install Cython (needed for some dependencies)
+RUN pip install --no-cache-dir cython
 
-# Install basicsr with the --user flag
-RUN python basicsr/setup.py develop --user
+# Install basicsr
+RUN python basicsr/setup.py develop
 
-# Install specific Cython version known to work with cupy
-RUN pip install cython
-# Install additional Python packages
-RUN pip install ffmpeg-python
-RUN pip install dlib
+# Install additional packages
+RUN pip install --no-cache-dir ffmpeg-python dlib
 
-# Set the environment variable to specify the GPU device
+# Set CUDA device settings
 ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
 ENV CUDA_VISIBLE_DEVICES=0
 
-# Command to run your application
+# Command to run the application
 CMD ["python", "app.py"]