Upload folder using huggingface_hub (part 3)

.gitattributes

@@ -53,3 +53,4 @@ model/Protenix-pxd/assets/mini_performance.png filter=lfs diff=lfs merge=lfs -te
 model/Protenix-pxd/assets/protenix_predictions.gif filter=lfs diff=lfs merge=lfs -text
 model/Protenix-pxd/assets/v0.5.0_metrics.png filter=lfs diff=lfs merge=lfs -text
 model/pxdesign/pxd_server/TimesNewRoman.ttf filter=lfs diff=lfs merge=lfs -text
+weight/release_data/ccd_cache/clusters-by-entity-40.txt filter=lfs diff=lfs merge=lfs -text

model/pxdesign/utils/infer.py

@@ -0,0 +1,535 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import hashlib
+import logging
+import os
+import sys
+import urllib
+from collections import defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable
+
+import torch
+from ml_collections.config_dict import ConfigDict
+from protenix.config import parse_configs
+from protenix.data.parser import DistillationMMCIFParser
+from protenix.utils.file_io import dump_gzip_pickle
+from pxdbench.pxd_configs.eval import eval_configs
+
+from pxdesign.configs.configs_base import configs as configs_base
+from pxdesign.configs.configs_data import data_configs
+from pxdesign.configs.configs_infer import inference_configs
+from pxdesign.data.utils import pdb_to_cif
+
+URL = {
+    "pxdesign_v0.1.0": "https://pxdesign.tos-cn-beijing.volces.com/release_model/pxdesign_v0.1.0.pt",
+    "protenix_base_default_v0.5.0": "https://pxdesign.tos-cn-beijing.volces.com/release_model/protenix_base_default_v0.5.0.pt",
+    "protenix_mini_default_v0.5.0": "https://pxdesign.tos-cn-beijing.volces.com/release_model/protenix_mini_default_v0.5.0.pt",
+    "protenix_mini_tmpl_v0.5.0": "https://pxdesign.tos-cn-beijing.volces.com/release_model/protenix_mini_tmpl_v0.5.0.pt",
+    "ccd_components_file": "https://pxdesign.tos-cn-beijing.volces.com/release_data/components.v20240608.cif",
+    "ccd_components_rdkit_mol_file": "https://pxdesign.tos-cn-beijing.volces.com/release_data/components.v20240608.cif.rdkit_mol.pkl",
+    "pdb_cluster_file": "https://pxdesign.tos-cn-beijing.volces.com/release_data/clusters-by-entity-40.txt",
+}
+
+ALIASES = {
+    "N_sample": "sample_diffusion.N_sample",
+    "N_step": "sample_diffusion.N_step",
+    "eta_type": "sample_diffusion.eta_schedule.type",
+    "eta_min": "sample_diffusion.eta_schedule.min",
+    "eta_max": "sample_diffusion.eta_schedule.max",
+    "gamma0": "sample_diffusion.gamma0",
+    "gamma_min": "sample_diffusion.gamma_min",
+    "sample_diffusion_chunk_size": "infer_setting.sample_diffusion_chunk_size",
+}
+
+logger = logging.getLogger(__name__)
+
+
+def download_inference_cache(configs) -> None:
+    def progress_callback(block_num, block_size, total_size):
+        downloaded = block_num * block_size
+        percent = min(100, downloaded * 100 / total_size)
+        bar_length = 30
+        filled_length = int(bar_length * percent // 100)
+        bar = "=" * filled_length + "-" * (bar_length - filled_length)
+
+        status = f"\r[{bar}] {percent:.1f}%"
+        print(status, end="", flush=True)
+
+        if downloaded >= total_size:
+            print()
+
+    def download_from_url(tos_url, checkpoint_path, check_weight=True):
+        urllib.request.urlretrieve(
+            tos_url, checkpoint_path, reporthook=progress_callback
+        )
+        if check_weight:
+            try:
+                ckpt = torch.load(checkpoint_path)
+                del ckpt
+            except:
+                os.remove(checkpoint_path)
+                raise RuntimeError(
+                    "Download model checkpoint failed, please download by yourself with "
+                    f"wget {tos_url} -O {checkpoint_path}"
+                )
+
+    for cache_name in (
+        "ccd_components_file",
+        "ccd_components_rdkit_mol_file",
+        "pdb_cluster_file",
+    ):
+        cur_cache_fpath = configs["data"][cache_name]
+        if not os.path.exists(cur_cache_fpath):
+            os.makedirs(os.path.dirname(cur_cache_fpath), exist_ok=True)
+            tos_url = URL[cache_name]
+            assert os.path.basename(tos_url) == os.path.basename(cur_cache_fpath), (
+                f"{cache_name} file name is incorrect, `{tos_url}` and "
+                f"`{cur_cache_fpath}`. Please check and try again."
+            )
+            logger.info(
+                f"Downloading data cache from\n {tos_url}... to {cur_cache_fpath}"
+            )
+            download_from_url(tos_url, cur_cache_fpath, check_weight=False)
+
+    checkpoint_path = os.path.join(
+        configs.load_checkpoint_dir, f"{configs.model_name}.pt"
+    )
+    if not os.path.exists(checkpoint_path):
+        os.makedirs(configs.load_checkpoint_dir, exist_ok=True)
+        tos_url = URL[configs.model_name]
+        logger.info(
+            f"Downloading model checkpoint from\n {tos_url}... to {checkpoint_path}"
+        )
+        download_from_url(tos_url, checkpoint_path)
+
+    # download protenix checkpoints
+    for model_name in [
+        "protenix_base_default_v0.5.0",
+        "protenix_mini_default_v0.5.0",
+        "protenix_mini_tmpl_v0.5.0",
+    ]:
+        checkpoint_path = os.path.join(configs.load_checkpoint_dir, f"{model_name}.pt")
+        if not os.path.exists(checkpoint_path):
+            tos_url = URL[model_name]
+            logger.info(
+                f"Downloading model checkpoint from\n {tos_url}... to {checkpoint_path}"
+            )
+            download_from_url(tos_url, checkpoint_path)
+
+    # set checkpoint dir for ptx tools in PXDesignBench
+    if hasattr(configs, "eval"):
+        configs.eval.binder.tools.ptx.load_checkpoint_dir = configs.load_checkpoint_dir
+        configs.eval.binder.tools.ptx_mini.load_checkpoint_dir = (
+            configs.load_checkpoint_dir
+        )
+
+
+def remap_arg_key(key: str) -> str:
+    if key.startswith("--"):
+        name = key[2:]
+        mapped = ALIASES.get(name, name)
+        return "--" + mapped
+    return key
+
+
+def parse_sys_args(argv=None):
+    if argv is None:
+        argv = sys.argv[1:]
+    remapped = []
+
+    i = 0
+    while i < len(argv):
+        k = argv[i]
+        # if k starts with "--", check whether it matches alias
+        if k.startswith("--") and i + 1 < len(argv):
+            remapped.append(remap_arg_key(k))
+            remapped.append(argv[i + 1])
+            i += 2
+        else:
+            remapped.append(k)
+            i += 1
+
+    return " ".join(remapped)
+
+
+def get_configs(argv=None) -> ConfigDict:
+    configs = {
+        **configs_base,
+        **{"data": data_configs},
+        **inference_configs,
+        **{"eval": eval_configs},
+    }
+    configs = parse_configs(
+        configs=configs,
+        arg_str=parse_sys_args(argv),
+        fill_required_with_null=True,
+    )
+    return configs
+
+
+class DisableLogging:
+    def __enter__(self):
+        logging.disable(logging.WARNING)
+
+    def __exit__(self, exc_type, exc, tb):
+        logging.disable(logging.NOTSET)
+
+
+# -------------------------
+# Handling PDB input
+# -------------------------
+
+
+def parse_ranges(range_str: str) -> list[tuple[int, int]]:
+    """
+    Parse "1-30,40-50,66" -> [(1,30),(40,50),(66,66)]
+    """
+    ranges: list[tuple[int, int]] = []
+    for part in range_str.split(","):
+        part = part.strip()
+        if not part:
+            continue
+        if "-" in part:
+            a, b = part.split("-")
+            ranges.append((int(a), int(b)))
+        else:
+            x = int(part)
+            ranges.append((x, x))
+    return ranges
+
+
+def format_ranges(ints: Iterable[int]) -> str:
+    """
+    Compress sorted integers into "a-b,c,d-e".
+    """
+    xs = sorted(set(int(x) for x in ints))
+    if not xs:
+        return ""
+    out: list[str] = []
+    s = e = xs[0]
+    for x in xs[1:]:
+        if x == e + 1:
+            e = x
+        else:
+            out.append(f"{s}-{e}" if s != e else f"{s}")
+            s = e = x
+    out.append(f"{s}-{e}" if s != e else f"{s}")
+    return ",".join(out)
+
+
+# -------------------------
+# Chain mapping
+# -------------------------
+
+
+def build_chain_mapping(
+    old_ids: Iterable[str],
+    new_ids: Iterable[str],
+    *,
+    keep_chains: Iterable[str] | None = None,
+    err_hint: str = "Please consider using a CIF structure file in your JSON file.",
+) -> dict[str, str]:
+    """
+    Build mapping old_id -> new_id with consistency check.
+
+    keep_chains:
+      - None: keep all chains
+      - otherwise: only build mapping for chains in keep_chains
+    """
+    old_ids = list(old_ids)
+    new_ids = list(new_ids)
+    if len(old_ids) != len(new_ids):
+        raise ValueError("old_ids and new_ids must have the same length.")
+
+    keep = set(keep_chains) if keep_chains is not None else None
+
+    mapping: dict[str, str] = {}
+    for old, new in zip(old_ids, new_ids):
+        if keep is not None and old not in keep:
+            continue
+        if old not in mapping:
+            mapping[old] = new
+        elif mapping[old] != new:
+            raise ValueError(
+                f"Inconsistent mapping: chain '{old}' maps to both "
+                f"'{mapping[old]}' and '{new}'. It will raise ambiguity. {err_hint}"
+            )
+    return mapping
+
+
+# -------------------------
+# Residue mapping (res_id <-> auth_res_id)
+# -------------------------
+
+
+@dataclass(frozen=True)
+class ResidueMaps:
+    # (chain_id, res_id) -> (auth_asym_id, auth_res_id)
+    resid2auth: dict[tuple[str, int], tuple[str, int]]
+    # (auth_asym_id, auth_res_id) -> (chain_id, res_id)
+    auth2resid: dict[tuple[str, int], tuple[str, int]]
+
+
+def build_residue_maps(
+    atom_array,
+    *,
+    strict_bijective: bool = True,
+    err_hint: str = "Please consider using a CIF structure file in your JSON file.",
+) -> ResidueMaps:
+    """
+    Build residue-level mapping with uniqueness checks.
+
+    Ensures:
+      - each (chain_id, res_id) maps to a single (auth_asym_id, auth_res_id)
+      - (optional) bijection: each (auth_asym_id, auth_res_id) maps back to a single (chain_id, res_id)
+    """
+    chain_id = atom_array.chain_id
+    res_id = atom_array.res_id
+    auth_asym_id = atom_array.auth_asym_id
+    auth_res_id = atom_array.auth_res_id
+
+    resid2auth: dict[tuple[str, int], tuple[str, int]] = {}
+    auth2resid: dict[tuple[str, int], tuple[str, int]] = {}
+
+    for c, r, ac, ar in zip(chain_id, res_id, auth_asym_id, auth_res_id):
+        key = (str(c), int(r))
+        val = (str(ac), int(ar))
+
+        if key in resid2auth and resid2auth[key] != val:
+            raise ValueError(
+                "Non-unique mapping detected: same (chain_id, res_id) maps to multiple "
+                f"(auth_asym_id, auth_res_id).\n  key={key}\n  first={resid2auth[key]}\n  new={val}\n"
+                f"{err_hint}"
+            )
+        resid2auth.setdefault(key, val)
+
+        if strict_bijective:
+            if val in auth2resid and auth2resid[val] != key:
+                raise ValueError(
+                    "Non-unique mapping detected: same (auth_asym_id, auth_res_id) maps to multiple "
+                    f"(chain_id, res_id).\n  val={val}\n  first={auth2resid[val]}\n  new={key}\n"
+                    f"{err_hint}"
+                )
+            auth2resid.setdefault(val, key)
+
+    return ResidueMaps(resid2auth=resid2auth, auth2resid=auth2resid)
+
+
+# -------------------------
+# Converters: crop / hotspot
+# -------------------------
+
+
+def convert_crop_auth_to_new(
+    crop_dict: dict[str, str],
+    residue_maps: ResidueMaps,
+    *,
+    strict_mapping: bool = True,
+) -> dict[str, str]:
+    result: dict[str, list[int]] = defaultdict(list)
+
+    for auth_chain, range_str in crop_dict.items():
+        for start, end in parse_ranges(range_str):
+            for auth_r in range(start, end + 1):
+                key = (auth_chain, int(auth_r))
+                if key not in residue_maps.auth2resid:
+                    if strict_mapping:
+                        raise KeyError(
+                            f"Requested auth residue not found in atom_array: {key}"
+                        )
+                    else:
+                        continue
+                new_c, new_r = residue_maps.auth2resid[key]
+                result[new_c].append(new_r)
+
+    return {new_c: format_ranges(rs) for new_c, rs in result.items() if rs}
+
+
+def convert_hotspot_auth_to_new(
+    hotspot_dict: dict[str, list[int]],
+    residue_maps: ResidueMaps,
+    *,
+    strict_mapping: bool = True,
+) -> dict[str, list[int]]:
+    """
+    {chain_id: [11,22]} -> {auth_asym_id: [auth_res_id,...]} (sorted unique)
+    """
+    result: dict[str, list[int]] = defaultdict(list)
+
+    for chain, res_list in hotspot_dict.items():
+        for r in res_list:
+            key = (chain, int(r))
+            if key not in residue_maps.auth2resid:
+                if strict_mapping:
+                    raise KeyError(f"Hotspot residue not found in atom_array: {key}")
+                else:
+                    continue
+            ac, ar = residue_maps.auth2resid[key]
+            result[ac].append(ar)
+
+    return {ac: sorted(set(ars)) for ac, ars in result.items() if ars}
+
+
+# -------------------------
+# Apply filter rewrite (chain_id / crop / msa / hotspot)
+# -------------------------
+
+
+def rewrite_input_dict_inplace(
+    input_dict: dict,
+    *,
+    chain_mapping: dict[str, str],  # old chain_id -> new chain_id (e.g. PDB->CIF)
+    residue_maps: ResidueMaps | None,
+) -> None:
+    """
+    Rewrite cond_dict['filter'] in-place using chain_mapping and residue_maps.
+    """
+    cond_dict = input_dict["condition"]
+    filt = cond_dict.get("filter", {})
+    if filt:
+        # chain_id list
+        if "chain_id" in filt and filt["chain_id"]:
+            filt["chain_id"] = [chain_mapping[c] for c in filt["chain_id"]]
+
+        # crop dict: {chain_id: "ranges"}
+        if "crop" in filt and filt["crop"]:
+            if residue_maps is None:
+                raise ValueError(
+                    "filter.crop requires residue_maps (atom_array) to convert res_id -> auth_res_id."
+                )
+            filt["crop"] = convert_crop_auth_to_new(filt["crop"], residue_maps)
+        cond_dict["filter"] = filt
+
+    # msa dict: {chain_id: ...}
+    if "msa" in cond_dict and cond_dict["msa"]:
+        cond_dict["msa"] = {chain_mapping[k]: v for k, v in cond_dict["msa"].items()}
+
+    input_dict["condition"] = cond_dict
+    # hotspot dict: {chain_id: [res_ids]}
+    if "hotspot" in input_dict and input_dict["hotspot"]:
+        if residue_maps is None:
+            raise ValueError(
+                "filter.hotspot requires residue_maps (atom_array) to convert res_id -> auth_res_id."
+            )
+        input_dict["hotspot"] = convert_hotspot_auth_to_new(
+            input_dict["hotspot"], residue_maps
+        )
+
+
+# -------------------------
+# Main entry: convert_to_bioassembly_dict
+# -------------------------
+
+
+def convert_to_bioassembly_dict(input_dict: dict, out_dir: str | None = None):
+    """
+    Returns:
+      - if input is already .pkl.gz: (str_file)  (kept as your original behavior)
+      - else: (out_path, chain_mapping)
+    """
+    assert "condition" in input_dict, "input_dict must have 'condition' key"
+    cond_dict = input_dict["condition"]
+    str_file = cond_dict["structure_file"]
+    if out_dir is None:
+        out_dir = os.path.dirname(str_file)
+
+    if str_file.endswith(".pkl.gz"):
+        return str_file
+
+    chain_mapping: dict[str, str] = {}
+    residue_maps: ResidueMaps | None = None
+    atom_array = None
+
+    if str_file.endswith(".cif"):
+        parser = DistillationMMCIFParser(str_file)
+        d = parser.get_structure_dict()
+
+    elif str_file.endswith(".pdb"):
+        cif_file = os.path.join(out_dir, os.path.basename(str_file)[:-4] + ".cif")
+        atom_array = pdb_to_cif(str_file, cif_file)
+
+        filter_chains = cond_dict.get("filter", {}).get("chain_id", [])
+        chain_mapping = build_chain_mapping(
+            atom_array.auth_asym_id,
+            atom_array.chain_id,
+            keep_chains=filter_chains if filter_chains else None,
+        )
+        residue_maps = build_residue_maps(atom_array)
+
+        rewrite_input_dict_inplace(
+            input_dict,
+            chain_mapping=chain_mapping,
+            residue_maps=residue_maps,
+        )
+        parser = DistillationMMCIFParser(cif_file)
+        d = parser.get_structure_dict()
+
+    else:
+        raise ValueError(f"Unsupported structure file! {str_file}")
+
+    out_path = Path(out_dir) / f"{Path(str_file).stem}.pkl.gz"
+    assert str(out_path).endswith(".pkl.gz"), "Bioassembly dict should end with .pkl.gz"
+    dump_gzip_pickle(d, out_path)
+    input_dict["condition"]["structure_file"] = str(out_path)
+
+    return d
+
+
+def configure_runtime_env(
+    use_fast_ln: bool = False, use_deepspeed_evo: bool = False
+) -> None:
+    """
+    Independent runtime knobs:
+      - use_fast_ln -> LAYERNORM_TYPE
+      - use_deepspeed_evo  -> DEEPSPEED_EVO (+ CUTLASS dependency)
+    """
+
+    # LayerNorm
+    if use_fast_ln:
+        os.environ["LAYERNORM_TYPE"] = "fast_layernorm"
+
+    # DeepSpeed Evo: fully independent
+    os.environ["DEEPSPEED_EVO"] = "true" if use_deepspeed_evo else "false"
+
+    if not use_deepspeed_evo:
+        return
+
+    if "CUTLASS_PATH" in os.environ and os.environ["CUTLASS_PATH"]:
+        cutlass_path = Path(os.environ["CUTLASS_PATH"]).expanduser()
+    else:
+        cutlass_path = Path.home() / "cutlass"
+        os.environ["CUTLASS_PATH"] = str(cutlass_path)
+
+    if not cutlass_path.is_dir():
+        print("")
+        print(f"[WARNING] CUTLASS not found at: {cutlass_path}")
+        print(
+            "  PXDesign uses DeepSpeed Evo kernels which require NVIDIA CUTLASS v3.5.1."
+        )
+        print("  To install:")
+        print(
+            '    git clone -b v3.5.1 https://github.com/NVIDIA/cutlass.git "$HOME/cutlass"'
+        )
+        print('    export CUTLASS_PATH="$HOME/cutlass"')
+        print("")
+
+
+def derive_seed(base_seed: int, rank: int = 0, digits: int = 6) -> int:
+    mod = 10**digits
+    msg = f"pxdesign|{base_seed}|{rank}".encode()
+    h = hashlib.blake2b(msg, digest_size=8).digest()
+    return int.from_bytes(h, "little") % mod

model/pxdesign/utils/inputs.py

@@ -0,0 +1,304 @@
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+import yaml
+from protenix.utils.file_io import load_gzip_pickle
+
+from pxdesign.data.utils import CIFWriter
+from pxdesign.utils.infer import convert_to_bioassembly_dict
+
+
+class NpEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, (np.integer,)):
+            return int(obj)
+        elif isinstance(obj, (np.floating,)):
+            return float(obj)
+        elif isinstance(obj, (np.ndarray,)):
+            return obj.tolist()
+        return super(NpEncoder, self).default(obj)
+
+
+def parse_yaml_to_json(yaml_path, json_path=None):
+    """
+    Parses the YAML config and converts it to the
+    JSON structure required by PXDesign model.
+    """
+    yaml_path = os.path.abspath(yaml_path)
+    if not os.path.exists(yaml_path):
+        raise FileNotFoundError(f"YAML config file not found: {yaml_path}")
+
+    with open(yaml_path, "r") as f:
+        try:
+            cfg = yaml.safe_load(f)
+        except yaml.YAMLError as e:
+            raise ValueError(f"Error parsing YAML file: {e}")
+
+    # --- 1. Top Level Fields ---
+    # Default task name uses filename if not provided
+    default_name = os.path.splitext(os.path.basename(yaml_path))[0]
+    task_name = cfg.get("task_name", default_name)
+
+    # Binder length (Required)
+    if "binder_length" not in cfg:
+        raise ValueError("Missing required field: 'binder_length'")
+    binder_length = int(cfg["binder_length"])
+
+    # --- 2. Target Parsing ---
+    target_cfg = cfg.get("target", {})
+    if "file" not in target_cfg:
+        raise ValueError("Missing required field: 'target.file'")
+
+    target_file_path = target_cfg["file"]
+    if not os.path.exists(target_file_path):
+        raise FileNotFoundError(f"Target structure file not found: {target_file_path}")
+
+    # Initialize containers
+    chain_ids = []
+    crop_dict = {}
+    hotspot_dict = {}
+    msa_dict_per_chain = {}
+
+    # --- 3. Chains Parsing ---
+    chains_cfg = target_cfg.get("chains", {})
+    if not chains_cfg:
+        raise ValueError("Missing required field: 'target.chains'")
+
+    for chain_id, props in chains_cfg.items():
+        chain_id = str(chain_id)
+        chain_ids.append(chain_id)
+
+        # Handle "A: all" or "A: null" shorthand
+        if props is None or (
+            isinstance(props, str) and props.lower() in ["all", "full"]
+        ):
+            props = {}
+
+        # --- Crop Logic ---
+        # User YAML: ["1-50", "80-100"] OR "1-100" OR "all"
+        # Internal JSON: "1-50,80-100" OR None
+        if "crop" in props:
+            raw_crop = props["crop"]
+            crop_val = None
+
+            if isinstance(raw_crop, list):
+                # Join list into comma-separated string
+                crop_val = ",".join(str(x) for x in raw_crop)
+            elif isinstance(raw_crop, str):
+                if raw_crop.lower() in ["all", "full"]:
+                    crop_val = None
+                else:
+                    crop_val = raw_crop
+
+            if crop_val:
+                crop_dict[chain_id] = crop_val
+
+        # --- Hotspot Logic ---
+        if "hotspots" in props:
+            # YAML list is already a Python list
+            hotspot_dict[chain_id] = props["hotspots"]
+
+        # --- MSA Logic ---
+        if "msa" in props and props["msa"]:
+            msa_path = props["msa"]
+            for fname in ["pairing.a3m", "non_pairing.a3m"]:
+                if not os.path.exists(os.path.join(msa_path, fname)):
+                    raise FileNotFoundError(
+                        f"MSA file not found: {os.path.join(msa_path, fname)}"
+                    )
+            msa_config = {
+                "precomputed_msa_dir": msa_path,  # Default to None (Auto)
+                "pairing_db": "uniref100",
+            }
+
+            msa_dict_per_chain[chain_id] = msa_config
+
+    # --- 4. Construct Internal JSON Structure ---
+    json_task = {
+        "name": task_name,
+        "condition": {
+            "structure_file": target_file_path,
+            "filter": {
+                "chain_id": chain_ids,
+                "crop": crop_dict,
+            },
+            "msa": msa_dict_per_chain,
+        },
+        "hotspot": hotspot_dict,
+        "generation": [
+            {
+                "type": "protein",
+                "length": binder_length,
+                "count": 1,
+            }
+        ],
+    }
+
+    if json_path is not None:
+        os.makedirs(os.path.dirname(json_path), exist_ok=True)
+        with open(json_path, "w") as f:
+            json.dump([json_task], f, cls=NpEncoder)
+
+    return [json_task]
+
+
+def check_yaml_file(yaml_path: str):
+    print(f"Checking YAML file: {yaml_path}...")
+    result = parse_yaml_to_json(yaml_path, None)
+    print("✅ YAML file is valid.")
+
+
+def process_input_file(input_path: str, out_dir: str = None) -> str:
+    """
+    Process the input file path to ensure it has the correct extension.
+    """
+    input_path = os.path.abspath(input_path)
+    if not os.path.exists(input_path):
+        raise FileNotFoundError(f"Input file not found: {input_path}")
+
+    # Check file extension
+    ext = os.path.splitext(input_path)[1].lower()
+    if ext not in [".json", ".yaml"]:
+        raise ValueError(
+            f"Unsupported input file format: {ext}. "
+            "Supported formats are: JSON, YAML."
+        )
+
+    # Convert YAML to JSON if necessary
+    if ext == ".yaml":
+        base, _ = os.path.splitext(os.path.basename(input_path))
+        out_dir = out_dir or os.path.dirname(input_path)
+        json_path = os.path.join(out_dir, f"{base}.json")
+
+        parse_yaml_to_json(input_path, json_path)
+        input_path = json_path
+
+    return input_path
+
+
+def dump_bioassembly_to_cif(
+    bio_dict: Union[str, Path, dict],
+    output_cif: Union[str, Path],
+    dump_unresolved_atoms: bool = False,
+):
+    """
+    Dump a bioassembly dict to CIF.
+    """
+    if isinstance(bio_dict, str) or isinstance(bio_dict, Path):
+        bio_dict = load_gzip_pickle(bio_dict)
+    atom_array = bio_dict["atom_array"]
+
+    if not dump_unresolved_atoms:
+        mask = atom_array.is_resolved
+    else:
+        mask = None
+
+    entity_poly_type = bio_dict["entity_poly_type"]
+    writer = CIFWriter(
+        atom_array=atom_array,
+        entity_poly_type=entity_poly_type,
+        atom_array_output_mask=mask,
+    )
+    writer.save_to_cif(
+        output_cif,
+        entry_id=Path(output_cif).stem.split(".")[0],
+        include_bonds=False,
+    )
+
+    return
+
+
+def generate_pml_from_json_input(cif_file_path: str, json_file_path: str) -> dict:
+    """
+    generate PML script from input_dict. for visualizing in pymol.
+    cif_file_path: path to the cif file of the target structure.
+    json_file_path: path to the json file of the input dict.
+    """
+    cif_fname = os.path.basename(cif_file_path)
+    cif_file_dir = os.path.dirname(cif_file_path)
+    pymol_cmds = [
+        "load " + cif_fname,
+        "hide",
+        "show cartoon",
+        'cmd.util.cbc(selection="(elem C)")',
+    ]
+    json_task_dict = json.load(open(json_file_path, "r"))[0]
+
+    cond_dict = json_task_dict.get("condition", {})
+    if "filter" in cond_dict:
+        crop_region = cond_dict["filter"].get("crop", {})
+        crop_sele_list = []
+        for chain_id, chain_crop in crop_region.items():
+            chain_crop_sele = (
+                f"(chain {chain_id} and resi " + chain_crop.replace(",", "+") + ")"
+            )
+            crop_sele_list.append(chain_crop_sele)
+        if len(crop_sele_list) > 0:
+            crop_sele_str = "select crop, " + " OR ".join(crop_sele_list)
+            pymol_cmds.append(crop_sele_str)
+            pymol_cmds.append("color marine, crop and elem C")
+
+    hotspot_residues = json_task_dict.get("hotspot", {})
+    hotspot_sele_list = []
+    for chain_id, residues in hotspot_residues.items():
+        hotspot_sele = (
+            f"(chain {chain_id} and resi " + "+".join([str(x) for x in residues]) + ")"
+        )
+        hotspot_sele_list.append(hotspot_sele)
+    if len(hotspot_sele_list) > 0:
+        hotspot_sele_str = "select hotspot, " + " OR ".join(hotspot_sele_list)
+        pymol_cmds.append(hotspot_sele_str)
+        pymol_cmds.extend(["color pink, hotspot and elem C", "show sticks, hotspot"])
+
+    pymol_cmds.append("color grey70, not (hotspot OR crop)")
+    pml_script_path = os.path.join(cif_file_dir, f"{Path(cif_file_path).stem}.pml")
+    with open(pml_script_path, "w") as f:
+        f.write("\n".join(pymol_cmds))
+    return
+
+
+def dump_target_cif_from_input_file(file_path: str, out_dir: str) -> dict:
+    """
+    Parse target structure from input_dict.
+    """
+    if os.path.splitext(file_path)[1].lower() == ".json":
+        json_path = file_path
+    else:
+        assert (
+            os.path.splitext(file_path)[1].lower() == ".yaml"
+        ), f"Input file must be JSON or YAML, but got {os.path.splitext(file_path)[1]}"
+        json_path = os.path.join(out_dir, "tmp", f"{Path(file_path).stem}.json")
+        parse_yaml_to_json(file_path, json_path)
+    with open(json_path, "r") as f:
+        json_task_dict = json.load(f)[0]
+    bioassembly_dict = convert_to_bioassembly_dict(
+        json_task_dict, os.path.join(out_dir, "tmp")
+    )
+    if isinstance(bioassembly_dict, str):
+        bioassembly_dict = load_gzip_pickle(bioassembly_dict)
+
+    output_cif = os.path.join(out_dir, f"{Path(file_path).stem}_parsed_target.cif")
+    dump_bioassembly_to_cif(bioassembly_dict, output_cif)
+    generate_pml_from_json_input(output_cif, json_path)
+    return
+
+
+# --- CLI Wrapper for Debugging ---
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--yaml_file", help="Path to input YAML file")
+    parser.add_argument(
+        "--output_json_file", help="Path to output JSON file", default=None
+    )
+    args = parser.parse_args()
+
+    try:
+        result = parse_yaml_to_json(args.yaml_file, args.output_json_file)
+    except Exception as e:
+        print(f"Error: {e}")
+        sys.exit(1)

