FROM nvidia/cuda:12.2.2-cudnn8-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get install -y \
    python3-pip \
    python3-dev \
    git \
    build-essential \
    ninja-build \
    && useradd -m user \
    && rm -rf /var/lib/apt/lists/*

USER user
WORKDIR /home/user/app

COPY --chown=user:user . .

RUN pip install --upgrade pip setuptools wheel packaging

RUN pip install torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu124

RUN pip install \
    gradio \
    pillow \
    matplotlib \
    numpy \
    requests \
    peft==0.15.2 \
    accelerate==1.4.0 \
    'transformers @ git+https://github.com/huggingface/transformers@ccf2ca162e33f381e454cdb74bf4b41a51ab976d'

# Install flash-attn matching dev environment
RUN MAX_JOBS=1 pip install flash-attn==2.7.4.post1 --no-build-isolation

EXPOSE 7860

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