Restore partial Windows compatibility by not importing jax unless it's actually needed #131
Description
Was using dalle-flow for months with AmericanPresidentJimmyCarter's Discord bot when a few weeks back it broke after a git pull. I shelved looking at why until now.
dalle-flow checks if flax is available, and if it is, will try to import jax. The problem is it then fails with jaxlib required. jaxlib is not available for Windows.
The solution:
pip uninstall jax
pip uninstall flax
And boom now stable-diffusion is working with dalle-flow again. Not sure which executor actually needs jax/flax, but please add a more sane check to not import jax unless it's actually needed. --enable-stable-diffusion does not need either flax or jax