onnx-web-worker-starter

onnx-web-worker-starter is an efficient starter kit designed for running ONNX models entirely on-device using a Web Worker, optimized for seamless integration with Next.js applications. This project aims to enhance the performance of AI model inference in web applications by alleviating the main UI thread from computational loads, ensuring that the user experience remains smooth and responsive, especially on mobile devices.

Key Features

• Minimal Setup: The project includes essential components such as

  • public/onnx-worker/onnxWorker.js: The classic worker script responsible for loading the ONNX Runtime.
  • src/lib/OnnxWorkerClient.js: A promise-based wrapper that facilitates communication via postMessage.
  • src/hooks/useOnnxWorker.js: An optional React hook for simplified usage within React applications.
  • example/DemoComponent.jsx: A straightforward example demonstrating how to utilize the library.
  • test/OnnxWorkerClient.test.js: Protocol tests to ensure robust client-to-worker interaction.

• Responsive UI: Transferring the model loading and inference process to a worker minimizes disruption to the main UI thread, offering an enhanced user experience without lag or stutter.

• Flexible Model Deployment: Users can bring their own models and inputs as needed, making this starter kit adaptable to various AI projects.

Usage Examples

React Example

Integrate the ONNX worker in a React component with ease:

"use client";

import { useOnnxWorker } from "./src/hooks/useOnnxWorker.js";

export function MyInferenceButton() {
  const { status, progress, error, run } = useOnnxWorker({
    modelUrl: "/models/my-model.onnx",
    workerUrl: "/onnx-worker/onnxWorker.js",
    runtimeUrl: "/onnx-worker/ort.min.js",
    wasmPaths: "/onnx-worker/",
  });

  async function handleClick() {
    const outputs = await run({
      input: {
        type: "float32",
        dims: [1, 3, 224, 224],
        data: new Float32Array(1 * 3 * 224 * 224),
      },
    });

    console.log(outputs);
  }

  return (
    <button onClick={handleClick} disabled={status !== "ready"}>
      {status === "running" ? "Running..." : "Run inference"}
      {progress ? ` ${progress.percent}%` : ""}
      {error ? ` ${error}` : ""}
    </button>
  );
}

Plain JavaScript Example

For non-React projects, use the following approach:

import { OnnxWorkerClient } from "./src/lib/OnnxWorkerClient.js";

const client = new OnnxWorkerClient({
  workerUrl: "/onnx-worker/onnxWorker.js",
  onProgress: (progress) => console.log(progress),
});

await client.preload({
  modelUrl: "/models/my-model.onnx",
  runtimeUrl: "/onnx-worker/ort.min.js",
  wasmPaths: "/onnx-worker/",
});

const outputs = await client.run({
  modelUrl: "/models/my-model.onnx",
  runtimeUrl: "/onnx-worker/ort.min.js",
  wasmPaths: "/onnx-worker/",
  inputs: {
    input: {
      type: "float32",
      dims: [1, 4],
      data: new Float32Array([1, 2, 3, 4]),
    },
  },
});

Additional Information

The preload() method initializes both the ONNX Runtime and the model prior to the first inference run, while the run() method processes tensor-like inputs in the worker and retrieves serialized output tensors. The worker efficiently manages requests and allows for a cached session for performance optimization.

Note that while leveraging a worker keeps inference off the UI thread, it does not automatically mitigate the computational demands of complex models, especially when dealing with high-resolution images. Hence, prior image processing such as cropping and resizing is advised to optimize performance.

This starter pattern was developed from production-level code implemented in Lumli, a privacy-focused AI studio that runs models locally on user devices.