How to Run ESMFold2 Locally Using Hugging Face with Custom Model Weights

You can run ESMFold2 locally by loading the model from the Hugging Face Hub with ESMFold2Model.from_pretrained() or by loading custom weights directly with torch.load(), then processing inputs through ESMFold2InputBuilder and decoding the results.

Running ESMFold2 from the Biohub/esm repository locally allows you to predict protein structures without relying on cloud APIs. The pipeline supports both the official biohub/ESMFold2 checkpoint from Hugging Face and custom fine-tuned weights. This guide covers the complete implementation path from model loading to structure prediction, including how to handle the Chemical Component Dictionary (CCD) cache for offline environments.

Architecture of the ESMFold2 Inference Pipeline

The local inference workflow consists of three distinct layers implemented in the Biohub/esm codebase. Understanding these layers helps debug custom weight loading and optimize performance for specific hardware.

Model Loading Layer

The ESMFold2Model class handles checkpoint management. When calling from_pretrained("biohub/ESMFold2"), the class downloads the official weights from the Hugging Face Hub and initializes the diffusion-based architecture. For custom weights, you bypass this method and load the state dict directly using torch.load().

Input Preparation Layer

The ESMFold2InputBuilder class, defined in esm/models/esmfold2/processor.py, converts high-level biological inputs into batched tensors. This builder handles:

• Normalization of chain-break symbols
• Grouping of identical sequences
• Construction of StructurePredictionInput objects containing proteins, DNA, RNA, or ligands

The prepare_input() method returns feature tensors and chain metadata required by the diffusion model.

Post-Processing and Decoding

After inference, the builder's decode() method reassembles raw diffusion outputs into a MolecularComplexResult. This process attaches pLDDT, pTM, and ipTM confidence metrics. The conversion relies on CCD conformer data stored in ccd.pkl, which the system lazily downloads from biohub/ESMFold2/ccd.pkl via the logic in esm/models/esmfold2/conformers.py.

Running ESMFold2 with Official Hugging Face Weights

First install the Biohub/esm package:

pip install "esm@git+https://github.com/Biohub/esm.git@main"

Then run the complete inference pipeline:

import torch
from esm.models.esmfold2.processor import ESMFold2InputBuilder
from esm.models.esmfold2.types import StructurePredictionInput, ProteinInput

# 1️⃣ Load the public model from the Hugging Face Hub

model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval()

# 2️⃣ Build the input specification

HHAI_SEQ = (
    "MSHHWGYGKHNGPEHWHKDFPIAKGERQSPVDIDTHTAKYDPSLKPLSVSYDQATSLRILNNGHAFNVEFDD"
    "SQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVL"
    "GIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPPLLECVTWIVLKEP"
    "ISVSSEQVLKFRKLNFNGEGEPEELMVDNWRPAQPLKNRQIKASFK"
)
spi = StructurePredictionInput(sequences=[ProteinInput(id="A", sequence=HHAI_SEQ)])

# 3️⃣ Prepare tensors and run the model

builder = ESMFold2InputBuilder()
features, chain_infos = builder.prepare_input(spi, seed=0, device=model.device)
output = model(**features, num_loops=3, num_sampling_steps=50)

# 4️⃣ Decode the prediction

result = builder.decode(output, features, chain_infos)

print(f"pLDDT mean: {float(result.plddt.mean()):.3f}, pTM: {float(result.ptm):.3f}, ipTM: {float(result.iptm):.3f}")

# Optional: write a CIF file

with open("prediction.cif", "w") as f:
    f.write(result.complex.to_mmcif())

The num_loops and num_sampling_steps parameters control the diffusion sampling process. Higher values increase accuracy but require more compute time.

Loading Custom Model Weights Locally

To use fine-tuned checkpoints instead of the official Hub weights, load the PyTorch state dict directly. Ensure your custom checkpoint matches the ESMFold2Model architecture defined in the source repository.

import torch
from esm.models.esmfold2.processor import ESMFold2InputBuilder
from esm.models.esmfold2.types import StructurePredictionInput, ProteinInput

# Path to your fine-tuned checkpoint

ckpt_path = "/path/to/my_esmfold2.pt"
model = torch.load(ckpt_path, map_location="cpu")
model.eval()
model.cuda()  # Move to GPU if available

# Initialize builder with optional local CCD cache to avoid Hub downloads

builder = ESMFold2InputBuilder(ccd_cache="/my_cache/ccd/")

# Prepare input identical to the official workflow

spi = StructurePredictionInput(sequences=[ProteinInput(id="A", sequence="MALWMRLLPLL...")])
features, chain_infos = builder.prepare_input(spi, device=model.device)

# Run inference with custom sampling parameters

output = model(**features, num_loops=3, num_sampling_steps=200)
result = builder.decode(output, features, chain_infos)

print(f"Result pLDDT: {result.plddt.mean():.3f}")

When loading custom weights, pass the ccd_cache parameter to ESMFold2InputBuilder to point to a pre-downloaded ccd.pkl file. This prevents automatic downloads from Hugging Face and enables fully offline inference.

Key Source Files

Reference these files in the Biohub/esm repository when customizing the pipeline:

• esm/models/esmfold2/processor.py – Contains ESMFold2InputBuilder with prepare_input() and decode() methods
• esm/models/esmfold2/conformers.py – Handles CCD conformer downloads and lookups for ligand handling
• esm/models/esmfold2/types.py – Defines dataclasses for StructurePredictionInput, ProteinInput, and other molecular types
• esm/models/esmfold2/constants.py – Stores residue-to-CCD mappings and chemical constants

Summary

• Load official weights using ESMFold2Model.from_pretrained("biohub/ESMFold2") for standard inference
• Load custom weights using torch.load() after ensuring the state dict matches the official architecture
• Process inputs through ESMFold2InputBuilder.prepare_input() to convert biological sequences to model tensors
• Decode outputs using builder.decode() to obtain MolecularComplexResult with confidence metrics and structural coordinates
• Enable offline mode by specifying ccd_cache when instantiating the input builder

Frequently Asked Questions

What is the difference between ESMFold and ESMFold2?

ESMFold2 is the updated protein structure prediction model shipped with the Biohub/esm repository, featuring improved diffusion-based architecture and support for complexes with DNA, RNA, and ligands. The original ESMFold used a different architecture and had more limited complex prediction capabilities.

Can I run ESMFold2 without an internet connection?

Yes. First run the model once with internet access to cache the CCD conformers, or manually download the ccd.pkl file from biohub/ESMFold2. Then instantiate ESMFold2InputBuilder(ccd_cache="/local/path/") to point to the cached file, preventing automatic downloads from Hugging Face during inference.

How do I fine-tune ESMFold2 on my own dataset?

Fine-tuning requires training the ESMFold2Model architecture on your specific protein structures, then saving the resulting state dict. The checkpoint must maintain the same layer names and dimensions as the original model. Load your trained weights locally using torch.load() as shown in the custom weights example above.

What hardware is required for local ESMFold2 inference?

ESMFold2 requires a CUDA-capable GPU with sufficient VRAM for the diffusion sampling process. While the code supports CPU inference via map_location="cpu", practical use requires GPU acceleration due to the computational demands of the num_sampling_steps loop in the diffusion model.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}