FROM pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime

WORKDIR /app

# System dependencies
RUN apt-get update && apt-get install -y \
    git wget bc libgl1-mesa-glx libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

# Clone all repos
RUN git clone https://github.com/kuai-lab/iccv25_faceshield.git /app/faceshield_src && \
    git clone https://github.com/choi403/DiffusionGuard.git /app/diffusionguard_src && \
    git clone https://github.com/py85252876/MMVGM.git /app/vgmshield_src

# Python dependencies (unified)
RUN pip install --no-cache-dir \
    accelerate \
    diffusers==0.31.0 \
    transformers==4.46.3 \
    omegaconf \
    pytorch-lightning \
    torch-mtcnn \
    insightface \
    opencv-python-headless \
    Pillow \
    scikit-image \
    scipy \
    einops \
    safetensors \
    tqdm \
    hydra-core \
    lpips \
    fire \
    av \
    open-clip-torch \
    torchmetrics \
    numpy

# Copy framework
COPY modules/ /app/modules/
COPY run_pipeline.py /app/run_pipeline.py

ENV HF_HOME=/root/.cache/huggingface
ENV PYTHONPATH=/app:/app/faceshield_src:/app/diffusionguard_src:/app/vgmshield_src

ENTRYPOINT ["python"]
CMD ["run_pipeline.py", "--help"]