Update FastPLMs runtime and model cards

README.md

@@ -40,9 +40,8 @@ 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
@@ -153,6 +152,7 @@ with torch.inference_mode():
     output = model.infer(
         "MKTLLILAVVAAALA",
         num_recycles=4,
+        verbose=False,
     )
 
 print(output["mean_plddt"])
@@ -185,7 +185,7 @@ folding requests raise.
 - 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/esmfold/modeling_fast_esmfold.py

@@ -16,10 +16,13 @@ from __future__ import annotations
 
 import torch
 import torch.nn as nn
+from collections.abc import Iterator
 from contextvars import ContextVar
 from dataclasses import dataclass
+from contextlib import contextmanager
 from typing import Any
 from einops import rearrange
+from tqdm.auto import tqdm
 from torch.nn import functional as F
 from transformers.modeling_outputs import ModelOutput
 from transformers.models.esm.configuration_esm import EsmConfig
@@ -142,6 +145,78 @@ _ESMFOLD_CAPTURED_ATTENTIONS: ContextVar[
 )
 
 
+@contextmanager
+def _folding_progress(
+    model: nn.Module,
+    *,
+    num_recycles: int | None,
+    verbose: bool,
+) -> Iterator[None]:
+    """Report actual embedding, recycling, and confidence stages."""
+
+    if not verbose:
+        yield
+        return
+
+    trunk = getattr(model, "trunk", None)
+    blocks = getattr(trunk, "blocks", None)
+    structure_module = getattr(trunk, "structure_module", None)
+    embedding = getattr(model, "esm", None)
+    confidence_heads = tuple(
+        getattr(model, name, None)
+        for name in ("distogram_head", "lm_head", "lddt_head", "ptm_head")
+    )
+    if not isinstance(blocks, nn.ModuleList) or not isinstance(structure_module, nn.Module):
+        raise RuntimeError("ESMFold progress requires the standard folding trunk modules.")
+
+    if num_recycles is None:
+        recycle_passes = getattr(getattr(trunk, "config", None), "max_recycles", 0)
+    else:
+        recycle_passes = num_recycles + 1
+    confidence_modules = tuple(
+        head for head in confidence_heads if isinstance(head, nn.Module)
+    )
+    total = (
+        int(isinstance(embedding, nn.Module))
+        + max(int(recycle_passes), 0) * (len(blocks) + 1)
+        + len(confidence_modules)
+    )
+    progress = tqdm(total=total, desc="ESMFold embeddings", unit="stage")
+
+    def update_progress(stage: str, *_args: Any) -> None:
+        progress.set_description(f"ESMFold {stage}")
+        progress.update(1)
+
+    handles = []
+    try:
+        if isinstance(embedding, nn.Module):
+            handles.append(
+                embedding.register_forward_hook(
+                    lambda *_args: update_progress("embeddings", *_args)
+                )
+            )
+        for block in blocks:
+            handles.append(
+                block.register_forward_hook(
+                    lambda *_args: update_progress("recycling", *_args)
+                )
+            )
+        handles.append(
+            structure_module.register_forward_hook(
+                lambda *_args: update_progress("recycling", *_args)
+            )
+        )
+        handles.extend(
+            head.register_forward_hook(lambda *_args: update_progress("confidence", *_args))
+            for head in confidence_modules
+        )
+        yield
+    finally:
+        for handle in handles:
+            handle.remove()
+        progress.close()
+
+
 def _align_internal_esm_attentions(
     attentions: tuple[torch.Tensor, ...],
     residue_mask: torch.Tensor,
@@ -674,8 +749,13 @@ class FastEsmForProteinFolding(FastPLMsAttentionMixin, EsmForProteinFolding):
         output_attentions: bool | None = None,
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
+        verbose: bool = False,
     ) -> FastEsmForProteinFoldingOutput | tuple[Any, ...]:
-        """Run folding with Meta ESMFold's 0-to-100 pLDDT convention."""
+        """Run folding with Meta ESMFold's 0-to-100 pLDDT convention.
+
+        Set ``verbose=True`` to display progress for embeddings, recycling, and
+        confidence heads. The default keeps the call silent.
+        """
 
         config = getattr(self, "config", None)
         resolved_attentions = (
@@ -698,14 +778,19 @@ class FastEsmForProteinFolding(FastPLMsAttentionMixin, EsmForProteinFolding):
         capture_token = _ESMFOLD_CAPTURED_ATTENTIONS.set(None)
         captured_attentions: tuple[torch.Tensor, ...] | None = None
         try:
-            output = super().forward(
-                input_ids,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                masking_pattern=masking_pattern,
+            with _folding_progress(
+                self,
                 num_recycles=num_recycles,
-                output_hidden_states=resolved_hidden_states,
-            )
+                verbose=verbose,
+            ):
+                output = super().forward(
+                    input_ids,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    masking_pattern=masking_pattern,
+                    num_recycles=num_recycles,
+                    output_hidden_states=resolved_hidden_states,
+                )
             captured_attentions = _ESMFOLD_CAPTURED_ATTENTIONS.get()
         finally:
             _ESMFOLD_CAPTURED_ATTENTIONS.reset(capture_token)
@@ -743,12 +828,14 @@ class FastEsmForProteinFolding(FastPLMsAttentionMixin, EsmForProteinFolding):
         num_recycles: int | None = None,
         residue_index_offset: int | None = 512,
         chain_linker: str | None = "G" * 25,
+        verbose: bool = False,
     ):
         """Fold raw sequences through Meta ESMFold's public input contract.
 
         Transformers v5 narrows ``infer`` even though ``forward`` retains the
         required controls. This adapter restores recycle selection, explicit
         residue indices, masking, and colon-delimited multimer preparation.
+        Set ``verbose=True`` to display folding progress.
         """
 
         sequence_batch = [sequences] if isinstance(sequences, str) else sequences
@@ -816,6 +903,7 @@ class FastEsmForProteinFolding(FastPLMsAttentionMixin, EsmForProteinFolding):
             position_ids=residx,
             masking_pattern=masking_pattern,
             num_recycles=num_recycles,
+            verbose=verbose,
         )
         output["atom37_atom_exists"] = output["atom37_atom_exists"] * linker_mask.unsqueeze(2)
         output["mean_plddt"] = (output["plddt"] * output["atom37_atom_exists"]).sum(

modeling_fastplms.py

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