Bootstrap: docker
From: nvcr.io/nvidia/jax:24.04-py3

%post
    python3 -m pip install tensorflow jax[cuda12]
    python3 -m pip install --no-warn-conflicts 'colabfold[alphafold] @ git+https://github.com/sokrypton/ColabFold'

%environment
    export PATH=${HOME}/.local/bin:$PATH
    unset XLA_FLAGS

%runscript
    conda deactivate >/dev/null 2>&1 || true
    exec /bin/bash -i