Hugging Face's logo

Transformers documentation

ESMFold2

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.15.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

This model was contributed to Hugging Face Transformers on 2026-08-19.

ESMFold2

Overview

ESMFold2 is an all-atom protein structure prediction model. It predicts 3D coordinates and per-residue confidence (pLDDT, PAE, PDE) directly from an amino-acid sequence, using the ESMC protein language model as its backbone. The architecture combines a sliding-window atom encoder with 3D rotary position embeddings, a pairwise folding trunk applied iteratively, a diffusion-based structure head, and a confidence head.

The model checkpoints are available on the Hugging Face Hub at biohub/ESMFold2 and biohub/ESMFold2-Fast

Usage example

import torch

from transformers import EsmFold2Model

# The ESMC backbone is bundled in the checkpoint and loaded with the model.
# bf16 is the recommended inference precision.
model = EsmFold2Model.from_pretrained("biohub/ESMFold2", dtype=torch.bfloat16, device_map="auto")

pdb_string = model.infer_protein_as_pdb("MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQ")
print(pdb_string)

infer_protein returns the raw outputs (atom coordinates, distogram logits and confidence metrics) as an EsmFold2Output if you need them instead of a PDB string. You may get slightly different predictions if you run the same sequence multiple times. Set a manual seed if you want exactly reproducible structures.

ESMFold2 draws config.structure_head.num_diffusion_samples structures per fold. infer_protein_as_pdb renders the best-ranked one (highest pTM); pass sample_idx to pick a specific sample instead. The PDB carries per-residue pLDDT in the b-factor column, on the same 0-1 scale as the plddt output.

forward vs fold

A structure prediction has two halves. EsmFold2Model.forward is the first: it runs the folding trunk over the featurized inputs and returns the refined pair representation plus the distogram, as an EsmFold2TrunkOutput. It does not produce 3D coordinates — ESMFold2 gets those by iterative denoising, and that sampling loop (the noise schedule, Kabsch alignment and the ODE/SDE update) lives in EsmFold2FoldingMixin along with the confidence head call:

MethodUse it for
infer_protein_as_pdb(sequence)a PDB string, straight from an amino-acid sequence
infer_protein(sequence)the raw EsmFold2Output
fold(**features)pre-featurized inputs (what infer_protein calls)
forward(**features)the trunk alone — a distogram and pair representation, no sampling

Call fold or infer_protein for an actual structure. Reach for forward when you only need the distogram, or when you want to drive the diffusion sampler yourself: fold calls forward once and then hands its output to EsmFold2DiffusionModule, whose own forward is the single denoising step.

Faster inference with a fused kernel

The folding trunk’s dominant cost is the triangle-multiplication update. Passing use_kernels=True to from_pretrained() swaps it for a fused Triton kernel loaded from the Hub via the kernels library, leaving the prediction unchanged. It is inference-only and CUDA-only; on CPU or without the kernel installed the model transparently falls back to the pure-PyTorch implementation. Make sure the model is on a CUDA device when kernelization happens (e.g. with device_map).

import torch

from transformers import EsmFold2Model

model = EsmFold2Model.from_pretrained(
    "biohub/ESMFold2", dtype=torch.bfloat16, device_map="cuda", use_kernels=True
)

pdb_string = model.infer_protein_as_pdb("MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQ")

EsmFold2Config

class transformers.EsmFold2Config

< >

( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: typing.Union[str, ForwardRef('torch.dtype'), NoneType] = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: typing.Optional[typing.Literal['regression', 'single_label_classification', 'multi_label_classification']] = Nonehidden_size: int | None = 384pairwise_hidden_size: int | None = 256single_inputs_size: int | None = 451pair_transition_intermediate_size: int | None = 1024sliding_window: int | None = 128chunk_size: int | None = 64num_relative_residx_bins: int | None = 32num_relative_chain_bins: int | None = 2num_loops: int | None = 10num_res_types: int | None = 33max_atomic_number: int | None = 128char_vocab_size: int | None = 64max_chars: int | None = 4max_atoms_per_token: int | None = 23atom_feature_dim: int | None = Nonefolding_trunk_num_hidden_layers: int | None = 24parcae_num_coda_layers: int | None = 2atom_encoder: dict | transformers.models.esmfold2.configuration_esmfold2.EsmFold2AtomEncoderConfig | None = Nonestructure_head: dict | transformers.models.esmfold2.configuration_esmfold2.EsmFold2StructureHeadConfig | None = Noneconfidence_head: dict | transformers.models.esmfold2.configuration_esmfold2.EsmFold2ConfidenceHeadConfig | None = Nonemsa_encoder: dict | transformers.models.esmfold2.configuration_esmfold2.EsmFold2MsaEncoderConfig | None = Nonelm_encoder: dict | transformers.models.esmfold2.configuration_esmfold2.EsmFold2LmEncoderConfig | None = Noneesmc_config: dict | transformers.models.esmc.configuration_esmc.EsmcConfig | None = None )

Parameters

  • hidden_size (int, optional, defaults to 384) — Single-representation width.
  • pairwise_hidden_size (int, optional, defaults to 256) — Pair-representation width.
  • single_inputs_size (int, optional, defaults to 451) — Width of the concatenated single-input features fed to the trunk and diffusion conditioning: atom_encoder.output_dim plus two residue-type one-hots and the profile scalar.
  • pair_transition_intermediate_size (int, optional, defaults to 1024) — SwiGLU width of the pair-stream transitions.
  • sliding_window (int, optional, defaults to 128) — Sliding-window size (token-index distance) for the atom-stack attention, as the total window width rather than the radius.
  • chunk_size (int, optional, defaults to 64) — Chunk size for the memory-heavy pair-/MSA-stream ops. None disables chunking.
  • num_relative_residx_bins (int, optional, defaults to 32) — Number of relative residue-index bins in the relative-position encoding.
  • num_relative_chain_bins (int, optional, defaults to 2) — Number of relative chain-index bins in the relative-position encoding.
  • num_loops (int, optional, defaults to 10) — Number of trunk refinement loops.
  • num_res_types (int, optional, defaults to 33) — Number of residue types.
  • max_atomic_number (int, optional, defaults to 128) — Size of the element one-hot in the atom features.
  • char_vocab_size (int, optional, defaults to 64) — Character-vocabulary size for the encoded atom names.
  • max_chars (int, optional, defaults to 4) — Number of characters per encoded atom name.
  • max_atoms_per_token (int, optional, defaults to 23) — Maximum number of atoms per token.
  • atom_feature_dim (int, optional) — Atom feature width: xyz, charge and mask, plus the element and atom-name-char one-hots. Derived from max_atomic_number, char_vocab_size and max_chars if unset.
  • folding_trunk_num_hidden_layers (int, optional, defaults to 24) — Number of pair-update blocks in the folding trunk.
  • parcae_num_coda_layers (int, optional, defaults to 2) — Number of pair-update blocks in the parcae coda.
  • atom_encoder (EsmFold2AtomEncoderConfig, optional) — Configuration for the inputs-embedder SWA atom encoder.
  • structure_head (EsmFold2StructureHeadConfig, optional) — Configuration for the diffusion structure-prediction head (holds the diffusion_module sub-config).
  • confidence_head (EsmFold2ConfidenceHeadConfig, optional) — Configuration for the confidence head.
  • msa_encoder (EsmFold2MsaEncoderConfig, optional) — Configuration for the MSA encoder.
  • lm_encoder (EsmFold2LmEncoderConfig, optional) — Configuration for the language-model hidden-state encoder.
  • esmc_config (EsmcConfig, optional) — Configuration for the bundled ESMC backbone.

This is the configuration class to store the configuration of a Esmfold2Model. It is used to instantiate a Esmfold2 model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the biohub/ESMFold2

Configuration objects inherit from PreTrainedConfig and can be used to control the model outputs. Read the documentation from PreTrainedConfig for more information.

validate_architecture

< >

( )

