Design protein binders at scale with ESMFold2 and ESMC

Protein folding was a landmark breakthrough in computational biology. But for many applications, we don’t just want to predict the structures of existing proteins — we want to design new proteins that can modulate biology.

One of the most important ways to do that is through binding. Protein-protein interactions drive much of biological function, and the ability to design molecules that bind specific targets opens the door to new research tools and therapeutics. Recent AI approaches have tackled binder design by inverting structure prediction models via an iterative optimization process:

Fold a candidate binder together with the target protein.
Score the resulting structure based on how well the binder folds and binds.
Take a step in sequence space that improves the score.
Repeat.

In this example, we’ll demonstrate how implement this process on Modal using ESMFold2 and ESMC, state-of-the-art models developed at Biohub that can predict the stucture of biomolecular complexes. Check out their technical report to see how the models were developed and used to design and experimentally validate binders against therapeutically relevant targets.

We’ll start by building a Modal Function that designs a single binder; then with only a few more lines of code, we’ll write an orchestrator function that executes a large-scale search powered by Modal’s autoscaling infrastructure and global GPU capacity.

Setup

from pathlib import Path
from typing import Optional

import modal

MINUTES = 60  # seconds
HOURS = 60 * MINUTES

app = modal.App(
    name="example-esmfold2-binder-design",
)

We’ll use Image.micromamba as our base image because a few of the packages we need are only available via Conda. We’ll also install the esm library from CZ Biohub (which pulls in a custom fork of transformers) and a few other helpful libraries for working with protein sequences.

We set CUBLAS_WORKSPACE_CONFIG which allows us to ensure reproducibility by calling torch.use_deterministic_algorithms(True) at the top of our remote code.

ESM_REVISION = (
    "f652b471d29da828b31e9b7a9cf7d0a7803240f5"  # see https://github.com/Biohub/esm
)

image = (
    modal.Image.micromamba(python_version="3.12")
    .run_commands("apt update && apt install -y git build-essential")
    .micromamba_install(
        "anarci=2024.05.21-0",
        channels=["conda-forge", "bioconda"],
    )
    .pip_install(
        f"esm @ git+https://github.com/Biohub/esm.git@{ESM_REVISION}",
        "abnumber==0.4.4",
        "pyarrow==18.1.0",
    )
    .env(
        {
            "HF_HOME": "/models",
            "HF_XET_HIGH_PERFORMANCE": "1",  # speed up Hugging Face downloads
            "XFORMERS_IGNORE_FLASH_VERSION_CHECK": "1",
            # required for torch.use_deterministic_algorithms(True)
            "CUBLAS_WORKSPACE_CONFIG": ":4096:8",
        }
    )
)

ESMFold2 builds on the 6B-parameter ESMC encoder; together with the four critic models used for final scoring, the model weights come in around ~50 GB. We cache them on a Modal Volume which delivers much better performance at cold-start time than re-downloading from Hugging Face each time.

models_volume = modal.Volume.from_name("esmfold2-models", create_if_missing=True)
models_dir = Path("/models")

A second Volume will store our results.

results_volume = modal.Volume.from_name(
    "esmfold2-binder-design-results", create_if_missing=True
)
results_dir = Path("/results")

To run binder design on Modal, we define a BinderDesignService class and wrap it with the @app.cls decorator. The decorator takes arguments that describe the infrastructure our code needs: the Image and both Volumes we defined, plus an H100 GPU which has enough memory for the 6B-parameter ESMC encoder and the four ESMFold2 “hero” critic models.

Inside the class, the @modal.enter() lifecycle hook downloads and initializes those models once per container start, so subsequent design calls on the same container reuse the loaded weights.

We decorate our design method with @modal.method() to enable remote execution. We’ll see it called both via .remote() (single design) and via .spawn() + modal.FunctionCall.gather (parallel sweep) further below. The class itself is a thin wrapper around ESMFold2Designer from the helper package, which handles the actual model loading and the gradient-guided optimization loop (design_binder in binder_design.design).

@app.cls(
    image=image,
    volumes={models_dir: models_volume},
    gpu="H100",
    timeout=1 * HOURS,
)
class BinderDesignService:
    """Modal entry point for ESMFold2-driven binder design.

    Set ``use_scaling_critics=True`` to also load the 15-checkpoint
    scaling-experiment ensemble (distogram binding confidence only).
    """

    use_scaling_critics: bool = modal.parameter(default=False)

    @modal.enter()
    def load(self):
        from .binder_design import ESMFold2Designer

        self._designer = ESMFold2Designer()
        self._designer.load(self.use_scaling_critics)

    @modal.method()
    def design(
        self,
        target_name: Optional[str] = None,
        target_sequence: Optional[str] = None,
        binder_name: Optional[str] = None,
        binder_sequence: Optional[str] = None,
        is_antibody: Optional[bool] = None,
        seed: int = 0,
        batch_size: int = 1,
    ):
        return self._designer.design(
            target_name=target_name,
            target_sequence=target_sequence,
            binder_name=binder_name,
            binder_sequence=binder_sequence,
            is_antibody=is_antibody,
            seed=seed,
            batch_size=batch_size,
        )

Fanning out a sweep with selection

