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.