Checks the width relations that span sub-configs; each sub-config checks its own.

EsmFold2PreTrainedModel

class transformers.EsmFold2PreTrainedModel

< >

( config: PreTrainedConfig*inputs**kwargs )

Parameters

  • config (PreTrainedConfig) — Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the from_pretrained() method to load the model weights.

This model inherits from PreTrainedModel. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.)

This model is also a PyTorch torch.nn.Module subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior.

EsmFold2Model

class transformers.EsmFold2Model

< >

( config: EsmFold2Config )

Parameters

  • config (EsmFold2Config) — Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the from_pretrained() method to load the model weights.

ESMFold2 all-atom protein structure predictor with a bundled ESMC protein-language-model backbone. This is the standard released ESMFold2 architecture, whose trunk is a linear-recurrent stack (internally referred to as “parcae”).

This model inherits from PreTrainedModel. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.)

This model is also a PyTorch torch.nn.Module subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior.

forward

< >

( token_index: Tensorresidue_index: Tensorasym_id: Tensorsym_id: Tensorentity_id: Tensormol_type: Tensorres_type: Tensortoken_bonds: Tensorattention_mask: Tensoratom_inputs: EsmFold2AtomInputsdeletion_mean: typing.Optional[torch.Tensor] = Nonemsa: typing.Optional[torch.Tensor] = Nonehas_deletion: typing.Optional[torch.Tensor] = Nonedeletion_value: typing.Optional[torch.Tensor] = Nonemsa_attention_mask: typing.Optional[torch.Tensor] = Noneinput_ids: typing.Optional[torch.Tensor] = Nonelm_hidden_states: typing.Optional[torch.Tensor] = Nonenum_loops: int | None = None**kwargs: Unpack ) → [~models.esmfold2.modeling_esmfold2.EsmFold2TrunkOutput] or tuple(torch.FloatTensor)

Parameters

  • token_index (torch.Tensor of shape (batch_size, num_tokens)) — Per-token positional index within the full complex; feeds the relative-position encoding.
  • residue_index (torch.Tensor of shape (batch_size, num_tokens)) — Residue index within each chain; feeds the relative-position encoding.
  • asym_id (torch.Tensor of shape (batch_size, num_tokens)) — Asymmetric-unit (chain) ID for each token.
  • sym_id (torch.Tensor of shape (batch_size, num_tokens)) — Symmetry-copy ID distinguishing identical chains of a homomer.
  • entity_id (torch.Tensor of shape (batch_size, num_tokens)) — Entity ID grouping tokens that belong to the same molecular entity.
  • mol_type (torch.Tensor of shape (batch_size, num_tokens)) — Molecule-type code for each token (0 = protein).
  • res_type (torch.Tensor of shape (batch_size, num_tokens)) — Residue-type (amino-acid identity) index for each token.
  • token_bonds (torch.Tensor of shape (batch_size, num_tokens, num_tokens, 1)) — Pairwise inter-token covalent-bond feature.
  • attention_mask (torch.Tensor of shape (batch_size, num_tokens)) — Mask marking valid tokens (1) versus padding (0). Inputs must be right-padded.
  • atom_inputs (~models.esmfold2.modeling_esmfold2.EsmFold2AtomInputs) — The raw reference-conformer atom tensors, bundled: positions, atomic numbers, formal charges, encoded atom names, per-atom group IDs, the atom padding mask and the atom->token map.
  • deletion_mean (torch.Tensor of shape (batch_size, num_tokens), optional) — Mean MSA deletion count per column. Defaults to zeros (no MSA).
  • msa (torch.Tensor of shape (batch_size, msa_depth, num_tokens), optional) — MSA residue-type tokens (row 0 is the query sequence). Defaults to a single-sequence MSA.
  • has_deletion (torch.Tensor of shape (batch_size, msa_depth, num_tokens), optional) — Boolean flag marking MSA positions preceded by a deletion.
  • deletion_value (torch.Tensor of shape (batch_size, msa_depth, num_tokens), optional) — Per-position MSA deletion counts.
  • msa_attention_mask (torch.Tensor of shape (batch_size, msa_depth, num_tokens), optional) — Validity mask for the MSA rows/columns.
  • input_ids (torch.Tensor of shape (batch_size, num_tokens), optional) — ESMC-vocabulary token ids for the sequence. Fed to the bundled ESMC backbone to produce lm_hidden_states when those are not passed directly; ignored when lm_hidden_states is given.
  • lm_hidden_states (torch.Tensor of shape (batch_size, num_tokens, num_esmc_layers + 1, esmc_hidden_size), optional) — Precomputed ESMC backbone hidden states, one per backbone layer plus the embeddings. When provided, the backbone is not run and input_ids is unused.
  • num_loops (int, optional) — Number of trunk refinement loops. Defaults to config.num_loops.

