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

ARG HYPIR_COMMIT=b61d107c6cef38f01a93c7833558869731cfa8c1

RUN apt-get update && apt-get install -y --no-install-recommends git ca-certificates python3.12-venv && \
    git clone https://github.com/XPixelGroup/HYPIR.git /opt/HYPIR && \
    cd /opt/HYPIR && git checkout "${HYPIR_COMMIT}" && \
    rm -rf /opt/HYPIR/.git /var/lib/apt/lists/* && \
    python -m venv --system-site-packages /opt/restore-venv && \
    /opt/restore-venv/bin/pip install --no-cache-dir \
      accelerate==1.4.0 diffusers==0.32.2 peft==0.14.0 \
      transformers==4.49.0 einops==0.8.1 omegaconf==2.3.0 \
      opencv-python-headless==4.11.0.86 safetensors pillow && \
    useradd --system --uid 10002 --home /nonexistent --shell /usr/sbin/nologin restoration-worker

COPY restoration_worker.py /app/restoration_worker.py
USER 10002:10002
WORKDIR /opt/HYPIR
ENV HF_HOME=/tmp/huggingface
ENTRYPOINT ["/opt/restore-venv/bin/python", "/app/restoration_worker.py"]
