Merge branch 'main' into torch-main-dep

setup.py

@@ -100,7 +100,6 @@ _deps = [
     "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",
@@ -137,6 +136,7 @@ _deps = [
     "requests",
     "tensorboard",
     "tiktoken>=0.7.0",
+    "flax>=0.4.1",
     "torch>=1.4",
     "torchvision",
     "transformers>=4.41.2",
@@ -252,6 +252,7 @@ if os.name == "nt":  # windows
 else:
     extras["flax"] = deps_list("jax", "jaxlib", "flax")
 
+
 extras["dev"] = (
     extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"] + extras["flax"]
 )
@@ -265,6 +266,7 @@ install_requires = [
     deps["requests"],
     deps["safetensors"],
     deps["Pillow"],
+    deps["torch"],
 ]
 
 version_range_max = max(sys.version_info[1], 10) + 1

src/diffusers/dependency_versions_table.py

@@ -7,7 +7,6 @@ deps = {
     "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",
@@ -44,6 +43,7 @@ deps = {
     "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",