Returns

[~models.esmfold2.modeling_esmfold2.EsmFold2TrunkOutput] or tuple(torch.FloatTensor)

A [~models.esmfold2.modeling_esmfold2.EsmFold2TrunkOutput] or a tuple of torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various elements depending on the configuration ([EsmFold2Config]) and inputs.

Run the folding trunk: featurize the inputs, embed them into a pair representation, refine it over num_loops recycling iterations, and read off the distogram. This is the deterministic half of a structure prediction; the diffusion sampler that turns the returned pair representation into 3D coordinates lives in EsmFold2FoldingMixin — call [~EsmFold2Model.fold] or [~EsmFold2Model.infer_protein] for an end-to-end prediction.

  • distogram_logits (torch.FloatTensor of shape (batch_size, num_tokens, num_tokens, distogram_bins)) — Predicted distance-distribution logits over residue pairs.
  • pair_states (torch.FloatTensor of shape (batch_size, num_tokens, num_tokens, pairwise_hidden_size)) — The trunk’s final pair representation, in fp32.
  • single_inputs (torch.FloatTensor of shape (batch_size, num_tokens, single_inputs_size)) — Concatenated single-input features built by the inputs embedder.
  • relative_position_encoding (torch.FloatTensor of shape (batch_size, num_tokens, num_tokens, pairwise_hidden_size)) — Relative-position pair encoding.
  • token_bonds_encoding (torch.FloatTensor of shape (batch_size, num_tokens, num_tokens, pairwise_hidden_size)) — Embedded inter-token covalent-bond feature.
  • atom_inputs (~models.esmfold2.modeling_esmfold2.EsmFold2AtomInputs, optional) — The featurized reference-conformer atom inputs, reused by the diffusion atom stack.

fold

< >

( attention_mask: Tensorasym_id: Tensormol_type: Tensordistogram_atom_idx: Tensornum_diffusion_samples: int | None = Nonenum_sampling_steps: int | None = None**trunk_kwargs )

Predict a structure end-to-end from featurized inputs: run the trunk ([EsmFold2Model.forward], which documents the feature arguments), sample coordinates from the diffusion structure head, and score them with the confidence head.

Only the arguments the sampler and the confidence head need are named here; the rest are forwarded to the trunk untouched, so fold takes the same arguments as forward.

attention_mask (torch.Tensor of shape (batch_size, num_tokens)): Mask marking valid tokens (1) versus padding (0). Also forwarded to the trunk. asym_id (torch.Tensor of shape (batch_size, num_tokens)): Asymmetric-unit (chain) ID for each token. Also forwarded to the trunk. mol_type (torch.Tensor of shape (batch_size, num_tokens)): Molecule-type code for each token (0 = protein). Also forwarded to the trunk. distogram_atom_idx (torch.Tensor of shape (batch_size, num_tokens)): Index of the representative atom (Cβ, or Cα for glycine) of each token. Used by the confidence head; the trunk does not need it. num_diffusion_samples (int, optional): Number of parallel structure samples to draw; the confidence head re-runs once per sample. Defaults to config.structure_head.num_diffusion_samples. num_sampling_steps (int, optional): Number of diffusion sampling steps. Defaults to config.structure_head.inference_num_steps. trunk_kwargs: The remaining featurized inputs (num_loops included, plus the raw per-atom tensors bundled below into atom_inputs), forwarded to [EsmFold2Model.forward], which documents them.

