Update Dockerfile

Dockerfile

@@ -1,73 +1,75 @@
-# Use the NVIDIA CUDA runtime as a base image
+# Use the NVIDIA CUDA runtime as a base image (MATCHING PyTorch)
 FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04
 
 # Set environment variables
 ENV DEBIAN_FRONTEND=noninteractive
 
+# Create a non-root user
 RUN useradd -m -u 1000 user
 
-USER user
-
-ENV HOME=/home/user \
-	PATH=/home/user/.local/bin:$PATH \
-    PYTHONPATH=$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 \
+    ninja-build \
     build-essential \
     libgl1-mesa-glx \
     libglib2.0-0 \
     ffmpeg \
-    python3.8 \
+    python3.9 \
     python3-pip \
-    python3.8-dev \
+    python3.9-dev \
+    python3.9-venv \
     && 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
+# Set Python 3.9 as the default python and pip versions
+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
 
+# 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} \
+    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
+
+# Upgrade pip and install numpy<2 first, then PyTorch (with correct CUDA version)
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir "numpy<2" && \
+    pip install --no-cache-dir torch==1.10.0 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113 && \
+    python -c "import torch; print('PyTorch version:', torch.__version__); print('CUDA Available:', torch.cuda.is_available()); print('Device:', torch.cuda.get_device_name(0))"
+
+# 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
-RUN pip install https://files.pythonhosted.org/packages/40/80/f4b6a9c656daf97f1c4acfd4ae1449f4c7c4e34f1092c8ba10f29b4bc7a8/cupy_cuda11x-11.0.0-cp38-cp38-manylinux1_x86_64.whl
-
-USER root
-# Install basicsr (assuming setup.py is in basicsr directory)
-RUN python basicsr/setup.py develop
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements_HF.txt
+RUN pip install --no-cache-dir gradio ffmpeg-python dlib-
+RUN pip install basicsr
 
-USER user
-# Install additional Python packages
-RUN pip install ffmpeg-python
-RUN pip install dlib-bin
-
-# 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"]