FROM nvidia/cuda:11.8.0-devel-ubuntu22.04

RUN apt-get update && apt-get install -y \
    python3.10 python3-pip git wget libgl1-mesa-glx libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

RUN useradd -m -u 1000 user
WORKDIR /app
COPY . /app

ENV TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6"
ENV FORCE_CUDA=1

RUN pip3 install torch==2.3.1 torchvision==0.18.1 \
    --index-url https://download.pytorch.org/whl/cu118
RUN pip3 install -r requirements.txt
RUN pip3 install gradio

RUN cd models/ops && python3 setup.py build install && cd ../..

RUN mkdir -p checkpoints && \
    wget -q -O checkpoints/r2g_res256_ep0849.pth \
    "https://huggingface.co/AGLO-AI/raster2seq-checkpoints/resolve/main/r2g_res256_ep0849.pth"

RUN chown -R user:user /app
USER user

EXPOSE 7860
CMD ["python3", "app.py"]