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

[YanxinLu]

I encountered the error while training the UNet.

Training the UNet...
Traceback (most recent call last):
File "/content/diffusers/examples/dreambooth/train_dreambooth.py", line 18, in
from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
File "/usr/local/lib/python3.10/dist-packages/diffusers/init.py", line 27, in
from .pipelines import OnnxRuntimeModel
File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/init.py", line 18, in
from .dance_diffusion import DanceDiffusionPipeline
File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/dance_diffusion/init.py", line 1, in
from .pipeline_dance_diffusion import DanceDiffusionPipeline
File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py", line 21, in
from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline
File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/pipeline_utils.py", line 39, in
from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/init.py", line 31, in
from .unet_2d_condition_flax import FlaxUNet2DConditionModel
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/unet_2d_condition_flax.py", line 25, in
from .modeling_flax_utils import FlaxModelMixin
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/modeling_flax_utils.py", line 45, in
class FlaxModelMixin:
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/modeling_flax_utils.py", line 192, 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'
Traceback (most recent call last):
File "/usr/local/bin/accelerate", line 8, in
sys.exit(main())
File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/accelerate_cli.py", line 43, in main
args.func(args)
File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/launch.py", line 837, in launch_command
simple_launcher(args)
File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/launch.py", line 354, in simple_launcher
raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
subprocess.CalledProcessError: Command '['/usr/bin/python3', '/content/diffusers/examples/dreambooth/train_dreambooth.py', '--image_captions_filename', '--train_only_unet', '--save_starting_step=500', '--save_n_steps=500', '--Session_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/t', '--pretrained_model_name_or_path=/content/stable-diffusion-custom', '--instance_data_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/t/instance_images', '--output_dir=/content/models/t', '--captions_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/t/captions', '--instance_prompt=', '--seed=805484', '--resolution=512', '--mixed_precision=fp16', '--train_batch_size=1', '--gradient_accumulation_steps=1', '--use_8bit_adam', '--learning_rate=2e-06', '--lr_scheduler=linear', '--lr_warmup_steps=0', '--max_train_steps=1500']' returned non-zero exit status 1.
Something went wrong

[Lime-tones]

Getting similar error, while loading session.

1.5 Detected
Session found, loading the trained model ...
Traceback (most recent call last):
File "/content/diffusers/scripts/convert_original_stable_diffusion_to_diffusers.py", line 19, in
from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt
File "/usr/local/lib/python3.10/dist-packages/diffusers/init.py", line 27, in
from .pipelines import OnnxRuntimeModel
File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/init.py", line 18, in
from .dance_diffusion import DanceDiffusionPipeline
File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/dance_diffusion/init.py", line 1, in
from .pipeline_dance_diffusion import DanceDiffusionPipeline
File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py", line 21, in
from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline
File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/pipeline_utils.py", line 39, in
from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/init.py", line 31, in
from .unet_2d_condition_flax import FlaxUNet2DConditionModel
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/unet_2d_condition_flax.py", line 25, in
from .modeling_flax_utils import FlaxModelMixin
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/modeling_flax_utils.py", line 45, in
class FlaxModelMixin:
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/modeling_flax_utils.py", line 192, 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'
Conversion error, if the error persists, remove the CKPT file from the current session folder

[martindubo]

Same problem, did anyone solve it?

[jeblbej]

Same problem here, problem with jax

[Geekphysique]

Also seeing the same. AttributeError: module 'jax.random' has no attribute 'KeyArray'

[vajayattila]

Same problem here.

[BlenderProxy]

