RuntimeError: Failed to import diffusers.pipelines.pipeline_utils because of the following error (look up to see its traceback): Failed to import diffusers.models.autoencoders.autoencoder_kl because of the following error (look up to see its traceback): cannot import name 'linear_util' from 'jax' (/usr/local/lib/python3.10/dist-packages/jax/__init__.py)

[LTYinfinite]

After I run the second cell of this colab notebook, it show ups this:

<ipython-input-3-fdfb929e2844>:26: DeprecationWarning: setDaemon() is deprecated, set the daemon attribute instead
  new_thread.setDaemon(True)
Your url is: https://moscow-correctly-stamp-tu.trycloudflare.com/
2024-07-29 11:08:48.336536: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-07-29 11:08:48.336587: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-07-29 11:08:48.337962: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-07-29 11:08:48.345098: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-07-29 11:08:49.478088: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/diffusers/utils/import_utils.py", line 808, in _get_module
    return importlib.import_module("." + module_name, self.__name__)
  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1050, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1027, in _find_and_load
  File "<frozen importlib._bootstrap>", line 992, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "<frozen importlib._bootstrap>", line 1050, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1027, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1006, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 688, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 883, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/autoencoders/__init__.py", line 1, in <module>
    from .autoencoder_asym_kl import AsymmetricAutoencoderKL
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/autoencoders/autoencoder_asym_kl.py", line 23, in <module>
    from .vae import DecoderOutput, DiagonalGaussianDistribution, Encoder, MaskConditionDecoder
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/autoencoders/vae.py", line 25, in <module>
    from ..unets.unet_2d_blocks import (
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/unets/__init__.py", line 18, in <module>
    from .unet_2d_condition_flax import FlaxUNet2DConditionModel
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/unets/unet_2d_condition_flax.py", line 16, in <module>
    import flax
  File "/usr/local/lib/python3.10/dist-packages/flax/__init__.py", line 18, in <module>
    from . import core as core
  File "/usr/local/lib/python3.10/dist-packages/flax/core/__init__.py", line 15, in <module>
    from .axes_scan import broadcast as broadcast
  File "/usr/local/lib/python3.10/dist-packages/flax/core/axes_scan.py", line 21, in <module>
    from jax import linear_util as lu
ImportError: cannot import name 'linear_util' from 'jax' (/usr/local/lib/python3.10/dist-packages/jax/__init__.py)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/diffusers/utils/import_utils.py", line 808, in _get_module
    return importlib.import_module("." + module_name, self.__name__)
  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1050, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1027, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1006, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 688, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 883, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/pipeline_utils.py", line 44, in <module>
    from ..models import AutoencoderKL
  File "<frozen importlib._bootstrap>", line 1075, in _handle_fromlist
  File "/usr/local/lib/python3.10/dist-packages/diffusers/utils/import_utils.py", line 798, in __getattr__
    module = self._get_module(self._class_to_module[name])
  File "/usr/local/lib/python3.10/dist-packages/diffusers/utils/import_utils.py", line 810, in _get_module
    raise RuntimeError(
RuntimeError: Failed to import diffusers.models.autoencoders.autoencoder_kl because of the following error (look up to see its traceback):
cannot import name 'linear_util' from 'jax' (/usr/local/lib/python3.10/dist-packages/jax/__init__.py)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/content/dalle-playground/backend/app.py", line 10, in <module>
    from stable_diffusion_wrapper import StableDiffusionWrapper
  File "/content/dalle-playground/backend/stable_diffusion_wrapper.py", line 1, in <module>
    from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
  File "<frozen importlib._bootstrap>", line 1075, in _handle_fromlist
  File "/usr/local/lib/python3.10/dist-packages/diffusers/utils/import_utils.py", line 799, in __getattr__
    value = getattr(module, name)
  File "/usr/local/lib/python3.10/dist-packages/diffusers/utils/import_utils.py", line 798, in __getattr__
    module = self._get_module(self._class_to_module[name])
  File "/usr/local/lib/python3.10/dist-packages/diffusers/utils/import_utils.py", line 810, in _get_module
    raise RuntimeError(
RuntimeError: Failed to import diffusers.pipelines.pipeline_utils because of the following error (look up to see its traceback):
Failed to import diffusers.models.autoencoders.autoencoder_kl because of the following error (look up to see its traceback):
cannot import name 'linear_util' from 'jax' (/usr/local/lib/python3.10/dist-packages/jax/__init__.py)

And it eventually halted
I know this thing was stop maintain by the author for years, so is there a solution?