I'm can't start training.

[Bonjonbon]

I'm using the default model of 1.5 and using external captions for training. Here is what shows up after clicking run. Also I've tried different models too.

`Training` the UNet...
Traceback (most recent call last):
  File "/content/diffusers/examples/dreambooth/train_dreambooth.py", line 18, in <module>
    from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
  File "/usr/local/lib/python3.10/dist-packages/diffusers/__init__.py", line 27, in <module>
    from .pipelines import OnnxRuntimeModel
  File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/__init__.py", line 18, in <module>
    from .dance_diffusion import DanceDiffusionPipeline
  File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/dance_diffusion/__init__.py", line 1, in <module>
    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 <module>
    from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline
  File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/pipeline_utils.py", line 39, in <module>
    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 <module>
    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 <module>
    from .modeling_flax_utils import FlaxModelMixin
  File "/usr/local/lib/python3.10/dist-packages/diffusers/models/modeling_flax_utils.py", line 45, in <module>
    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 <module>
    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', '--external_captions', '--image_captions_filename', '--train_only_unet', '--save_starting_step=500', '--save_n_steps=0', '--Session_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/Nevermore_Beta_2', '--pretrained_model_name_or_path=/content/stable-diffusion-v1-5', '--instance_data_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/Nevermore_Beta_2/instance_images', '--output_dir=/content/models/Nevermore_Beta_2', '--captions_dir=/content/gdrive/MyDrive/Fast-Dreambooth/Sessions/Nevermore_Beta_2/captions', '--instance_prompt=', '--seed=415965', '--resolution=512', '--mixed_precision=fp16', '--train_batch_size=1', '--gradient_accumulation_steps=1', '--use_8bit_adam', '--learning_rate=7e-06', '--lr_scheduler=linear', '--lr_warmup_steps=0', '--max_train_steps=3750']' returned non-zero exit status 1.
Something went wrong

[TallisHe]

I meet the same problem, it doesn't work in colab and in the runpod, I can not find the code of the author, any one can help ? TT

[nohholmez]

Exact same error for me, hopefully there is a fix for it! Connected to a T4 machine if that's helpful info.

[TheLastBen]

fixed, use the latest notebook

[Bonjonbon]

Thank you! Seriously, if you don't get thanks enough for what you do and provide, then thank you again!