[Bug]:Incompatibility between jaxlib and jax packages

[CYDOPS]

Checklist

• The issue exists after disabling all extensions
• The issue exists on a clean installation of webui
• The issue is caused by an extension, but I believe it is caused by a bug in the webui
• The issue exists in the current version of the webui
• The issue has not been reported before recently
• The issue has been reported before but has not been fixed yet

What happened?

txt2img not working as expected

Steps to reproduce the problem

After the WIN10 computer restarts, it is estimated that the system update will make the SD unusable. It can be used before restarting.

What should have happened?

Run normally and open weibe UI

What browsers do you use to access the UI ?

Google Chrome

Sysinfo

SD-WebUI Launcher Diagnostic File

Date: 2024-05-20 14:44:24
Launcher Version: 2.8.4.332
Data File Version: 2024-05-18 17:21
SD-WebUI Version: bef51ae (2024-03-02 12:03:13)
Working Directory: C:\sd-webui-aki-v4

System Information:
OS: Microsoft Windows NT 10.0.19045.0
CPU: 16 cores
Memory Size: 32768 MB
Page File Size: 30408 MB

NVIDIA Management Library:
NVIDIA Driver Version: 552.44
NVIDIA Management Library Version: 12.552.44

CUDA Driver:
Version: 12040
Devices:
00000000:01:00.0 0: NVIDIA GeForce RTX 3060 [86] 11 GB

NvApi:
Version: 55244 r552_19

HIP Driver:
Not Available

DirectML Driver:
Devices:
9476 0: NVIDIA GeForce RTX 3060 11 GB
18050 1: Intel(R) UHD Graphics 730 0 GB
9476 2: NVIDIA GeForce RTX 3060 11 GB

Intel Level Zero Driver:
Drivers:
C1670301A0F83C000000000000000000 Driver 010367C1 API 00010003
Devices:
Intel(R) UHD Graphics 730
TYP 1
VEN 00008086
DEV 00004682
FLG 5
SDI 00000000
CLK 1550
ALC 4294959104
UUID 868082460C0000000002000000000000
Memories:
DDR (13646823424) CLK 0 BUS 64