model/pxdesign/utils/pipeline.py

@@ -0,0 +1,365 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# -*- coding: utf-8 -*-
+"""
+Utility helpers & constants for the design pipeline.
+"""
+
+import ast
+import os
+import sys
+from glob import glob
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from biotite.structure import get_residues
+from protenix.data.constants import PROT_STD_RESIDUES_ONE_TO_THREE
+from protenix.data.parser import MMCIFParser
+
+from pxdesign.data.utils import CIFWriter, pdb_to_cif
+
+# -------- Global subdir constants (used everywhere) --------
+ORIG_SUBDIR = os.path.join("orig_designed")
+AF2_SUBDIR = os.path.join("passing-AF2-IG-easy")
+PTX_SUBDIR = os.path.join("passing-Protenix-basic")
+
+
+# -------- Small general helpers --------
+
+
+def convert_strlist_col(df: pd.DataFrame) -> pd.DataFrame:
+    """Convert stringified lists like '[0.3, 0.6]' to their numeric mean."""
+
+    def _safe_mean(v):
+        if isinstance(v, list):
+            return float(np.mean(v)) if len(v) > 0 else np.nan
+        if isinstance(v, str):
+            try:
+                parsed = ast.literal_eval(v)
+                if isinstance(parsed, list) and len(parsed) > 0:
+                    return float(np.mean(parsed))
+            except Exception:
+                pass
+        return v
+
+    out = df.copy()
+    for c in out.columns:
+        out[c] = out[c].apply(_safe_mean)
+    return out
+
+
+# -------- File post-processing --------
+
+
+def add_terms_to_cif(input_cif: str, output_cif: str) -> None:
+    """Insert license/disclaimer blocks into a CIF file."""
+    terms1 = [
+        "# By using this file you agree to the legally binding terms of use found at https://protenix-server.com/terms-of-service\n",
+        "# Version: 0.1\n",
+    ]
+    terms2 = [
+        "#\n",
+        "loop_\n",
+        "_audit_author.name\n",
+        "_audit_author.pdbx_ordinal\n",
+        '"ByteDance Seed" 1\n',
+        "#\n",
+        "loop_\n",
+        "_pdbx_data_usage.details\n",
+        "_pdbx_data_usage.id\n",
+        "_pdbx_data_usage.type\n",
+        "_pdbx_data_usage.url\n",
+        ";NON-COMMERCIAL USE ONLY, BY USING THIS FILE YOU AGREE TO THE TERMS OF USE FOUND\n",
+        "AT https://protenix-server.com/terms-of-service.\n",
+        ";\n",
+        "1 license    ?\n",
+        ";THE INFORMATION IS NOT INTENDED FOR, HAS NOT BEEN VALIDATED FOR, AND IS NOT\n",
+        "APPROVED FOR CLINICAL USE. IT SHOULD NOT BE USED FOR CLINICAL PURPOSE OR RELIED\n",
+        "ON FOR MEDICAL OR OTHER PROFESSIONAL ADVICE. IT IS THEORETICAL MODELLING ONLY\n",
+        'AND CAUTION SHOULD BE EXERCISED IN ITS USE. IT IS PROVIDED "AS-IS" WITHOUT ANY\n',
+        "WARRANTY OF ANY KIND, WHETHER EXPRESSED OR IMPLIED. NO WARRANTY IS GIVEN THAT\n",
+        "USE OF THE INFORMATION SHALL NOT INFRINGE THE RIGHTS OF ANY THIRD PARTY.\n",
+        ";\n",
+        "2 disclaimer ?\n",
+    ]
+    with open(input_cif, "r") as f:
+        lines = f.readlines()
+    out, inserted = [], False
+    out.extend(terms1)
+    for line in lines:
+        out.append(line)
+        if not inserted and line.strip().startswith("_entry.id"):
+            out.extend(terms2)
+            inserted = True
+    with open(output_cif, "w") as f:
+        f.writelines(out)
+
+
+def replace_last_xpb_chain_sequence(
+    input_cif_path: str, out_cif_path: str, new_sequence: str
+) -> str:
+    """Replace the sequence on the LAST chain where all residues are 'xpb'."""
+    parser = MMCIFParser(input_cif_path)
+    atom_array = parser.get_structure(
+        altloc="first", model=1, bond_lenth_threshold=None
+    )
+
+    uniq_chain_ids, first_idx = np.unique(atom_array.chain_id, return_index=True)
+    ordered_chain_ids = uniq_chain_ids[np.argsort(first_idx)]
+
+    candidates = []
+    for ch in ordered_chain_ids:
+        ch_mask = atom_array.chain_id == ch
+        if np.any(ch_mask) and np.all(atom_array.res_name[ch_mask] == "xpb"):
+            candidates.append(ch)
+    if not candidates:
+        raise ValueError("No chain found where all residues are 'xpb'.")
+
+    target_chain = candidates[-1]
+    ch_mask = atom_array.chain_id == target_chain
+    res_ids, _ = get_residues(atom_array[ch_mask])
+    if len(res_ids) != len(new_sequence):
+        raise ValueError(
+            f"Length mismatch on chain {target_chain}: structure={len(res_ids)} vs seq={len(new_sequence)}"
+        )
+
+    for resid, one in zip(res_ids, new_sequence):
+        if one not in PROT_STD_RESIDUES_ONE_TO_THREE:
+            raise ValueError(f"Invalid amino acid letter '{one}' in new_sequence.")
+        res3 = PROT_STD_RESIDUES_ONE_TO_THREE[one]
+        atom_array.res_name[
+            (atom_array.chain_id == target_chain) & (atom_array.res_id == resid)
+        ] = res3
+
+    writer = CIFWriter(atom_array=atom_array, entity_poly_type=parser.entity_poly_type)
+    writer.save_to_cif(out_cif_path, include_bonds=True)
+    return out_cif_path
+
+
+# -------- Structure writers (project-specific) --------
+
+
+def save_design_cif(
+    task,
+    base_dir: str,
+    output_dir: str,
+    output_subdir: str = "",
+    rank_col: str = "rank",
+) -> str:
+    """Write original designed CIF and inject license blocks."""
+    task_name = task["task_name"]
+    src = os.path.join(
+        base_dir,
+        f"global_run_{task['run_idx']}",
+        task_name,
+        f"seed_{task['seed']}",
+        "predictions",
+        task["name"] + ".cif",
+    )
+    dst = os.path.join(output_dir, output_subdir, f"rank_{task[rank_col]}.cif")
+    os.makedirs(os.path.dirname(dst), exist_ok=True)
+    replace_last_xpb_chain_sequence(src, dst, task["sequence"])
+    add_terms_to_cif(dst, dst)
+    return dst
+
+
+def save_af2_docked(
+    task,
+    base_dir: str,
+    output_dir: str,
+    output_subdir: str = "",
+) -> str:
+    """Copy AF2-docked PDB to <output>/<task>/<output_subdir>/rank_<k>.pdb"""
+    task_name, sample_name, rank = task["task_name"], task["name"], task["rank"]
+    matches = glob(
+        os.path.join(
+            base_dir,
+            f"global_run_{task['run_idx']}",
+            task_name,
+            f"seed_{task['seed']}",
+            "predictions",
+            "af2_pred",
+            f"*{sample_name}_seq{task['seq_idx']}_model*.pdb",
+        )
+    )
+    if len(matches) != 1:
+        raise FileNotFoundError(
+            f"Expect 1 AF2 PDB for {sample_name}, got {len(matches)}"
+        )
+    src = matches[0]
+    dst = os.path.join(output_dir, output_subdir, f"rank_{rank}.cif")
+    os.makedirs(os.path.dirname(dst), exist_ok=True)
+    pdb_to_cif(src, dst, reset_res_id=False, pad_chain_id=True)
+    add_terms_to_cif(dst, dst)
+    return dst
+
+
+def save_ptx_docked(
+    task,
+    base_dir: str,
+    output_dir: str,
+    output_subdir: str = "",
+    is_large: bool = True,
+) -> str:
+    """Write Protenix-docked CIF (with hotspot-based chain permutation)."""
+    sample_name, rank = task["name"], task["rank"]
+    ptx_redocked_name = f"run_{task['run_idx']}_{task['name']}_seq{task['seq_idx']}"
+    matches = glob(
+        os.path.join(
+            base_dir,
+            "ptx_pred" if is_large else "ptx_mini_pred",
+            ptx_redocked_name + "_seq0",
+            "seed_*",
+            "predictions",
+            f"{ptx_redocked_name}*.cif",
+        )
+    )
+    if len(matches) != 1:
+        raise FileNotFoundError(
+            f"Expect 1 PTX CIF for {sample_name}, got {len(matches)}"
+        )
+    src = matches[0]
+    dst = os.path.join(output_dir, output_subdir, f"rank_{rank}.cif")
+    os.makedirs(os.path.dirname(dst), exist_ok=True)
+    add_terms_to_cif(src, dst)
+    return dst
+
+
+# -------- Misc helpers --------
+
+
+def trim_summary_df(df: pd.DataFrame) -> pd.DataFrame:
+    """Keep a curated subset of columns and rename AF2 metrics if present."""
+    base_metrics = [
+        "plddt",
+        "ptm_binder",
+        "ptm_target",
+        "iptm",
+        "ptm",
+        "iptm_binder",
+        "pred_design_rmsd",
+    ]
+    ptx_cols = ["ptx_" + m for m in base_metrics]
+    ptx_mini_cols = ["ptx_mini_" + m for m in base_metrics]
+    columns_to_keep = (
+        [
+            "rank",
+            "task_name",
+            "sequence",
+            "af2_easy_success",
+            "af2_opt_success",
+            "ptx_success",
+            "ptx_basic_success",
+            "pLDDT",
+            "pTM",
+            "i_pTM",
+            "pAE",
+            "unscaled_i_pAE",
+            "pLDDT_MONOMER",
+            "pTM_MONOMER",
+            "pAE_MONOMER",
+            "bound_unbound_RMSD",
+            "af2_binder_pred_design_rmsd",
+            "af2_complex_pred_design_rmsd",
+        ]
+        + ptx_cols
+        + ptx_mini_cols
+        + ["alpha", "beta", "loop", "Rg", "chosen_struct_type", "chosen_struct_path"]
+    )
+    existing_cols = [c for c in columns_to_keep if c in df.columns]
+    df = df[existing_cols].copy()
+
+    # ---- Rename AF2 columns to canonical names ----
+    df = df.rename(
+        columns={
+            "pLDDT": "af2_plddt",
+            "pTM": "af2_ptm",
+            "i_pTM": "af2_iptm",
+            "pAE": "af2_pAE",
+            "unscaled_i_pAE": "af2_ipAE",
+            "pLDDT_MONOMER": "af2_monomer_plddt",
+            "pTM_MONOMER": "af2_monomer_ptm",
+            "pAE_MONOMER": "af2_monomer_pAE",
+            "bound_unbound_RMSD": "af2_bound_unbound_RMSD",
+            "af2_easy_success": "AF2-IG-easy-success",
+            "af2_opt_success": "AF2-IG-success",
+            "ptx_success": "Protenix-success",
+            "ptx_basic_success": "Protenix-basic-success",
+        }
+    )
+    return df
+
+
+def parse_hotspot_json(json_list):
+    """Parse hotspots from the input design JSON for downstream PTX permutation."""
+    result = {}
+    for entry in json_list:
+        name = entry["name"]
+        if "condition" in entry:  # structure input
+            chain_ids = entry["condition"]["filter"]["chain_id"]
+            hotspot = entry.get("hotspot", {})
+            chain_mapping = {
+                chain_id: f"{chr(65 + i)}0" for i, chain_id in enumerate(chain_ids)
+            }
+            new_hotspot = {
+                chain_mapping[chain_id]: residues
+                for chain_id, residues in hotspot.items()
+                if chain_id in chain_mapping
+            }
+        else:  # sequence input
+            assert "sequences" in entry
+            new_hotspot = {}
+            for i, seq_entity in enumerate(entry["sequences"]):
+                entity_type = list(seq_entity.keys())[0]
+                new_hotspot[f"{chr(65 + i)}0"] = seq_entity[entity_type].get(
+                    "hotspot", {}
+                )
+        result[name] = new_hotspot
+    return result
+
+
+def check_tool_weights() -> None:
+    """
+    Sanity check for required tool weights.
+    Equivalent to the original shell script.
+    """
+    root = os.environ.get("TOOL_WEIGHTS_ROOT")
+    if not root:
+        raise RuntimeError(
+            "Environment variable TOOL_WEIGHTS_ROOT is not set.\n"
+            "Please set TOOL_WEIGHTS_ROOT or run download_tool_weights.sh."
+        )
+
+    root = Path(root)
+
+    required_files = [
+        # ---- AF2 ----
+        root / "af2" / "params_model_1.npz",
+        root / "af2" / "params_model_1_ptm.npz",
+    ]
+
+    print(f"Checking tool weights in: {root}")
+
+    missing = [p for p in required_files if not p.is_file()]
+    if missing:
+        print()
+        print("Missing required tool weights:")
+        for p in missing:
+            print(f"   {p}")
+        print()
+        print("Please run:")
+        print("   bash download_tool_weights.sh")
+        sys.exit(1)

