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

[cseti007]

Thanks for providing the example notebook, but when I try to run it with default settings I get the following error while running the training cell. Can you please help what the problem is?

/content/animatediff
HEAD is now at 4735142 FreeU added
From https://github.com/tumurzakov/AnimateDiff

• branch main -> FETCH_HEAD
  Already up to date.
  configs/bunny.yaml
  The following values were not passed to accelerate launch and had defaults used instead:
  --num_processes was set to a value of 1
--num_machines was set to a value of 1
--mixed_precision was set to a value of 'no'
--dynamo_backend was set to a value of 'no'
  To avoid this warning pass in values for each of the problematic parameters or run accelerate config.
  Traceback (most recent call last):
  File "/content/animatediff/train.py", line 15, in
  import diffusers
  File "/usr/local/lib/python3.10/dist-packages/diffusers/init.py", line 46, in
  from .pipeline_utils import DiffusionPipeline
  File "/usr/local/lib/python3.10/dist-packages/diffusers/pipeline_utils.py", line 38, in
  from .schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME
  File "/usr/local/lib/python3.10/dist-packages/diffusers/schedulers/init.py", line 51, in
  from .scheduling_ddpm_flax import FlaxDDPMScheduler
  File "/usr/local/lib/python3.10/dist-packages/diffusers/schedulers/scheduling_ddpm_flax.py", line 80, in
  class FlaxDDPMScheduler(FlaxSchedulerMixin, ConfigMixin):
  File "/usr/local/lib/python3.10/dist-packages/diffusers/schedulers/scheduling_ddpm_flax.py", line 216, in FlaxDDPMScheduler
  key: random.KeyArray,
  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 48, in main
  args.func(args)
  File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/launch.py", line 1106, in launch_command
  simple_launcher(args)
  File "/usr/local/lib/python3.10/dist-packages/accelerate/commands/launch.py", line 704, in simple_launcher
  raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
  subprocess.CalledProcessError: Command '['/usr/bin/python3', 'train.py', '--config=configs/bunny.yaml']' returned non-zero exit status 1.