Feat: Add MLX backend and testing the DDPM schedular with MLX

This commit adds the mlx backend for diffusers library, the idea is to make it easy to run Diffusion models
like JAX backend for TPUs.

Author: guptaaryan16

setup.py

@@ -108,6 +108,7 @@
     "isort>=5.5.4",
     "jax>=0.4.1",
     "jaxlib>=0.4.1",
+    "mlx",
     "Jinja2",
     "k-diffusion>=0.0.12",
     "torchsde",
@@ -235,8 +236,13 @@ def run(self):
 else:
     extras["flax"] = deps_list("jax", "jaxlib", "flax")
 
+if sys.platform == "darwin" and os.uname().machine == "arm64":  # Apple Silicon
+    extras["mlx"] = deps_list("mlx")
+else:
+    extras["mlx"] = []
+
 extras["dev"] = (
-    extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"] + extras["flax"]
+    extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"] + extras["flax"] + extras["mlx"]
 )
 
 install_requires = [
@@ -255,7 +261,7 @@ def run(self):
 setup(
     name="diffusers",
     version="0.31.0.dev0",  # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
-    description="State-of-the-art diffusion in PyTorch and JAX.",
+    description="State-of-the-art diffusion in PyTorch, JAX and MLX.",
     long_description=open("README.md", "r", encoding="utf-8").read(),
     long_description_content_type="text/markdown",
     keywords="deep learning diffusion jax pytorch stable diffusion audioldm",

src/diffusers/__init__.py

@@ -7,6 +7,7 @@
     OptionalDependencyNotAvailable,
     _LazyModule,
     is_flax_available,
+    is_mlx_available,
     is_k_diffusion_available,
     is_librosa_available,
     is_note_seq_available,
@@ -35,6 +36,7 @@
     "utils": [
         "OptionalDependencyNotAvailable",
         "is_flax_available",
+        "is_mlx_available",
         "is_inflect_available",
         "is_invisible_watermark_available",
         "is_k_diffusion_available",
@@ -519,6 +521,26 @@
         ]
     )
 
+try:
+    if not is_mlx_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    from .utils import dummy_mlx_objects  # noqa F403
+
+    _import_structure["utils.dummy_mlx_objects"] = [
+        name for name in dir(dummy_mlx_objects) if not name.startswith("_")
+    ]
+
+
+else:
+    _import_structure["schedulers"].extend(
+        [
+            "MLXDDPMScheduler",
+            "MLXEulerDiscreteScheduler",
+        ]
+    )
+
+
 try:
     if not (is_note_seq_available()):
         raise OptionalDependencyNotAvailable()

src/diffusers/schedulers/__init__.py

@@ -20,6 +20,7 @@
     _LazyModule,
     get_objects_from_module,
     is_flax_available,
+    is_mlx_available,
     is_scipy_available,
     is_torch_available,
     is_torchsde_available,
@@ -99,6 +100,19 @@
         "broadcast_to_shape_from_left",
     ]
 
+try:
+    if not is_mlx_available():
+        raise OptionalDependencyNotAvailable()
+
+except OptionalDependencyNotAvailable:
+    from ..utils import dummy_mlx_objects  # noqa F403
+
+    _dummy_modules.update(get_objects_from_module(dummy_mlx_objects))
+
+else:
+    _import_structure["scheduling_euler_discrete_mlx"] = ["MLXEulerDiscreteScheduler"]
+    _import_structure["scheduling_ddpm_mlx"] = ["MLXDDPMScheduler"]
+
 
 try:
     if not (is_torch_available() and is_scipy_available()):
@@ -127,6 +141,7 @@
     from ..utils import (
         OptionalDependencyNotAvailable,
         is_flax_available,
+        is_mlx_available,
         is_scipy_available,
         is_torch_available,
         is_torchsde_available,
@@ -197,6 +212,20 @@
         )
 
     try:
+        if not is_mlx_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ..utils.dummy_mlx_objects import *  # noqa F403
+    else:
+        from .scheduling_euler_discrete_mlx import MLXEulerDiscreteScheduler
+        from .scheduling_ddpm_mlx import MLXDDPMScheduler
+        from .scheduling_utils_mlx import (
+            MLXKarrasDiffusionSchedulers,
+            MLXSchedulerMixin,
+            MLXSchedulerOutput
+        )
+
+    try:        
         if not (is_torch_available() and is_scipy_available()):
             raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable: