Synthyra/Boltz2

NOTE

The GitHub with the implementation and requirements can be found here.

Boltz2 AutoModel (Inference-only)

This is a barebones Huggingface AutoModel compatible implementation of Boltz2 focused on fast inference workflows.

The implementation is located in fastplms/boltz/ and exposes:

• Boltz2Config
• Boltz2Model
• predict_structure(amino_acid_sequence, ...)
• save_as_cif(structure_output, output_path, ...)

Design goals

• Inference-only (no training hooks, no Lightning trainer usage).
• Lightweight runtime around torch + transformers (plus numpy).
• AutoModel remote-code compatibility via trust_remote_code=True.
• Confidence outputs included in prediction outputs (plddt, ptm, iptm, and derived confidence score when available).

Runtime note

This implementation is self-contained inside fastplms/boltz/ and does not require the original cloned boltz package at runtime.

Use with transformers

Load from an exported directory

import torch
from transformers import AutoModel

model = AutoModel.from_pretrained(
    "Synthyra/Boltz2",
    trust_remote_code=True,
    dtype=torch.float32,
).eval()

Predict structure from sequence

out = model.predict_structure(
    amino_acid_sequence="MSTNPKPQRKTKRNTNRRPQDVKFPGG",
    recycling_steps=3,
    num_sampling_steps=200,
    diffusion_samples=1,
)

print(out.sample_atom_coords.shape)
print(None if out.plddt is None else out.plddt.shape)

Save CIF

model.save_as_cif(out, "prediction.cif")

Convert Boltz checkpoint to HF export

Use:

py -m fastplms.boltz.get_weights --checkpoint_path fastplms/boltz/weights/boltz2_conf.ckpt --output_dir boltz2_automodel_export

The export directory contains:

• config.json
• pytorch_model.bin
• modeling_boltz2.py
• minimal_featurizer.py
• minimal_structures.py
• cif_writer.py
• vb_*.py (self-contained vendored Boltz2 inference modules/constants)

Output object fields

predict_structure(...) returns Boltz2StructureOutput with:

• sample_atom_coords
• atom_pad_mask
• plddt
• complex_plddt
• ptm
• iptm
• confidence_score (derived when available)
• raw_output

Limitations

• Current featurization path is protein-only and minimal.
• This implementation is meant for practical inference and export workflows, not full Boltz training parity.

Docker-first compliance testing

Build the container at repo root:

docker build -t fastplms-test -f Dockerfile .

Launch a test shell:

docker run --rm --gpus all -it -v ${PWD}:/workspace fastplms-test bash

Inside the container, run Boltz2 compliance against pip boltz:

python -m testing.run_boltz2_compliance --device cuda --dtype float32 --seed 42 --num-sequences 3 --recycling-steps 3 --num-sampling-steps 200 --diffusion-samples 1 --pass-coord-metric aligned --write-cif-artifacts

Artifacts are written to testing/results/<timestamp>/boltz2_compliance/ by default:

• metrics.json
• metrics.csv
• summary.txt
• structures/seq_<idx>/ours_seq<idx>.cif
• structures/seq_<idx>/ref_seq<idx>.cif

Coordinate metrics now include both raw and rigid-aligned variants:

• coord_mae, coord_rmse, coord_max_abs (raw frame-dependent deltas)
• coord_mae_aligned, coord_rmse_aligned, coord_max_abs_aligned (Kabsch aligned)
• pairwise_dist_mae (frame-invariant pairwise-distance delta)

Pass/fail uses --pass-coord-metric aligned by default. Set --pass-coord-metric raw to use the raw coordinate thresholds.

Citations

@misc{FastPLMs,
  author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.},
  title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.},
  year={2024},
  url={https://huggingface.co/Synthyra/ESMplusplus_small},
  DOI={10.57967/hf/3726},
  publisher={Hugging Face}
}

@article{passaro2025boltz2,
  title={Boltz-2: Exploring the Frontiers of Biomolecular Prediction},
  author={Passaro, Saro and Corso, Gabriele and Wohlwend, Jeremy and Reveiz, Mateo and Bordes, Florian and Wicky, Basile and Dayan, Peter and Jing, Bowen},
  journal={bioRxiv},
  year={2025}
}

@article{wohlwend2024boltz1,
  title={Boltz-1: Democratizing Biomolecular Interaction Modeling},
  author={Wohlwend, Jeremy and Corso, Gabriele and Passaro, Saro and Reveiz, Mateo and Leidal, Ken and Swanson, Wojtek and Kher, Gilmer and Lember, Tommi and Jaakkola, Tommi},
  journal={bioRxiv},
  year={2024}
}

@inproceedings{paszke2019pytorch,
  title={PyTorch: An Imperative Style, High-Performance Deep Learning Library},
  author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith},
  booktitle={Advances in Neural Information Processing Systems 32},
  year={2019}
}