Apptainer

Bootstrap: docker
From: nvcr.io/nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04

# The post section is equivalent to the Docker RUN instruction.
%post
    apt-get -q update
    DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \
        git \
        python3.10 \
        python3-pip
    python3.10 -m pip install -q -U --no-cache-dir pip
    rm -rf /var/lib/apt/lists/*
    apt-get autoremove -y
    apt-get clean
    
    # Clone the RFdiffusion repository
    git clone https://github.com/WEHI-ResearchComputing/RFdiffusion_WEHI.git /app/RFdiffusion

    # Installing the provided versions of pip packages
    pip install -q --no-cache-dir \
        dgl==1.0.2+cu116 -f https://data.dgl.ai/wheels/cu116/repo.html \
        torch==2.0.1 \
        e3nn==0.3.3 \
        wandb==0.12.0 \
        pynvml==11.0.0 \
        git+https://github.com/NVIDIA/dllogger#egg=dllogger \
        decorator==4.4.2 \
        hydra-core==1.3.2 \
        pyrsistent==0.19.3 \
        icecream==2.1.3 \
        jax==0.4.16 \
        dm-haiku==0.0.10 \
        dm-tree==0.1.8 \
        biopython==1.81 \
        ml-collections==0.1.1 \
        tensorflow==2.13.0 \
        /app/RFdiffusion/env/SE3Transformer

    # Install JAXlib using the new installation protocol
    pip install -q --no-cache-dir https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.4.16+cuda11.cudnn86-cp310-cp310-manylinux2014_x86_64.whl

    # Install from provided pip command
    pip install -q --no-cache-dir git+https://github.com/sokrypton/ColabDesign.git

    # Installing RFdiffusion
    pip install --no-cache-dir /app/RFdiffusion --no-deps

# Setting environment variables
%environment
    DGLBACKEND="pytorch"

# Entry point for the container when run.
%runscript
    exec python3.10 /app/RFdiffusion/scripts/run_inference.py "$@"