Upload folder using huggingface_hub

cif_writer.py

@@ -0,0 +1,114 @@
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+
+from .minimal_structures import ProteinStructureTemplate
+
+
+def _confidence_per_atom(
+    plddt: Optional[torch.Tensor],
+    atom_to_residue: list[int],
+    num_atoms: int,
+    sample_index: int,
+) -> np.ndarray:
+    if plddt is None:
+        return np.ones((num_atoms,), dtype=np.float32) * 100.0
+
+    values = plddt.detach().cpu()
+    if values.ndim == 1:
+        values = values.unsqueeze(0)
+    assert values.ndim == 2, "Expected pLDDT with shape [samples, tokens/atoms]."
+    assert sample_index < values.shape[0], "sample_index out of range for pLDDT."
+
+    selected = values[sample_index]
+    if selected.shape[0] == num_atoms:
+        return (selected.numpy() * 100.0).astype(np.float32)
+
+    num_residues = max(atom_to_residue) + 1
+    if selected.shape[0] == num_residues:
+        expanded = np.zeros((num_atoms,), dtype=np.float32)
+        selected_np = selected.numpy()
+        for atom_idx, residue_idx in enumerate(atom_to_residue):
+            expanded[atom_idx] = selected_np[residue_idx] * 100.0
+        return expanded
+
+    return np.ones((num_atoms,), dtype=np.float32) * 100.0
+
+
+def write_cif(
+    structure_template: ProteinStructureTemplate,
+    atom_coords: torch.Tensor,
+    atom_mask: torch.Tensor,
+    output_path: str,
+    plddt: Optional[torch.Tensor] = None,
+    sample_index: int = 0,
+) -> str:
+    coords = atom_coords.detach().cpu()
+    if coords.ndim == 2:
+        coords = coords.unsqueeze(0)
+    assert coords.ndim == 3, "Expected coordinates with shape [samples, atoms, 3]."
+    assert sample_index < coords.shape[0], "sample_index out of range."
+    selected_coords = coords[sample_index].numpy()
+
+    mask = atom_mask.detach().cpu()
+    if mask.ndim == 2:
+        mask = mask[0]
+    assert mask.ndim == 1, "Expected atom mask with shape [atoms]."
+    assert mask.shape[0] == selected_coords.shape[0], "Atom mask/coord size mismatch."
+
+    b_iso = _confidence_per_atom(
+        plddt=plddt,
+        atom_to_residue=structure_template.atom_residue_index,
+        num_atoms=structure_template.num_atoms,
+        sample_index=sample_index,
+    )
+    assert b_iso.shape[0] == structure_template.num_atoms
+
+    lines = [
+        "data_boltz2_prediction",
+        "#",
+        "loop_",
+        "_atom_site.group_PDB",
+        "_atom_site.id",
+        "_atom_site.type_symbol",
+        "_atom_site.label_atom_id",
+        "_atom_site.label_comp_id",
+        "_atom_site.label_asym_id",
+        "_atom_site.label_seq_id",
+        "_atom_site.Cartn_x",
+        "_atom_site.Cartn_y",
+        "_atom_site.Cartn_z",
+        "_atom_site.occupancy",
+        "_atom_site.B_iso_or_equiv",
+        "_atom_site.pdbx_PDB_model_num",
+    ]
+
+    atom_id = 1
+    for idx in range(structure_template.num_atoms):
+        if mask[idx] <= 0:
+            continue
+
+        residue_idx = structure_template.atom_residue_index[idx]
+        residue_name = structure_template.residue_names[residue_idx]
+        atom_name = structure_template.atom_names[idx]
+        element = structure_template.atom_elements[idx]
+        chain_id = structure_template.atom_chain_id[idx]
+        x_val, y_val, z_val = selected_coords[idx].tolist()
+        b_factor = float(b_iso[idx])
+
+        line = (
+            f"ATOM {atom_id} {element} {atom_name} {residue_name} {chain_id} "
+            f"{residue_idx + 1} {x_val:.3f} {y_val:.3f} {z_val:.3f} 1.00 {b_factor:.2f} 1"
+        )
+        lines.append(line)
+        atom_id += 1
+
+    lines.append("#")
+    text = "\n".join(lines) + "\n"
+
+    out_path = Path(output_path)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_text(text, encoding="utf-8")
+    return str(out_path)

config.json

@@ -0,0 +1,138 @@
+{
+  "architectures": [
+    "Boltz2Model"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_boltz2.Boltz2Config",
+    "AutoModel": "modeling_boltz2.Boltz2Model"
+  },
+  "core_kwargs": {
+    "alpha_pae": 1,
+    "atom_feature_dim": 388,
+    "atom_s": 128,
+    "atom_z": 16,
+    "atoms_per_window_keys": 128,
+    "atoms_per_window_queries": 32,
+    "bond_type_feature": true,
+    "conditioning_cutoff_max": 20.0,
+    "conditioning_cutoff_min": 4.0,
+    "confidence_model_args": {
+      "add_s_input_to_s": true,
+      "add_s_to_z_prod": true,
+      "add_z_input_to_z": true,
+      "confidence_args": {
+        "num_pae_bins": 64,
+        "num_pde_bins": 64,
+        "num_plddt_bins": 50,
+        "relative_confidence": "none",
+        "use_separate_heads": true
+      },
+      "max_dist": 22,
+      "no_trunk_feats": false,
+      "num_dist_bins": 64,
+      "pairformer_args": {
+        "dropout": 0.25,
+        "num_blocks": 8,
+        "num_heads": 16,
+        "v2": true
+      },
+      "use_gaussian": false,
+      "use_miniformer": false,
+      "use_s_diffusion": false
+    },
+    "confidence_prediction": true,
+    "cyclic_pos_enc": true,
+    "diffusion_process_args": {
+      "P_mean": -1.2,
+      "P_std": 1.5,
+      "alignment_reverse_diff": true,
+      "coordinate_augmentation": true,
+      "gamma_0": 0.8,
+      "gamma_min": 1.0,
+      "mse_rotational_alignment": true,
+      "noise_scale": 1.0,
+      "rho": 7,
+      "sigma_data": 16.0,
+      "sigma_max": 160.0,
+      "sigma_min": 0.0004,
+      "step_scale": 1.0,
+      "step_scale_random": [
+        1.0,
+        1.1,
+        1.2,
+        1.3,
+        1.4,
+        1.5
+      ],
+      "synchronize_sigmas": false
+    },
+    "embedder_args": {
+      "add_cyclic_flag": true,
+      "add_method_conditioning": true,
+      "add_modified_flag": true,
+      "add_mol_type_feat": true,
+      "atom_encoder_depth": 3,
+      "atom_encoder_heads": 4
+    },
+    "fix_sym_check": true,
+    "msa_args": {
+      "activation_checkpointing": true,
+      "miniformer_blocks": false,
+      "msa_blocks": 4,
+      "msa_dropout": 0.15,
+      "msa_s": 64,
+      "pairwise_head_width": 32,
+      "pairwise_num_heads": 4,
+      "use_paired_feature": true,
+      "use_trifast": true,
+      "z_dropout": 0.25
+    },
+    "num_bins": 64,
+    "pairformer_args": {
+      "activation_checkpointing": true,
+      "dropout": 0.25,
+      "num_blocks": 64,
+      "num_heads": 16,
+      "post_layer_norm": false,
+      "use_trifast": true
+    },
+    "run_trunk_and_structure": true,
+    "score_model_args": {
+      "activation_checkpointing": true,
+      "atom_decoder_depth": 3,
+      "atom_decoder_heads": 4,
+      "atom_encoder_depth": 3,
+      "atom_encoder_heads": 4,
+      "conditioning_transition_layers": 2,
+      "dim_fourier": 256,
+      "sigma_data": 16,
+      "token_transformer_depth": 24,
+      "token_transformer_heads": 16,
+      "transformer_post_ln": false
+    },
+    "skip_run_structure": false,
+    "steering_args": {
+      "contact_guidance_update": false,
+      "fk_lambda": 4.0,
+      "fk_resampling_interval": 3,
+      "fk_steering": false,
+      "num_gd_steps": 16,
+      "num_particles": 3,
+      "physical_guidance_update": false
+    },
+    "token_level_confidence": true,
+    "token_s": 384,
+    "token_z": 128,
+    "use_atom_backbone_feat": false,
+    "use_kernels": false,
+    "use_no_atom_char": false,
+    "use_residue_feats_atoms": false
+  },
+  "default_diffusion_samples": 5,
+  "default_recycling_steps": 3,
+  "default_sampling_steps": 200,
+  "dtype": "float32",
+  "model_type": "boltz2_automodel",
+  "num_bins": 64,
+  "transformers_version": "5.1.0"
+}

