Transpose the code for more efficient VRAM and faster inference?

[remybonnav]

I was wondering if it would be possible that someone transpose the code to use xformers/JAX/Transformers/diffusers in order to optimize VRAM necessary and to speed up inference? (as used in other generative models such as stable diffusion or else).

Maybe some people from https://github.com/OpenBioML could help?

[remybonnav]

I think I wrote to fast, I did not know that SE3 Transformers was actually already a huge optimization and even provide a low memory mode.
I will try low memory mode on my 4GB VRAM laptop.