model/setup.py

@@ -0,0 +1,63 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+from pathlib import Path
+
+from setuptools import find_packages, setup
+
+this_directory = Path(__file__).parent
+long_description = (this_directory / "README.md").read_text()
+# Read requirements from the file
+# with open("requirements.txt") as f:
+#     install_requires = f.read().splitlines()
+
+# Check if the user specified the CPU option
+# if "--cpu" in sys.argv:
+#     # Remove the gpu packages
+#     try:
+#         to_drop = [x for x in install_requires if "nvidia" in x or "cuda" in x]
+#         for x in to_drop:
+#             install_requires.remove(x)
+#     except ValueError:
+#         pass
+#     # Remove the --cpu option from sys.argv so setuptools doesn't get confused
+#     sys.argv.remove("--cpu")
+
+setup(
+    name="pxdesign",
+    python_requires=">=3.10",
+    version="0.1.0",
+    description="PXDesign: a Fast, Modular, and Accurate Model For De Novo Protein Binder Design",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    author="Bytedance Inc.",
+    url="https://github.com/bytedance/PXDesign",
+    author_email="ai4s-bio@bytedance.com",
+    packages=find_packages(
+        exclude=(
+            "assets",
+            "benchmark",
+            "*.egg-info",
+        )
+    ),
+    # install_requires=install_requires,
+    license="Apache 2.0 License",
+    platforms="manylinux1",
+    entry_points={
+        "console_scripts": [
+            "pxdesign = pxdesign.runner.cli:cli",
+        ],
+    },
+)

