Dreambooth in flax?

[droidXrobot]

Describe the bug

I’ve tried just about everything to get dreambooth running with flax but get errors each time. Has anyone got it working with flax?

Reproduction

I’ve tried an A100 and TPU in Colab and followed the directions in the readme

Logs

No response

System Info

Colab, using requirements in your dreambooth readme