FROM pytorch/pytorch:2.11.0-cuda12.8-cudnn9-runtime

ARG DIFFUSERS_VERSION=0.40.0
ARG TRANSFORMERS_VERSION=5.15.1
ARG ACCELERATE_VERSION=1.14.0
ARG HF_HUB_VERSION=1.28.0

RUN python -m venv --system-site-packages /opt/flux-venv && \
    /opt/flux-venv/bin/pip install --no-cache-dir \
      "diffusers==${DIFFUSERS_VERSION}" \
      "transformers==${TRANSFORMERS_VERSION}" \
      "accelerate==${ACCELERATE_VERSION}" \
      "huggingface-hub==${HF_HUB_VERSION}" \
      sentencepiece protobuf safetensors pillow && \
    useradd --system --uid 10002 --home /nonexistent --shell /usr/sbin/nologin flux

COPY flux_worker.py /app/flux_worker.py
USER 10002:10002
ENTRYPOINT ["/opt/flux-venv/bin/python", "/app/flux_worker.py"]