requirements.txt

@@ -0,0 +1,33 @@
+PyYaml
+scipy
+ml_collections
+tqdm
+pandas
+optree
+rdkit
+biopython==1.83
+modelcif==0.7
+biotite==1.0.1
+scikit-learn
+scikit-learn-extra
+deepspeed>=0.15.1
+protobuf==3.20.2
+icecream
+ipdb
+wandb
+numpy==1.26.3
+matplotlib==3.9.2
+ipywidgets
+py3Dmol
+nvidia-cublas-cu12
+torch==2.3.1
+fair-esm
+posix_ipc
+einops
+transformers==4.51.3
+optax==0.2.5
+dm-haiku==0.0.13
+natsort
+dm-tree
+protenix>=0.1.0
+pxdbench>=0.1.0

scripts/Dockerfile

@@ -0,0 +1,15 @@
+FROM ai4s-cn-beijing.cr.volces.com/infra/protenix:v0.0.3
+
+# Install Protenix
+RUN pip --no-cache-dir install git+https://github.com/bytedance/Protenix.git@v0.5.0+pxd
+
+# Install PXDesignBench
+RUN pip install git+https://github.com/sokrypton/ColabDesign.git --no-deps
+RUN pip install posix_ipc einops transformers==4.51.3 optax==0.2.5 dm-haiku==0.0.13
+RUN pip install "jax[cuda]==0.4.29" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+RUN pip install numpy==1.26.3 natsort dm-tree
+RUN pip install git+https://github.com/bytedance/PXDesignBench.git@v0.1.2 --no-deps
+
+# CUTLASS (for DeepSpeed Evo attention)
+RUN git clone -b v3.5.1 https://github.com/NVIDIA/cutlass.git /opt/cutlass
+ENV CUTLASS_PATH=/opt/cutlass

