[BUG REPORT] Launch error on Google Colab.

[JustinsKittens]

Describe the bug When I download the google colab notebook file, open it in Google Colab, follow the instructions and run each of the things, I get the following error when everything installs on the "Launching" part:

Checking installs and venv + autodebug checks
Python version:  3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0]
Webui version: 7eb428accafc56d3836ee6b399fa9f852ac75f89 - Fri Feb 16 19:25:22 2024 +0100
Found extensions: 
Installed wheel!             
Installed PyTorch!             
Installed diffusers!             
Installed gradio!             
Installed huggingface-hub!             
Installed audio2numpy!             
Installed bark!             
Installed sox!             
Installed audiolm!             
Installed praat-parselmouth!             
Installed pyworld!             
Installed faiss!             
Installed torchcrepe!             
Installed ffmpeg-python!             
Installed noisereduce!             
Installed demucs!             
Installed pytube!             
Installed openai-whisper!             
Installed audiocraft!             
Installed beartype!             
Done installing/checking installs.
Activating extensions
Preparing
Traceback (most recent call last):
  File "/content/audio-webui/main.py", line 45, in <module>
    autodebug.catcher(e)
  File "/content/audio-webui/autodebug/autodebug.py", line 79, in catcher
    raise e
  File "/content/audio-webui/main.py", line 27, in <module>
    from webui.modules.implementations.tts_monkeypatching import patch as patch1
  File "/content/audio-webui/webui/modules/implementations/__init__.py", line 1, in <module>
    import webui.modules.implementations.ttsmodels as tts
  File "/content/audio-webui/webui/modules/implementations/ttsmodels.py", line 15, in <module>
    from webui.modules.implementations.patches.bark_custom_voices import wav_to_semantics, generate_fine_from_wav, \
  File "/content/audio-webui/webui/modules/implementations/patches/bark_custom_voices.py", line 8, in <module>
    from webui.modules.implementations.patches.bark_generation import generate_text_semantic_new, generate_coarse_new, generate_fine_new
  File "/content/audio-webui/webui/modules/implementations/patches/bark_generation.py", line 7, in <module>
    from webui.ui.tabs import settings
  File "/content/audio-webui/webui/ui/tabs/__init__.py", line 5, in <module>
    from .audioldm import create_tab as audioldm_tab
  File "/content/audio-webui/webui/ui/tabs/audioldm/__init__.py", line 2, in <module>
    from .audioldm import audioldm_tab
  File "/content/audio-webui/webui/ui/tabs/audioldm/audioldm.py", line 2, in <module>
    import webui.modules.implementations.audioldm as aldm
  File "/content/audio-webui/webui/modules/implementations/audioldm.py", line 4, in <module>
    import diffusers
  File "/usr/local/lib/python3.10/dist-packages/diffusers/__init__.py", line 38, in <module>
    from .models import (
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/__init__.py", line 36, in <module>
    from .controlnet_flax import FlaxControlNetModel
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/controlnet_flax.py", line 25, in <module>
    from .modeling_flax_utils import FlaxModelMixin
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/modeling_flax_utils.py", line 46, in <module>
    class FlaxModelMixin(PushToHubMixin):
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/modeling_flax_utils.py", line 195, in FlaxModelMixin
    def init_weights(self, rng: jax.random.KeyArray) -> Dict:
  File "/usr/local/lib/python3.10/dist-packages/jax/_src/deprecations.py", line 54, in getattr
    raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax.random' has no attribute 'KeyArray'

To Reproduce Steps to reproduce the behavior:

1. Open up the Google Colab notebook
2. Follow the instructions.
3. Run each code block in order.
4. See the error.

Expected behavior I expected there to not be an error.

Screenshots

Additional context Add any other context about the problem here.