Dockerfile

# 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

# Install system dependencies as 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.9 \
    python3-pip \
    python3.9-dev \
    python3.9-venv \
    && rm -rf /var/lib/apt/lists/*

# 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

ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6"
ENV TCNN_CUDA_ARCHITECTURES=86;80;75;70;61;60
ENV FORCE_CUDA=1

ENV CUDA_HOME=/usr/local/cuda
ENV PATH=${CUDA_HOME}/bin:/home/${USER_NAME}/.local/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
ENV LIBRARY_PATH=${CUDA_HOME}/lib64/stubs:${LIBRARY_PATH}

# Set home to the user's home directory
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 \
	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 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 application files
COPY app.py .
COPY requirements_HF.txt .

# Install Python dependencies
RUN pip install --no-cache-dir -r requirements_HF.txt
RUN pip install --no-cache-dir gradio ffmpeg-python dlib-bin
RUN pip install basicsr

# 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"]