TeaCache added

README.md

@@ -53,9 +53,23 @@ simplefold \
     --plddt \                           # output pLDDT
     --fasta_path [FASTA_PATH] \         # path to the target fasta directory or file
     --output_dir [OUTPUT_DIR] \         # path to the output directory
-    --backend [mlx, torch]              # choose from MLX and PyTorch for inference backend 
+    --backend [mlx, torch] \            # choose from MLX and PyTorch for inference backend
+    --teacache 0.0                      # optional TeaCache threshold (0 disables; see below)
 ```
 
+### TeaCache acceleration
+
+SimpleFold supports optional **TeaCache** acceleration, a training-free technique that caches and reuses velocity predictions across timesteps where the model output is changing slowly, trading a small amount of quality for a large speedup in the diffusion sampling loop. The implementation is adapted from [TeaCache](https://github.com/ali-vilab/TeaCache) (*"Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model"*) and works with both the PyTorch and MLX backends.
+
+Enable it by passing a positive threshold via `--teacache`:
+
+- `--teacache 0.1` — conservative (higher quality, smaller speedup)
+- `--teacache 0.15` — balanced default
+- `--teacache 0.2` — aggressive (more speed, larger quality loss)
+- `--teacache 0.0` — disabled (use the standard Euler–Maruyama sampler)
+
+Lower thresholds force the sampler to recompute more often; higher thresholds allow more cache hits. A short warm-up period at the start of sampling always runs at full precision before caching kicks in.
+
 ## Evaluation
 
 We provide predicted structures from SimpleFold of different model sizes:

src/simplefold/cli.py

@@ -28,6 +28,7 @@ def main():
     parser.add_argument("--output_format", type=str, default="mmcif", choices=["pdb", "mmcif"], help="Output file format.")
     parser.add_argument("--backend", type=str, default='torch', choices=['torch', 'mlx'], help="Backend to run inference either torch or mlx")
     parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility.")
+    parser.add_argument("--teacache", type=float, default=0.0, help="Enable TeaCache acceleration with the given threshold (e.g. 0.1 for quality, 0.2 for speed). 0 disables.")
     parser.add_argument(
         "--version",
         action="version",

src/simplefold/inference.py

@@ -17,6 +17,8 @@
 
 from model.flow import LinearPath
 from model.torch.sampler import EMSampler
+from model.torch.teacache import TeaCacheSampler as TeaCacheSamplerTorch
+from model.torch.teacache import TeaCacheConfig as TeaCacheConfigTorch
 
 from processor.protein_processor import ProteinDataProcessor
 from utils.datamodule_utils import process_one_inference_structure
@@ -31,6 +33,7 @@
     from mlx.utils import tree_unflatten, tree_flatten
     from model.mlx.sampler import EMSampler as EMSamplerMLX
     from model.mlx.esm_network import ESM2 as ESM2MLX
+    from model.mlx.teacache import TeaCacheSampler, TeaCacheConfig
     from utils.mlx_utils import map_torch_to_mlx, map_plddt_torch_to_mlx
     MLX_AVAILABLE = True
 except:
@@ -216,18 +219,44 @@ def initialize_others(args, device):
     # define flow process and sampler
     flow = LinearPath()
 
+    teacache_threshold = getattr(args, 'teacache', 0.0)
+
     if args.backend == "torch":
-        sampler_cls = EMSampler
+        if teacache_threshold > 0:
+            config = TeaCacheConfigTorch(threshold=teacache_threshold)
+            sampler = TeaCacheSamplerTorch(
+                num_timesteps=args.num_steps,
+                t_start=1e-4,
+                tau=args.tau,
+                config=config,
+            )
+            print(f"TeaCache enabled (threshold={teacache_threshold})")
+        else:
+            sampler = EMSampler(
+                num_timesteps=args.num_steps,
+                t_start=1e-4,
+                tau=args.tau,
+                log_timesteps=True,
+                w_cutoff=0.99,
+            )
     elif args.backend == "mlx":
-        sampler_cls = EMSamplerMLX
-
-    sampler = sampler_cls(
-        num_timesteps=args.num_steps,
-        t_start=1e-4,
-        tau=args.tau,
-        log_timesteps=True,
-        w_cutoff=0.99,
-    )
+        if teacache_threshold > 0:
+            config = TeaCacheConfig(threshold=teacache_threshold)
+            sampler = TeaCacheSampler(
+                num_timesteps=args.num_steps,
+                t_start=1e-4,
+                tau=args.tau,
+                config=config,
+            )
+            print(f"TeaCache enabled (threshold={teacache_threshold})")
+        else:
+            sampler = EMSamplerMLX(
+                num_timesteps=args.num_steps,
+                t_start=1e-4,
+                tau=args.tau,
+                log_timesteps=True,
+                w_cutoff=0.99,
+            )
     return tokenizer, featurizer, processor, flow, sampler