remove flax from setup.py

Author: sayakpaul

setup.py

@@ -84,7 +84,6 @@
     you need to go back to main before executing this.
 """
 
-import os
 import re
 import sys
 
@@ -100,7 +99,6 @@
     "compel==0.1.8",
     "datasets",
     "filelock",
-    "flax>=0.4.1",
     "hf-doc-builder>=0.3.0",
     "huggingface-hub>=0.34.0",
     "requests-mock==1.10.0",
@@ -245,14 +243,8 @@ 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"] + extras["flax"]
-)
+extras["dev"] = extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"]
 
 install_requires = [
     deps["importlib_metadata"],

src/diffusers/dependency_versions_table.py

@@ -7,7 +7,6 @@
     "compel": "compel==0.1.8",
     "datasets": "datasets",
     "filelock": "filelock",
-    "flax": "flax>=0.4.1",
     "hf-doc-builder": "hf-doc-builder>=0.3.0",
     "huggingface-hub": "huggingface-hub>=0.34.0",
     "requests-mock": "requests-mock==1.10.0",