search

hako-mikan / sd-webui-traintrain

LoRA training extention for Stable Diffusion Web-UI
GNU Affero General Public License v3.0
142 stars 5 forks source link

Fix error: AttributeError: module 'jax.random' has no attribute 'KeyArray' (Update diffusers version) #15

Closed shinshin86 closed 5 months ago

shinshin86 commented 5 months ago

I encountered the following error when trying to install and use the Stable Diffusion web UI in a Google Colab environment: (The error occurred when I attempted to launch the Stable Diffusion web UI after installing TrainTrain.)

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

Here are the more detailed error logs:

*** Error loading script: traintrain.py
    Traceback (most recent call last):
      File "/content/stable-diffusion-webui/modules/scripts.py", line 527, in load_scripts
        script_module = script_loading.load_module(scriptfile.path)
      File "/content/stable-diffusion-webui/modules/script_loading.py", line 10, in load_module
        module_spec.loader.exec_module(module)
      File "<frozen importlib._bootstrap_external>", line 883, in exec_module
      File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
      File "/content/stable-diffusion-webui/extensions/sd-webui-traintrain/scripts/traintrain.py", line 12, in <module>
        from trainer import train, trainer, gen
      File "/content/stable-diffusion-webui/extensions/sd-webui-traintrain/trainer/train.py", line 14, in <module>
        from trainer import trainer, dataset
      File "/content/stable-diffusion-webui/extensions/sd-webui-traintrain/trainer/trainer.py", line 6, in <module>
        from diffusers import StableDiffusionPipeline, DDPMScheduler, StableDiffusionXLPipeline
      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'

I believe this issue originates from the dependencies in diffusers. Upgrading the dependent version to 0.25.0 resolves this problem, as I have confirmed that LoRA can be successfully created after installing TrainTrain on the Stable Diffusion web UI launched on Google Colab.

If you find this acceptable, could you please merge this PR? I would appreciate your consideration.

Thank you very much.

hako-mikan commented 5 months ago

Thanks!