infer_protein

< >

( sequence: str**forward_kwargs )

Featurize sequence and fold it. forward_kwargs go to [~EsmFold2Model.fold].

infer_protein_as_pdb

< >

( sequence: strsample_idx: int | None = None**forward_kwargs )

Fold sequence and render the prediction as a PDB string.

sample_idx picks which diffusion sample to render; by default the best-ranked one.

EsmFold2Output

class transformers.models.esmfold2.modeling_esmfold2.EsmFold2Output

< >

( distogram_logits: typing.Optional[torch.Tensor] = Nonesample_atom_coords: typing.Optional[torch.Tensor] = Noneplddt_logits: typing.Optional[torch.Tensor] = Noneplddt: typing.Optional[torch.Tensor] = Noneplddt_per_atom: typing.Optional[torch.Tensor] = Noneplddt_ca: typing.Optional[torch.Tensor] = Nonecomplex_plddt: typing.Optional[torch.Tensor] = Nonecomplex_iplddt: typing.Optional[torch.Tensor] = Nonepae_logits: typing.Optional[torch.Tensor] = Nonepae: typing.Optional[torch.Tensor] = Nonepde_logits: typing.Optional[torch.Tensor] = Nonepde: typing.Optional[torch.Tensor] = Noneresolved_logits: typing.Optional[torch.Tensor] = Noneptm: typing.Optional[torch.Tensor] = Noneiptm: typing.Optional[torch.Tensor] = Nonepair_chains_iptm: typing.Optional[torch.Tensor] = None )

Parameters

  • distogram_logits (torch.FloatTensor of shape (batch_size, num_tokens, num_tokens, distogram_bins)) — Predicted distance-distribution logits over residue pairs (RNG-independent; no diffusion sampling).
  • sample_atom_coords (torch.FloatTensor of shape (num_diffusion_samples, num_atoms, 3)) — Predicted all-atom Cartesian coordinates for each diffusion sample.
  • plddt_logits (torch.FloatTensor of shape (num_diffusion_samples, num_atoms, num_plddt_bins)) — Per-atom pLDDT bin logits.
  • plddt (torch.FloatTensor of shape (num_diffusion_samples, num_tokens)) — Per-residue predicted lDDT confidence.
  • plddt_per_atom (torch.FloatTensor of shape (num_diffusion_samples, num_atoms)) — Per-atom predicted lDDT confidence.
  • plddt_ca (torch.FloatTensor of shape (num_diffusion_samples, num_tokens)) — Predicted lDDT at the representative (Cα) atom of each token.
  • complex_plddt (torch.FloatTensor of shape (num_diffusion_samples,)) — Mean pLDDT over all atoms of the complex.
  • complex_iplddt (torch.FloatTensor of shape (num_diffusion_samples,)) — Interface-weighted complex pLDDT.
  • pae_logits (torch.FloatTensor of shape (num_diffusion_samples, num_tokens, num_tokens, num_pae_bins)) — Predicted-aligned-error bin logits.
  • pae (torch.FloatTensor of shape (num_diffusion_samples, num_tokens, num_tokens)) — Expected predicted aligned error (Å) for each residue pair.
  • pde_logits (torch.FloatTensor of shape (num_diffusion_samples, num_tokens, num_tokens, num_pde_bins)) — Predicted-distance-error bin logits.
  • pde (torch.FloatTensor of shape (num_diffusion_samples, num_tokens, num_tokens)) — Expected predicted distance error (Å) for each residue pair.
  • resolved_logits (torch.FloatTensor of shape (num_diffusion_samples, num_atoms, 2)) — Per-atom resolved/unresolved logits.
  • ptm (torch.FloatTensor of shape (num_diffusion_samples,)) — Predicted TM-score for each sample.
  • iptm (torch.FloatTensor of shape (num_diffusion_samples,)) — Predicted interface TM-score for each sample.
  • pair_chains_iptm (torch.FloatTensor of shape (num_diffusion_samples, num_chains, num_chains)) — Predicted interface TM-score for each ordered chain pair.

