Update FastPLMs runtime and model cards

README.md

@@ -36,9 +36,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
@@ -113,6 +112,7 @@ output = model.predict_structure(
     num_sampling_steps=50,
     diffusion_samples=1,
     seed=7,
+    verbose=False,
 )
 model.save_as_cif(output, "prediction.cif")
 
@@ -149,7 +149,7 @@ continues independently of the ESM++ and ESMFold2 release gates.
 - 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/boltz/modeling_boltz2.py

@@ -9,6 +9,7 @@ from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from typing import Any, ClassVar
 from torch import Tensor
+from tqdm.auto import tqdm
 from transformers import PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import ModelOutput
 
@@ -572,10 +573,19 @@ class Boltz2InferenceCore(nn.Module):
         max_parallel_samples: int | None = None,
         run_confidence_sequentially: bool = True,
         detach_confidence: bool = True,
+        verbose: bool = False,
     ) -> dict[str, Tensor]:
         # b is the batch size, t the token count, a the padded atom count,
         # d_s the token width, d_z the pair width, and m the diffusion multiplicity.
-        s_inputs = self.input_embedder(feats)  # (b, t, d_s)
+        with tqdm(
+            total=1,
+            desc="Boltz2: Input embeddings",
+            disable=not verbose,
+            dynamic_ncols=True,
+            leave=False,
+        ) as progress:
+            s_inputs = self.input_embedder(feats)  # (b, t, d_s)
+            progress.update()
         s_init = self.s_init(s_inputs)  # (b, t, d_s)
 
         z_init = (
@@ -598,7 +608,14 @@ class Boltz2InferenceCore(nn.Module):
         pair_mask = mask[:, :, None] * mask[:, None, :]  # (b, t, t)
 
         if self.run_trunk_and_structure:
-            for _ in range(recycling_steps + 1):
+            for _ in tqdm(
+                range(recycling_steps + 1),
+                desc="Boltz2: Recycling",
+                unit="cycle",
+                disable=not verbose,
+                dynamic_ncols=True,
+                leave=False,
+            ):
                 s = s_init + self.s_recycle(self.s_norm(s))  # (b, t, d_s)
                 z = z_init + self.z_recycle(self.z_norm(z))  # (b, t, t, d_z)
                 z = z + self.msa_module(
@@ -650,6 +667,7 @@ class Boltz2InferenceCore(nn.Module):
                     max_parallel_samples=max_parallel_samples,
                     steering_args=self.steering_args,
                     diffusion_conditioning=diffusion_conditioning,
+                    verbose=verbose,
                 )  # tensor values include sample_atom_coords: (b * m, a, 3)
             output.update(struct_out)
 
@@ -676,19 +694,27 @@ class Boltz2InferenceCore(nn.Module):
                 x_pred_c = x_pred  # (b * m, ..., a, 3)
                 pdist_c = output["pdistogram"][:, :, :, 0]  # (b, t, t, n_bin)
 
-            output.update(
-                self.confidence_module(
-                    s_inputs=s_inputs_c,
-                    s=s_c,
-                    z=z_c,
-                    x_pred=x_pred_c,
-                    feats=feats,
-                    pred_distogram_logits=pdist_c,
-                    multiplicity=diffusion_samples,
-                    run_sequentially=run_confidence_sequentially,
-                    use_kernels=self.use_kernels,
+            with tqdm(
+                total=1,
+                desc="Boltz2: Confidence",
+                disable=not verbose,
+                dynamic_ncols=True,
+                leave=False,
+            ) as progress:
+                output.update(
+                    self.confidence_module(
+                        s_inputs=s_inputs_c,
+                        s=s_c,
+                        z=z_c,
+                        x_pred=x_pred_c,
+                        feats=feats,
+                        pred_distogram_logits=pdist_c,
+                        multiplicity=diffusion_samples,
+                        run_sequentially=run_confidence_sequentially,
+                        use_kernels=self.use_kernels,
+                    )
                 )
-            )
+                progress.update()
 
         return output  # named tensors retain the shapes traced above
 
@@ -836,6 +862,7 @@ class Boltz2Model(PreTrainedModel):
         output_attentions: bool | None = None,
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
+        verbose: bool = False,
     ) -> Boltz2ModelOutput | tuple[Any, ...]:
         # Feature shapes follow build_boltz2_features; b and t are batch and token counts.
         output_attentions = (
@@ -866,6 +893,7 @@ class Boltz2Model(PreTrainedModel):
             max_parallel_samples=max_parallel_samples,
             run_confidence_sequentially=run_confidence_sequentially,
             detach_confidence=detach_confidence,
+            verbose=verbose,
         )  # named tensors use the Boltz2InferenceCore.forward shapes
         token_state = raw_output.get("s")  # (b, t, d_s) or None
         pair_state = raw_output.get("z")  # (b, t, t, d_z) or None
@@ -907,6 +935,7 @@ class Boltz2Model(PreTrainedModel):
         run_confidence_sequentially: bool = True,
         float_dtype: torch.dtype = torch.float32,
         seed: int | None = None,
+        verbose: bool = False,
     ) -> Boltz2StructureOutput:
         if float_dtype != torch.float32:
             raise ValueError(
@@ -923,16 +952,24 @@ class Boltz2Model(PreTrainedModel):
             )
 
         with _seed_context(seed):
