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

[Vect-r]

I am getting this error while trying v.3.0.1
is this version outdated or am i doing something wrong.

`2024-10-09 11:17:06.599485: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-10-09 11:17:06.619755: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-10-09 11:17:06.626391: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-10-09 11:17:07.858085: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Traceback (most recent call last):
File "/content/InvokeAI/scripts/invokeai-model-install.py", line 1, in
from invokeai.frontend.install.model_install import main
File "/content/InvokeAI/invokeai/frontend/install/init.py", line 4, in
from .invokeai_configure import main as invokeai_configure
File "/content/InvokeAI/invokeai/frontend/install/invokeai_configure.py", line 4, in
from ...backend.install.invokeai_configure import main
File "/content/InvokeAI/invokeai/backend/init.py", line 4, in
from .generator import InvokeAIGeneratorBasicParams, InvokeAIGenerator, InvokeAIGeneratorOutput, Img2Img, Inpaint
File "/content/InvokeAI/invokeai/backend/generator/init.py", line 4, in
from .base import (
File "/content/InvokeAI/invokeai/backend/generator/base.py", line 9, in
import diffusers
File "/usr/local/lib/python3.10/dist-packages/diffusers/init.py", line 38, in
from .models import (
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/init.py", line 35, in
from .controlnet_flax import FlaxControlNetModel
File "/usr/local/lib/python3.10/dist-packages/diffusers/models/controlnet_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 194, in FlaxModelMixin
def init_weights(self, rng: jax.random.KeyArray) -> Dict:
File "/usr/local/lib/python3.10/dist-packages/jax/_src/deprecations.py", line 55, in getattr
raise AttributeError(f"module {module!r} has no attribute {name!r}")
AttributeError: module 'jax.random' has no attribute 'KeyArray'

CalledProcessError Traceback (most recent call last)
in <cell line: 26>()
24 ssh_name = "id_rsa"
25 ssh_path = Path(os.path.dirname(os.getcwd())) / ssh_name
---> 26 gen_key(ssh_path)
27
28 import threading

1 frames
/usr/lib/python3.10/subprocess.py in run(input, capture_output, timeout, check, *popenargs, **kwargs)
524 retcode = process.poll()
525 if check and retcode:
--> 526 raise CalledProcessError(retcode, process.args,
527 output=stdout, stderr=stderr)
528 return CompletedProcess(process.args, retcode, stdout, stderr)

CalledProcessError: Command '['ssh-keygen', '-t', 'rsa', '-b', '4096', '-N', '', '-q', '-f', '/content/id_rsa']' returned non-zero exit status 1.`