Same problem here! :(

Training the UNet...
Traceback (most recent call last):
File "/content/diffusers/examples/dreambooth/train_dreambooth.py", line 18, in
from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
File "/usr/local/lib/python3.10/dist-packages/diffusers/init.py", line 27, in
from .pipelines import OnnxRuntimeModel
File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/init.py", line 18, in
from .dance_diffusion import DanceDiffusionPipeline
File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/dance_diffusion/init.py", line 1, in
from .pipeline_dance_diffusion import DanceDiffusionPipeline
File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py", line 21, in
from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline
File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/pipeline_utils.py", line 39, in
from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/init.py", line 31, in
from .unet_2d_condition_flax import FlaxUNet2DConditionModel
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/unet_2d_condition_flax.py", line 25, in
from .modeling_flax_utils import FlaxModelMixin
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/modeling_flax_utils.py", line 45, in
class FlaxModelMixin:
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/modeling_flax_utils.py", line 192, 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'
Traceback (most recent call last):
File "/usr/local/bin/accelerate", line 8, in
sys.exit(main())
File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/accelerate_cli.py", line 43, in main
args.func(args)
File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/launch.py", line 837, in launch_command
simple_launcher(args)
File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/launch.py", line 354, in simple_launcher
raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
subprocess.CalledProcessError: Command '['/usr/bin/python3', '/content/diffusers/examples/dreambooth/train_dreambooth.py', '--image_captions_filename', '--train_only_unet', '--save_starting_step=500', '--save_n_steps=0', '--Session_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/fkvnd-pro2', '--pretrained_model_name_or_path=/content/stable-diffusion-v2-512', '--instance_data_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/fkvnd-pro2/instance_images', '--output_dir=/content/models/fkvnd-pro2', '--captions_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/fkvnd-pro2/captions', '--instance_prompt=', '--seed=71560', '--resolution=512', '--mixed_precision=fp16', '--train_batch_size=1', '--gradient_accumulation_steps=1', '--use_8bit_adam', '--learning_rate=2e-06', '--lr_scheduler=linear', '--lr_warmup_steps=0', '--max_train_steps=1500']' returned non-zero exit status 1.
Something went wrong

[unexpectedly]

Ditto

[DjNastyMagic]

jax.random.KeyArray was removed in JAX v0.4.24. current diffusers only work with JAX v0.4.23 or earlier.

add under requirements:

!pip install "jax[cuda12_pip]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

[Geekphysique]

jax.random.KeyArray was removed in JAX v0.4.24. current diffusers only work with JAX v0.4.23 or earlier.

add under requirements:

!pip install "jax[cuda12_pip]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Worked like a charm. Thank you

[nohholmez]

Can someone please help me with where specifically I need to put !pip install "jax[cuda12_pip]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html I am very new to all this, muddled my way through the colab workbook hit the error and came here. Thank you in advance!

[BosmansR]

Can someone please help me with where specifically I need to put !pip install "jax[cuda12_pip]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html I am very new to all this, muddled my way through the colab workbook hit the error and came here. Thank you in advance!

You have to click "show code" in the dependencies cell.
Then you just paste it under the "!pip install gradio..."

That made it work for me :)

[nohholmez]

Can someone please help me with where specifically I need to put !pip install "jax[cuda12_pip]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html I am very new to all this, muddled my way through the colab workbook hit the error and came here. Thank you in advance!

You have to click "show code" in the dependencies cell. Then you just paste it under the "!pip install gradio..."

That made it work for me :)

Legend thank you will try it now :D Edit - works a treat!

[Talionic]

Yeah, just started having this problem. DjNastyMagic fix worked.

[jeblbej]

I applied fix by DJNastyMagic.
But have now this error:

Training the UNet...
Traceback (most recent call last):
File "/content/diffusers/examples/dreambooth/train_dreambooth.py", line 803, in
main()
File "/content/diffusers/examples/dreambooth/train_dreambooth.py", line 512, in main
vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae")
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/modeling_utils.py", line 558, in from_pretrained
raise ValueError(
ValueError: Cannot load <class 'diffusers.models.autoencoder_kl.AutoencoderKL'> from /content/stable-diffusion-custom because the following keys are missing:
decoder.mid_block.attentions.0.key.bias, encoder.mid_block.attentions.0.value.weight, decoder.mid_block.attentions.0.proj_attn.bias, decoder.mid_block.attentions.0.value.bias, decoder.mid_block.attentions.0.query.bias, decoder.mid_block.attentions.0.query.weight, encoder.mid_block.attentions.0.proj_attn.bias, encoder.mid_block.attentions.0.query.weight, encoder.mid_block.attentions.0.key.weight, encoder.mid_block.attentions.0.value.bias, decoder.mid_block.attentions.0.proj_attn.weight, encoder.mid_block.attentions.0.query.bias, decoder.mid_block.attentions.0.value.weight, encoder.mid_block.attentions.0.proj_attn.weight, encoder.mid_block.attentions.0.key.bias, decoder.mid_block.attentions.0.key.weight.
Please make sure to pass low_cpu_mem_usage=False and device_map=None if you want to randomely initialize those weights or else make sure your checkpoint file is correct.
Traceback (most recent call last):
File "/usr/local/bin/accelerate", line 8, in
sys.exit(main())
File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/accelerate_cli.py", line 43, in main
args.func(args)
File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/launch.py", line 837, in launch_command
simple_launcher(args)
File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/launch.py", line 354, in simple_launcher
raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
subprocess.CalledProcessError: Command '['/usr/bin/python3', '/content/diffusers/examples/dreambooth/train_dreambooth.py', '--image_captions_filename', '--train_only_unet', '--save_starting_step=500', '--save_n_steps=0', '--Session_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/Mysession', '--pretrained_model_name_or_path=/content/stable-diffusion-custom', '--instance_data_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/Mysession/instance_images', '--output_dir=/content/models/Mysession', '--captions_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/Mysession/captions', '--instance_prompt=', '--seed=601321', '--resolution=512', '--mixed_precision=fp16', '--train_batch_size=1', '--gradient_accumulation_steps=1', '--use_8bit_adam', '--learning_rate=2e-06', '--lr_scheduler=linear', '--lr_warmup_steps=0', '--max_train_steps=1800']' returned non-zero exit status 1.
Something went wrong

[LIQUIDMIND111]

ERROR INSTALLING CUDA

aceback (most recent call last):
File "/content/diffusers/examples/dreambooth/train_dreambooth.py", line 803, in
main()
File "/content/diffusers/examples/dreambooth/train_dreambooth.py", line 535, in main
import bitsandbytes as bnb
File "/usr/local/lib/python3.10/dist-packages/bitsandbytes/init.py", line 6, in
from .autograd._functions import (
File "/usr/local/lib/python3.10/dist-packages/bitsandbytes/autograd/_functions.py", line 5, in
import bitsandbytes.functional as F
File "/usr/local/lib/python3.10/dist-packages/bitsandbytes/functional.py", line 13, in
from .cextension import COMPILED_WITH_CUDA, lib
File "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cextension.py", line 41, in
lib = CUDALibrary_Singleton.get_instance().lib
File "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cextension.py", line 37, in get_instance
cls._instance.initialize()
File "/usr/local/lib/python3.10/dist-packages/bitsandbytes/cextension.py", line 27, in initialize
raise Exception('CUDA SETUP: Setup Failed!')
Exception: CUDA SETUP: Setup Failed!
Traceback (most recent call last):
File "/usr/local/bin/accelerate", line 8, in
sys.exit(main())
File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/accelerate_cli.py", line 43, in main
args.func(args)
File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/launch.py", line 837, in launch_command
simple_launcher(args)
File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/launch.py", line 354, in simple_launcher
raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
subprocess.CalledProcessError: Command '['/usr/bin/python3', '/content/diffusers/examples/dreambooth/train_dreambooth.py', '--image_captions_filename', '--train_only_unet', '--save_starting_step=500', '--save_n_steps=0', '--Session_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/bEBA', '--pretrained_model_name_or_path=/content/stable-diffusion-v1-5', '--instance_data_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/bEBA/instance_images', '--output_dir=/content/models/bEBA', '--captions_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/bEBA/captions', '--instance_prompt=', '--seed=247601', '--resolution=768', '--mixed_precision=fp16', '--train_batch_size=1', '--gradient_accumulation_steps=1', '--use_8bit_adam', '--learning_rate=2e-06', '--lr_scheduler=linear', '--lr_warmup_steps=0', '--max_train_steps=1500']' returned non-zero exit status 1.
Something went wrong