Add basic support for mps (Apple Silicon) as an accelerator

Dockerfile

@@ -0,0 +1,32 @@
+FROM pytorch/pytorch:2.7.0-cuda12.6-cudnn9-runtime
+
+LABEL author="Colby T. Ford <colby@tuple.xyz>"
+
+## Environment Settings
+ENV DEBIAN_FRONTEND=noninteractive
+
+## Install Basic Dependencies
+RUN apt-get clean && \
+    apt-get update && \
+    apt-get -y install \
+        sudo \
+        git \
+        curl \
+        wget \
+        g++
+
+## Install Boltz (+ dependencies)
+# RUN pip install boltz==2.1.1 -U
+RUN pip install boltz[cuda]==2.2.1 -U
+
+## Download CCD data and model weights 
+RUN mkdir ~/.boltz && \
+    cd ~/.boltz && \
+    wget https://huggingface.co/boltz-community/boltz-2/resolve/main/mols.tar && \
+    tar -xvf mols.tar && \
+    rm mols.tar && \
+    wget https://model-gateway.boltz.bio/boltz2_conf.ckpt && \
+    wget https://model-gateway.boltz.bio/boltz2_aff.ckpt
+
+
+CMD ["boltz"]

README.md

@@ -72,6 +72,18 @@ To encourage reproducibility and facilitate comparison with other models, on top
 If you're interested in retraining the model, currently for Boltz-1 but soon for Boltz-2, see our [training instructions](docs/training.md).
 
 
+## Docker
+
+Build locally:
+```bash
+docker build --platform linux/amd64 -t boltz .
+```
+
+Run locally with GPUs and mounted volume:
+```bash
+docker run --gpus all -v ./:/mnt/ --name boltz --rm -it boltz /bin/bash
+```
+
 ## Contributing
 
 We welcome external contributions and are eager to engage with the community. Connect with us on our [Slack channel](https://boltz.bio/join-slack) to discuss advancements, share insights, and foster collaboration around Boltz-2.

docs/prediction.md

@@ -145,7 +145,7 @@ Examples of common options include:
 | `--cache`                | `PATH`          | `~/.boltz`                  | The directory where to download the data and model. Will use environment variable `BOLTZ_CACHE` as an absolute path if set                                                          |
 | `--checkpoint`           | `PATH`          | None                        | An optional checkpoint. Uses the provided Boltz-2 model by default.                                                                                                                 |
 | `--devices`              | `INTEGER`       | `1`                         | The number of devices to use for prediction.                                                                                                                                        |
-| `--accelerator`          | `[gpu,cpu,tpu]` | `gpu`                       | The accelerator to use for prediction.                                                                                                                                              |
+| `--accelerator`          | `[gpu,cpu,tpu,mps]` | `gpu`                       | The accelerator to use for prediction.                                                                                                                                              |
 | `--recycling_steps`      | `INTEGER`       | `3`                         | The number of recycling steps to use for prediction.                                                                                                                                |
 | `--sampling_steps`       | `INTEGER`       | `200`                       | The number of sampling steps to use for prediction.                                                                                                                                 |
 | `--diffusion_samples`    | `INTEGER`       | `1`                         | The number of diffusion samples to use for prediction.                                                                                                                              |

src/boltz/main.py

@@ -845,7 +845,7 @@ def cli() -> None:
 )
 @click.option(
     "--accelerator",
-    type=click.Choice(["gpu", "cpu", "tpu"]),
+    type=click.Choice(["gpu", "cpu", "tpu", "mps"]),
     help="The accelerator to use for prediction. Default is gpu.",
     default="gpu",
 )