Update FastPLMs runtime and model cards

README.md

@@ -10,6 +10,47 @@ tags:
 
 # ESMFold2
 
+## Quick start
+
+Load the published model, fold two protein chains together, and write an mmCIF
+file. The example omits `num_sampling_steps` and uses the model default.
+
+```python
+from pathlib import Path
+
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "Synthyra/ESMFold2",
+    trust_remote_code=True,
+    dtype=torch.float32,
+    device_map="cuda",
+    esmc_precision="bf16",
+    attn_implementation="sdpa",
+).eval()
+model.set_chunk_size(32)
+
+types = model.input_types
+complex_input = types.StructurePredictionInput(
+    sequences=[
+        types.ProteinInput(id="A", sequence="MSTNPKPQRKTKRNT"),
+        types.ProteinInput(id="B", sequence="MKTIIALSYIFCLVFA"),
+    ]
+)
+with torch.inference_mode():
+    result = model.fold(
+        complex_input,
+        num_loops=3,
+        num_diffusion_samples=1,
+        seed=17,
+        verbose=True,
+    )
+Path("complex.cif").write_text(model.result_to_cif(result), encoding="utf-8")
+```
+
+Set `verbose=False` to silence the folding progress display. This variant has an enabled confidence head and returns confidence fields.
+
 ## Model overview
 
 `Synthyra/ESMFold2` packages the `biohub/ESMFold2` checkpoint with the FastPLMs
@@ -41,30 +82,13 @@ This model requires Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13.
 
 The artifact requirements include the structure dependencies.
 
-The release contract requires a CUDA device. The current validated target is
-the exact NVIDIA GH200 on Linux aarch64. Linux x86-64, CPU-only, Windows, and
-macOS structure runs are not release evidence.
+Validation runs in Docker on any compatible CUDA device. Record the container,
+hardware, precision, and inputs; no GPU product or workstation is required.
 
 The Hub quick start needs network access for the first download. For an
 air-gapped run, build the manifest-pinned local artifact first and use the
 offline example.
 
-## Quick start
-
-```python
-from transformers import AutoModel
-
-model_id = "Synthyra/ESMFold2"
-model = AutoModel.from_pretrained(
-    model_id,
-    trust_remote_code=True,
-    attn_implementation="sdpa",
-).eval()
-```
-
-For offline validation, replace `model_id` with the manifest-built
-`dist/hub/ESMFold2` path. Pass `local_files_only=True`.
-
 ## Attention backends
 
 The quick start uses `sdpa`.
