FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04

ENV DEBIAN_FRONTEND=noninteractive

# Create 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 \
    software-properties-common \
    git \
    cmake \
    ninja-build \
    build-essential \
    libgl1-mesa-glx \
    libglib2.0-0 \
    ffmpeg

# Add deadsnakes PPA and install Python 3.10 and related packages
RUN add-apt-repository ppa:deadsnakes/ppa && apt-get update && \
    apt-get install -y python3.10 python3.10-dev python3.10-distutils python3.10-venv

# Install curl (and any other missing tools)
RUN apt-get update && apt-get install -y --no-install-recommends curl

# Install pip for Python 3.10 using get-pip.py
RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10

# Set Python 3.10 as the default and update pip alternatives.
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1 && \
    update-alternatives --install /usr/bin/pip pip /usr/local/bin/pip3.10 1

# Switch to non-root user
USER user

# Set CUDA and environment variables
ENV CUDA_HOME=/usr/local/cuda
ENV PATH=${CUDA_HOME}/bin:${PATH}
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6"
ENV FORCE_CUDA=1
ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
ENV CUDA_VISIBLE_DEVICES=0

ENV HOME=/home/user \
    PYTHONPATH=/home/user/app \
    PYTHONUNBUFFERED=1 \
    GRADIO_ALLOW_FLAGGING=never \
    GRADIO_NUM_PORTS=1 \
    GRADIO_SERVER_NAME=0.0.0.0 \
    GRADIO_THEME=huggingface \
    SYSTEM=spaces

WORKDIR $HOME/app

# Upgrade pip and install numpy, then PyTorch
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())"

# Clone repository and install dependencies
RUN git clone --recursive https://github.com/jnjaby/KEEP.git .
COPY app.py .
COPY requirements_HF.txt .

RUN pip install --no-cache-dir -r requirements_HF.txt && \
    pip install --no-cache-dir gradio ffmpeg-python dlib-bin basicsr 

RUN pip install --no-cache-dir diffusers==0.11.0
RUN pip install --no-cache-dir huggingface-hub==0.25.2
    
CMD ["python", "app.py"]