@sconedev/ai_toolkit
Version:
Simplify AI integration in web apps with local and offline model support
370 lines (369 loc) • 15.1 kB
JavaScript
import { getConfig } from './config';
import { getCached, setCached } from './cache';
import { OnnxModelError, OnnxRuntimeError, runOnnxInference, validateInputData } from "./onnxRuntime";
// Dynamic import for onnxruntime-web only when in browser
let ortPromise = null;
/**
* Load ONNX Runtime Web asynchronously
*/
async function loadOrt() {
if (!ortPromise) {
ortPromise = import('onnxruntime-web').catch(err => {
console.error('Failed to load onnxruntime-web:', err);
throw new OnnxRuntimeError('ONNX Runtime Web could not be loaded. Make sure it is properly installed.');
});
}
return ortPromise;
}
/**
* Run inference using local model file in browser
*
* @param modelFile File from file input or drop
* @param modelConfig Model configuration
* @param inputData Input data for inference
* @param options Additional options
*/
export async function runModelFromFile(modelFile, modelConfig, inputData, options) {
if (!modelFile) {
throw new OnnxModelError('No model file provided');
}
try {
// Validate input data
validateInputData(inputData, modelConfig);
const config = getConfig();
const modelName = modelConfig.name || modelFile.name.replace(/\.[^/.]+$/, "");
// Create cache keys with file name and last modified date for uniqueness
const fileId = `${modelFile.name}-${modelFile.lastModified}`;
const resultCacheKey = `onnx:result:${modelName}:${fileId}:${JSON.stringify(inputData)}`;
const modelCacheKey = `onnx:model:${modelName}:${fileId}`;
if (options?.progressCallback) {
options.progressCallback(5); // Starting
}
// Check for cached results
if (options?.cacheResults !== false) {
const cached = getCached(resultCacheKey);
if (cached) {
if (options?.progressCallback) {
options.progressCallback(100);
}
return cached;
}
}
// Get or create session
let session = getCached(modelCacheKey);
if (!session) {
try {
const ortLib = await loadOrt();
if (options?.progressCallback) {
options.progressCallback(20);
}
// Convert to ArrayBuffer
const arrayBuffer = await modelFile.arrayBuffer();
if (options?.progressCallback) {
options.progressCallback(40);
}
// Log request if configured
if (config.onRequest) {
config.onRequest({
provider: 'onnx',
function: 'runModelFromFile',
input: { model: modelName, fileSize: modelFile.size }
});
}
// Combine default options with user-provided options
const inferenceOptions = {
executionProviders: ['webgpu', 'wasm'],
graphOptimizationLevel: 'all',
...(options?.inferenceOptions || {})
};
// Create session
session = await ortLib.InferenceSession.create(arrayBuffer, inferenceOptions);
if (options?.progressCallback) {
options.progressCallback(60);
}
// Cache if requested
if (options?.cacheModel !== false) {
setCached(modelCacheKey, session);
}
}
catch (error) {
throw new OnnxModelError(`Failed to load model from file: ${error instanceof Error ? error.message : String(error)}`);
}
}
else if (options?.progressCallback) {
options.progressCallback(50); // Using cached model
}
// Run inference
const ortLib = await loadOrt();
const feeds = {};
// Prepare inputs
for (const inputName of modelConfig.inputNames) {
if (!inputData[inputName]) {
throw new OnnxModelError(`Missing required input '${inputName}'`);
}
const data = inputData[inputName];
// Handle different types of input data
if (Array.isArray(data)) {
feeds[inputName] = new ortLib.Tensor('float32', new Float32Array(data.flat()), data.length > 0 ?
(Array.isArray(data[0]) ? [data.length, data[0].length] : [data.length]) : [0]);
}
else if (data instanceof Float32Array) {
feeds[inputName] = new ortLib.Tensor('float32', data, [data.length]);
}
else if (data instanceof Uint8Array && inputName.toLowerCase().includes('image')) {
const dims = inputData[`${inputName}_dims`] || [1, 3, 224, 224];
feeds[inputName] = new ortLib.Tensor('uint8', data, dims);
}
else if (data instanceof Int32Array) {
feeds[inputName] = new ortLib.Tensor('int32', data, [data.length]);
}
else if (data instanceof Uint8Array) {
const dims = inputData[`${inputName}_dims`] || [data.length];
feeds[inputName] = new ortLib.Tensor('uint8', data, dims);
}
else {
throw new OnnxModelError(`Unsupported data type for input '${inputName}'`);
}
}
if (options?.progressCallback) {
options.progressCallback(75); // Inputs prepared
}
// Run the model
const results = await session.run(feeds);
if (options?.progressCallback) {
options.progressCallback(90); // Inference completed
}
// Process outputs
const outputs = {};
for (const outputName of modelConfig.outputNames) {
if (results[outputName]) {
outputs[outputName] = results[outputName].data;
}
}
// Cache results if enabled
if (options?.cacheResults !== false) {
setCached(resultCacheKey, outputs);
}
if (options?.progressCallback) {
options.progressCallback(100); // All done
}
return outputs;
}
catch (error) {
if (error instanceof OnnxModelError) {
throw error;
}
throw new OnnxRuntimeError(`Error running model from file: ${error instanceof Error ? error.message : String(error)}`);
}
}
/**
* Load a model from URL with caching support
* Useful for offline-first applications
*/
export async function runOnnxInferenceBrowser(modelConfig, inputData, options) {
// This is just a browser-specific wrapper for runOnnxInference
// Validate and prepare inputs for the browser environment
validateInputData(inputData, modelConfig);
// Process browser-specific logic then call common implementation
// Note: You may need to adapt this based on your actual implementation
return runOnnxInference(modelConfig, inputData, options);
}
export async function loadModelFromUrl(url, options) {
try {
// Check if the model exists in the cache API
if (options?.checkCache && 'caches' in window) {
try {
const cache = await caches.open('onnx-models');
const cachedResponse = await cache.match(url);
if (cachedResponse) {
if (options?.progressCallback) {
options.progressCallback(50);
}
const buffer = await cachedResponse.arrayBuffer();
if (options?.progressCallback) {
options.progressCallback(100);
}
return buffer;
}
}
catch (err) {
console.warn('Failed to check cache:', err);
}
}
// Check IndexedDB
if (options?.checkCache) {
try {
const dbName = 'onnx_models_db';
const storeName = 'model_urls';
const urlHash = await hashString(url);
const modelFromDB = await new Promise((resolve, reject) => {
const request = indexedDB.open(dbName, 1);
request.onupgradeneeded = (event) => {
const db = event.target.result;
if (!db.objectStoreNames.contains(storeName)) {
db.createObjectStore(storeName);
}
};
request.onsuccess = (event) => {
const db = event.target.result;
try {
const transaction = db.transaction(storeName, 'readonly');
const store = transaction.objectStore(storeName);
const getRequest = store.get(urlHash);
getRequest.onsuccess = () => {
resolve(getRequest.result || null);
};
getRequest.onerror = () => {
resolve(null);
};
}
catch (err) {
resolve(null);
}
};
request.onerror = () => {
resolve(null);
};
});
if (modelFromDB) {
if (options?.progressCallback) {
options.progressCallback(100);
}
return modelFromDB;
}
}
catch (err) {
console.warn('Failed to check IndexedDB:', err);
}
}
// Download the model
if (options?.progressCallback) {
options.progressCallback(10); // Starting download
}
// Create a controller to potentially abort the download
const controller = new AbortController();
const signal = controller.signal;
// Download with progress tracking
const response = await fetch(url, { signal });
if (!response.ok) {
throw new Error(`Failed to download model: ${response.status} ${response.statusText}`);
}
// Get total size for progress calculation
const totalSize = parseInt(response.headers.get('Content-Length') || '0', 10);
let loadedSize = 0;
// Create a reader to stream the response with progress updates
const reader = response.body.getReader();
const chunks = [];
while (true) {
const { done, value } = await reader.read();
if (done) {
break;
}
chunks.push(value);
loadedSize += value.length;
// Report download progress
if (options?.progressCallback && totalSize > 0) {
const progress = Math.min(10 + Math.floor((loadedSize / totalSize) * 70), 80);
options.progressCallback(progress);
}
}
// Combine chunks into a single buffer
const totalLength = chunks.reduce((sum, chunk) => sum + chunk.length, 0);
const modelBuffer = new Uint8Array(totalLength);
let position = 0;
for (const chunk of chunks) {
modelBuffer.set(chunk, position);
position += chunk.length;
}
if (options?.progressCallback) {
options.progressCallback(85); // Download complete
}
// Store in Cache API if requested
if (options?.saveToCache && 'caches' in window) {
try {
const cache = await caches.open('onnx-models');
const responseToCache = new Response(modelBuffer);
await cache.put(url, responseToCache);
if (options?.progressCallback) {
options.progressCallback(90);
}
}
catch (err) {
console.warn('Failed to save to cache:', err);
}
}
// Store in IndexedDB if requested
if (options?.saveToIndexedDB) {
try {
const dbName = 'onnx_models_db';
const storeName = 'model_urls';
const urlHash = await hashString(url);
const request = indexedDB.open(dbName, 1);
request.onupgradeneeded = (event) => {
const db = event.target.result;
if (!db.objectStoreNames.contains(storeName)) {
db.createObjectStore(storeName);
}
};
request.onsuccess = (event) => {
const db = event.target.result;
const transaction = db.transaction(storeName, 'readwrite');
const store = transaction.objectStore(storeName);
store.put(modelBuffer.buffer, urlHash);
if (options?.progressCallback) {
options.progressCallback(95);
}
};
}
catch (err) {
console.warn('Failed to save to IndexedDB:', err);
}
}
if (options?.progressCallback) {
options.progressCallback(100); // All done
}
return modelBuffer.buffer;
}
catch (error) {
throw new OnnxModelError(`Failed to load model from URL: ${error instanceof Error ? error.message : String(error)}`);
}
}
/**
* Generate a hash from a string (URL)
* Used for IndexedDB storage keys
*/
async function hashString(str) {
// Use Web Crypto API if available
if (window.crypto && window.crypto.subtle) {
const encoder = new TextEncoder();
const data = encoder.encode(str);
const hashBuffer = await window.crypto.subtle.digest('SHA-256', data);
const hashArray = Array.from(new Uint8Array(hashBuffer));
return hashArray.map(b => b.toString(16).padStart(2, '0')).join('');
}
else {
// Simple fallback if crypto API is not available
let hash = 0;
for (let i = 0; i < str.length; i++) {
const char = str.charCodeAt(i);
hash = ((hash << 5) - hash) + char;
hash = hash & hash; // Convert to 32bit integer
}
return hash.toString(36);
}
}
/**
* Check if WebGPU is supported in the current browser
*/
export function isWebGPUSupported() {
return typeof navigator !== 'undefined' &&
navigator.gpu !== undefined;
}
/**
* Check if WebAssembly is supported in the current browser
*/
export function isWasmSupported() {
return (typeof WebAssembly === 'object' &&
WebAssembly !== null &&
typeof WebAssembly.instantiate === 'function');
}