Update Dockerfile

Dockerfile

@@ -1,4 +1,4 @@
-# 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
@@ -37,8 +37,7 @@ 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} \
+    LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH \
     PYTHONUNBUFFERED=1 \
     GRADIO_ALLOW_FLAGGING=never \
     GRADIO_NUM_PORTS=1 \
@@ -50,11 +49,11 @@ ENV HOME=/home/user \
 # Set working directory
 WORKDIR $HOME/app
 
-# Upgrade pip and install numpy<2 first, then PyTorch
+# 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.8.0 torchvision --extra-index-url https://download.pytorch.org/whl/cu113 && \
-    python -c "import torch; print('PyTorch version:', torch.__version__)"
+    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 .
@@ -65,23 +64,12 @@ COPY requirements_HF.txt .
 
 # Install Python dependencies
 RUN pip install --no-cache-dir -r requirements_HF.txt
-RUN pip install --no-cache-dir gradio
-
-# Install Cython (needed for some dependencies)
-#RUN pip install --no-cache-dir cython==0.29.36
-
-# Install basicsr with selective warnings
-#RUN BASICSR_EXT=True CFLAGS="-Wno-deprecated-declarations" python basicsr/setup.py develop
-
+RUN pip install --no-cache-dir gradio ffmpeg-python dlib-
 RUN pip install basicsr
-#RUN pip install gfpgan --use-pep517
-
-# Install additional packages
-RUN pip install --no-cache-dir ffmpeg-python dlib-bin
 
 # Set CUDA device settings
 ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
 ENV CUDA_VISIBLE_DEVICES=0
 
 # Command to run the application
-CMD ["python", "app.py"]
+CMD ["python", "app.py"]