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