AttributeError: module 'jax.random' has no attribute 'KeyArray'

[amrzv]

Hello. When running google colab notebook this error is raised on the line from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler

The issue is that jax library needs to be downgraded to 0.4.23, so adding the following into the first cell resolves the error

jax==0.4.23 jaxlib==0.4.23

stackoverflow

Should the requirements.txt be patched? I can create an MR.

[eliphatfs]

Thank you for your report! Could you specify a little bit more -- are you using jax for inference or you are using torch but importing diffusers causes a problem?

[amrzv]

The second scenario, as it is in colab.