Files

24 lines
958 B
Docker

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 apt-get update && apt-get install -y --no-install-recommends python3.12-venv && \
rm -rf /var/lib/apt/lists/* && \
python -m venv --system-site-packages /opt/image-venv && \
/opt/image-venv/bin/pip install --no-cache-dir \
"diffusers==${DIFFUSERS_VERSION}" \
"transformers==${TRANSFORMERS_VERSION}" \
"accelerate==${ACCELERATE_VERSION}" \
"huggingface-hub==${HF_HUB_VERSION}" \
"nvidia-modelopt==0.46.0" bitsandbytes \
sentencepiece protobuf safetensors pillow && \
useradd --system --uid 10002 --home /nonexistent --shell /usr/sbin/nologin image-worker
COPY image_worker.py /app/image_worker.py
COPY image_worker_9b.py /app/image_worker_9b.py
USER 10002:10002
ENTRYPOINT ["/opt/image-venv/bin/python", "/app/image_worker_9b.py"]