update readme and cli

Author: zhangyuxuann

README.md

@@ -34,6 +34,7 @@ Protenix is built for high-accuracy structure prediction. It serves as an initia
 - **[Protenix-Dock](https://github.com/bytedance/Protenix-Dock)**: Our implementation of a classical protein-ligand docking framework that leverages empirical scoring functions. Without using deep neural networks, Protenix-Dock delivers competitive performance in rigid docking tasks.
 
 ## 🎉 Updates
+- 2025-11-05: [**Protenix-v0.7.0**](./assets/inference_time_vs_ntoken.png) is now open-sourced, with new options for faster diffusion inference: shared variable caching, efficient bias fusion, and TF32 acceleration.
 - 2025-07-17: **Protenix-Mini released!**: Lightweight model variants with significantly reduced inference cost are now available. Users can choose from multiple configurations to balance speed and accuracy based on deployment needs. See our [paper](https://arxiv.org/abs/2507.11839) and [model configs](./configs/configs_model_type.py) for more information. 
 - 2025-07-17: [***New constraint feature***](docs/infer_json_format.md#constraint) is released! Now supports **atom-level contact** and **pocket** constraints, significantly improving performance in our evaluations.
 - 2025-05-30: **Protenix-v0.5.0** is now available! You may try Protenix-v0.5.0 by accessing the [server](https://protenix-server.com), or upgrade to the latest version using pip.

runner/batch_inference.py

@@ -170,6 +170,9 @@ def get_default_runner(
     use_msa: bool = True,
     trimul_kernel="cuequivariance",
     triatt_kernel="triattention",
+    enable_cache=True,
+    enable_fusion=True,
+    enable_tf32=True,
 ) -> InferenceRunner:
     inference_configs["model_name"] = model_name
     configs = {**configs_base, **{"data": data_configs}, **inference_configs}
@@ -194,7 +197,14 @@ def get_default_runner(
     configs.use_msa = use_msa
     configs.triangle_multiplicative = trimul_kernel
     configs.triangle_attention = triatt_kernel
+    configs.enable_diffusion_shared_vars_cache = enable_cache
+    configs.enable_efficient_fusion = enable_fusion
+    configs.enable_tf32 = enable_tf32
 
+    logger.info(
+        f"enable_diffusion_shared_vars_cache: {configs.enable_diffusion_shared_vars_cache}, "
+        + f"enable_efficient_fusion: {configs.enable_efficient_fusion}, enable_tf32: {configs.enable_tf32}"
+    )
     download_infercence_cache(configs)
     return InferenceRunner(configs)
 
@@ -210,6 +220,9 @@ def inference_jsons(
     model_name: str = "protenix_base_default_v0.5.0",
     trimul_kernel="cuequivariance",
     triatt_kernel="triattention",
+    enable_cache=True,
+    enable_fusion=True,
+    enable_tf32=True,
     msa_server_mode: str = "protenix",
 ) -> None:
     """
@@ -246,6 +259,9 @@ def inference_jsons(
         use_msa,
         trimul_kernel,
         triatt_kernel,
+        enable_cache,
+        enable_fusion,
+        enable_tf32,
     )
     configs = runner.configs
     for idx, infer_json in enumerate(tqdm.tqdm(infer_jsons)):
@@ -266,15 +282,16 @@ def protenix_cli():
 
 
 @click.command()
-@click.option("--input", type=str, help="json files or dir for inference")
-@click.option("--out_dir", default="./output", type=str, help="infer result dir")
+@click.option("-i", "--input", type=str, help="json files or dir for inference")
+@click.option("-o", "--out_dir", default="./output", type=str, help="infer result dir")
 @click.option(
-    "--seeds", type=str, default="101", help="the inference seed, split by comma"
+    "-s", "--seeds", type=str, default="101", help="the inference seed, split by comma"
 )
-@click.option("--cycle", type=int, default=10, help="pairformer cycle number")
-@click.option("--step", type=int, default=200, help="diffusion step")
-@click.option("--sample", type=int, default=5, help="sample number")
+@click.option("-c", "--cycle", type=int, default=10, help="pairformer cycle number")
+@click.option("-p", "--step", type=int, default=200, help="diffusion step")
+@click.option("-e", "--sample", type=int, default=5, help="sample number")
 @click.option(
+    "-n",
     "--model_name",
     type=str,
     default="protenix_base_default_v0.5.0",
@@ -301,6 +318,24 @@ def protenix_cli():
     default="triattention",
     help="Kernel to use for triangle attention. Options: 'triattention', 'cuequivariance', 'deepspeed', 'torch'.",
 )
+@click.option(
+    "--enable_cache",
+    type=bool,
+    default=True,
+    help="The diffusion module precomputes and caches pair_z, p_lm, and c_l (which are shareable across the N_sample and N_step dimensions)",
+)
+@click.option(
+    "--enable_fusion",
+    type=bool,
+    default=True,
+    help="The diffusion transformer consists of 24 transformer blocks, and the biases in these blocks can be pre-transformed in terms of dimensionality and normalization",
+)
+@click.option(
+    "--enable_tf32",
+    type=bool,
+    default=True,
+    help="When the diffusion module uses FP32 computation, enabling enable_tf32 reduces the matrix multiplication precision from FP32 to TF32.",
+)
 @click.option(
     "--msa_server_mode",
     type=str,
@@ -319,6 +354,9 @@ def predict(
     use_default_params,
     trimul_kernel,
     triatt_kernel,
+    enable_cache,
+    enable_fusion,
+    enable_tf32,
     msa_server_mode,
 ):
     """
@@ -380,6 +418,9 @@ def predict(
         model_name=model_name,
         trimul_kernel=trimul_kernel,
         triatt_kernel=triatt_kernel,
+        enable_cache=enable_cache,
+        enable_fusion=enable_fusion,
+        enable_tf32=enable_tf32,
         msa_server_mode=msa_server_mode,
     )
 

runner/inference.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import logging
 import os
+import time
 import traceback
 import urllib.request
 from argparse import Namespace
@@ -307,10 +308,13 @@ def infer_predict(runner: InferenceRunner, configs: Any) -> None:
         return
 
     num_data = len(dataloader.dataset)
+    t0_start = time.time()
     for seed in configs.seeds:
         seed_everything(seed=seed, deterministic=configs.deterministic)
+        t1_start = time.time()
         for batch in dataloader:
             try:
+                t2_start = time.time()
                 data, atom_array, data_error_message = batch[0]
                 sample_name = data["sample_name"]
 
@@ -338,9 +342,9 @@ def infer_predict(runner: InferenceRunner, configs: Any) -> None:
                     atom_array=atom_array,
                     entity_poly_type=data["entity_poly_type"],
                 )
-
+                t2_end = time.time()
                 logger.info(
-                    f"[Rank {DIST_WRAPPER.rank}] {data['sample_name']} succeeded.\n"
+                    f"[Rank {DIST_WRAPPER.rank}] {data['sample_name']} succeeded. Model forward time: {t2_end-t2_start}s.\n"
                     f"Results saved to {configs.dump_dir}"
                 )
                 torch.cuda.empty_cache()
@@ -352,6 +356,14 @@ def infer_predict(runner: InferenceRunner, configs: Any) -> None:
                     f.write(error_message)
                 if hasattr(torch.cuda, "empty_cache"):
                     torch.cuda.empty_cache()
+        t1_end = time.time()
+        logger.info(
+            f"[Rank {DIST_WRAPPER.rank}] seed {seed} succeeded. Total task time: {t1_end-t1_start}s.\n"
+        )
+    t0_end = time.time()
+    logger.info(
+        f"[Rank {DIST_WRAPPER.rank}] job succeeded. Total job time: {t0_end-t0_start}s.\n"
+    )
 
 
 def main(configs: Any) -> None:

setup.py

@@ -38,7 +38,7 @@
 setup(
     name="protenix",
     python_requires=">=3.11",
-    version="0.6.3",
+    version="0.7.0",
     description="A trainable PyTorch reproduction of AlphaFold 3.",
     long_description=long_description,
     long_description_content_type="text/markdown",