minimal_featurizer.py

@@ -0,0 +1,469 @@
+import math
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+from torch.nn.functional import one_hot
+
+from .minimal_structures import ProteinStructureTemplate
+from .runtime import ensure_boltz_importable
+
+
+ensure_boltz_importable()
+from boltz.data import const  # noqa: E402
+
+
+_ELEMENT_TO_Z = {
+    "H": 1,
+    "C": 6,
+    "N": 7,
+    "O": 8,
+    "P": 15,
+    "S": 16,
+}
+
+
+def _normalize_sequence(sequence: str) -> str:
+    seq = sequence.strip().upper()
+    assert len(seq) > 0, "Amino acid sequence must be non-empty."
+    for aa in seq:
+        assert aa in const.prot_letter_to_token, f"Unsupported residue code '{aa}'."
+    return seq
+
+
+def _atom_name_to_element(atom_name: str) -> str:
+    name = atom_name.strip().upper()
+    if len(name) == 0:
+        return "C"
+    if name[0].isdigit():
+        name = name[1:]
+    if len(name) >= 2 and name[0:2] in ("CL", "BR", "FE", "MG", "ZN", "NA", "CA"):
+        return name[0]
+    return name[0]
+
+
+def _atom_name_to_codes(atom_name: str) -> torch.Tensor:
+    clipped = atom_name.strip()[:4]
+    vals = [ord(ch) - 32 for ch in clipped]
+    while len(vals) < 4:
+        vals.append(0)
+    out = torch.tensor(vals, dtype=torch.long)
+    assert torch.all(out >= 0) and torch.all(out < 64), (
+        f"Invalid atom-name encoding for '{atom_name}'."
+    )
+    return out
+
+
+def _default_atom_offset(atom_name: str, atom_idx: int) -> np.ndarray:
+    if atom_name == "N":
+        return np.array([-1.20, 0.80, 0.00], dtype=np.float32)
+    if atom_name == "CA":
+        return np.array([0.00, 0.00, 0.00], dtype=np.float32)
+    if atom_name == "C":
+        return np.array([1.35, 0.65, 0.00], dtype=np.float32)
+    if atom_name == "O":
+        return np.array([2.25, 0.10, 0.00], dtype=np.float32)
+    if atom_name == "CB":
+        return np.array([0.15, -1.20, 1.00], dtype=np.float32)
+    angle = (atom_idx + 1) * 0.7
+    radius = 1.4 + 0.03 * atom_idx
+    return np.array(
+        [
+            radius * math.cos(angle),
+            radius * math.sin(angle),
+            0.1 * ((atom_idx % 5) - 2),
+        ],
+        dtype=np.float32,
+    )
+
+
+def _build_template(
+    sequence: str,
+) -> Tuple[
+    ProteinStructureTemplate,
+    List[str],
+    List[int],
+    List[int],
+    List[int],
+    List[np.ndarray],
+    List[int],
+]:
+    residue_names: List[str] = []
+    residue_token_ids: List[int] = []
+    atom_names: List[str] = []
+    atom_elements: List[str] = []
+    atom_residue_index: List[int] = []
+    atom_chain_id: List[str] = []
+    atom_positions: List[np.ndarray] = []
+    residue_center_atom_idx: List[int] = []
+    residue_disto_atom_idx: List[int] = []
+    residue_frame_atom_idx: List[int] = []
+
+    global_atom_idx = 0
+    for res_idx, aa in enumerate(sequence):
+        token_name = const.prot_letter_to_token[aa]
+        residue_names.append(token_name)
+        residue_token_ids.append(const.token_ids[token_name])
+
+        residue_atoms = const.ref_atoms[token_name]
+        assert len(residue_atoms) > 0, f"No reference atoms for residue {token_name}."
+        center_atom_name = const.res_to_center_atom[token_name]
+        disto_atom_name = const.res_to_disto_atom[token_name]
+
+        base = np.array([3.8 * res_idx, 0.0, 0.0], dtype=np.float32)
+        center_idx = -1
+        disto_idx = -1
+        n_idx = -1
+        ca_idx = -1
+        c_idx = -1
+
+        for local_idx, atom_name in enumerate(residue_atoms):
+            atom_names.append(atom_name)
+            element = _atom_name_to_element(atom_name)
+            atom_elements.append(element)
+            atom_residue_index.append(res_idx)
+            atom_chain_id.append("A")
+
+            atom_pos = base + _default_atom_offset(atom_name, local_idx)
+            atom_positions.append(atom_pos)
+
+            if atom_name == center_atom_name:
+                center_idx = global_atom_idx
+            if atom_name == disto_atom_name:
+                disto_idx = global_atom_idx
+            if atom_name == "N":
+                n_idx = global_atom_idx
+            if atom_name == "CA":
+                ca_idx = global_atom_idx
+            if atom_name == "C":
+                c_idx = global_atom_idx
+            global_atom_idx += 1
+
+        if center_idx == -1:
+            center_idx = global_atom_idx - len(residue_atoms)
+        if disto_idx == -1:
+            disto_idx = center_idx
+        if n_idx == -1:
+            n_idx = center_idx
+        if ca_idx == -1:
+            ca_idx = center_idx
+        if c_idx == -1:
+            c_idx = center_idx
+
+        residue_center_atom_idx.append(center_idx)
+        residue_disto_atom_idx.append(disto_idx)
+        residue_frame_atom_idx.extend([n_idx, ca_idx, c_idx])
+
+    template = ProteinStructureTemplate(
+        sequence=sequence,
+        residue_names=residue_names,
+        atom_names=atom_names,
+        atom_elements=atom_elements,
+        atom_residue_index=atom_residue_index,
+        atom_chain_id=atom_chain_id,
+    )
+
+    return (
+        template,
+        residue_names,
+        residue_token_ids,
+        residue_center_atom_idx,
+        residue_disto_atom_idx,
+        atom_positions,
+        residue_frame_atom_idx,
+    )
+
+
+def build_boltz2_features(
+    amino_acid_sequence: str,
+    num_bins: int = 64,
+    atoms_per_window_queries: int = 32,
+) -> Tuple[Dict[str, torch.Tensor], ProteinStructureTemplate]:
+    sequence = _normalize_sequence(amino_acid_sequence)
+    (
+        template,
+        residue_names,
+        residue_token_ids,
+        residue_center_atom_idx,
+        residue_disto_atom_idx,
+        atom_positions_np,
+        residue_frame_atom_idx_flat,
+    ) = _build_template(sequence)
+
+    num_tokens = len(residue_names)
+    num_atoms = len(atom_positions_np)
+    assert num_tokens > 0 and num_atoms > 0
+
+    atom_positions = torch.tensor(np.asarray(atom_positions_np), dtype=torch.float32)
+    atom_positions = atom_positions - atom_positions.mean(dim=0, keepdim=True)
+
+    token_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0)
+    residue_index = torch.arange(num_tokens, dtype=torch.long).unsqueeze(0)
+    asym_id = torch.zeros((1, num_tokens), dtype=torch.long)
+    entity_id = torch.zeros((1, num_tokens), dtype=torch.long)
+    sym_id = torch.zeros((1, num_tokens), dtype=torch.long)
+    mol_type = torch.full(
+        (1, num_tokens),
+        fill_value=const.chain_type_ids["PROTEIN"],
+        dtype=torch.long,
+    )
+
+    res_type_ids = torch.tensor(residue_token_ids, dtype=torch.long)
+    res_type = one_hot(res_type_ids, num_classes=const.num_tokens).float().unsqueeze(0)
+
+    token_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.float32)
+    type_bonds = torch.zeros((num_tokens, num_tokens), dtype=torch.long)
+    if "COVALENT" in const.bond_type_ids:
+        covalent_bond_id = const.bond_type_ids["COVALENT"]
+    else:
+        covalent_bond_id = 1
+    for idx in range(num_tokens - 1):
+        token_bonds[idx, idx + 1] = 1.0
+        token_bonds[idx + 1, idx] = 1.0
+        type_bonds[idx, idx + 1] = covalent_bond_id
+        type_bonds[idx + 1, idx] = covalent_bond_id
+    token_bonds = token_bonds.unsqueeze(0).unsqueeze(-1)
+    type_bonds = type_bonds.unsqueeze(0)
+
+    token_pad_mask = torch.ones((1, num_tokens), dtype=torch.float32)
+    token_resolved_mask = torch.ones((1, num_tokens), dtype=torch.float32)
+    token_disto_mask = torch.ones((1, num_tokens), dtype=torch.float32)
+
+    num_contact_classes = len(const.contact_conditioning_info)
+    unspecified_id = const.contact_conditioning_info["UNSPECIFIED"]
+    contact_ids = torch.full(
+        (num_tokens, num_tokens),
+        fill_value=unspecified_id,
+        dtype=torch.long,
+    )
+    contact_conditioning = one_hot(
+        contact_ids,
+        num_classes=num_contact_classes,
+    ).float().unsqueeze(0)
+    contact_threshold = torch.zeros((1, num_tokens, num_tokens), dtype=torch.float32)
+
+    assert "x-ray diffraction" in const.method_types_ids
+    method_feature = torch.full(
+        (1, num_tokens),
+        fill_value=const.method_types_ids["x-ray diffraction"],
+        dtype=torch.long,
+    )
+    modified = torch.zeros((1, num_tokens), dtype=torch.long)
+    cyclic_period = torch.zeros((1, num_tokens), dtype=torch.float32)
+    affinity_token_mask = torch.zeros((1, num_tokens), dtype=torch.float32)
+
+    ref_pos = atom_positions.unsqueeze(0)
+    atom_pad_mask = torch.ones((1, num_atoms), dtype=torch.float32)
+    atom_resolved_mask = torch.ones((1, num_atoms), dtype=torch.float32)
+
+    atom_name_codes = torch.stack(
+        [_atom_name_to_codes(atom_name) for atom_name in template.atom_names],
+        dim=0,
+    )
+    ref_atom_name_chars = one_hot(atom_name_codes, num_classes=64).float().unsqueeze(0)
+
+    atomic_numbers = []
+    for element in template.atom_elements:
+        if element in _ELEMENT_TO_Z:
+            z_value = _ELEMENT_TO_Z[element]
+        else:
+            z_value = _ELEMENT_TO_Z["C"]
+        assert z_value < const.num_elements
+        atomic_numbers.append(z_value)
+    ref_element = one_hot(
+        torch.tensor(atomic_numbers, dtype=torch.long),
+        num_classes=const.num_elements,
+    ).float().unsqueeze(0)
+
+    ref_charge = torch.zeros((1, num_atoms), dtype=torch.float32)
+    ref_chirality = torch.zeros((1, num_atoms), dtype=torch.long)
+    ref_space_uid = torch.tensor(template.atom_residue_index, dtype=torch.long).unsqueeze(0)
+
+    atom_to_token = one_hot(
+        torch.tensor(template.atom_residue_index, dtype=torch.long),
+        num_classes=num_tokens,
+    ).float().unsqueeze(0)
+    token_to_rep_atom = one_hot(
+        torch.tensor(residue_disto_atom_idx, dtype=torch.long),
+        num_classes=num_atoms,
+    ).float().unsqueeze(0)
+    token_to_center_atom = one_hot(
+        torch.tensor(residue_center_atom_idx, dtype=torch.long),
+        num_classes=num_atoms,
+    ).float().unsqueeze(0)
+    r_set_to_rep_atom = token_to_center_atom.clone()
+
+    num_backbone_classes = (
+        1
+        + len(const.protein_backbone_atom_index)
+        + len(const.nucleic_backbone_atom_index)
+    )
+    backbone_ids = []
+    for atom_name in template.atom_names:
+        if atom_name in const.protein_backbone_atom_index:
+            backbone_ids.append(const.protein_backbone_atom_index[atom_name] + 1)
+        else:
+            backbone_ids.append(0)
+    atom_backbone_feat = one_hot(
+        torch.tensor(backbone_ids, dtype=torch.long),
+        num_classes=num_backbone_classes,
+    ).float().unsqueeze(0)
+
+    coords = ref_pos.unsqueeze(1).contiguous()
+    disto_coords = torch.stack(
+        [atom_positions[idx] for idx in residue_disto_atom_idx],
+        dim=0,
+    )
+    disto_coords_ensemble = disto_coords.unsqueeze(0).unsqueeze(0).contiguous()
+
+    bfactor = torch.zeros((1, num_atoms), dtype=torch.float32)
+    atom_plddt = torch.ones((1, num_atoms), dtype=torch.float32)
+
+    assert atoms_per_window_queries > 0
+    pad_atoms = (
+        ((num_atoms - 1) // atoms_per_window_queries + 1) * atoms_per_window_queries
+        - num_atoms
+    )
+    if pad_atoms > 0:
+        ref_pos = torch.nn.functional.pad(ref_pos, (0, 0, 0, pad_atoms), value=0.0)
+        atom_pad_mask = torch.nn.functional.pad(atom_pad_mask, (0, pad_atoms), value=0.0)
+        atom_resolved_mask = torch.nn.functional.pad(
+            atom_resolved_mask,
+            (0, pad_atoms),
+            value=0.0,
+        )
+        ref_atom_name_chars = torch.nn.functional.pad(
+            ref_atom_name_chars,
+            (0, 0, 0, 0, 0, pad_atoms),
+            value=0.0,
+        )
+        ref_element = torch.nn.functional.pad(ref_element, (0, 0, 0, pad_atoms), value=0.0)
+        ref_charge = torch.nn.functional.pad(ref_charge, (0, pad_atoms), value=0.0)
+        ref_chirality = torch.nn.functional.pad(ref_chirality, (0, pad_atoms), value=0)
+        atom_backbone_feat = torch.nn.functional.pad(
+            atom_backbone_feat,
+            (0, 0, 0, pad_atoms),
+            value=0.0,
+        )
+        ref_space_uid = torch.nn.functional.pad(ref_space_uid, (0, pad_atoms), value=0)
+        coords = torch.nn.functional.pad(coords, (0, 0, 0, pad_atoms), value=0.0)
+        atom_to_token = torch.nn.functional.pad(atom_to_token, (0, 0, 0, pad_atoms), value=0.0)
+        token_to_rep_atom = torch.nn.functional.pad(
+            token_to_rep_atom,
+            (0, pad_atoms),
+            value=0.0,
+        )
+        token_to_center_atom = torch.nn.functional.pad(
+            token_to_center_atom,
+            (0, pad_atoms),
+            value=0.0,
+        )
+        r_set_to_rep_atom = torch.nn.functional.pad(
+            r_set_to_rep_atom,
+            (0, pad_atoms),
+            value=0.0,
+        )
+        bfactor = torch.nn.functional.pad(bfactor, (0, pad_atoms), value=0.0)
+        atom_plddt = torch.nn.functional.pad(atom_plddt, (0, pad_atoms), value=0.0)
+
+    frames_idx = torch.tensor(
+        residue_frame_atom_idx_flat,
+        dtype=torch.long,
+    ).reshape(num_tokens, 3)
+    frames_idx = frames_idx.unsqueeze(0).unsqueeze(1)
+    frame_resolved_mask = torch.ones((1, 1, num_tokens), dtype=torch.float32)
+
+    msa = torch.tensor(residue_token_ids, dtype=torch.long).unsqueeze(0).unsqueeze(0)
+    msa_paired = torch.ones((1, 1, num_tokens), dtype=torch.float32)
+    deletion_value = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
+    has_deletion = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
+    msa_mask = torch.ones((1, 1, num_tokens), dtype=torch.float32)
+    deletion_mean = torch.zeros((1, num_tokens), dtype=torch.float32)
+    profile = one_hot(
+        torch.tensor(residue_token_ids, dtype=torch.long),
+        num_classes=const.num_tokens,
+    ).float().unsqueeze(0)
+
+    template_restype = one_hot(
+        torch.zeros((1, 1, num_tokens), dtype=torch.long),
+        num_classes=const.num_tokens,
+    ).float()
+    template_frame_rot = torch.zeros((1, 1, num_tokens, 3, 3), dtype=torch.float32)
+    template_frame_t = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
+    template_cb = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
+    template_ca = torch.zeros((1, 1, num_tokens, 3), dtype=torch.float32)
+    template_mask_cb = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
+    template_mask_frame = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
+    template_mask = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
+    query_to_template = torch.zeros((1, 1, num_tokens), dtype=torch.long)
+    visibility_ids = torch.zeros((1, 1, num_tokens), dtype=torch.float32)
+
+    disto_target = torch.zeros(
+        (1, num_tokens, num_tokens, 1, num_bins),
+        dtype=torch.float32,
+    )
+    disto_center = torch.stack(
+        [atom_positions[idx] for idx in residue_disto_atom_idx],
+        dim=0,
+    ).unsqueeze(0)
+
+    features: Dict[str, torch.Tensor] = {
+        "token_index": token_index,
+        "residue_index": residue_index,
+        "asym_id": asym_id,
+        "entity_id": entity_id,
+        "sym_id": sym_id,
+        "mol_type": mol_type,
+        "res_type": res_type,
+        "disto_center": disto_center,
+        "token_bonds": token_bonds,
+        "type_bonds": type_bonds,
+        "token_pad_mask": token_pad_mask,
+        "token_resolved_mask": token_resolved_mask,
+        "token_disto_mask": token_disto_mask,
+        "contact_conditioning": contact_conditioning,
+        "contact_threshold": contact_threshold,
+        "method_feature": method_feature,
+        "modified": modified,
+        "cyclic_period": cyclic_period,
+        "affinity_token_mask": affinity_token_mask,
+        "ref_pos": ref_pos,
+        "atom_resolved_mask": atom_resolved_mask,
+        "ref_atom_name_chars": ref_atom_name_chars,
+        "ref_element": ref_element,
+        "ref_charge": ref_charge,
+        "ref_chirality": ref_chirality,
+        "atom_backbone_feat": atom_backbone_feat,
+        "ref_space_uid": ref_space_uid,
+        "coords": coords,
+        "atom_pad_mask": atom_pad_mask,
+        "atom_to_token": atom_to_token,
+        "token_to_rep_atom": token_to_rep_atom,
+        "r_set_to_rep_atom": r_set_to_rep_atom,
+        "token_to_center_atom": token_to_center_atom,
+        "disto_target": disto_target,
+        "disto_coords_ensemble": disto_coords_ensemble,
+        "bfactor": bfactor,
+        "plddt": atom_plddt,
+        "frames_idx": frames_idx,
+        "frame_resolved_mask": frame_resolved_mask,
+        "msa": msa,
+        "msa_paired": msa_paired,
+        "deletion_value": deletion_value,
+        "has_deletion": has_deletion,
+        "deletion_mean": deletion_mean,
+        "profile": profile,
+        "msa_mask": msa_mask,
+        "template_restype": template_restype,
+        "template_frame_rot": template_frame_rot,
+        "template_frame_t": template_frame_t,
+        "template_cb": template_cb,
+        "template_ca": template_ca,
+        "template_mask_cb": template_mask_cb,
+        "template_mask_frame": template_mask_frame,
+        "template_mask": template_mask,
+        "query_to_template": query_to_template,
+        "visibility_ids": visibility_ids,
+    }
+
+    return features, template

minimal_structures.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from typing import List
+
+
+@dataclass
+class ProteinStructureTemplate:
+    sequence: str
+    residue_names: List[str]
+    atom_names: List[str]
+    atom_elements: List[str]
+    atom_residue_index: List[int]
+    atom_chain_id: List[str]
+
+    @property
+    def num_atoms(self) -> int:
+        return len(self.atom_names)
+
+    @property
+    def num_residues(self) -> int:
+        return len(self.residue_names)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da4a19b5c0afb0643b6895cd283b9e11e576a41d5605e5207727b564b8d14649
+size 2083451028

modeling_boltz2.py

@@ -0,0 +1,813 @@
+import copy
+import inspect
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import torch
+import torch._dynamo
+import torch.nn as nn
+from torch import Tensor
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import ModelOutput
+
+from .cif_writer import write_cif
+from .minimal_featurizer import build_boltz2_features
+from .minimal_structures import ProteinStructureTemplate
+from .runtime import ensure_boltz_importable
+
+
+ensure_boltz_importable()
+import boltz.model.layers.initialize as init  # noqa: E402
+from boltz.data import const  # noqa: E402
+from boltz.model.layers.pairformer import PairformerModule  # noqa: E402
+from boltz.model.modules.confidencev2 import ConfidenceModule  # noqa: E402
+from boltz.model.modules.diffusion_conditioning import DiffusionConditioning  # noqa: E402
+from boltz.model.modules.diffusionv2 import AtomDiffusion, DiffusionModule  # noqa: E402
+from boltz.model.modules.encodersv2 import RelativePositionEncoder  # noqa: E402
+from boltz.model.modules.trunkv2 import (  # noqa: E402
+    ContactConditioning,
+    DistogramModule,
+    InputEmbedder,
+    MSAModule,
+)
+
+
+def _default_steering_args() -> dict[str, Any]:
+    return {
+        "fk_steering": False,
+        "num_particles": 3,
+        "fk_lambda": 4.0,
+        "fk_resampling_interval": 3,
+        "physical_guidance_update": False,
+        "contact_guidance_update": False,
+        "num_gd_steps": 16,
+    }
+
+
+def _require_key(mapping: dict[str, Any], key: str) -> Any:
+    assert key in mapping, f"Missing required key '{key}' in checkpoint hyperparameters."
+    return mapping[key]
+
+
+def _state_dict_without_wrappers(state_dict: dict[str, Tensor]) -> dict[str, Tensor]:
+    cleaned: dict[str, Tensor] = {}
+    for key, value in state_dict.items():
+        if key.startswith("ema."):
+            continue
+        new_key = key
+        if new_key.startswith("model."):
+            new_key = new_key[len("model.") :]
+        if new_key.startswith("module."):
+            new_key = new_key[len("module.") :]
+        cleaned[new_key] = value
+    return cleaned
+
+
+def _to_cpu_detached(value: Any) -> Any:
+    if torch.is_tensor(value):
+        return value.detach().cpu()
+    if isinstance(value, dict):
+        out: dict[Any, Any] = {}
+        for key, nested_value in value.items():
+            out[key] = _to_cpu_detached(nested_value)
+        return out
+    if isinstance(value, list):
+        return [_to_cpu_detached(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_to_cpu_detached(item) for item in value)
+    return value
+
+
+def _to_plain_python(value: Any) -> Any:
+    if isinstance(value, Mapping):
+        out: dict[Any, Any] = {}
+        for key, nested_value in value.items():
+            out[key] = _to_plain_python(nested_value)
+        return out
+    if isinstance(value, list):
+        return [_to_plain_python(item) for item in value]
+    if isinstance(value, tuple):
+        return [_to_plain_python(item) for item in value]
+    if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
+        return [_to_plain_python(item) for item in value]
+    return value
+
+
+def _filtered_kwargs(target: Any, kwargs: dict[str, Any]) -> dict[str, Any]:
+    signature = inspect.signature(target.__init__)
+    allowed = set(signature.parameters.keys())
+    allowed.discard("self")
+    filtered: dict[str, Any] = {}
+    for key, value in kwargs.items():
+        if key in allowed:
+            filtered[key] = value
+    return filtered
+
+
+@dataclass
+class Boltz2StructureOutput(ModelOutput):
+    sample_atom_coords: Optional[torch.Tensor] = None
+    atom_pad_mask: Optional[torch.Tensor] = None
+    plddt: Optional[torch.Tensor] = None
+    confidence_score: Optional[torch.Tensor] = None
+    complex_plddt: Optional[torch.Tensor] = None
+    iptm: Optional[torch.Tensor] = None
+    ptm: Optional[torch.Tensor] = None
+    sequence: Optional[str] = None
+    structure_template: Optional[ProteinStructureTemplate] = None
+    raw_output: Optional[dict[str, torch.Tensor]] = None
+
+
+class Boltz2Config(PretrainedConfig):
+    model_type = "boltz2_automodel"
+
+    def __init__(
+        self,
+        core_kwargs: Optional[dict[str, Any]] = None,
+        num_bins: int = 64,
+        default_recycling_steps: int = 3,
+        default_sampling_steps: int = 200,
+        default_diffusion_samples: int = 1,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        if core_kwargs is None:
+            core_kwargs = {}
+        self.core_kwargs = core_kwargs
+        self.num_bins = num_bins
+        self.default_recycling_steps = default_recycling_steps
+        self.default_sampling_steps = default_sampling_steps
+        self.default_diffusion_samples = default_diffusion_samples
+
+    @classmethod
+    def from_hyperparameters(
+        cls,
+        hparams: dict[str, Any],
+        use_kernels: bool = False,
+        default_recycling_steps: Optional[int] = None,
+        default_sampling_steps: Optional[int] = None,
+        default_diffusion_samples: Optional[int] = None,
+    ) -> "Boltz2Config":
+        assert isinstance(hparams, dict), "Expected checkpoint hyperparameters as a dictionary."
+        required = [
+            "atom_s",
+            "atom_z",
+            "token_s",
+            "token_z",
+            "num_bins",
+            "embedder_args",
+            "msa_args",
+            "pairformer_args",
+            "score_model_args",
+            "diffusion_process_args",
+        ]
+        for key in required:
+            _require_key(hparams, key)
+
+        core_kwargs: dict[str, Any] = {
+            "atom_s": hparams["atom_s"],
+            "atom_z": hparams["atom_z"],
+            "token_s": hparams["token_s"],
+            "token_z": hparams["token_z"],
+            "num_bins": hparams["num_bins"],
+            "embedder_args": _to_plain_python(copy.deepcopy(hparams["embedder_args"])),
+            "msa_args": _to_plain_python(copy.deepcopy(hparams["msa_args"])),
+            "pairformer_args": _to_plain_python(copy.deepcopy(hparams["pairformer_args"])),
+            "score_model_args": _to_plain_python(copy.deepcopy(hparams["score_model_args"])),
+            "diffusion_process_args": _to_plain_python(
+                copy.deepcopy(hparams["diffusion_process_args"])
+            ),
+            "use_kernels": use_kernels,
+        }
+
+        if "confidence_model_args" in hparams:
+            core_kwargs["confidence_model_args"] = _to_plain_python(
+                copy.deepcopy(hparams["confidence_model_args"])
+            )
+        else:
+            core_kwargs["confidence_model_args"] = None
+
+        if "confidence_prediction" in hparams:
+            core_kwargs["confidence_prediction"] = hparams["confidence_prediction"]
+        else:
+            core_kwargs["confidence_prediction"] = True
+
+        if "token_level_confidence" in hparams:
+            core_kwargs["token_level_confidence"] = hparams["token_level_confidence"]
+        else:
+            core_kwargs["token_level_confidence"] = True
+
+        if "alpha_pae" in hparams:
+            core_kwargs["alpha_pae"] = hparams["alpha_pae"]
+        else:
+            core_kwargs["alpha_pae"] = 0.0
+
+        if "atoms_per_window_queries" in hparams:
+            core_kwargs["atoms_per_window_queries"] = hparams["atoms_per_window_queries"]
+        else:
+            core_kwargs["atoms_per_window_queries"] = 32
+
+        if "atoms_per_window_keys" in hparams:
+            core_kwargs["atoms_per_window_keys"] = hparams["atoms_per_window_keys"]
+        else:
+            core_kwargs["atoms_per_window_keys"] = 128
+
+        if "atom_feature_dim" in hparams:
+            core_kwargs["atom_feature_dim"] = hparams["atom_feature_dim"]
+        else:
+            core_kwargs["atom_feature_dim"] = 128
+
+        if "bond_type_feature" in hparams:
+            core_kwargs["bond_type_feature"] = hparams["bond_type_feature"]
+        else:
+            core_kwargs["bond_type_feature"] = False
+
+        if "run_trunk_and_structure" in hparams:
+            core_kwargs["run_trunk_and_structure"] = hparams["run_trunk_and_structure"]
+        else:
+            core_kwargs["run_trunk_and_structure"] = True
+
+        if "skip_run_structure" in hparams:
+            core_kwargs["skip_run_structure"] = hparams["skip_run_structure"]
+        else:
+            core_kwargs["skip_run_structure"] = False
+
+        if "fix_sym_check" in hparams:
+            core_kwargs["fix_sym_check"] = hparams["fix_sym_check"]
+        else:
+            core_kwargs["fix_sym_check"] = False
+
+        if "cyclic_pos_enc" in hparams:
+            core_kwargs["cyclic_pos_enc"] = hparams["cyclic_pos_enc"]
+        else:
+            core_kwargs["cyclic_pos_enc"] = False
+
+        if "use_no_atom_char" in hparams:
+            core_kwargs["use_no_atom_char"] = hparams["use_no_atom_char"]
+        else:
+            core_kwargs["use_no_atom_char"] = False
+
+        if "use_atom_backbone_feat" in hparams:
+            core_kwargs["use_atom_backbone_feat"] = hparams["use_atom_backbone_feat"]
+        else:
+            core_kwargs["use_atom_backbone_feat"] = False
+
+        if "use_residue_feats_atoms" in hparams:
+            core_kwargs["use_residue_feats_atoms"] = hparams["use_residue_feats_atoms"]
+        else:
+            core_kwargs["use_residue_feats_atoms"] = False
+
+        if "conditioning_cutoff_min" in hparams:
+            core_kwargs["conditioning_cutoff_min"] = hparams["conditioning_cutoff_min"]
+        else:
+            core_kwargs["conditioning_cutoff_min"] = 4.0
+
+        if "conditioning_cutoff_max" in hparams:
+            core_kwargs["conditioning_cutoff_max"] = hparams["conditioning_cutoff_max"]
+        else:
+            core_kwargs["conditioning_cutoff_max"] = 20.0
+
+        if "steering_args" in hparams and hparams["steering_args"] is not None:
+            core_kwargs["steering_args"] = _to_plain_python(
+                copy.deepcopy(hparams["steering_args"])
+            )
+        else:
+            core_kwargs["steering_args"] = _default_steering_args()
+
+        if "validation_args" in hparams:
+            validation_args = hparams["validation_args"]
+            assert isinstance(validation_args, Mapping), (
+                "Expected 'validation_args' in checkpoint hyperparameters to be a mapping."
+            )
+            if default_recycling_steps is None and "recycling_steps" in validation_args:
+                default_recycling_steps = validation_args["recycling_steps"]
+            if default_sampling_steps is None and "sampling_steps" in validation_args:
+                default_sampling_steps = validation_args["sampling_steps"]
+            if default_diffusion_samples is None and "diffusion_samples" in validation_args:
+                default_diffusion_samples = validation_args["diffusion_samples"]
+
+        if default_recycling_steps is None:
+            default_recycling_steps = 3
+        if default_sampling_steps is None:
+            default_sampling_steps = 200
+        if default_diffusion_samples is None:
+            default_diffusion_samples = 1
+
+        return cls(
+            core_kwargs=core_kwargs,
+            num_bins=hparams["num_bins"],
+            default_recycling_steps=default_recycling_steps,
+            default_sampling_steps=default_sampling_steps,
+            default_diffusion_samples=default_diffusion_samples,
+        )
+
+
+class Boltz2InferenceCore(nn.Module):
+    def __init__(
+        self,
+        atom_s: int,
+        atom_z: int,
+        token_s: int,
+        token_z: int,
+        num_bins: int,
+        embedder_args: dict[str, Any],
+        msa_args: dict[str, Any],
+        pairformer_args: dict[str, Any],
+        score_model_args: dict[str, Any],
+        diffusion_process_args: dict[str, Any],
+        confidence_model_args: Optional[dict[str, Any]] = None,
+        atom_feature_dim: int = 128,
+        confidence_prediction: bool = True,
+        token_level_confidence: bool = True,
+        alpha_pae: float = 0.0,
+        atoms_per_window_queries: int = 32,
+        atoms_per_window_keys: int = 128,
+        run_trunk_and_structure: bool = True,
+        skip_run_structure: bool = False,
+        bond_type_feature: bool = False,
+        fix_sym_check: bool = False,
+        cyclic_pos_enc: bool = False,
+        use_no_atom_char: bool = False,
+        use_atom_backbone_feat: bool = False,
+        use_residue_feats_atoms: bool = False,
+        conditioning_cutoff_min: float = 4.0,
+        conditioning_cutoff_max: float = 20.0,
+        use_kernels: bool = False,
+        steering_args: Optional[dict[str, Any]] = None,
+    ) -> None:
+        super().__init__()
+        self.use_kernels = use_kernels
+        self.confidence_prediction = confidence_prediction
+        self.token_level_confidence = token_level_confidence
+        self.alpha_pae = alpha_pae
+        self.run_trunk_and_structure = run_trunk_and_structure
+        self.skip_run_structure = skip_run_structure
+        self.bond_type_feature = bond_type_feature
+        self.steering_args = steering_args if steering_args is not None else _default_steering_args()
+
+        full_embedder_args = {
+            "atom_s": atom_s,
+            "atom_z": atom_z,
+            "token_s": token_s,
+            "token_z": token_z,
+            "atoms_per_window_queries": atoms_per_window_queries,
+            "atoms_per_window_keys": atoms_per_window_keys,
+            "atom_feature_dim": atom_feature_dim,
+            "use_no_atom_char": use_no_atom_char,
+            "use_atom_backbone_feat": use_atom_backbone_feat,
+            "use_residue_feats_atoms": use_residue_feats_atoms,
+            **embedder_args,
+        }
+        full_embedder_args = _filtered_kwargs(InputEmbedder, full_embedder_args)
+        self.input_embedder = InputEmbedder(**full_embedder_args)
+
+        self.s_init = nn.Linear(token_s, token_s, bias=False)
+        self.z_init_1 = nn.Linear(token_s, token_z, bias=False)
+        self.z_init_2 = nn.Linear(token_s, token_z, bias=False)
+        self.rel_pos = RelativePositionEncoder(
+            token_z,
+            fix_sym_check=fix_sym_check,
+            cyclic_pos_enc=cyclic_pos_enc,
+        )
+        self.token_bonds = nn.Linear(1, token_z, bias=False)
+        if self.bond_type_feature:
+            self.token_bonds_type = nn.Embedding(len(const.bond_types) + 1, token_z)
+
+        self.contact_conditioning = ContactConditioning(
+            token_z=token_z,
+            cutoff_min=conditioning_cutoff_min,
+            cutoff_max=conditioning_cutoff_max,
+        )
+        self.s_norm = nn.LayerNorm(token_s)
+        self.z_norm = nn.LayerNorm(token_z)
+
+        self.s_recycle = nn.Linear(token_s, token_s, bias=False)
+        self.z_recycle = nn.Linear(token_z, token_z, bias=False)
+        init.gating_init_(self.s_recycle.weight)
+        init.gating_init_(self.z_recycle.weight)
+
+        torch._dynamo.config.cache_size_limit = 512  # noqa: SLF001
+        torch._dynamo.config.accumulated_cache_size_limit = 512  # noqa: SLF001
+
+        msa_kwargs = _filtered_kwargs(MSAModule, {"token_z": token_z, "token_s": token_s, **msa_args})
+        self.msa_module = MSAModule(**msa_kwargs)
+
+        pairformer_kwargs = _filtered_kwargs(
+            PairformerModule,
+            {"token_s": token_s, "token_z": token_z, **pairformer_args},
+        )
+        assert "token_s" in pairformer_kwargs and "token_z" in pairformer_kwargs
+        pairformer_token_s = pairformer_kwargs.pop("token_s")
+        pairformer_token_z = pairformer_kwargs.pop("token_z")
+        self.pairformer_module = PairformerModule(
+            pairformer_token_s,
+            pairformer_token_z,
+            **pairformer_kwargs,
+        )
+
+        diffusion_conditioning_kwargs = {
+            "token_s": token_s,
+            "token_z": token_z,
+            "atom_s": atom_s,
+            "atom_z": atom_z,
+            "atoms_per_window_queries": atoms_per_window_queries,
+            "atoms_per_window_keys": atoms_per_window_keys,
+            "atom_encoder_depth": score_model_args["atom_encoder_depth"],
+            "atom_encoder_heads": score_model_args["atom_encoder_heads"],
+            "token_transformer_depth": score_model_args["token_transformer_depth"],
+            "token_transformer_heads": score_model_args["token_transformer_heads"],
+            "atom_decoder_depth": score_model_args["atom_decoder_depth"],
+            "atom_decoder_heads": score_model_args["atom_decoder_heads"],
+            "atom_feature_dim": atom_feature_dim,
+            "conditioning_transition_layers": score_model_args["conditioning_transition_layers"],
+            "use_no_atom_char": use_no_atom_char,
+            "use_atom_backbone_feat": use_atom_backbone_feat,
+            "use_residue_feats_atoms": use_residue_feats_atoms,
+        }
+        diffusion_conditioning_kwargs = _filtered_kwargs(
+            DiffusionConditioning,
+            diffusion_conditioning_kwargs,
+        )
+        self.diffusion_conditioning = DiffusionConditioning(**diffusion_conditioning_kwargs)
+
+        structure_score_model_args = {
+            "token_s": token_s,
+            "atom_s": atom_s,
+            "atoms_per_window_queries": atoms_per_window_queries,
+            "atoms_per_window_keys": atoms_per_window_keys,
+            **score_model_args,
+        }
+        structure_score_model_args = _filtered_kwargs(
+            DiffusionModule,
+            structure_score_model_args,
+        )
+        structure_module_kwargs = {
+            "score_model_args": structure_score_model_args,
+            "compile_score": False,
+            **diffusion_process_args,
+        }
+        structure_module_kwargs = _filtered_kwargs(AtomDiffusion, structure_module_kwargs)
+        self.structure_module = AtomDiffusion(**structure_module_kwargs)
+        self.distogram_module = DistogramModule(token_z, num_bins)
+
+        if self.confidence_prediction:
+            assert confidence_model_args is not None, (
+                "confidence_prediction=True requires confidence_model_args in config."
+            )
+            confidence_kwargs = {
+                "token_s": token_s,
+                "token_z": token_z,
+                "token_level_confidence": token_level_confidence,
+                "bond_type_feature": bond_type_feature,
+                "fix_sym_check": fix_sym_check,
+                "cyclic_pos_enc": cyclic_pos_enc,
+                "conditioning_cutoff_min": conditioning_cutoff_min,
+                "conditioning_cutoff_max": conditioning_cutoff_max,
+                **confidence_model_args,
+            }
+            confidence_kwargs = _filtered_kwargs(ConfidenceModule, confidence_kwargs)
+            self.confidence_module = ConfidenceModule(**confidence_kwargs)
+
+    def forward(
+        self,
+        feats: dict[str, Tensor],
+        recycling_steps: int = 3,
+        num_sampling_steps: Optional[int] = None,
+        diffusion_samples: int = 1,
+        max_parallel_samples: Optional[int] = None,
+        run_confidence_sequentially: bool = True,
+    ) -> dict[str, Tensor]:
+        s_inputs = self.input_embedder(feats)
+        s_init = self.s_init(s_inputs)
+
+        z_init = self.z_init_1(s_inputs)[:, :, None] + self.z_init_2(s_inputs)[:, None, :]
+        relative_position_encoding = self.rel_pos(feats)
+        z_init = z_init + relative_position_encoding
+        z_init = z_init + self.token_bonds(feats["token_bonds"].float())
+        if self.bond_type_feature:
+            z_init = z_init + self.token_bonds_type(feats["type_bonds"].long())
+        z_init = z_init + self.contact_conditioning(feats)
+
+        s = torch.zeros_like(s_init)
+        z = torch.zeros_like(z_init)
+        mask = feats["token_pad_mask"].float()
+        pair_mask = mask[:, :, None] * mask[:, None, :]
+
+        if self.run_trunk_and_structure:
+            for _ in range(recycling_steps + 1):
+                s = s_init + self.s_recycle(self.s_norm(s))
+                z = z_init + self.z_recycle(self.z_norm(z))
+                z = z + self.msa_module(
+                    z,
+                    s_inputs,
+                    feats,
+                    use_kernels=self.use_kernels,
+                )
+                s, z = self.pairformer_module(
+                    s,
+                    z,
+                    mask=mask,
+                    pair_mask=pair_mask,
+                    use_kernels=self.use_kernels,
+                )
+
+        pdistogram = self.distogram_module(z)
+        output: dict[str, Tensor] = {
+            "pdistogram": pdistogram,
+            "s": s,
+            "z": z,
+        }
+
+        if self.run_trunk_and_structure and (not self.skip_run_structure):
+            q, c, to_keys, atom_enc_bias, atom_dec_bias, token_trans_bias = (
+                self.diffusion_conditioning(
+                    s_trunk=s,
+                    z_trunk=z,
+                    relative_position_encoding=relative_position_encoding,
+                    feats=feats,
+                )
+            )
+            diffusion_conditioning = {
+                "q": q,
+                "c": c,
+                "to_keys": to_keys,
+                "atom_enc_bias": atom_enc_bias,
+                "atom_dec_bias": atom_dec_bias,
+                "token_trans_bias": token_trans_bias,
+            }
+            with torch.autocast("cuda", enabled=False):
+                struct_out = self.structure_module.sample(
+                    s_trunk=s.float(),
+                    s_inputs=s_inputs.float(),
+                    feats=feats,
+                    num_sampling_steps=num_sampling_steps,
+                    atom_mask=feats["atom_pad_mask"].float(),
+                    multiplicity=diffusion_samples,
+                    max_parallel_samples=max_parallel_samples,
+                    steering_args=self.steering_args,
+                    diffusion_conditioning=diffusion_conditioning,
+                )
+            output.update(struct_out)
+
+        if self.confidence_prediction:
+            if self.skip_run_structure:
+                x_pred = feats["coords"].repeat_interleave(diffusion_samples, 0)
+            else:
+                assert "sample_atom_coords" in output, (
+                    "Structure sampling did not produce sample_atom_coords."
+                )
+                x_pred = output["sample_atom_coords"]
+
+            output.update(
+                self.confidence_module(
+                    s_inputs=s_inputs.detach(),
+                    s=s.detach(),
+                    z=z.detach(),
+                    x_pred=x_pred.detach(),
+                    feats=feats,
+                    pred_distogram_logits=output["pdistogram"][:, :, :, 0].detach(),
+                    multiplicity=diffusion_samples,
+                    run_sequentially=run_confidence_sequentially,
+                    use_kernels=self.use_kernels,
+                )
+            )
+
+        return output
+
+
+class Boltz2Model(PreTrainedModel):
+    config_class = Boltz2Config
+    base_model_prefix = "core"
+
+    def __init__(self, config: Boltz2Config) -> None:
+        super().__init__(config)
+        assert isinstance(config.core_kwargs, dict), "config.core_kwargs must be a dictionary."
+        self.core = Boltz2InferenceCore(**config.core_kwargs)
+
+    def _init_weights(self, module: nn.Module) -> None:  # noqa: ARG002
+        return
+
+    def _detied_state_dict(self) -> dict[str, Tensor]:
+        raw_state = self.state_dict()
+        seen_ptrs: dict[int, str] = {}
+        out: dict[str, Tensor] = {}
+        for key, tensor in raw_state.items():
+            if torch.is_tensor(tensor):
+                ptr = tensor.untyped_storage().data_ptr()
+                if ptr in seen_ptrs:
+                    out[key] = tensor.clone()
+                else:
+                    seen_ptrs[ptr] = key
+                    out[key] = tensor
+            else:
+                out[key] = tensor
+        return out
+
+    def save_pretrained(self, save_directory: str, **kwargs: Any) -> None:
+        if "safe_serialization" not in kwargs:
+            kwargs["safe_serialization"] = False
+        if "state_dict" not in kwargs:
+            kwargs["state_dict"] = self._detied_state_dict()
+        super().save_pretrained(save_directory, **kwargs)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    @classmethod
+    def from_boltz_checkpoint(
+        cls,
+        checkpoint_path: str,
+        map_location: str | torch.device = "cpu",
+        use_kernels: bool = False,
+        default_recycling_steps: Optional[int] = None,
+        default_sampling_steps: Optional[int] = None,
+        default_diffusion_samples: Optional[int] = None,
+    ) -> "Boltz2Model":
+        # Boltz Lightning checkpoints include OmegaConf objects and require full unpickling.
+        checkpoint = torch.load(
+            checkpoint_path,
+            map_location=map_location,
+            weights_only=False,
+        )
+        assert isinstance(checkpoint, dict), "Checkpoint must deserialize to a dictionary."
+        _require_key(checkpoint, "hyper_parameters")
+        _require_key(checkpoint, "state_dict")
+
+        hparams = checkpoint["hyper_parameters"]
+        assert isinstance(hparams, dict), "Checkpoint hyper_parameters must be a dictionary."
+        state_dict = checkpoint["state_dict"]
+        assert isinstance(state_dict, dict), "Checkpoint state_dict must be a dictionary."
+
+        config = Boltz2Config.from_hyperparameters(
+            hparams,
+            use_kernels=use_kernels,
+            default_recycling_steps=default_recycling_steps,
+            default_sampling_steps=default_sampling_steps,
+            default_diffusion_samples=default_diffusion_samples,
+        )
+        model = cls(config)
+        cleaned = _state_dict_without_wrappers(state_dict)
+        target_keys = set(model.core.state_dict().keys())
+        filtered: dict[str, Tensor] = {}
+        for key, value in cleaned.items():
+            if key in target_keys:
+                filtered[key] = value
+
+        missing = sorted(target_keys.difference(filtered.keys()))
+        allowed_missing_substrings = [
+            ".attention.norm_s.weight",
+            ".attention.norm_s.bias",
+        ]
+        disallowed_missing: list[str] = []
+        for key in missing:
+            is_allowed = False
+            for token in allowed_missing_substrings:
+                if token in key:
+                    is_allowed = True
+            if not is_allowed:
+                disallowed_missing.append(key)
+        assert len(disallowed_missing) == 0, (
+            "Checkpoint is missing required parameters for Boltz2 inference core. "
+            f"Missing keys (first 20): {disallowed_missing[:20]}"
+        )
+
+        load_result = model.core.load_state_dict(filtered, strict=False)
+        loaded_missing = sorted(load_result.missing_keys)
+        loaded_disallowed_missing: list[str] = []
+        for key in loaded_missing:
+            is_allowed = False
+            for token in allowed_missing_substrings:
+                if token in key:
+                    is_allowed = True
+            if not is_allowed:
+                loaded_disallowed_missing.append(key)
+        assert len(loaded_disallowed_missing) == 0, (
+            "Model has unexpected missing keys after load_state_dict. "
+            f"Missing keys (first 20): {loaded_disallowed_missing[:20]}"
+        )
+        assert len(load_result.unexpected_keys) == 0
+        model.eval()
+        return model
+
+    def forward(
+        self,
+        feats: dict[str, Tensor],
+        recycling_steps: Optional[int] = None,
+        num_sampling_steps: Optional[int] = None,
+        diffusion_samples: Optional[int] = None,
+        max_parallel_samples: Optional[int] = None,
+        run_confidence_sequentially: bool = True,
+    ) -> dict[str, Tensor]:
+        if recycling_steps is None:
+            recycling_steps = self.config.default_recycling_steps
+        if num_sampling_steps is None:
+            num_sampling_steps = self.config.default_sampling_steps
+        if diffusion_samples is None:
+            diffusion_samples = self.config.default_diffusion_samples
+        return self.core(
+            feats=feats,
+            recycling_steps=recycling_steps,
+            num_sampling_steps=num_sampling_steps,
+            diffusion_samples=diffusion_samples,
+            max_parallel_samples=max_parallel_samples,
+            run_confidence_sequentially=run_confidence_sequentially,
+        )
+
+    def _to_model_device(
+        self,
+        feats: dict[str, Tensor],
+        float_dtype: torch.dtype,
+    ) -> dict[str, Tensor]:
+        moved: dict[str, Tensor] = {}
+        for key, value in feats.items():
+            if torch.is_tensor(value):
+                if value.is_floating_point():
+                    moved[key] = value.to(device=self.device, dtype=float_dtype)
+                else:
+                    moved[key] = value.to(device=self.device)
+            else:
+                moved[key] = value
+        return moved
+
+    def predict_structure(
+        self,
+        amino_acid_sequence: str,
+        recycling_steps: Optional[int] = None,
+        num_sampling_steps: Optional[int] = None,
+        diffusion_samples: Optional[int] = None,
+        max_parallel_samples: Optional[int] = None,
+        run_confidence_sequentially: bool = True,
+        float_dtype: Optional[torch.dtype] = None,
+    ) -> Boltz2StructureOutput:
+        if float_dtype is None:
+            float_dtype = torch.float32
+
+        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,
+        )
+        feats = self._to_model_device(feats, float_dtype=float_dtype)
+
+        with torch.no_grad():
+            output = self.forward(
+                feats=feats,
+                recycling_steps=recycling_steps,
+                num_sampling_steps=num_sampling_steps,
+                diffusion_samples=diffusion_samples,
+                max_parallel_samples=max_parallel_samples,
+                run_confidence_sequentially=run_confidence_sequentially,
+            )
+
+        sample_atom_coords = output["sample_atom_coords"].detach().cpu()
+        atom_pad_mask = feats["atom_pad_mask"][0].detach().cpu()
+        plddt = output["plddt"].detach().cpu() if "plddt" in output else None
+        complex_plddt = output["complex_plddt"].detach().cpu() if "complex_plddt" in output else None
+        iptm = output["iptm"].detach().cpu() if "iptm" in output else None
+        ptm = output["ptm"].detach().cpu() if "ptm" in output else None
+
+        confidence_score = None
+        if (complex_plddt is not None) and (iptm is not None) and (ptm is not None):
+            if torch.allclose(iptm, torch.zeros_like(iptm)):
+                confidence_score = (4 * complex_plddt + ptm) / 5
+            else:
+                confidence_score = (4 * complex_plddt + iptm) / 5
+
+        return Boltz2StructureOutput(
+            sample_atom_coords=sample_atom_coords,
+            atom_pad_mask=atom_pad_mask,
+            plddt=plddt,
+            confidence_score=confidence_score,
+            complex_plddt=complex_plddt,
+            iptm=iptm,
+            ptm=ptm,
+            sequence=template.sequence,
+            structure_template=template,
+            raw_output={key: _to_cpu_detached(val) for key, val in output.items()},
+        )
+
+    def save_as_cif(
+        self,
+        structure_output: Boltz2StructureOutput,
+        output_path: str,
+        sample_index: int = 0,
+    ) -> str:
+        assert structure_output.structure_template is not None, (
+            "structure_output.structure_template is required for CIF export."
+        )
+        assert structure_output.sample_atom_coords is not None, (
+            "structure_output.sample_atom_coords is required for CIF export."
+        )
+        assert structure_output.atom_pad_mask is not None, (
+            "structure_output.atom_pad_mask is required for CIF export."
+        )
+        return write_cif(
+            structure_template=structure_output.structure_template,
+            atom_coords=structure_output.sample_atom_coords,
+            atom_mask=structure_output.atom_pad_mask,
+            output_path=output_path,
+            plddt=structure_output.plddt,
+            sample_index=sample_index,
+        )

runtime.py

@@ -0,0 +1,32 @@
+import os
+import sys
+from pathlib import Path
+
+
+def ensure_boltz_importable() -> None:
+    candidates: list[Path] = []
+
+    if "BOLTZ_SRC_DIR" in os.environ:
+        env_override = os.environ["BOLTZ_SRC_DIR"]
+        if len(env_override.strip()) > 0:
+            candidates.append(Path(env_override.strip()))
+
+    module_dir = Path(__file__).resolve().parent
+    for parent in [module_dir, *module_dir.parents]:
+        candidates.append(parent / "boltz" / "src")
+
+    cwd_dir = Path.cwd().resolve()
+    for parent in [cwd_dir, *cwd_dir.parents]:
+        candidates.append(parent / "boltz" / "src")
+
+    for candidate in candidates:
+        if candidate.exists():
+            candidate_str = str(candidate)
+            if candidate_str not in sys.path:
+                sys.path.insert(0, candidate_str)
+            return
+
+    assert False, (
+        "Could not locate cloned Boltz source directory ('boltz/src'). "
+        "Set BOLTZ_SRC_DIR explicitly if needed."
+    )