Console logs

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\transformers\utils\import_utils.py", line 1086, in _get_module
    return importlib.import_module("." + module_name, self.__name__)
  File "importlib\__init__.py", line 126, in import_module
  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 "C:\sd-webui-aki-v4\py310\lib\site-packages\transformers\modeling_utils.py", line 85, in <module>
    from accelerate import __version__ as accelerate_version
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\accelerate\__init__.py", line 3, in <module>
    from .accelerator import Accelerator
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\accelerate\accelerator.py", line 41, in <module>
    from .tracking import LOGGER_TYPE_TO_CLASS, GeneralTracker, filter_trackers
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\accelerate\tracking.py", line 43, in <module>
    from torch.utils import tensorboard
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\torch\utils\tensorboard\__init__.py", line 12, in <module>
    from .writer import FileWriter, SummaryWriter  # noqa: F401
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\torch\utils\tensorboard\writer.py", line 16, in <module>
    from ._embedding import (
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\torch\utils\tensorboard\_embedding.py", line 9, in <module>
    _HAS_GFILE_JOIN = hasattr(tf.io.gfile, "join")
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorboard\lazy.py", line 65, in __getattr__
    return getattr(load_once(self), attr_name)
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorboard\lazy.py", line 97, in wrapper
    cache[arg] = f(arg)
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorboard\lazy.py", line 50, in load_once
    module = load_fn()
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorboard\compat\__init__.py", line 45, in tf
    import tensorflow
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorflow\__init__.py", line 54, in <module>
    from tensorflow._api.v2 import compat
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorflow\_api\v2\compat\__init__.py", line 8, in <module>
    from tensorflow._api.v2.compat import v1
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorflow\_api\v2\compat\v1\__init__.py", line 32, in <module>
    from tensorflow._api.v2.compat.v1 import compat
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorflow\_api\v2\compat\v1\compat\__init__.py", line 8, in <module>
    from tensorflow._api.v2.compat.v1.compat import v1
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorflow\_api\v2\compat\v1\compat\v1\__init__.py", line 49, in <module>
    from tensorflow._api.v2.compat.v1 import lite
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorflow\_api\v2\compat\v1\lite\__init__.py", line 9, in <module>
    from tensorflow._api.v2.compat.v1.lite import experimental
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorflow\_api\v2\compat\v1\lite\experimental\__init__.py", line 8, in <module>
    from tensorflow._api.v2.compat.v1.lite.experimental import authoring
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorflow\_api\v2\compat\v1\lite\experimental\authoring\__init__.py", line 8, in <module>
    from tensorflow.lite.python.authoring.authoring import compatible # line: 265
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorflow\lite\python\authoring\authoring.py", line 43, in <module>
    from tensorflow.lite.python import convert
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorflow\lite\python\convert.py", line 29, in <module>
    from tensorflow.lite.python import util
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\tensorflow\lite\python\util.py", line 51, in <module>
    from jax import xla_computation as _xla_computation
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\jax\__init__.py", line 39, in <module>
    from jax import config as _config_module
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\jax\config.py", line 15, in <module>
    from jax._src.config import config as _deprecated_config  # noqa: F401
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\jax\_src\config.py", line 28, in <module>
    from jax._src import lib
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\jax\_src\lib\__init__.py", line 75, in <module>
    version = check_jaxlib_version(
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\jax\_src\lib\__init__.py", line 67, in check_jaxlib_version
    raise RuntimeError(
RuntimeError: jaxlib version 0.4.28 is newer than and incompatible with jax version 0.4.24. Please update your jax and/or jaxlib packages.

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

Traceback (most recent call last):
  File "C:\sd-webui-aki-v4\launch.py", line 48, in <module>
    main()
  File "C:\sd-webui-aki-v4\launch.py", line 44, in main
    start()
  File "C:\sd-webui-aki-v4\modules\launch_utils.py", line 465, in start
    import webui
  File "C:\sd-webui-aki-v4\webui.py", line 13, in <module>
    initialize.imports()
  File "C:\sd-webui-aki-v4\modules\initialize.py", line 17, in imports
    import pytorch_lightning  # noqa: F401
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\pytorch_lightning\__init__.py", line 35, in <module>
    from pytorch_lightning.callbacks import Callback  # noqa: E402
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\pytorch_lightning\callbacks\__init__.py", line 14, in <module>
    from pytorch_lightning.callbacks.batch_size_finder import BatchSizeFinder
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\pytorch_lightning\callbacks\batch_size_finder.py", line 24, in <module>
    from pytorch_lightning.callbacks.callback import Callback
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\pytorch_lightning\callbacks\callback.py", line 25, in <module>
    from pytorch_lightning.utilities.types import STEP_OUTPUT
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\pytorch_lightning\utilities\types.py", line 27, in <module>
    from torchmetrics import Metric
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\torchmetrics\__init__.py", line 14, in <module>
    from torchmetrics import functional  # noqa: E402
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\torchmetrics\functional\__init__.py", line 77, in <module>
    from torchmetrics.functional.text.bleu import bleu_score
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\torchmetrics\functional\text\__init__.py", line 30, in <module>
    from torchmetrics.functional.text.bert import bert_score  # noqa: F401
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\torchmetrics\functional\text\bert.py", line 24, in <module>
    from torchmetrics.functional.text.helper_embedding_metric import (
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\torchmetrics\functional\text\helper_embedding_metric.py", line 26, in <module>
    from transformers import AutoModelForMaskedLM, AutoTokenizer, PreTrainedModel, PreTrainedTokenizerBase
  File "<frozen importlib._bootstrap>", line 1075, in _handle_fromlist
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\transformers\utils\import_utils.py", line 1076, in __getattr__
    module = self._get_module(self._class_to_module[name])
  File "C:\sd-webui-aki-v4\py310\lib\site-packages\transformers\utils\import_utils.py", line 1088, in _get_module
    raise RuntimeError(
RuntimeError: Failed to import transformers.modeling_utils because of the following error (look up to see its traceback):
jaxlib version 0.4.28 is newer than and incompatible with jax version 0.4.24. Please update your jax and/or jaxlib packages.
[Exited, code 1 (0x00000001)]
------------------------
Fault Traceback: 
Not Available

Additional information

1. Uninstall and reinstall
2. Clear cache files and try reinstalling.
3. I also tried lowering the version.
   I have tried all the above methods and still get an error!

[serick4126]

I encountered the same issue on my PC, and upon investigating the versions of jax and jaxlib in the WebUI's virtual environment, I found a version mismatch. The version of jax was 0.4.26, so I resolved the issue by installing the matching version of jaxlib.

Here is how I fixed it:

1. Open a console in the WebUI directory.
2. Activate the virtual environment with the following command:

.\venv\Scripts\activate

3. Check the versions of jax and jaxlib:

pip show jax jaxlib

4. Ensure that the versions of jax and jaxlib match. If they do not, install the version of jaxlib that matches the version of jax. For example, if the version of jax is 0.4.24, install jaxlib version 0.4.24:

pip install jaxlib==0.4.24

[LucianoDaluz]

I encountered the same issue on my PC, and upon investigating the versions of jax and jaxlib in the WebUI's virtual environment, I found a version mismatch. The version of jax was 0.4.26, so I resolved the issue by installing the matching version of jaxlib.

Here is how I fixed it:

1. Open a console in the WebUI directory.

2. Activate the virtual environment with the following command:

.\venv\Scripts\activate

3. Check the versions of jax and jaxlib:

pip show jax jaxlib

4. Ensure that the versions of jax and jaxlib match. If they do not, install the version of jaxlib that matches the version of jax. For example, if the version of jax is 0.4.24, install jaxlib version 0.4.24:

pip install jaxlib==0.4.24

Thanks a lot

[Maverous]

I recently had this issue too, so before I went crazy, I just deleted the jax folders in the site-packages folder

\venv\Lib\site-packages

Then I added git pull to my webui-user.bat

git pull

Once I ran the bat shortcut, SD reinstalled all required dependencies and I have not had the issue again yet.
There is most likely a better way than deleting folders, but that was the fastest way to remove corrupted files I could think of. Hope this helps