scripts/download_tool_weights.sh

@@ -0,0 +1,125 @@
+#!/usr/bin/env bash
+# download_model_weights.sh
+#
+# Usage:
+#   bash download_model_weights.sh               # Download into ./tool_weights
+#   bash download_model_weights.sh /path/to/dir  # Custom download root directory
+#
+# The final directory structure will look like:
+#   MODELS_ROOT/
+#     af2/
+#       params_model_{1..5}.npz
+#       params_model_{1..5}_ptm.npz
+#       params_model_{1..5}_multimer_v3.npz
+#       LICENSE
+#     mpnn/
+#       ca_model_weights/
+#       soluble_model_weights/
+#       vanilla_model_weights/
+#
+# After downloading, update pxdbench/globals.py to point to these locations.
+
+set -euo pipefail
+
+# Set root directory for storing all model weights
+MODELS_ROOT="${1:-$(pwd)/tool_weights}"
+
+AF2_DIR="${MODELS_ROOT}/af2"
+MPNN_DIR="${MODELS_ROOT}/mpnn"
+
+echo "Model root directory: ${MODELS_ROOT}"
+mkdir -p "${AF2_DIR}" "${MPNN_DIR}"
+########################################
+# 1. AlphaFold2 parameters
+########################################
+echo "==> Downloading AlphaFold2 parameters ..."
+
+AF2_TAR="alphafold_params_2022-12-06.tar"
+AF2_URL="https://storage.googleapis.com/alphafold/${AF2_TAR}"
+
+# If AF2 parameters already exist, skip download
+if compgen -G "${AF2_DIR}/params_model_1*.npz" > /dev/null; then
+  echo "  AlphaFold2 params appear to already exist — skipping download."
+else
+  tmp_tar="${MODELS_ROOT}/${AF2_TAR}"
+  echo "  Downloading from: ${AF2_URL}"
+  curl -L "${AF2_URL}" -o "${tmp_tar}"
+
+  echo "  Extracting to: ${AF2_DIR}"
+  tar -xf "${tmp_tar}" -C "${AF2_DIR}"
+  rm -f "${tmp_tar}"
+
+  echo "  AlphaFold2 parameters downloaded successfully."
+fi
+
+
+########################################
+# 2. ProteinMPNN weights
+########################################
+echo "==> Downloading ProteinMPNN weights ..."
+
+TMP_DIR="$(mktemp -d)"
+cleanup() {
+  rm -rf "${TMP_DIR}"
+}
+trap cleanup EXIT
+
+echo "  Cloning dauparas/ProteinMPNN (shallow clone)..."
+git clone --depth 1 https://github.com/dauparas/ProteinMPNN.git "${TMP_DIR}"
+
+# Copy each weight directory
+for subdir in ca_model_weights soluble_model_weights vanilla_model_weights; do
+  src="${TMP_DIR}/${subdir}"
+  dst="${MPNN_DIR}/${subdir}"
+
+  if [ -d "${dst}" ]; then
+    echo "  ${subdir} already exists — skipping."
+  else
+    echo "  Copying ${subdir} → ${dst}"
+    mkdir -p "${MPNN_DIR}"
+    cp -r "${src}" "${dst}"
+  fi
+done
+
+echo "  ProteinMPNN weights are ready in: ${MPNN_DIR}"
+
+########################################
+# 3. CCD cache (PXDesign release_data)
+########################################
+echo "==> Downloading CCD cache ..."
+
+CCD_DIR="${1:-$(pwd)/release_data/ccd_cache}"
+mkdir -p "${CCD_DIR}"
+
+CCD_COMPONENTS_URL="https://pxdesign.tos-cn-beijing.volces.com/release_data/components.v20240608.cif"
+CCD_RDKIT_URL="https://pxdesign.tos-cn-beijing.volces.com/release_data/components.v20240608.cif.rdkit_mol.pkl"
+PDB_CLUSTER_URL="https://pxdesign.tos-cn-beijing.volces.com/release_data/clusters-by-entity-40.txt"
+
+CCD_COMPONENTS_FILE="${CCD_DIR}/components.v20240608.cif"
+CCD_RDKIT_FILE="${CCD_DIR}/components.v20240608.cif.rdkit_mol.pkl"
+PDB_CLUSTER_FILE="${CCD_DIR}/clusters-by-entity-40.txt"
+
+download_if_missing() {
+  local url="$1"
+  local out="$2"
+  if [ -f "$out" ]; then
+    echo "  $(basename "$out") already exists — skipping."
+  else
+    echo "  Downloading $(basename "$out")"
+    curl -L -C - "$url" -o "$out"
+  fi
+}
+
+download_if_missing "${CCD_COMPONENTS_URL}" "${CCD_COMPONENTS_FILE}"
+download_if_missing "${CCD_RDKIT_URL}"     "${CCD_RDKIT_FILE}"
+download_if_missing "${PDB_CLUSTER_URL}"   "${PDB_CLUSTER_FILE}"
+
+echo "  CCD cache is ready in: ${CCD_DIR}"
+
+########################################
+
+echo "==> All downloads completed."
+echo "Model weight directories:"
+echo "  AF2:      ${AF2_DIR}"
+echo "  MPNN:     ${MPNN_DIR}"
+echo "CCD cache:  ${CCD_DIR}"

