For people having jax error on colab in animatediff mode, run this once, in a new line:

!pip uninstall jax -y