Add MLX backend support and optimize inference for Apple Silicon

cellpose/__main__.py

@@ -198,7 +198,9 @@ def _evaluate_cellposemodel_cli(args, logger, imf, device, pretrained_model, nor
             ">>>> running cellpose on %d images using all channels" % nimg)
 
     # handle built-in model exceptions
-    model = models.CellposeModel(device=device, pretrained_model=pretrained_model,)
+    use_mlx = getattr(args, 'use_mlx', False)
+    model = models.CellposeModel(device=device, pretrained_model=pretrained_model,
+                                  use_mlx=use_mlx)
 
     tqdm_out = utils.TqdmToLogger(logger, level=logging.INFO)
 

cellpose/cli.py

@@ -30,6 +30,11 @@ def get_arg_parser():
     hardware_args.add_argument(
         "--gpu_device", required=False, default="0", type=str,
         help="which gpu device to use, use an integer for torch, or mps for M1")
+    hardware_args.add_argument(
+        "--use_mlx", nargs="?", const=True, default=False,
+        help="use MLX backend for Apple Silicon acceleration. "
+             "Pass without value to enable, or pass 'auto' to auto-detect "
+             "(requires macOS + Apple Silicon + mlx package)")
 
     # settings for locating and formatting images
     input_img_args = parser.add_argument_group("Input Image Arguments")

cellpose/core.py

@@ -10,12 +10,18 @@
 
 TORCH_ENABLED = True
 
+try:
+    import mlx.core as mx
+    MLX_AVAILABLE = True
+except ImportError:
+    MLX_AVAILABLE = False
+
 core_logger = logging.getLogger(__name__)
 tqdm_out = utils.TqdmToLogger(core_logger, level=logging.INFO)
 
 
 def use_gpu(gpu_number=0, use_torch=True):
-    """ 
+    """
     Check if GPU is available for use.
 
     Args:
@@ -34,6 +40,23 @@ def use_gpu(gpu_number=0, use_torch=True):
         raise ValueError("cellpose only runs with PyTorch now")
 
 
+def use_mlx():
+    """Check if MLX is available and running on Apple Silicon.
+
+    Returns:
+        bool: True if MLX is available.
+    """
+    if not MLX_AVAILABLE:
+        return False
+    try:
+        _ = mx.zeros((1, 1))
+        mx.eval(_)
+        core_logger.info("** MLX available on Apple Silicon. **")
+        return True
+    except Exception:
+        return False
+
+
 def _use_gpu_torch(gpu_number=0):
     """
     Checks if CUDA or MPS is available and working with PyTorch.
@@ -109,6 +132,22 @@ def assign_device(use_torch=True, gpu=False, device=0):
     return device, gpu
 
 
+def _forward_mlx(net, x):
+    """Converts images to MLX arrays, runs the network model, and returns numpy arrays.
+
+    Args:
+        net: The MLX TransformerMLX model.
+        x (numpy.ndarray): The input images of shape (N, C, H, W).
+
+    Returns:
+        Tuple[numpy.ndarray, numpy.ndarray]: The output predictions and style features.
+    """
+    X = mx.array(x)
+    y, style = net(X)
+    mx.eval(y, style)
+    return np.array(y, copy=False), np.array(style, copy=False)
+
+
 def _to_device(x, device, dtype=torch.float32):
     """
     Converts the input tensor or numpy array to the specified device.
@@ -163,10 +202,10 @@ def _forward(net, x):
 
 
 def run_net(net, imgi, batch_size=8, augment=False, tile_overlap=0.1, bsize=224,
-            rsz=None):
-    """ 
+            rsz=None, use_mlx_backend=False):
+    """
     Run network on stack of images.
-    
+
     (faster if augment is False)
 
     Args:
