slim_colab_rfdiffusion.def

#    You must also download the weights to pass through as you run the script.
#
#    Example:
#    inference.model_directory_path=./models
#    inference.output_prefix=./outputs/test 
#    inference.schedule_directory_path=$TMPDIR/RFdiffusion 
#    inference.input_pdb=./inputs/5TPN.pdb
#    
#    wget http://files.ipd.uw.edu/pub/RFdiffusion/6f5902ac237024bdd0c176cb93063dc4/Base_ckpt.pt &
#    wget http://files.ipd.uw.edu/pub/RFdiffusion/e29311f6f1bf1af907f9ef9f44b8328b/Complex_base_ckpt.pt &
#    wget http://files.ipd.uw.edu/pub/RFdiffusion/60f09a193fb5e5ccdc4980417708dbab/Complex_Fold_base_ckpt.pt &
#    wget http://files.ipd.uw.edu/pub/RFdiffusion/74f51cfb8b440f50d70878e05361d8f0/InpaintSeq_ckpt.pt &
#    wget http://files.ipd.uw.edu/pub/RFdiffusion/76d00716416567174cdb7ca96e208296/InpaintSeq_Fold_ckpt.pt &
#    wget http://files.ipd.uw.edu/pub/RFdiffusion/5532d2e1f3a4738decd58b19d633b3c3/ActiveSite_ckpt.pt &
#    wget http://files.ipd.uw.edu/pub/RFdiffusion/12fc204edeae5b57713c5ad7dcb97d39/Base_epoch8_ckpt.pt &
#    # Optional:
#    wget http://files.ipd.uw.edu/pub/RFdiffusion/f572d396fae9206628714fb2ce00f72e/Complex_beta_ckpt.pt &
#    # original structure prediction weights
#    wget http://files.ipd.uw.edu/pub/RFdiffusion/1befcb9b28e2f778f53d47f18b7597fa/RF_structure_prediction_weights.pt &
#    wait

Bootstrap: docker
From: nvidia/cuda:11.8.0-base-ubuntu22.04

%post
    apt update && apt install --no-install-recommends -y wget git python3.10 python3-pip 
    apt autoremove -y && apt clean
    python3.10 -m pip install -q -U --no-cache-dir pip
    rm -rf /var/lib/apt/lists/*

    # Clone the RFdiffusion repository
    git clone --depth 1 https://github.com/RosettaCommons/RFdiffusion.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

    apt remove --purge -y wget git && apt autoremove -y && apt clean -y
    rm -rf /var/lib/apt/lists/* /root/.cache
    
%environment
    export DGLBACKEND="pytorch"
    export PYTHONPATH="/app/RFdiffusion/scripts:$PYTHONPATH"

%runscript
    python3.10 /app/RFdiffusion/scripts/run_inference.py "$@"