adapt for esmfold

[JinyuanSun]

Made small changes to the esm.py file to meet the requirement of esmfold code.

[pengzhangzhi]

Thanks Jinyuan! I like ur implementation, efficient and elegant! A couple of things we should do here：

1. Add a test function to compare the errors and efficiency.
2. Visualize nice structures for the README.

Happy to take the PR if we finish these two.

[hmms117]

This is great! Could you continue this plot to 1000aa or even 2000aa? Memory and runtime issues with default ESMFold only starts becoming an issue for very long proteins (without chunking) so the gains should be even more for >1000aa. Moreover, default ESMFold needs chunking to work for very large proteins, making faesmfold faster (add a plot line for esmfold with chunking?)
If you dont have the VMEM, I would be happy to run the benchmark.

[JinyuanSun]

After carefully check the memory usage, I found this improvement is likely a bug of the benchmark script😅. I mannuly checked a 560aa sequence. The raw version EsmFold used 16073 MB vmem, and replace the esm with faesm only reduced the memory consumption to 15986 MB (~0.5% smaller). I will try to modifiy the EsmFoldAttention see if can do better.