@@ -177,39 +216,42 @@ def run_net(net, imgi, batch_size=8, augment=False, tile_overlap=0.1, bsize=224,
         augment (bool, optional): Tiles image with overlapping tiles and flips overlapped regions to augment. Defaults to False.
         tile_overlap (float, optional): Fraction of overlap of tiles when computing flows. Defaults to 0.1.
         bsize (int, optional): Size of tiles to use in pixels [bsize x bsize]. Defaults to 224.
+        use_mlx_backend (bool, optional): Use MLX backend for inference. Defaults to False.
 
     Returns:
         Tuple[numpy.ndarray, numpy.ndarray]: outputs of network y and style. If tiled `y` is averaged in tile overlaps. Size of [Ly x Lx x 3] or [Lz x Ly x Lx x 3].
-            y[...,0] is Y flow; y[...,1] is X flow; y[...,2] is cell probability. 
+            y[...,0] is Y flow; y[...,1] is X flow; y[...,2] is cell probability.
             style is a 1D array of size 256 summarizing the style of the image, if tiled `style` is averaged over tiles.
     """
+    forward_fn = _forward_mlx if use_mlx_backend else _forward
+
     # run network
-    Lz, Ly0, Lx0, nchan = imgi.shape 
+    Lz, Ly0, Lx0, nchan = imgi.shape
     if rsz is not None:
         if not isinstance(rsz, list) and not isinstance(rsz, np.ndarray):
             rsz = [rsz, rsz]
         Lyr, Lxr = int(Ly0 * rsz[0]), int(Lx0 * rsz[1])
     else:
         Lyr, Lxr = Ly0, Lx0
-    
+
     ly, lx = bsize, bsize
     ypad1, ypad2, xpad1, xpad2 = transforms.get_pad_yx(Lyr, Lxr, min_size=(bsize, bsize))
     Ly, Lx = Lyr + ypad1 + ypad2, Lxr + xpad1 + xpad2
     pads = np.array([[0, 0], [ypad1, ypad2], [xpad1, xpad2]])
-    
+
     if augment:
         ny = max(2, int(np.ceil(2. * Ly / bsize)))
         nx = max(2, int(np.ceil(2. * Lx / bsize)))
     else:
         ny = 1 if Ly <= bsize else int(np.ceil((1. + 2 * tile_overlap) * Ly / bsize))
         nx = 1 if Lx <= bsize else int(np.ceil((1. + 2 * tile_overlap) * Lx / bsize))
-    
-    
+
+
     # run multiple slices at the same time
     ntiles = ny * nx
     nimgs = max(1, batch_size // ntiles) # number of imgs to run in the same batch
     niter = int(np.ceil(Lz / nimgs))
-    ziterator = (trange(niter, file=tqdm_out, mininterval=30) 
+    ziterator = (trange(niter, file=tqdm_out, mininterval=30)
                     if niter > 10 or Lz > 1 else range(niter))
     for k in ziterator:
         inds = np.arange(k * nimgs, min(Lz, (k + 1) * nimgs))
@@ -221,13 +263,13 @@ def run_net(net, imgi, batch_size=8, augment=False, tile_overlap=0.1, bsize=224,
             IMG, ysub, xsub, Lyt, Lxt = transforms.make_tiles(
                 imgb, bsize=bsize, augment=augment,
                 tile_overlap=tile_overlap)
-            IMGa[i * ntiles : (i+1) * ntiles] = np.reshape(IMG, 
+            IMGa[i * ntiles : (i+1) * ntiles] = np.reshape(IMG,
                                             (ny * nx, nchan, ly, lx))
-        
+
         # run network
         for j in range(0, IMGa.shape[0], batch_size):
             bslc = slice(j, min(j + batch_size, IMGa.shape[0]))
-            ya0, stylea0 = _forward(net, IMGa[bslc])
+            ya0, stylea0 = forward_fn(net, IMGa[bslc])
             if j == 0:
                 nout = ya0.shape[1]
                 ya = np.zeros((IMGa.shape[0], nout, ly, lx), "float32")
@@ -252,13 +294,13 @@ def run_net(net, imgi, batch_size=8, augment=False, tile_overlap=0.1, bsize=224,
             # styles[b] = stylei
     # slices from padding
     yf = yf[:, :, ypad1 : Ly-ypad2, xpad1 : Lx-xpad2]
-    yf = yf.transpose(0,2,3,1)   
+    yf = yf.transpose(0,2,3,1)
     return yf, np.array(styles)
 
 
 def run_3D(net, imgs, batch_size=8, augment=False,
            tile_overlap=0.1, bsize=224, net_ortho=None,
-           progress=None):
+           progress=None, use_mlx_backend=False):
     """ 
     Run network on image z-stack.
 
@@ -293,9 +335,9 @@ def run_3D(net, imgs, batch_size=8, augment=False,
         core_logger.info("running %s: %d planes of size (%d, %d)" %
                          (sstr[p], shape[pm[p][0]], shape[pm[p][1]], shape[pm[p][2]]))
         y, style = run_net(net,
-                           xsl, batch_size=batch_size, augment=augment, 
-                           bsize=bsize, tile_overlap=tile_overlap, 
-                           rsz=None)
+                           xsl, batch_size=batch_size, augment=augment,
+                           bsize=bsize, tile_overlap=tile_overlap,
+                           rsz=None, use_mlx_backend=use_mlx_backend)
         yf[..., -1] += y[..., -1].transpose(ipm[p])
         for j in range(2):
             yf[..., cp[p][j]] += y[..., cpy[p][j]].transpose(ipm[p])