add flax back.

Author: sayakpaul

setup.py

@@ -84,6 +84,7 @@
     you need to go back to main before executing this.
 """
 
+import os
 import re
 import sys
 
@@ -134,6 +135,7 @@
     "requests",
     "tensorboard",
     "tiktoken>=0.7.0",
+    "flax>=0.4.1",
     "torch>=1.4",
     "torchvision",
     "transformers>=4.41.2",
@@ -243,6 +245,11 @@ def run(self):
 extras["optimum_quanto"] = deps_list("optimum_quanto", "accelerate")
 extras["torchao"] = deps_list("torchao", "accelerate")
 
+if os.name == "nt":  # windows
+    extras["flax"] = []  # jax is not supported on windows
+else:
+    extras["flax"] = deps_list("jax", "jaxlib", "flax")
+
 
 extras["dev"] = extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"]
 

src/diffusers/dependency_versions_table.py

@@ -42,6 +42,7 @@
     "requests": "requests",
     "tensorboard": "tensorboard",
     "tiktoken": "tiktoken>=0.7.0",
+    "flax": "flax>=0.4.1",
     "torch": "torch>=1.4",
     "torchvision": "torchvision",
     "transformers": "transformers>=4.41.2",