-            feats, template = build_boltz2_features(
-                amino_acid_sequence=amino_acid_sequence,
-                num_bins=self.config.num_bins,
-                atoms_per_window_queries=(
-                    self.core.input_embedder.atom_encoder.atoms_per_window_queries
-                ),
-            )  # feature shapes are traced in build_boltz2_features
-            feats = self._to_model_device(
-                feats, float_dtype=torch.float32
-            )  # unchanged feature shapes
+            with tqdm(
+                total=1,
+                desc="Boltz2: Prepare features",
+                disable=not verbose,
+                dynamic_ncols=True,
+                leave=False,
+            ) as progress:
+                feats, template = build_boltz2_features(
+                    amino_acid_sequence=amino_acid_sequence,
+                    num_bins=self.config.num_bins,
+                    atoms_per_window_queries=(
+                        self.core.input_embedder.atom_encoder.atoms_per_window_queries
+                    ),
+                )  # feature shapes are traced in build_boltz2_features
+                feats = self._to_model_device(
+                    feats, float_dtype=torch.float32
+                )  # unchanged feature shapes
+                progress.update()
             autocast_context = (
                 torch.autocast(device_type="cuda", dtype=torch.bfloat16)
                 if self.device.type == "cuda"
@@ -947,6 +984,7 @@ class Boltz2Model(PreTrainedModel):
                     max_parallel_samples=max_parallel_samples,
                     run_confidence_sequentially=run_confidence_sequentially,
                     return_dict=True,
+                    verbose=verbose,
                 )  # named tensors use Boltz2InferenceCore.forward shapes
 
         sample_atom_coords_value = output.get(

fastplms/models/boltz/vb_modules_diffusionv2.py

@@ -11,6 +11,7 @@ import torch.nn.functional as F
 from einops import rearrange
 from torch import nn
 from torch.nn import Module
+from tqdm.auto import tqdm
 
 from . import vb_const as const
 from . import vb_layers_initialize as init
@@ -294,6 +295,7 @@ class AtomDiffusion(Module):
         multiplicity=1,
         max_parallel_samples=None,
         steering_args=None,
+        verbose: bool = False,
         **network_condition_kwargs,
     ):
         if steering_args is not None and (
@@ -339,7 +341,16 @@ class AtomDiffusion(Module):
         atom_coords_denoised = None
 
         # gradually denoise
-        for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(sigmas_and_gammas):
+        for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(
+            tqdm(
+                sigmas_and_gammas,
+                desc="Boltz2: Diffusion sampling",
+                unit="step",
+                disable=not verbose,
+                dynamic_ncols=True,
+                leave=False,
+            )
+        ):
             random_R, random_tr = compute_random_augmentation(
                 multiplicity, device=atom_coords.device, dtype=atom_coords.dtype
             )

modeling_fastplms.py

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