Jax Attributes error

[AdvaithVellanki]

The Jax version that gets installed according to requirements.txt has several missing attributes that are needed for the rest of the code.

Hoping to get some clarity on this. Here is the error message:
AttributeError: module 'jax.random' has no attribute 'KeyArray'

[rohitgandikota]

Hey - we haven't tested this repo on jax. Let me know if you would like to open a pull request for a Jax version. Would be happy to collaborate

[RedParisArts]

They're referring to the colab notebook...it doesn't work because the error related jax.random

[kordc]

You just need to downgrade JAX -- refer to https://stackoverflow.com/questions/78302031/stable-diffusion-attributeerror-module-jax-random-has-no-attribute-keyarray