Output of EsmFold2Model. All confidence scores are on a 0-1 scale; per-sample tensors have a leading num_diffusion_samples axis.

EsmFold2TrunkOutput

class transformers.models.esmfold2.modeling_esmfold2.EsmFold2TrunkOutput

< >

( distogram_logits: typing.Optional[torch.Tensor] = Nonepair_states: typing.Optional[torch.Tensor] = Nonesingle_inputs: typing.Optional[torch.Tensor] = Nonerelative_position_encoding: typing.Optional[torch.Tensor] = Nonetoken_bonds_encoding: typing.Optional[torch.Tensor] = Noneatom_inputs: transformers.models.esmfold2.modeling_esmfold2.EsmFold2AtomInputs | None = None )

Parameters

  • distogram_logits (torch.FloatTensor of shape (batch_size, num_tokens, num_tokens, distogram_bins)) — Predicted distance-distribution logits over residue pairs.
  • pair_states (torch.FloatTensor of shape (batch_size, num_tokens, num_tokens, pairwise_hidden_size)) — The trunk’s final pair representation, in fp32.
  • single_inputs (torch.FloatTensor of shape (batch_size, num_tokens, single_inputs_size)) — Concatenated single-input features built by the inputs embedder.
  • relative_position_encoding (torch.FloatTensor of shape (batch_size, num_tokens, num_tokens, pairwise_hidden_size)) — Relative-position pair encoding.
  • token_bonds_encoding (torch.FloatTensor of shape (batch_size, num_tokens, num_tokens, pairwise_hidden_size)) — Embedded inter-token covalent-bond feature.
  • atom_inputs (~models.esmfold2.modeling_esmfold2.EsmFold2AtomInputs, optional) — The featurized reference-conformer atom inputs, reused by the diffusion atom stack.

Output of EsmFold2Model.forward(): the folding trunk’s pair representation, the distogram read off it, and the conditioning tensors that the structure and confidence heads consume. Everything here is deterministic given the inputs apart from the trunk’s random initial pair state.

EsmFold2AtomInputs

class transformers.models.esmfold2.modeling_esmfold2.EsmFold2AtomInputs

< >

( ref_pos: Tensorref_charge: Tensoratom_attention_mask: Tensorref_element: Tensorref_atom_name_chars: Tensorref_space_uid: Tensoratom_to_token: Tensor )

Parameters

  • ref_pos — (batch_size, num_atoms, 3) reference-conformer Cartesian coordinates.
  • ref_charge — (batch_size, num_atoms) formal charge of each atom.
  • atom_attention_mask — (batch_size, num_atoms) valid atoms (1) versus padding (0).
  • ref_element — (batch_size, num_atoms, max_atomic_number) one-hot atomic number.
  • ref_atom_name_chars — (batch_size, num_atoms, max_chars, char_vocab_size) one-hot atom name.
  • ref_space_uid — (batch_size, num_atoms) per-atom group ID (the atom’s token index).
  • atom_to_token — (batch_size, num_atoms) index of the token each atom belongs to.

Featurized reference-conformer atom inputs, bundled so the atom stack takes one argument.

Taken as one argument by [EsmFold2Model.forward], holding the raw featurizer tensors. EsmFold2Model._prepare_features then one-hot encodes the categorical fields and zeroes the padding, and forward swaps those in with dataclasses.replace, so every consumer downstream of that point sees the encoded forms documented below. Returned on [EsmFold2TrunkOutput] so the diffusion sampler and the confidence head reuse the trunk’s featurization rather than redoing it.

A plain dataclass rather than a [~utils.ModelOutput]: every field here is required and is read unconditionally downstream, and ModelOutput permits at most one required field (the rest must default to None), which would turn a missing tensor from a constructor error into a None surfacing deep in the atom encoder.

Update on GitHub

←ESMC EuroBERT→