@@ -166,28 +190,6 @@ cif_text = model.result_to_cif(result)
 print(result.ptm, result.plddt.mean().item())
 ```
 
-No target structure is required. For complexes, construct the input from the
-types exposed by the loaded artifact:
-
-```python
-types = model.input_types
-complex_input = types.StructurePredictionInput(
-    sequences=[
-        types.ProteinInput(id="A", sequence="MSTNPKPQRKTKRNT"),
-        types.ProteinInput(id="B", sequence="MKTIIALSYIFCLVFA"),
-        types.DNAInput(id="C", sequence="ATGC"),
-        types.LigandInput(id="L", smiles="O"),
-    ]
-)
-complex_result = model.fold(
-    complex_input,
-    num_loops=1,
-    num_sampling_steps=200,
-    seed=7,
-)
-print(complex_result.ptm, complex_result.plddt.mean().item())
-```
-
 The typed interface also supports RNA, protein MSAs, modifications, and covalent
 bonds. The public schema recognizes `PocketConditioning` and
 `DistogramConditioning`, but the pinned official forward consumes neither. Its
@@ -296,7 +298,7 @@ adapter modules are excluded from checkpoint state. It is not a generic
 - Redistributable: `true`
 - Complete weight publication required: `false`
 
-## Validation and provenance
+## Validation and sources
 
 FastPLMs pins the checkpoint, upstream source revisions, state transformation,
 and required files in `models.toml`. Built artifacts record exact source

fastplms/models.toml

@@ -392,7 +392,7 @@ checkpoint_license = "MIT"
 hub_license = "mit"
 weights_publication_allowed = true
 state_transform = "identity"
-conversion_provenance = "Input: each pinned Biohub ESMFold2 checkpoint and its separately pinned ESMC checkpoint. Transformation: apply identity to preserve the folding checkpoint exactly, load its parameters in FP32 for CUDA BF16-autocast execution, retain canonical BF16 ESMC weights, and optionally rebuild exactly 80 ESMC attention output projections as transient Transformer Engine linears. Output: the corresponding pinned Synthyra ESMFold2 checkpoint plus its declared ESMC precision policy. Validation: release parity covers exact canonical state, learned projection, prepared features, and seeded BF16 folding; experimental FP8 validation covers strict unavailable-device behavior, all four variants, and three BF16-to-FP8 reload cycles on the standard variant. Limitation: only the four manifest-listed ESMFold2 variants are supported; FP8 is experimental, applies only to inference-time ESMC execution, and requires direct CUDA loading with Transformer Engine availability."
+conversion_provenance = "Input: each pinned Biohub ESMFold2 checkpoint and its separately pinned ESMC checkpoint. Transformation: apply identity to preserve the folding checkpoint exactly, load its parameters in FP32 for CUDA BF16-autocast execution, retain canonical BF16 ESMC weights, and optionally rebuild exactly 80 ESMC attention output projections as transient Transformer Engine linears. Output: the corresponding pinned Synthyra ESMFold2 checkpoint plus its declared ESMC precision policy. Validation: release parity covers exact canonical state, learned projection, prepared features, and seeded BF16 folding; experimental FP8 validation covers strict unavailable-device behavior, the four 6B-backbone variants, and three BF16-to-FP8 reload cycles on the standard variant. Limitation: only the six manifest-listed ESMFold2 variants are supported; FP8 is experimental, applies only to inference-time ESMC execution, and requires direct CUDA loading with Transformer Engine availability."
 representative = "esmfold2"
 documentation = "docs/esmfold2.md"
 test_tiers = ["check", "compliance", "structure", "feature", "artifact", "benchmark"]
@@ -1237,3 +1237,39 @@ official_files = [
   "model.safetensors=sha256:4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f",
 ]
 auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" }
+
+[[models]]
+id = "esmfold2_300"
+family = "esmfold2"
+size_category = "structure"
+generation_contract = "not_applicable"
+msa_conditioning = false
+publication_status = "published"
+fast_repo = "Synthyra/ESMFold2-300"
+fast_revision = "a38a62ae930d157484b331c2bf4241684573adba"
+fast_files = ["config.json=git-sha1:47ec20cf8b234c3b41d6f3ae1bdfe95d4eb4849e", "model.safetensors=sha256:44d6797c5efebf24753d502b40950e0874871c96ceea14f2d7f7e39cebac67fd"]
+official_repo = "biohub/ESMFold2-Experimental-Fast-base300M-step1500k"
+official_revision = "21531e59002c9205284715e28ee802dafb430637"
+official_files = ["config.json=git-sha1:8c9a04fe22b0e5fca77bc4e2861a12c9494ef4d4", "model.safetensors=sha256:44d6797c5efebf24753d502b40950e0874871c96ceea14f2d7f7e39cebac67fd"]
+notes = "Experimental Fast checkpoint with a frozen 300M ESM++ backbone, tensor-exact in BF16 with the pinned step-1500000 source, 24 folding blocks, no MSA conditioning, and no confidence head. BF16 execution uses FP32 folding parameters with CUDA autocast. FP8 is unsupported. Docker BF16 inference validation passed on the compact Protein G case."
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" }
+backbone_model = "esmc_small"
+backbone = { repo = "biohub/ESMC-300M-1500000", revision = "56803b6378b82e16c3b24aac49d1fce4445540b7", files = ["config.json=git-sha1:7fe728a0eb3fb81b24491d6cc1de816bf7797c27", "model.safetensors=sha256:8bd6cacf9b5a92d51954b64b20407f1f9f564a7e4849b8663470784d2a8b7ed2", "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c", "tokenizer_config.json=git-sha1:f49f57b24a1c93bd544974811e8ecbd61b7fae89", "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b"] }
+
+[[models]]
+id = "esmfold2_600"
+family = "esmfold2"
+size_category = "structure"
+generation_contract = "not_applicable"
+msa_conditioning = false
+publication_status = "published"
+fast_repo = "Synthyra/ESMFold2-600"
+fast_revision = "71c67d0b2b73dc245ea7c3cc0d0476439a882d08"
+fast_files = ["config.json=git-sha1:8e271837cbdada96c4974c8e543f84065e0f06f1", "model.safetensors=sha256:11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602"]
+official_repo = "biohub/ESMFold2-Experimental-Fast-base600M-step1500k"
+official_revision = "15cf2d6648692f6c17cee1297d8a285476fffa9b"
+official_files = ["config.json=git-sha1:95517e555f17a1eca4b68866c89033a1f6916a5d", "model.safetensors=sha256:11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602"]
+notes = "Experimental Fast checkpoint with a frozen 600M ESM++ backbone, tensor-exact in BF16 with the pinned step-1500000 source, 24 folding blocks, no MSA conditioning, and no confidence head. BF16 execution uses FP32 folding parameters with CUDA autocast. FP8 is unsupported. Configuration, weight identities, and artifact reload are verified; this model is not inference-validated."
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" }
+backbone_model = "esmc_large"
+backbone = { repo = "biohub/ESMC-600M-1500000", revision = "21af9cc429af76ebda6c48074fb624db4735aaaf", files = ["config.json=git-sha1:ec29f6009b21d710f64bf1c058f3a9710833d692", "model.safetensors=sha256:d6869f5ae0f11e5dc829b195e062e87cfcc2f851a08a5edbaf5d1083ae7f76cc", "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c", "tokenizer_config.json=git-sha1:f49f57b24a1c93bd544974811e8ecbd61b7fae89", "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b"] }

fastplms/models/esm_plusplus/modeling_esm_plusplus.py

@@ -1076,6 +1076,8 @@ class PreTrainedESMplusplusModel(FastPLMsAttentionMixin, PreTrainedModel):
         all other operations and all SAE weights remain BF16.
         """
 
+        if self.config.hidden_size != 2560 or self.config.num_hidden_layers != 80:
+            raise ValueError("FP8 is supported only for the ESMC-6B backbone.")
         if self._esmc_fp8:
             return self._esmc_precision_status
         if self.training:

fastplms/models/esmfold2/configuration_esmfold2.py

@@ -1,326 +1,336 @@
-# Copyright 2026 Biohub. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Configuration schema for release and experimental ESMFold2 checkpoints."""
-
-from __future__ import annotations
-
-from dataclasses import asdict, dataclass, field
-from typing import Any, TypeVar, cast
-
-from transformers.configuration_utils import PretrainedConfig
-
-from fastplms.attention import canonical_checkpoint_attention_backend
-
-
-_ESMC_ATTENTION_IMPLEMENTATIONS = frozenset({"eager", "flex_attention", "sdpa"})
-_ESMC_PRECISIONS = frozenset({"auto", "bf16", "fp32", "fp8"})
-
-
-def _esmc_backbone_checkpoint_ids() -> tuple[str, str]:
-    """Return the manifest-pinned official and FastPLMs ESMC repositories."""
-
-    from fastplms.registry import RegistryError, get_model_registry
-
-    registry = get_model_registry()
-    family = registry.families["esmfold2"]
-    if family.backbone_model is None:
-        raise RegistryError("families.esmfold2 must declare backbone_model.")
-    backbone = registry[family.backbone_model]
-    return backbone.official.repo_id, backbone.fast.repo_id
-
-
-def normalize_esmc_id(esmc_id: str) -> str:
-    """Resolve an official ESMC identifier to its FastPLMs checkpoint mirror."""
-
-    official_repo, fast_repo = _esmc_backbone_checkpoint_ids()
-    return fast_repo if esmc_id == official_repo else esmc_id
-
-
-def normalize_esmc_attention_implementation(
-    implementation: str | dict[str, str] | None,
-) -> str | None:
-    """Validate the ESMC backend and translate the historical ``flex`` name."""
-
-    if isinstance(implementation, dict):
-        if tuple(implementation) != ("",):
-            raise ValueError(
-                "ESMFold2 has one ESMC attention backbone; use a string or {'': implementation}."
-            )
-        implementation = implementation[""]
-    canonical = canonical_checkpoint_attention_backend(implementation)
-    if canonical is not None and canonical not in _ESMC_ATTENTION_IMPLEMENTATIONS:
-        expected = sorted(_ESMC_ATTENTION_IMPLEMENTATIONS)
-        raise ValueError(
-            f"Unsupported ESMFold2 attention implementation {canonical!r}; "
-            f"expected one of {expected}."
-        )
-    return canonical
-
-
-NestedConfig = TypeVar("NestedConfig")
-
-
-def _nested_config(value: Any, config_type: type[NestedConfig]) -> NestedConfig:
-    if isinstance(value, config_type):
-        return value
-    return config_type(**value) if isinstance(value, dict) else config_type()
-
-
-def _coerce_nested_field(
-    value: NestedConfig | dict[str, Any], config_type: type[NestedConfig]
-) -> NestedConfig:
-    """Convert serialized nested dictionaries while retaining supplied objects."""
-
-    return config_type(**value) if isinstance(value, dict) else value
-
-
-@dataclass
-class AtomAttentionConfig:
-    """Sliding-window atom attention and three-dimensional RoPE settings."""
-
-    d_atom: int = field(default=128)
-    d_token: int = field(default=768)
-    n_blocks: int = field(default=3)
-    n_heads: int = field(default=4)
-    swa_window_size: int = field(default=128)
-    expansion_ratio: int = field(default=2)
-    spatial_rope_base_frequency: float = field(default=20.0)
-    n_spatial_rope_pairs_per_axis: int = field(default=2)
-    n_uid_rope_pairs: int = field(default=10)
-    uid_rope_base_frequency: float = field(default=10000.0)
-
-
-@dataclass
-class DiffusionModuleConfig:
-    """Dimensions and depth of the coordinate diffusion network."""
-
-    sigma_data: float = field(default=16.0)
-    c_atom: int = field(default=128)
-    c_token: int = field(default=768)
-    c_z: int = field(default=256)
-    c_s_inputs: int = field(default=451)
-    fourier_dim: int = field(default=256)
-    relpos_r_max: int = field(default=32)
-    relpos_s_max: int = field(default=2)
-    atom_num_blocks: int = field(default=3)
-    atom_num_heads: int = field(default=4)
-    token_num_blocks: int = field(default=12)
-    token_num_heads: int = field(default=16)
-    transition_multiplier: int = field(default=2)
-
-
-@dataclass
-class FoldingTrunkConfig:
-    """Iterative pair/single trunk dimensions."""
-
-    n_layers: int = field(default=24)
-    n_heads: int = field(default=8)
-    dropout: float = field(default=0.0)
-
-
-@dataclass
-class InputsEmbedderConfig:
-    """Input feature width and atom encoder settings."""
-
-    d_inputs: int = field(default=451)
-    atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig)
-
-    def __post_init__(self) -> None:
-        self.atom_encoder = _coerce_nested_field(self.atom_encoder, AtomAttentionConfig)
-
-
-@dataclass
-class DiffusionStructureHeadConfig:
-    """Training and inference schedules for coordinate denoising."""
-
-    diffusion_module: DiffusionModuleConfig = field(default_factory=DiffusionModuleConfig)
-    distogram_bins: int = field(default=128)
-    train_noise_log_mean: float = field(default=-1.2)
-    train_noise_log_std: float = field(default=1.5)
-    gamma_0: float = field(default=0.605)
-    gamma_min: float = field(default=1.107)
-    noise_scale: float = field(default=0.0)
-    step_scale: float = field(default=1.0)
-    inference_s_max: float = field(default=160.0)
-    inference_s_min: float = field(default=4e-4)
-    inference_p: float = field(default=8.0)
-    inference_num_steps: int = field(default=68)
-
-    def __post_init__(self) -> None:
-        self.diffusion_module = _coerce_nested_field(self.diffusion_module, DiffusionModuleConfig)
-
-
-@dataclass
-class ConfidenceHeadConfig:
-    """Confidence-bin definitions and the compact confidence trunk."""
-
-    enabled: bool = field(default=True)
-    num_plddt_bins: int = field(default=50)
-    num_pde_bins: int = field(default=64)
-    num_pae_bins: int = field(default=64)
-    min_dist: float = field(default=2.0)
-    max_dist: float = field(default=52.0)
-    distogram_bins: int = field(default=128)
-    folding_trunk: FoldingTrunkConfig = field(
-        default_factory=lambda: FoldingTrunkConfig(n_layers=4)
-    )
-
-    def __post_init__(self) -> None:
-        self.folding_trunk = _coerce_nested_field(self.folding_trunk, FoldingTrunkConfig)
-
-
-@dataclass
-class MSAEncoderConfig:
-    """Optional multiple-sequence-alignment encoder settings."""
-
-    enabled: bool = field(default=False)
-    d_msa: int = field(default=128)
-    d_hidden: int = field(default=32)
-    n_layers: int = field(default=4)
-    n_heads_msa: int = field(default=8)
-    msa_head_width: int = field(default=32)
-
-
-@dataclass
-class LMEncoderConfig:
-    """Release-model pair encoder derived from language-model states."""
-
-    enabled: bool = field(default=True)
-    n_layers: int = field(default=4)
-    lm_dropout: float = field(default=0.25)
-    per_loop_lm_dropout: bool = field(default=True)
-
-
-@dataclass
-class ParcaeConfig:
-    """Release-model diffusion-loop scheduler settings."""
-
-    enabled: bool = field(default=True)
-    poisson_mean: float = field(default=3.0)
-    min_steps: int = field(default=1)
-    max_steps: int | None = field(default=6)
-    coda_n_layers: int = field(default=2)
-
-
-_SCALAR_DEFAULTS: tuple[tuple[str, Any], ...] = (
-    ("d_single", 384),
-    ("d_pair", 256),
-    ("n_relative_residx_bins", 32),
-    ("n_relative_chain_bins", 2),
-    ("num_loops", 10),
-    ("num_diffusion_samples", 8),
-    ("disable_msa_features", False),
-    ("lm_dropout", 0.0),
-    ("force_lm_dropout_during_inference", False),
-    ("lm_mask_pct", 0.0),
-    ("lm_d_model", 2560),
-    ("lm_num_layers", 80),
-)
-_NESTED_CONFIGS = (
-    ("inputs", InputsEmbedderConfig),
-    ("folding_trunk", FoldingTrunkConfig),
-    ("structure_head", DiffusionStructureHeadConfig),
-    ("confidence_head", ConfidenceHeadConfig),
-    ("msa_encoder", MSAEncoderConfig),
-    ("parcae", ParcaeConfig),
-    ("lm_encoder", LMEncoderConfig),
-)
-
-
-class ESMFold2Config(PretrainedConfig):
-    """Serializable ESMFold2 architecture, runtime, and precision settings."""
-
-    model_type = "esmfold2"
-    has_no_defaults_at_init = True
-
-    def __init__(self, **kwargs: Any) -> None:
-        legacy_backend = normalize_esmc_attention_implementation(kwargs.get("esmc_attn_backend"))
-        requested_backend = normalize_esmc_attention_implementation(
-            kwargs.get("attn_implementation")
-        )
-        resolved_backend = requested_backend or legacy_backend
-        kwargs["attn_implementation"] = resolved_backend
-        super().__init__(**kwargs)
-
-        self.type = kwargs.get("type", "release")
-        if self.type not in {"experimental", "release"}:
-            raise ValueError(
-                f"ESMFold2Config.type must be 'release' or 'experimental', got {self.type!r}"
-            )
-
-        for name, default in _SCALAR_DEFAULTS:
-            setattr(self, name, kwargs.get(name, default))
-
-        _official_esmc_repo, default_esmc_repo = _esmc_backbone_checkpoint_ids()
-        self.esmc_id = normalize_esmc_id(kwargs.get("esmc_id", default_esmc_repo))
-        self.esmc_attn_backend = resolved_backend
-        self.esmc_precision = str(kwargs.get("esmc_precision", "auto"))
-        if self.esmc_precision not in _ESMC_PRECISIONS:
-            raise ValueError(
-                "esmc_precision must be 'auto', 'bf16', 'fp32', or 'fp8', "
-                f"got {self.esmc_precision!r}."
-            )
-
-        for name, config_type in _NESTED_CONFIGS:
-            setattr(self, name, _nested_config(kwargs.get(name), config_type))
-        if not isinstance(self.msa_encoder.enabled, bool):
-            raise TypeError("msa_encoder.enabled must be a boolean.")
-        declared_msa_conditioning = kwargs.get("msa_conditioning")
-        if "msa_conditioning" in kwargs and not isinstance(declared_msa_conditioning, bool):
-            raise TypeError("msa_conditioning must be a boolean when provided.")
-        self.msa_conditioning = (
-            self.msa_encoder.enabled
-            if "msa_conditioning" not in kwargs
-            else declared_msa_conditioning
-        )
-        if self.msa_conditioning != self.msa_encoder.enabled:
-            raise ValueError(
-                "msa_conditioning must match msa_encoder.enabled; received "
-                f"{self.msa_conditioning!r} and {self.msa_encoder.enabled!r}."
-            )
-        self.msa_encoder_overwrite = bool(kwargs.get("msa_encoder_overwrite", True))
-
-        self.classifier_train_scope = str(kwargs.get("classifier_train_scope", "probe"))
-        if self.classifier_train_scope not in {"probe", "projection"}:
-            raise ValueError(
-                "classifier_train_scope must be 'probe' or 'projection', "
-                f"got {self.classifier_train_scope!r}."
-            )
-        self.classifier_probe_hidden_size = int(
-            kwargs.get("classifier_probe_hidden_size", 512)
-        )
-        self.classifier_probe_num_heads = int(kwargs.get("classifier_probe_num_heads", 4))
-        self.classifier_probe_dropout = float(kwargs.get("classifier_probe_dropout", 0.1))
-        self.classifier_hidden_size = int(kwargs.get("classifier_hidden_size", 4096))
-        self.classifier_dropout = float(kwargs.get("classifier_dropout", 0.2))
-        self.classifier_pooling_types = list(
-            kwargs.get("classifier_pooling_types", ["mean"])
-        )
-
-    def to_dict(self) -> dict[str, Any]:
-        output = cast(dict[str, Any], super().to_dict())
-        for name, _config_type in _NESTED_CONFIGS:
-            output[name] = asdict(getattr(self, name))
-        return output
-
-
-__all__ = [
-    "ESMFold2Config",
-    "LMEncoderConfig",
-    "MSAEncoderConfig",
-    "ParcaeConfig",
-    "normalize_esmc_attention_implementation",
-    "normalize_esmc_id",
-]
+# Copyright 2026 Biohub. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Configuration schema for release and experimental ESMFold2 checkpoints."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from typing import Any, TypeVar, cast
+
+from transformers.configuration_utils import PretrainedConfig
+
+from fastplms.attention import canonical_checkpoint_attention_backend
+
+
+_ESMC_ATTENTION_IMPLEMENTATIONS = frozenset({"eager", "flex_attention", "sdpa"})
+_ESMC_PRECISIONS = frozenset({"auto", "bf16", "fp32", "fp8"})
+
+
+def _esmc_backbone_checkpoint_ids() -> tuple[str, str]:
+    """Return the manifest-pinned official and FastPLMs ESMC repositories."""
+
+    from fastplms.registry import RegistryError, get_model_registry
+
+    registry = get_model_registry()
+    family = registry.families["esmfold2"]
+    if family.backbone_model is None:
+        raise RegistryError("families.esmfold2 must declare backbone_model.")
+    backbone = registry[family.backbone_model]
+    return backbone.official.repo_id, backbone.fast.repo_id
+
+
+def normalize_esmc_id(esmc_id: str) -> str:
+    """Resolve an official ESMC identifier to its FastPLMs checkpoint mirror."""
+
+    from fastplms.registry import get_model_registry
+
+    official_repo, fast_repo = _esmc_backbone_checkpoint_ids()
+    if esmc_id == official_repo:
+        return fast_repo
+    registry = get_model_registry()
+    for spec in registry.by_family("esmfold2"):
+        if spec.backbone is not None and spec.backbone_model is not None:
+            backbone = registry[spec.backbone_model]
+            if esmc_id in {spec.backbone.repo_id, backbone.official.repo_id}:
+                return backbone.fast.repo_id
+    return esmc_id
+
+
+def normalize_esmc_attention_implementation(
+    implementation: str | dict[str, str] | None,
+) -> str | None:
+    """Validate the ESMC backend and translate the historical ``flex`` name."""
+
+    if isinstance(implementation, dict):
+        if tuple(implementation) != ("",):
+            raise ValueError(
+                "ESMFold2 has one ESMC attention backbone; use a string or {'': implementation}."
+            )
+        implementation = implementation[""]
+    canonical = canonical_checkpoint_attention_backend(implementation)
+    if canonical is not None and canonical not in _ESMC_ATTENTION_IMPLEMENTATIONS:
+        expected = sorted(_ESMC_ATTENTION_IMPLEMENTATIONS)
+        raise ValueError(
+            f"Unsupported ESMFold2 attention implementation {canonical!r}; "
+            f"expected one of {expected}."
+        )
+    return canonical
+
+
+NestedConfig = TypeVar("NestedConfig")
+
+
+def _nested_config(value: Any, config_type: type[NestedConfig]) -> NestedConfig:
+    if isinstance(value, config_type):
+        return value
+    return config_type(**value) if isinstance(value, dict) else config_type()
+
+
+def _coerce_nested_field(
+    value: NestedConfig | dict[str, Any], config_type: type[NestedConfig]
+) -> NestedConfig:
+    """Convert serialized nested dictionaries while retaining supplied objects."""
+
+    return config_type(**value) if isinstance(value, dict) else value
+
+
+@dataclass
+class AtomAttentionConfig:
+    """Sliding-window atom attention and three-dimensional RoPE settings."""
+
+    d_atom: int = field(default=128)
+    d_token: int = field(default=768)
+    n_blocks: int = field(default=3)
+    n_heads: int = field(default=4)
+    swa_window_size: int = field(default=128)
+    expansion_ratio: int = field(default=2)
+    spatial_rope_base_frequency: float = field(default=20.0)
+    n_spatial_rope_pairs_per_axis: int = field(default=2)
+    n_uid_rope_pairs: int = field(default=10)
+    uid_rope_base_frequency: float = field(default=10000.0)
+
+
+@dataclass
+class DiffusionModuleConfig:
+    """Dimensions and depth of the coordinate diffusion network."""
+
+    sigma_data: float = field(default=16.0)
+    c_atom: int = field(default=128)
+    c_token: int = field(default=768)
+    c_z: int = field(default=256)
+    c_s_inputs: int = field(default=451)
+    fourier_dim: int = field(default=256)
+    relpos_r_max: int = field(default=32)
+    relpos_s_max: int = field(default=2)
+    atom_num_blocks: int = field(default=3)
+    atom_num_heads: int = field(default=4)
+    token_num_blocks: int = field(default=12)
+    token_num_heads: int = field(default=16)
+    transition_multiplier: int = field(default=2)
+
+
+@dataclass
+class FoldingTrunkConfig:
+    """Iterative pair/single trunk dimensions."""
+
+    n_layers: int = field(default=24)
+    n_heads: int = field(default=8)
+    dropout: float = field(default=0.0)
+
+
+@dataclass
+class InputsEmbedderConfig:
+    """Input feature width and atom encoder settings."""
+
+    d_inputs: int = field(default=451)
+    atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig)
+
+    def __post_init__(self) -> None:
+        self.atom_encoder = _coerce_nested_field(self.atom_encoder, AtomAttentionConfig)
+
+
+@dataclass
+class DiffusionStructureHeadConfig:
+    """Training and inference schedules for coordinate denoising."""
+
+    diffusion_module: DiffusionModuleConfig = field(default_factory=DiffusionModuleConfig)
+    distogram_bins: int = field(default=128)
+    train_noise_log_mean: float = field(default=-1.2)
+    train_noise_log_std: float = field(default=1.5)
+    gamma_0: float = field(default=0.605)
+    gamma_min: float = field(default=1.107)
+    noise_scale: float = field(default=0.0)
+    step_scale: float = field(default=1.0)
+    inference_s_max: float = field(default=160.0)
+    inference_s_min: float = field(default=4e-4)
+    inference_p: float = field(default=8.0)
+    inference_num_steps: int = field(default=68)
+
+    def __post_init__(self) -> None:
+        self.diffusion_module = _coerce_nested_field(self.diffusion_module, DiffusionModuleConfig)
+
+
+@dataclass
+class ConfidenceHeadConfig:
+    """Confidence-bin definitions and the compact confidence trunk."""
+
+    enabled: bool = field(default=True)
+    num_plddt_bins: int = field(default=50)
+    num_pde_bins: int = field(default=64)
+    num_pae_bins: int = field(default=64)
+    min_dist: float = field(default=2.0)
+    max_dist: float = field(default=52.0)
+    distogram_bins: int = field(default=128)
+    folding_trunk: FoldingTrunkConfig = field(
+        default_factory=lambda: FoldingTrunkConfig(n_layers=4)
+    )
+
+    def __post_init__(self) -> None:
+        self.folding_trunk = _coerce_nested_field(self.folding_trunk, FoldingTrunkConfig)
+
+
+@dataclass
+class MSAEncoderConfig:
+    """Optional multiple-sequence-alignment encoder settings."""
+
+    enabled: bool = field(default=False)
+    d_msa: int = field(default=128)
+    d_hidden: int = field(default=32)
+    n_layers: int = field(default=4)
+    n_heads_msa: int = field(default=8)
+    msa_head_width: int = field(default=32)
+
+
+@dataclass
+class LMEncoderConfig:
+    """Release-model pair encoder derived from language-model states."""
+
+    enabled: bool = field(default=True)
+    n_layers: int = field(default=4)
+    lm_dropout: float = field(default=0.25)
+    per_loop_lm_dropout: bool = field(default=True)
+
+
+@dataclass
+class ParcaeConfig:
+    """Release-model diffusion-loop scheduler settings."""
+
+    enabled: bool = field(default=True)
+    poisson_mean: float = field(default=3.0)
+    min_steps: int = field(default=1)
+    max_steps: int | None = field(default=6)
+    coda_n_layers: int = field(default=2)
+
+
+_SCALAR_DEFAULTS: tuple[tuple[str, Any], ...] = (
+    ("d_single", 384),
+    ("d_pair", 256),
+    ("n_relative_residx_bins", 32),
+    ("n_relative_chain_bins", 2),
+    ("num_loops", 10),
+    ("num_diffusion_samples", 8),
+    ("disable_msa_features", False),
+    ("lm_dropout", 0.0),
+    ("force_lm_dropout_during_inference", False),
+    ("lm_mask_pct", 0.0),
+    ("lm_d_model", 2560),
+    ("lm_num_layers", 80),
+)
+_NESTED_CONFIGS = (
+    ("inputs", InputsEmbedderConfig),
+    ("folding_trunk", FoldingTrunkConfig),
+    ("structure_head", DiffusionStructureHeadConfig),
+    ("confidence_head", ConfidenceHeadConfig),
+    ("msa_encoder", MSAEncoderConfig),
+    ("parcae", ParcaeConfig),
+    ("lm_encoder", LMEncoderConfig),
+)
+
+
+class ESMFold2Config(PretrainedConfig):
+    """Serializable ESMFold2 architecture, runtime, and precision settings."""
+
+    model_type = "esmfold2"
+    has_no_defaults_at_init = True
+
+    def __init__(self, **kwargs: Any) -> None:
+        legacy_backend = normalize_esmc_attention_implementation(kwargs.get("esmc_attn_backend"))
+        requested_backend = normalize_esmc_attention_implementation(
+            kwargs.get("attn_implementation")
+        )
+        resolved_backend = requested_backend or legacy_backend
+        kwargs["attn_implementation"] = resolved_backend
+        super().__init__(**kwargs)
+
+        self.type = kwargs.get("type", "release")
+        if self.type not in {"experimental", "release"}:
+            raise ValueError(
+                f"ESMFold2Config.type must be 'release' or 'experimental', got {self.type!r}"
+            )
+
+        for name, default in _SCALAR_DEFAULTS:
+            setattr(self, name, kwargs.get(name, default))
+
+        _official_esmc_repo, default_esmc_repo = _esmc_backbone_checkpoint_ids()
+        self.esmc_id = normalize_esmc_id(kwargs.get("esmc_id", default_esmc_repo))
+        self.esmc_attn_backend = resolved_backend
+        self.esmc_precision = str(kwargs.get("esmc_precision", "auto"))
+        if self.esmc_precision not in _ESMC_PRECISIONS:
+            raise ValueError(
+                "esmc_precision must be 'auto', 'bf16', 'fp32', or 'fp8', "
+                f"got {self.esmc_precision!r}."
+            )
+
+        for name, config_type in _NESTED_CONFIGS:
+            setattr(self, name, _nested_config(kwargs.get(name), config_type))
+        if not isinstance(self.msa_encoder.enabled, bool):
+            raise TypeError("msa_encoder.enabled must be a boolean.")
+        declared_msa_conditioning = kwargs.get("msa_conditioning")
+        if "msa_conditioning" in kwargs and not isinstance(declared_msa_conditioning, bool):
+            raise TypeError("msa_conditioning must be a boolean when provided.")
+        self.msa_conditioning = (
+            self.msa_encoder.enabled
+            if "msa_conditioning" not in kwargs
+            else declared_msa_conditioning
+        )
+        if self.msa_conditioning != self.msa_encoder.enabled:
+            raise ValueError(
+                "msa_conditioning must match msa_encoder.enabled; received "
+                f"{self.msa_conditioning!r} and {self.msa_encoder.enabled!r}."
+            )
+        self.msa_encoder_overwrite = bool(kwargs.get("msa_encoder_overwrite", True))
+
+        self.classifier_train_scope = str(kwargs.get("classifier_train_scope", "probe"))
+        if self.classifier_train_scope not in {"probe", "projection"}:
+            raise ValueError(
+                "classifier_train_scope must be 'probe' or 'projection', "
+                f"got {self.classifier_train_scope!r}."
+            )
+        self.classifier_probe_hidden_size = int(
+            kwargs.get("classifier_probe_hidden_size", 512)
+        )
+        self.classifier_probe_num_heads = int(kwargs.get("classifier_probe_num_heads", 4))
+        self.classifier_probe_dropout = float(kwargs.get("classifier_probe_dropout", 0.1))
+        self.classifier_hidden_size = int(kwargs.get("classifier_hidden_size", 4096))
+        self.classifier_dropout = float(kwargs.get("classifier_dropout", 0.2))
+        self.classifier_pooling_types = list(
+            kwargs.get("classifier_pooling_types", ["mean"])
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        output = cast(dict[str, Any], super().to_dict())
+        for name, _config_type in _NESTED_CONFIGS:
+            output[name] = asdict(getattr(self, name))
+        return output
+
+
+__all__ = [
+    "ESMFold2Config",
+    "LMEncoderConfig",
+    "MSAEncoderConfig",
+    "ParcaeConfig",
+    "normalize_esmc_attention_implementation",
+    "normalize_esmc_id",
+]

fastplms/models/esmfold2/esmfold2_mmcif_parsing.py

@@ -58,7 +58,7 @@ def round_mmcif_columns(cif_file: CIFFile) -> None:
         original = atom_site[name]
         values = original.as_array(np.float64)
         strings = np.asarray(
-            [f"{value:.{decimals}f}" for value in values],
+            ["?" if not np.isfinite(value) else f"{value:.{decimals}f}" for value in values],
             dtype=np.str_,
         )
         atom_site[name] = CIFColumn(

fastplms/models/esmfold2/esmfold2_output.py

@@ -79,7 +79,7 @@ class _ComplexRecords:
 
 def build_molecular_complex_from_features(
     coords: torch.Tensor,
-    plddt: torch.Tensor,
+    plddt: torch.Tensor | None,
     atom_mask: torch.Tensor,
     ref_element: torch.Tensor,
     ref_atom_name_chars: torch.Tensor,
@@ -96,7 +96,7 @@ def build_molecular_complex_from_features(
     X = coords.float().cpu().numpy()
     atom_names = ref_atom_name_chars.cpu().numpy()
     elements = ref_element.cpu().numpy()
-    confidence = plddt.float().cpu().numpy()
+    confidence = None if plddt is None else plddt.float().cpu().numpy()
     records = _ComplexRecords()
 
     def decode_atoms(tokens: Iterable[Any]):
@@ -117,8 +117,8 @@ def build_molecular_complex_from_features(
         if is_nonpolymer:
             mean_confidence = (
                 float(np.mean([confidence[token.token_index] for token in chain.tokens]))
-                if chain.tokens
-                else 0.0
+                if confidence is not None and chain.tokens
+                else float("nan")
             )
             records.add_token(
                 residue_name=chain.tokens[0].residue_name if chain.tokens else "LIG",
@@ -135,7 +135,11 @@ def build_molecular_complex_from_features(
             records.add_token(
                 residue_name=residue_tokens[0].residue_name,
                 asym_id=chain.asym_id,
-                plddt=float(np.mean([confidence[token.token_index] for token in residue_tokens])),
+                plddt=(
+                    float(np.mean([confidence[token.token_index] for token in residue_tokens]))
+                    if confidence is not None
+                    else float("nan")
+                ),
                 atoms=decode_atoms(residue_tokens),
                 hetero=False,
             )

fastplms/models/esmfold2/esmfold2_processor.py

@@ -10,6 +10,7 @@ from typing import Any
 import numpy as np
 import torch
 from torch import Tensor
+from tqdm.auto import tqdm
 
 from .esmfold2_conformers import load_ccd
 from .esmfold2_molecular_complex import MolecularComplexResult
@@ -240,15 +241,19 @@ class ESMFold2InputBuilder:
 
     def _decode_sample(
         self,
-        output: Mapping[str, Tensor],
+        output: Mapping[str, Tensor | None],
         features: dict[str, Tensor],
         chain_infos: list[ChainInfo],
         sample: int,
         complex_id: str,
     ) -> MolecularComplexResult:
-        plddt = output["plddt"][sample]
+        plddt_output = output.get("plddt")
+        plddt = None if plddt_output is None else plddt_output[sample]
+        coords_output = output.get("sample_atom_coords")
+        if coords_output is None:
+            raise ValueError("ESMFold2 output omitted sample_atom_coords.")
         molecular_complex = build_molecular_complex_from_features(
-            coords=output["sample_atom_coords"][sample],
+            coords=coords_output[sample],
             plddt=plddt,
             atom_mask=features["atom_attention_mask"][0],
             ref_element=features["ref_element"][0],
@@ -269,7 +274,7 @@ class ESMFold2InputBuilder:
         iptm = output.get("iptm")
         return MolecularComplexResult(
             complex=molecular_complex,
-            plddt=plddt.detach().cpu(),
+            plddt=None if plddt is None else plddt.detach().cpu(),
             ptm=float(ptm[sample].item()) if ptm is not None else None,
             iptm=float(iptm[sample].item()) if iptm is not None else None,
             pae=sample_tensor("pae"),
@@ -281,16 +286,19 @@ class ESMFold2InputBuilder:
 
     def decode(
         self,
-        output: Mapping[str, Tensor],
+        output: Mapping[str, Tensor | None],
         features: dict[str, Tensor],
         chain_infos: list[ChainInfo],
         *,
         num_diffusion_samples: int = 1,
         complex_id: str = "pred",
     ) -> MolecularComplexResult | list[MolecularComplexResult]:
+        coords_output = output.get("sample_atom_coords")
+        if coords_output is None:
+            raise ValueError("ESMFold2 output omitted sample_atom_coords.")
         results = [
             self._decode_sample(output, features, chain_infos, sample, complex_id)
-            for sample in range(output["sample_atom_coords"].shape[0])
+            for sample in range(coords_output.shape[0])
         ]
         return results[0] if num_diffusion_samples == 1 and len(results) == 1 else results
 
@@ -308,13 +316,24 @@ class ESMFold2InputBuilder:
         max_inference_sigma: int | None = None,
         early_exit: bool = False,
         complex_id: str = "pred",
+        verbose: bool = False,
     ) -> MolecularComplexResult | list[MolecularComplexResult]:
-        features, chain_infos = self.prepare_model_input(
-            model,
-            input,
-            seed=seed,
-            device=model.device,
-        )
+        if verbose:
+            with tqdm(total=1, desc="ESMFold2 typed features", unit="stage") as progress:
+                features, chain_infos = self.prepare_model_input(
+                    model,
+                    input,
+                    seed=seed,
+                    device=model.device,
+                )
+                progress.update()
+        else:
+            features, chain_infos = self.prepare_model_input(
+                model,
+                input,
+                seed=seed,
+                device=model.device,
+            )
         overrides = _sampler_overrides(noise_scale, step_scale, max_inference_sigma)
         with torch.no_grad(), seed_context(seed):
             output = model(
@@ -324,8 +343,20 @@ class ESMFold2InputBuilder:
                 num_diffusion_samples=num_diffusion_samples,
                 early_exit=early_exit,
                 return_dict=True,
+                verbose=verbose,
                 **overrides,
             )
+        if verbose:
+            with tqdm(total=1, desc="ESMFold2 decode", unit="stage") as progress:
+                result = self.decode(
+                    output,
+                    features,
+                    chain_infos,
+                    num_diffusion_samples=num_diffusion_samples,
+                    complex_id=complex_id,
+                )
+                progress.update()
+            return result
         return self.decode(
             output,
             features,

fastplms/models/esmfold2/modeling_esmfold2_common.py

@@ -19,6 +19,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 from torch.utils.checkpoint import checkpoint
+from tqdm.auto import tqdm
 
 from .configuration_esmfold2 import ESMFold2Config
 from .reproducibility import seed_context
@@ -2615,6 +2616,7 @@ class DiffusionStructureHead(nn.Module):
         return_atom_repr: bool = False,
         use_inference_cache: bool = True,
         denoising_early_exit_rmsd: float | None = None,
+        verbose: bool = False,
     ) -> dict[str, Tensor | None]:
         """Diffusion sampling (Algorithm 18).
 
@@ -2656,7 +2658,16 @@ class DiffusionStructureHead(nn.Module):
         step_pairs = list(zip(schedule[:-1], schedule[1:], gammas[1:], strict=True))
         num_steps = len(step_pairs)
 
-        for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(step_pairs):
+        step_iterator = step_pairs
+        if verbose:
+            step_iterator = tqdm(
+                step_pairs,
+                total=num_steps if denoising_early_exit_rmsd is None else None,
+                desc="ESMFold2 diffusion",
+                unit="step",
+            )
+
+        for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(step_iterator):
             x, x_denoised_prev = self._center_random_augmentation(
                 x, atom_mask, second_coords=x_denoised_prev
             )

fastplms/models/esmfold2/modeling_esmfold2_experimental.py

@@ -9,6 +9,7 @@ re-injection and a different confidence/MSA stack.
 from __future__ import annotations
 
 import gc
+from collections.abc import Mapping
 from pathlib import Path
 from typing import Any, ClassVar, cast
 
@@ -16,6 +17,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
+from tqdm.auto import tqdm
 from transformers.modeling_utils import PreTrainedModel
 
 from .attention import ESMFold2AttentionMixin
@@ -569,9 +571,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
             precision=precision,
             device=device,
             local_files_only=(
-                self._esmc_local_files_only
-                if local_files_only is None
-                else local_files_only
+                self._esmc_local_files_only if local_files_only is None else local_files_only
             ),
         )
 
@@ -627,6 +627,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         residue_index: Tensor,
         mol_type: Tensor,
         tok_mask: Tensor,
+        verbose: bool = False,
     ) -> Tensor:
         if self._esmc_fp8 and torch.is_grad_enabled():
             _reload_esmc_bf16_for_gradients(
@@ -640,6 +641,19 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
             raise RuntimeError("ESMFold2 language-model features require load_esmc=True.")
         pad_to = 16 if self._esmc_fp8 else None
         with _lm_precision_context(self._esmc_precision_status.resolved, self.device):
+            if verbose:
+                with tqdm(total=1, desc="ESMFold2 backbone", unit="stage") as progress:
+                    result = compute_lm_hidden_states(
+                        self._esmc,
+                        input_ids,
+                        asym_id,
+                        residue_index,
+                        mol_type,
+                        tok_mask,
+                        pad_to_multiple=pad_to,
+                    )
+                    progress.update()
+                return result
             return compute_lm_hidden_states(
                 self._esmc,
                 input_ids,
@@ -696,6 +710,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         frames_idx: Tensor | None = None,
         disto_cond: Tensor | None = None,
         disto_cond_mask: Tensor | None = None,
+        verbose: bool = False,
     ) -> ESMFold2Output | tuple[Any, ...]:
         output_hidden_states, return_dict = _resolve_structure_output_controls(
             self.config,
@@ -796,7 +811,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
 
             if lm_hidden_states is None and input_ids is not None and self._esmc is not None:
                 lm_hidden_states = self._compute_lm_hidden_states(
-                    input_ids, asym_id, residue_index, mol_type, tok_mask
+                    input_ids, asym_id, residue_index, mol_type, tok_mask, verbose=verbose
                 )
             if lm_hidden_states is not None:
                 lm_dropout = (
@@ -845,7 +860,16 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
             z = torch.zeros_like(z_init)
             prev_pair: Tensor | None = None
             prev_disto_probs: Tensor | None = None
-            for loop_num in range(n_loops + 1):
+            loop_iterator = range(n_loops + 1)
+            if verbose:
+                loop_iterator = tqdm(
+                    loop_iterator,
+                    total=n_loops + 1 if not early_exit else None,
+                    desc="ESMFold2 recycling",
+                    unit="loop",
+                )
+
+            for loop_num in loop_iterator:
                 z = z_init + self.pair_loop_proj(z)
                 if msa_kwargs is not None and self.msa_encoder is not None:
                     z = z + self.msa_encoder(x_pair=z, **msa_kwargs).to(z.dtype)
@@ -899,6 +923,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
                 step_scale=step_scale,
                 return_atom_repr=False,
                 denoising_early_exit_rmsd=(0.10 if early_exit else None),
+                verbose=verbose,
             )
         sample_coords = structure_output["sample_atom_coords"]
         if sample_coords is None:
@@ -924,21 +949,45 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
             "sample_atom_coords": sample_coords,
             "representative_atom_coords": representative_atom_coords,
         }
-        if calculate_confidence and self.confidence_head is not None:
-            confidence_output = self.confidence_head(
-                s_inputs=x_inputs.detach(),
-                z=z.detach().float(),
-                x_pred=sample_coords.detach(),
-                distogram_atom_idx=distogram_atom_idx,
-                token_attention_mask=tok_mask,
-                atom_to_token=atom_to_token,
-                atom_attention_mask=atm_mask,
-                asym_id=asym_id,
-                mol_type=mol_type,
-                num_diffusion_samples=n_samples,
-                relative_position_encoding=relative_position_encoding.detach(),
-                token_bonds_encoding=token_bonds_encoding.detach(),
-            )
+        confidence_config = self.config.confidence_head
+        confidence_enabled = (
+            confidence_config.get("enabled", True)
+            if isinstance(confidence_config, Mapping)
+            else getattr(confidence_config, "enabled", True)
+        )
+        if calculate_confidence and confidence_enabled and self.confidence_head is not None:
+            if verbose:
+                with tqdm(total=1, desc="ESMFold2 confidence", unit="stage") as progress:
+                    confidence_output = self.confidence_head(
+                        s_inputs=x_inputs.detach(),
+                        z=z.detach().float(),
+                        x_pred=sample_coords.detach(),
+                        distogram_atom_idx=distogram_atom_idx,
+                        token_attention_mask=tok_mask,
+                        atom_to_token=atom_to_token,
+                        atom_attention_mask=atm_mask,
+                        asym_id=asym_id,
+                        mol_type=mol_type,
+                        num_diffusion_samples=n_samples,
+                        relative_position_encoding=relative_position_encoding.detach(),
+                        token_bonds_encoding=token_bonds_encoding.detach(),
+                    )
+                    progress.update()
+            else:
+                confidence_output = self.confidence_head(
+                    s_inputs=x_inputs.detach(),
+                    z=z.detach().float(),
+                    x_pred=sample_coords.detach(),
+                    distogram_atom_idx=distogram_atom_idx,
+                    token_attention_mask=tok_mask,
+                    atom_to_token=atom_to_token,
+                    atom_attention_mask=atm_mask,
+                    asym_id=asym_id,
+                    mol_type=mol_type,
+                    num_diffusion_samples=n_samples,
+                    relative_position_encoding=relative_position_encoding.detach(),
+                    token_bonds_encoding=token_bonds_encoding.detach(),
+                )
             output.update(confidence_output)
         output["atom_pad_mask"] = atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask
         output["residue_index"] = residue_index
@@ -981,7 +1030,12 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
             raise ValueError(
                 "infer_protein always returns a mapping; return_dict=False is invalid."
             )
-        features = prepare_protein_features(seq)
+        if forward_kwargs.get("verbose", False):
+            with tqdm(total=1, desc="ESMFold2 features", unit="stage") as progress:
+                features = prepare_protein_features(seq)
+                progress.update()
+        else:
+            features = prepare_protein_features(seq)
         if not self.config.msa_conditioning:
             for name in MSA_CONDITIONING_INPUT_NAMES:
                 features.pop(name, None)
@@ -1011,6 +1065,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         max_inference_sigma: int | None = None,
         early_exit: bool = False,
         complex_id: str = "pred",
+        verbose: bool = False,
     ):
         return self.input_builder.fold(
             self,
@@ -1024,6 +1079,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
             max_inference_sigma=max_inference_sigma,
             early_exit=early_exit,
             complex_id=complex_id,
+            verbose=verbose,
         )
 
     def fold_protein(
@@ -1036,6 +1092,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         num_diffusion_samples: int = 1,
         seed: int | None = None,
         complex_id: str = "pred",
+        verbose: bool = False,
     ):
         from .esmfold2_types import ProteinInput, StructurePredictionInput
 
@@ -1047,6 +1104,7 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
             num_diffusion_samples=num_diffusion_samples,
             seed=seed,
             complex_id=complex_id,
+            verbose=verbose,
         )
 
     @staticmethod

modeling_fastplms.py

@@ -13,7 +13,7 @@ from zipfile import ZIP_DEFLATED, ZipFile
 
 from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH
 
-if RUNTIME_HASH != "c76f023da79d5db93093291f4f3858a0e6cde4d9b4fa61ade779cf941f2175da":
+if RUNTIME_HASH != "4f6a9d59126e9bee9e65b30d55161d648473a058b6aed36727f2ae2cd99d44d7":
     raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
 
 _RUNTIME_TEMPORARIES = []