A single design run gives you one candidate per batch slot. To recover the kind of hit rates reported in the paper, you want many seeds, several binder templates, and several targets, then a selection pass that ranks designs by a combined ipTM / distogram-ipTM-proxy score.

We orchestrate from inside a Modal Function so you don’t have to worry about keeping a long-running process alive locally or installing any local dependencies.

@app.function(
    image=image,
    volumes={results_dir: results_volume},
    gpu="H100",
    timeout=2 * HOURS,
)
def run_sweep(
    line_sweeps: dict[str, list],
    use_scaling_critics: bool = False,
    save_filename: str = "selection.parquet",
) -> bytes:
    """Fan a grid sweep across GPUs, gather results, select top designs, ave results + return parquet."""
    import io

    from .binder_design.sweep import expand_sweep, select_designs

    designer = BinderDesignService(use_scaling_critics=use_scaling_critics)
    configs = expand_sweep(line_sweeps)

    print(f"🧬 spawning {len(configs)} design jobs")
    calls = [designer.design.spawn(**cfg) for cfg in configs]
    raw_results = modal.FunctionCall.gather(*calls)

    df_select = select_designs(configs, raw_results)

    buf = io.BytesIO()
    df_select.to_parquet(buf, index=False)
    parquet_bytes = buf.getvalue()

    save_path = results_dir / save_filename
    save_path.write_bytes(parquet_bytes)
    results_volume.commit()
    print(f"🧬 saved {len(df_select)} selected designs to volume:{save_path}")

    return parquet_bytes

From the command line

main runs a single design. Override the target_name / binder_name to try one of the bundled targets (cd45, ctla4, egfr, pd-l1, pdgfr) and binder templates (minibinder, trastuzumab_framework_vhvl, atezolizumab_framework_vhvl, ocankitug_framework_vhvl), or pass an arbitrary target_sequence / binder_sequence directly.

modal run -m 06_gpu_and_ml.binder-design.esmfold2_binder_design::main \
    --target-name pd-l1 --binder-name minibinder

@app.local_entrypoint()
def main(
    target_name: Optional[str] = "pd-l1",
    target_sequence: Optional[str] = None,
    binder_name: Optional[str] = "minibinder",
    binder_sequence: Optional[str] = None,
    is_antibody: Optional[bool] = None,
    use_scaling_critics: bool = False,
    seed: int = 0,
    batch_size: int = 1,
):
    designer = BinderDesignService(use_scaling_critics=use_scaling_critics)
    seq, trajectory, results = designer.design.remote(
        target_name=target_name,
        target_sequence=target_sequence,
        binder_name=binder_name,
        binder_sequence=binder_sequence,
        is_antibody=is_antibody,
        seed=seed,
        batch_size=batch_size,
    )

    avg_final_loss = sum(r["final_loss"] for r in results) / len(results)
    print(f"🧬 designed sequence: {seq}")
    print(f"🧬 trajectory length: {len(trajectory)} steps")
    print(f"🧬 average final loss: {avg_final_loss:.4f}")

sweep runs a grid sweep across every (target, binder, seed) combination of the targets and binders you pass in, scaling design horizontally with Modal’s asynchronous job processing. The selection pass runs server-side and the resulting parquet is written to both the esmfold2-binder-design-results Volume and to a local file for inspection.

target_names and binder_names are passed as comma-separated strings because Modal’s local-entrypoint CLI doesn’t accept lists directly. The defaults sweep one target across two binder modalities — a minibinder and the trastuzumab_framework_vhvl antibody template — so a single command fans out across both at once:

modal run -m 06_gpu_and_ml.binder-design.esmfold2_binder_design::sweep \
    --target-names pd-l1,ctla4 \
    --binder-names minibinder,trastuzumab_framework_vhvl \
    --n-seeds 8

@app.local_entrypoint()
def sweep(
    target_names: str = "pd-l1",
    binder_names: str = "minibinder,trastuzumab_framework_vhvl",
    use_scaling_critics: bool = False,
    n_seeds: int = 8,
    output_path: Optional[str] = None,
):
    target_name_list = [
        name.strip() for name in target_names.split(",") if name.strip()
    ]
    binder_name_list = [
        name.strip() for name in binder_names.split(",") if name.strip()
    ]

    line_sweeps = {
        "target_name": target_name_list,
        "target_sequence": [None],
        "binder_name": binder_name_list,
        "binder_sequence": [None],
        "seed": list(range(n_seeds)),
        "batch_size": [1],
    }

    print(
        f"🧬 launching sweep: targets={target_name_list}, binders={binder_name_list}, "
        f"n_seeds={n_seeds}, use_scaling_critics={use_scaling_critics}"
    )
    parquet_bytes = run_sweep.remote(
        line_sweeps, use_scaling_critics=use_scaling_critics
    )

    if output_path is None:
        output_path = Path("/tmp") / "esmfold2_binder_design" / "selection.parquet"
    else:
        output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_bytes(parquet_bytes)
    print(f"🧬 wrote selection parquet to {output_path}")

Design protein binders at scale with ESMFold2 and ESMC

Setup

Defining our Modal Image

Caching weights and persisting results on Modal Volumes

Designing a binder on Modal

Fanning out a sweep with selection

From the command line

Try this on Modal!