Benchmarking MPS vs MLX encoding speed using SigLIP

[maxlund]

Aggregated totals by architecture and backend:
                    architecture backend  frame_count  duration_s    avg_fps
0           ViT-B-16-SigLIP2-256     MPS         1637       28.21  58.029068
1            ViT-L-16-SigLIP-256     MPS         1637       91.39  17.912244
2            ViT-L-16-SigLIP-384     MPS         1637      226.56   7.225459
3       ViT-SO400M-14-SigLIP-384     MPS         1637      483.48   3.385869
4       siglip-large-patch16-384     MLX         1637      180.88   9.050199
5  siglip-large-patch16-384-4bit     MLX         1637      208.80   7.840038
6      siglip-so400m-patch14-224     MLX         1637      106.37  15.389678
7      siglip-so400m-patch14-384     MLX         1637      341.62   4.791874

Was expecting a more significant speedup, perhaps I'm missing something? I'm taking extracted frames from videos and encoding them with batch size 32, using an MBP M1 Pro 16GB. Here's some snippets of my code:

 def process_batch(batch: torch.Tensor) -> torch.Tensor:
    # batch has shape [B, 3, H, W] on CPU (by default).
    if self.is_mlx_model:
        mx_in = self.mx.array(batch)
        dtype = (
            self.mlx_model.vision_model.vision_model.embeddings.patch_embedding.weight.dtype
        )
        mx_in = mx_in.transpose(0, 2, 3, 1).astype(dtype)
        return torch.from_numpy(np.array(self.mlx_model.get_image_features(pixel_values=mx_in, return_dict=False, output_attentions=False), copy=False))

I also avoid passing the pre-allocated tensors between CPU and MPS when running the MLX pipeline, where otherwise I would allocate them on CPU for the frame extraction and move the whole thing over to MPS when it's done in order to start computing embeddings. Maybe (hopefully) I'm missing something? Will be interesting to see how SigLIP2 performs on MLX regardless, as you can see I included ViT-B-16-SigLIP2-256 and it's by far the fastest.

[Blaizzy]

Thanks @maxlund for the detailed breakdown!

Could you share you full script so I can debug and see what's happening?

[maxlund]

import os
import time
from typing import Generator, Optional

import cv2
import mlx.core as mx
import numpy as np
import torch
from PIL import Image
from mlx_embeddings.utils import load as mlx_load

pretrained_model_dir_or_hf_repo = "path/to/model/or/repo"
path_to_video = "/path/to/dads-child-recordings.mp4"
image_size = 384  # or other resolution of the model
patch_size = 16  # or other patch size of the model
emb_size = 1024  # or emb size of your model
frames_to_process = 640  # or however many frames you want to extract from a video
mlx_model, mlx_processor = mlx_load(
    path_or_hf_repo=pretrained_model_dir_or_hf_repo,
    # note: this is using the changes from my PR here https://github.com/Blaizzy/mlx-embeddings/pull/22
    model_config={
        "image_size": image_size,
        "patch_size": patch_size,
        "num_labels": 0
    }
)


def open_video_with_orientation(video_path: str) -> cv2.VideoCapture | None:
    cap = cv2.VideoCapture(video_path)
    cap.set(cv2.CAP_PROP_ORIENTATION_AUTO, 1)
    return cap


def _validate_and_setup_video(source_path: str, start: int, end: Optional[int]):
    source_path = os.path.abspath(os.path.expanduser(source_path))
    video = open_video_with_orientation(source_path)
    if not video.isOpened():
        raise Exception(f"Could not open video at {source_path}")
    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    if end is not None and end > total_frames:
        raise Exception("_validate_and_setup_video Requested frames are outbound")
    start = max(start, 0)
    end = min(end, total_frames) if end is not None else total_frames
    video.set(cv2.CAP_PROP_POS_FRAMES, start)
    return video, start, end


def get_video_frames_generator(
        source_path: str,
        stride: int = 1,
        start: int = 0,
        end: Optional[int] = None
) -> Generator[np.ndarray, None, None]:
    source_path = os.path.abspath(os.path.expanduser(source_path))
    video, start, end = _validate_and_setup_video(source_path, start, end)
    frame_position = start
    while True:
        success, frame = video.read()
        if not success or frame_position >= end:
            break
        yield frame
        for _ in range(stride - 1):
            success = video.grab()
            if not success:
                break
        frame_position += stride
    video.release()


def preprocess_image(pil_img: Image.Image):
    proc = mlx_processor(images=pil_img, return_tensors="pt")
    tensor = proc["pixel_values"]  # shape [1, 3, H, W]
    return tensor[0]  # shape [3, H, W]


def process_batch(batch: torch.Tensor) -> torch.Tensor:
    # batch has shape [B, 3, H, W] on CPU (by default).
    # MLX path: transpose to [B, H, W, 3], cast dtype
    mx_in = mx.array(batch)
    dtype = (
        mlx_model.vision_model.vision_model.embeddings.patch_embedding.weight.dtype
    )
    mx_in = mx_in.transpose(0, 2, 3, 1).astype(dtype)
    return torch.from_numpy(np.array(
        mlx_model.get_image_features(
            pixel_values=mx_in,
            return_dict=False,
            output_attentions=False
        ),
        copy=False)
    )


print(f"pre-processing {frames_to_process=}")
preprocessed_frames = torch.zeros((frames_to_process, 3, image_size, image_size), device=torch.device('cpu'))
for i, frame in enumerate(get_video_frames_generator(path_to_video)):
    if i >= frames_to_process:
        break
    preprocessed_frames[i] = preprocess_image(Image.fromarray(frame, mode="RGB"))
print("pre-processing done")
features = torch.zeros((frames_to_process, emb_size), device=torch.device("cpu"))

batch_size = 32
print(f"computing embedding for {frames_to_process=} using {batch_size=}")
st = time.time()
for idx in range(0, frames_to_process, batch_size):
    batch_end = min(idx + batch_size, frames_to_process)
    batch_input = preprocessed_frames[idx:batch_end]
    emb = process_batch(batch_input)
    features[idx:batch_end] = emb
print(f"computed embeddings for {frames_to_process=} producing {features.shape=} in {time.time() - st:.2f} seconds")

Here is the video I used for the benchmarks in the plots:
https://drive.google.com/file/d/1tVZE8naya8FR9Je88UstV2i5qFmyGcW2/view?usp=drive_link

As you see in the screenshot it looks like we're fully using the GPU anyway