FROM nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04

LABEL org.opencontainers.image.source="https://github.com/poly-hammer/poly-hammer-portal"

ENV DEBIAN_FRONTEND=noninteractive \
    PIP_DISABLE_PIP_VERSION_CHECK=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1

RUN apt-get update && apt-get install -y --no-install-recommends \
        python3 python3-venv ca-certificates \
    && rm -rf /var/lib/apt/lists/*

RUN python3 -m venv /opt/ph-install
ENV PATH="/opt/ph-install/bin:$PATH"

COPY . /opt/worker-src
RUN --mount=type=cache,target=/root/.cache/pip \
    /opt/ph-install/bin/python -m pip install --only-binary=:all: pip setuptools wheel \
    && PH_IN_DOCKER=1 \
       PH_PYTHON=/opt/ph-install/bin/python \
       PH_WORKER_SRC=/opt/worker-src \
       bash /opt/worker-src/poly_hammer_worker/jobs/flux2_klein/recipes/linux_install.sh \
           --profile "flux2-klein/linux-x64/docker" --prefix /opt/ph-install \
           --wheelhouse /opt/worker-src/wheelhouse \
           --python-version 3.12 \
           --torch-index-url https://download.pytorch.org/whl/cu128 \
           --torch-version 2.7.1 \
    && rm -rf /opt/worker-src

ENV WEIGHTS_PATH=/weights \
    CUDA_CACHE_PATH=/weights/.nv/ComputeCache \
    PH_FLUX_QUANTIZATION=none \
    PH_FLUX_DTYPE=bfloat16 \
    PH_TEXT_ENCODER=local \
    PH_OFFLOAD=none \
    PH_PROFILE_ID=flux2-klein/linux-x64/docker-cu128 \
    PH_JOB_CLASS=poly_hammer_worker.jobs.flux2_klein.job.Flux2KleinJob

CMD ["ph-worker", "start"]
