FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04

RUN apt-get update && \
    apt-get install -y python3.10 python3.10-dev python3.10-distutils curl git && \
    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 && \
    curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10

RUN python3 -m pip install --upgrade pip

RUN pip install --no-cache-dir torch==2.4.0+cu121 torchvision==0.19.0 --extra-index-url https://download.pytorch.org/whl/cu121
RUN pip install --no-cache-dir \
    "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.3.11/flash_attn-2.8.0+cu121torch2.4-cp310-cp310-linux_x86_64.whl"
RUN pip install --no-cache-dir \
    diffusers transformers accelerate safetensors opencv-python datasets dataclasses \
    scikit-learn huggingface-hub tensorboard peft ftfy jinja2 bitsandbytes

WORKDIR /app
COPY . .

ENV PYTHONUNBUFFERED=1
ENV TQDM_POSITION=-1
ENV TQDM_MININTERVAL=1

ENV OPENBLAS_NUM_THREADS=1
ENV MKL_NUM_THREADS=1
ENV OMP_NUM_THREADS=1

RUN mkdir /.cache && chmod 777 /.cache && chmod 777 /app

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