scripts/install.sh

@@ -0,0 +1,369 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+############################################################
+# PXDesign One-Click Installation Script
+#
+# This script will:
+#   1. Create a dedicated conda/mamba/micromamba environment
+#   2. Install GPU PyTorch matching a specified CUDA version
+#   3. Install Protenix
+#   4. Install PXDesignBench dependencies
+#   5. Clone PXDesign repo and install
+#   6. Run basic import sanity checks
+#
+# Supported options:
+#   --env <name>           Conda/mamba environment name (default: pxdesign)
+#   --pkg_manager <tool>   conda | mamba | micromamba (default: conda)
+#   --cuda-version <ver>   CUDA version string, e.g. 12.1, 12.2, 12.4
+#                          Required. Must be >= 12.1.
+############################################################
+
+# Default configuration
+env_name="pxdesign"
+pkg_manager="conda"      # conda | mamba | micromamba
+cuda_version=""          # e.g. 12.1, 12.2, 12.4
+
+# ----------------------------------------------------------
+# Parse command-line options
+# ----------------------------------------------------------
+OPTIONS=e:p:c:
+LONGOPTIONS=env:,pkg_manager:,cuda-version:
+
+PARSED=$(getopt --options="${OPTIONS}" --longoptions="${LONGOPTIONS}" --name "$0" -- "$@") || {
+  echo "Error: failed to parse command line options."
+  exit 1
+}
+eval set -- "${PARSED}"
+
+while true; do
+  case "$1" in
+    -e|--env)
+      env_name="$2"
+      shift 2
+      ;;
+    -p|--pkg_manager)
+      pkg_manager="$2"
+      shift 2
+      ;;
+    -c|--cuda-version)
+      cuda_version="$2"
+      shift 2
+      ;;
+    --)
+      shift
+      break
+      ;;
+    *)
+      echo "Invalid option: $1" >&2
+      exit 1
+      ;;
+  esac
+done
+
+echo "=================================================="
+echo " PXDesign Installation"
+echo "   Environment name : ${env_name}"
+echo "   Package manager  : ${pkg_manager}"
+echo "   CUDA version     : ${cuda_version:-<not specified>}"
+echo "=================================================="
+
+SECONDS=0
+
+############################################################
+# CUDA version checks & PyTorch CUDA tag selection
+############################################################
+
+# Helper: check if version >= 12.1
+check_cuda_ge_12_1() {
+  local ver="$1"
+  local major="${ver%%.*}"
+  local rest="${ver#*.}"
+  local minor="${rest%%.*}"
+
+  if (( major > 12 )); then
+    return 0
+  elif (( major == 12 && minor >= 1 )); then
+    return 0
+  else
+    return 1
+  fi
+}
+
+if [ -z "${cuda_version}" ]; then
+  echo "Error: --cuda-version must be specified (e.g., --cuda-version 12.1)."
+  exit 1
+fi
+
+if ! check_cuda_ge_12_1 "${cuda_version}"; then
+  echo "Error: CUDA version must be >= 12.1, but got '${cuda_version}'."
+  exit 1
+fi
+
+# Decide PyTorch CUDA tag from CUDA version
+# Extend this mapping as needed.
+torch_tag=""
+torch_version="2.3.1"   # adjust if needed
+
+if [[ "${cuda_version}" == 12.1* || "${cuda_version}" == 12.2* ]]; then
+  torch_tag="cu121"
+elif [[ "${cuda_version}" == 12.4* || "${cuda_version}" == 12.5* ]]; then
+  torch_tag="cu124"
+else
+  echo "Error: unsupported CUDA version '${cuda_version}' for this installer."
+  echo "       Currently supported: 12.1, 12.2 (cu121) and 12.4+ (cu124)."
+  exit 1
+fi
+
+############################################################
+# Package manager detection and initialization
+############################################################
+
+case "${pkg_manager}" in
+  conda)
+    if ! command -v conda >/dev/null 2>&1; then
+      echo "Error: conda is not installed or not in PATH."
+      exit 1
+    fi
+    env_tool="conda"
+    ;;
+  mamba)
+    if ! command -v mamba >/dev/null 2>&1; then
+      echo "Error: mamba is not installed or not in PATH."
+      exit 1
+    fi
+    if ! command -v conda >/dev/null 2>&1; then
+      echo "Error: mamba is installed but conda is not available."
+      exit 1
+    fi
+    env_tool="mamba"
+    ;;
+  micromamba)
+    if ! command -v micromamba >/dev/null 2>&1; then
+      echo "Error: micromamba is not installed or not in PATH."
+      exit 1
+    fi
+    env_tool="micromamba"
+    ;;
+  *)
+    echo "Error: unsupported pkg_manager '${pkg_manager}'. Use 'conda', 'mamba', or 'micromamba'."
+    exit 1
+    ;;
+esac
+
+install_dir=$(pwd)
+echo "Install root    : ${install_dir}"
+
+############################################################
+# Create and activate environment
+############################################################
+
+if [ "${env_tool}" = "micromamba" ]; then
+  echo ">>> Using micromamba to manage environments"
+
+  # Ensure MAMBA_ROOT_PREFIX is defined to avoid "unbound variable" under `set -u`
+  export MAMBA_ROOT_PREFIX="${MAMBA_ROOT_PREFIX:-$HOME/micromamba}"
+
+  # Initialize micromamba shell hook for bash
+  eval "$(micromamba shell hook -s bash)"
+
+  echo ">>> Creating environment '${env_name}' (Python 3.11) with micromamba"
+  micromamba create -y -n "${env_name}" python=3.11 || {
+    echo "Error: failed to create environment ${env_name} with micromamba"
+    exit 1
+  }
+
+  echo ">>> Activating environment '${env_name}' (micromamba)"
+  micromamba activate "${env_name}" || {
+    echo "Error: failed to activate environment ${env_name} with micromamba"
+    exit 1
+  }
+
+else
+  echo ">>> Using ${env_tool} to manage environments"
+
+  CONDA_BASE=$(conda info --base 2>/dev/null) || {
+    echo "Error: conda is not installed or cannot be initialized."
+    exit 1
+  }
+
+  echo "Conda base      : ${CONDA_BASE}"
+
+  echo ">>> Creating environment '${env_name}' (Python 3.11) with ${env_tool}"
+  "${env_tool}" create -y -n "${env_name}" python=3.11 || {
+    echo "Error: failed to create environment ${env_name} with ${env_tool}"
+    exit 1
+  }
+
+  echo ">>> Activating environment '${env_name}' (${env_tool})"
+  # shellcheck disable=SC1090
+  source "${CONDA_BASE}/etc/profile.d/conda.sh"
+  conda activate "${env_name}" || {
+    echo "Error: failed to activate environment ${env_name}"
+    exit 1
+  }
+
+  if [ "${CONDA_DEFAULT_ENV:-}" != "${env_name}" ]; then
+    echo "Error: expected environment '${env_name}' to be active, but got '${CONDA_DEFAULT_ENV:-}'."
+    exit 1
+  fi
+fi
+
+echo "Environment '${env_name}' successfully activated."
+
+############################################################
+# Python package installation
+############################################################
+
+echo ">>> Upgrading pip"
+python -m pip install --upgrade pip
+
+# ----------------------------------------------------------
+# 1) Install GPU PyTorch first (matching CUDA version)
+# ----------------------------------------------------------
+echo ">>> Installing PyTorch (GPU, CUDA ${cuda_version}, tag ${torch_tag})"
+pip install --no-cache-dir \
+  "torch==${torch_version}" \
+  --index-url "https://download.pytorch.org/whl/${torch_tag}" \
+  || { echo "Error: failed to install PyTorch ${torch_version} with ${torch_tag} wheels."; exit 1; }
+
+python - << 'PYTORCH_CHECK'
+import torch
+print("PyTorch version:", torch.__version__)
+print("CUDA available :", torch.cuda.is_available())
+print("Torch CUDA     :", torch.version.cuda)
+if torch.cuda.is_available():
+    print("CUDA devices   :", torch.cuda.device_count())
+PYTORCH_CHECK
+
+# ----------------------------------------------------------
+# 2) Install Protenix & PXDesignBench
+# ----------------------------------------------------------
+
+echo ">>> Installing Protenix"
+pip install --no-cache-dir "git+https://github.com/bytedance/Protenix.git@v0.5.0+pxd" \
+  || { echo "Error: failed to install Protenix."; exit 1; }
+
+echo ">>> Installing PXDesignBench base dependencies"
+pip install --no-cache-dir \
+  einops \
+  natsort \
+  dm-tree \
+  posix_ipc \
+  "transformers==4.51.3" \
+  "dm-haiku==0.0.13" \
+  "optax==0.2.5" \
+  || { echo "Error: failed to install base Python dependencies."; exit 1; }
+
+echo ">>> Installing ColabDesign (without dependencies)"
+pip install --no-cache-dir git+https://github.com/sokrypton/ColabDesign.git --no-deps \
+  || { echo "Error: failed to install ColabDesign."; exit 1; }
+
+echo ">>> Installing JAX with CUDA support"
+pip install --no-cache-dir \
+  "jax[cuda]==0.4.29" \
+  -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
+  || { echo "Error: failed to install JAX (CUDA build)."; exit 1; }
+
+# downgrade numpy
+pip install --no-cache-dir \
+  "numpy==1.26.3" \
+  || { echo "Error: failed to install numpy 1.26.3."; exit 1; }
+
+echo ">>> Installing PXDesignBench"
+pip install --no-cache-dir git+https://github.com/bytedance/PXDesignBench.git@v0.1.2 --no-deps \
+  || { echo "Error: failed to install PXDesignBench."; exit 1; }
+
+echo ">>> Installing PXDesign"
+pip install -e .
+
+if [ "${env_tool}" = "micromamba" ]; then
+  micromamba install -c conda-forge cudnn -y || { echo "Error: failed to install cudnn with micromamba."; exit 1; }
+else
+  conda install -c conda-forge cudnn -y || { echo "Error: failed to install cudnn with conda."; exit 1; }
+fi
+
+
+# -------------------------------
+# 3) CUTLASS (for DeepSpeed Evo attention)
+# -------------------------------
+
+# Default to $HOME/cutlass if CUTLASS_PATH is not set by the user
+export CUTLASS_PATH="${CUTLASS_PATH:-$HOME/cutlass}"
+
+echo "[CUTLASS] Using CUTLASS_PATH=${CUTLASS_PATH}"
+
+if [ ! -d "${CUTLASS_PATH}" ]; then
+  echo "[CUTLASS] CUTLASS not found, cloning NVIDIA/cutlass v3.5.1 ..."
+  git clone -b v3.5.1 https://github.com/NVIDIA/cutlass.git "${CUTLASS_PATH}"
+else
+  echo "[CUTLASS] Existing CUTLASS directory detected, skipping clone."
+fi
+
+
+############################################################
+# Sanity checks
+############################################################
+
+echo ">>> Running sanity checks (import tests)"
+
+python - << 'PYCODE'
+import sys
+
+def check(mod):
+    try:
+        __import__(mod)
+        print(f"[OK] import {mod}")
+    except Exception as e:
+        print(f"[FAIL] import {mod}: {e}", file=sys.stderr)
+        raise
+
+modules = [
+    "torch",
+    "jax",
+    "jax.numpy",
+    "colabdesign",
+    "protenix",
+    "pxdbench",
+    "pxdesign"
+]
+
+for m in modules:
+    check(m)
+
+import jax
+print("JAX devices:", jax.devices())
+PYCODE
+
+echo "Sanity checks completed."
+
+############################################################
+# Cleanup and final message
+############################################################
+
+echo ">>> Cleaning up package manager caches"
+
+if [ "${env_tool}" = "micromamba" ]; then
+  micromamba clean -a -y || echo "Warning: failed to clean micromamba caches."
+  micromamba deactivate || true
+else
+  "${env_tool}" clean -a -y || echo "Warning: failed to clean ${env_tool} caches."
+  conda deactivate || true
+fi
+
+t=${SECONDS}
+echo "=================================================="
+echo " PXDesign environment setup done!"
+echo "   Environment name : ${env_name}"
+echo "   Package manager  : ${pkg_manager}"
+echo "   CUDA version     : ${cuda_version} (torch tag: ${torch_tag})"
+echo
+echo " Activate with:"
+if [ "${env_tool}" = "micromamba" ]; then
+  echo "   micromamba activate ${env_name}"
+else
+  echo "   conda activate ${env_name}"
+fi
+echo
+echo " Installation time: $((t / 3600))h $(((t / 60) % 60))m $((t % 60))s"
+echo "=================================================="

weight/release_data/ccd_cache/clusters-by-entity-40.txt

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

weight/release_data/ccd_cache/components.v20240608.cif

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

weight/release_data/ccd_cache/components.v20240608.cif.rdkit_mol.pkl

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

weight/release_data/checkpoint/protenix_base_default_v0.5.0.pt

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

weight/release_data/checkpoint/protenix_mini_default_v0.5.0.pt

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

weight/release_data/checkpoint/protenix_mini_tmpl_v0.5.0.pt

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

weight/release_data/checkpoint/pxdesign_v0.1.0.pt

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