perf: reduce GPU memory usage during inference

[haoyu-haoyu]

Summary

• Offload trajectory to CPU per step: Move px0, x_t, seq, plddt to CPU immediately after each denoising step instead of accumulating all steps on GPU. Peak GPU memory is reduced from O(T × L × 14 × 3) to O(L × 14 × 3) for trajectory storage.

• Remove unused ComputeAllAtomCoords() instantiations: get_next_ca() and get_next_pose() in inference/utils.py created a ComputeAllAtomCoords() module every timestep but never used it. Removed both instantiations and the now-unused import.

Why this is safe

x_t and seq_t are produced fresh by sampler.sample_step() each iteration — the active computation tensors remain on GPU. Only the previous step's saved copies (for trajectory visualization) are moved to CPU. The final torch.stack() and torch.flip() operate on CPU tensors, which is fine since they're only used for PDB writing.

Impact

For a 200-residue, 50-step design, this reduces GPU memory for trajectory storage from ~50× to 1× the single-step footprint. This is especially important for longer proteins (>300 residues) that can OOM on 16GB GPUs.

Test plan

• Compare output PDB coordinates with and without this change (should be identical)
• Profile GPU memory with torch.cuda.max_memory_allocated()
• Run tests/test_diffusion.py