@sconedev/ai_toolkit
Version:
Simplify AI integration in web apps with local and offline model support
872 lines (871 loc) • 37.1 kB
JavaScript
import { getConfig } from './config';
import { getCached, setCached } from './cache';
import { isBrowser } from '../utils/browser-check';
import * as fs from 'fs';
import * as path from 'path';
// Custom error types for better error handling
export class OnnxModelError extends Error {
constructor(message) {
super(message);
this.name = 'OnnxModelError';
}
}
export class OnnxRuntimeError extends Error {
constructor(message) {
super(message);
this.name = 'OnnxRuntimeError';
}
}
// Use CDN URL for ONNX runtime with fallback options
const ONNX_CDN_URLS = [
"https://cdn.jsdelivr.net/npm/onnxruntime-web@1.21.0-dev.20250206-d981b153d3/webgpu/+esm",
"https://cdn.jsdelivr.net/npm/onnxruntime-web@1.16.3/dist/esm/ort.js" // Fallback to stable version
];
// Dynamic import or script load for ONNX runtime
let ortPromise = null;
let ort = null; // For Node.js usage
// Load ONNX Runtime Web asynchronously when needed with improved error handling
async function loadOrt() {
if (isBrowser()) {
if (!ortPromise) {
ortPromise = new Promise(async (resolve, reject) => {
// Try each CDN URL in order
for (const cdnUrl of ONNX_CDN_URLS) {
try {
// Create script element to load from CDN
const script = document.createElement('script');
script.type = 'module';
script.async = true;
script.src = cdnUrl;
// Create a promise to track loading
const loadPromise = new Promise((loadResolve, loadReject) => {
script.onload = () => loadResolve();
script.onerror = () => loadReject(new Error(`Failed to load from ${cdnUrl}`));
});
// Add to document to begin loading
document.head.appendChild(script);
// Wait for script to load
await loadPromise;
// ESM import needs to be dynamically imported
const module = await import(cdnUrl);
resolve(module);
return; // Success, exit the loop
}
catch (err) {
console.warn(`Failed to load ONNX Runtime from ${cdnUrl}. Trying next option...`);
// Continue to next URL
}
}
// If we get here, all CDN URLs failed
reject(new OnnxRuntimeError('Failed to load ONNX Runtime from all CDN URLs'));
});
}
return ortPromise;
}
else {
// In Node.js environment, use the regular package
if (!ort) {
try {
ort = require('onnxruntime-node');
}
catch (err) {
throw new OnnxRuntimeError('ONNX Runtime Node.js could not be loaded. Make sure it is properly installed.');
}
}
return ort;
}
}
/**
* Validates the input data against the expected input shapes and types
* @param inputData The input data to validate
* @param modelConfig The model configuration
* @throws Error if validation fails
*/
function validateInputData(inputData, modelConfig) {
// Check that all required inputs are present
for (const inputName of modelConfig.inputNames) {
if (!inputData[inputName]) {
throw new OnnxModelError(`Missing required input '${inputName}' for ONNX model '${modelConfig.name}'`);
}
// If shapes are defined, validate them
if (modelConfig.inputShapes) {
const inputIndex = modelConfig.inputNames.indexOf(inputName);
if (inputIndex >= 0 && modelConfig.inputShapes[inputIndex]) {
const expectedShape = modelConfig.inputShapes[inputIndex];
const data = inputData[inputName];
// Check array length for simple arrays
if (Array.isArray(data) && expectedShape.length === 1) {
if (data.length !== expectedShape[0] && expectedShape[0] > 0) {
console.warn(`Expected input '${inputName}' to have length ${expectedShape[0]}, got ${data.length}`);
}
}
// Check tensor dimensions for TypedArrays
else if ((data instanceof Float32Array || data instanceof Int32Array || data instanceof Uint8Array)) {
// Calculate expected size from shape (product of all dimensions)
const expectedSize = expectedShape.reduce((a, b) => a * b, 1);
if (data.length !== expectedSize && expectedSize > 0) {
console.warn(`Expected input '${inputName}' to have ${expectedSize} elements, got ${data.length}`);
}
}
}
}
}
}
/**
* Run inference using a local ONNX model file
*
* @param modelPath Path to the local ONNX model file or File object in browser
* @param modelConfig Configuration for the model (input/output names)
* @param inputData Input data for inference
* @param options Additional options for running the model
* @returns Result of the model inference
*/
export async function runModelFromPath(modelPath, modelConfig, inputData, options) {
try {
// Validate input data before proceeding
validateInputData(inputData, modelConfig);
if (isBrowser()) {
if (modelPath instanceof File) {
return runModelFromBrowserFile(modelPath, modelConfig, inputData, options);
}
else if (modelPath instanceof ArrayBuffer) {
return runModelFromArrayBuffer(modelPath, modelConfig, inputData, options);
}
else {
throw new OnnxModelError('In browser environments, model path must be a File or ArrayBuffer');
}
}
else {
if (typeof modelPath === 'string') {
return runModelFromNodePath(modelPath, modelConfig, inputData, options);
}
else {
throw new OnnxModelError('In Node.js environments, model path must be a string');
}
}
}
catch (error) {
if (error instanceof OnnxModelError || error instanceof OnnxRuntimeError) {
throw error;
}
throw new OnnxRuntimeError(`Error running model from path: ${error instanceof Error ? error.message : String(error)}`);
}
}
/**
* Run inference using a local ONNX model from ArrayBuffer
* This is useful for models loaded from IndexedDB, fetch(), etc.
*/
async function runModelFromArrayBuffer(modelBuffer, modelConfig, inputData, options) {
const config = getConfig();
const modelName = modelConfig.name || 'model-from-buffer';
// Create unique identifier for this buffer (using hash of first few bytes)
const bytes = new Uint8Array(modelBuffer.slice(0, 64));
const hashVal = bytes.reduce((prev, curr) => ((prev << 5) - prev) + curr, 0);
const bufferId = `buffer-${hashVal}`;
// Create cache keys
const resultCacheKey = `onnx:result:${modelName}:${bufferId}:${JSON.stringify(inputData)}`;
const modelCacheKey = `onnx:model:${modelName}:${bufferId}`;
// Progress reporting if callback provided
if (options?.progressCallback) {
options.progressCallback(5); // 5% - Starting
}
// Check for cached results
if (options?.cacheResults !== false) {
const cached = getCached(resultCacheKey);
if (cached) {
if (options?.progressCallback) {
options.progressCallback(100); // Complete - Using cached result
}
return cached;
}
}
// Get or create session
let session = getCached(modelCacheKey);
if (!session) {
try {
// Load ONNX Runtime Web
const ortLib = await loadOrt();
if (options?.progressCallback) {
options.progressCallback(20); // 20% - ONNX Runtime loaded
}
// Notify about model loading if configured
if (config.onRequest) {
config.onRequest({
provider: 'onnx',
function: 'runModelFromArrayBuffer',
input: { model: modelName, bufferSize: modelBuffer.byteLength }
});
}
// Create inference options combining default with user options
const inferenceOptions = {
executionProviders: ['webgpu', 'wasm'],
graphOptimizationLevel: 'all',
...(options?.inferenceOptions || {})
};
// Create ONNX Session from array buffer
session = await ortLib.InferenceSession.create(modelBuffer, inferenceOptions);
if (options?.progressCallback) {
options.progressCallback(60); // 60% - Model loaded
}
// Cache the model if requested
if (options?.cacheModel !== false) {
setCached(modelCacheKey, session);
}
}
catch (error) {
throw new OnnxModelError(`Failed to load ONNX model from buffer: ${error instanceof Error ? error.message : String(error)}`);
}
}
else if (options?.progressCallback) {
options.progressCallback(50); // 50% - Using cached model
}
try {
// Load ONNX Runtime if not already loaded
const ortLib = await loadOrt();
// Prepare inputs
const feeds = {};
for (const inputName of modelConfig.inputNames) {
if (!inputData[inputName]) {
throw new OnnxModelError(`Missing required input '${inputName}' for ONNX model '${modelName}'`);
}
// Convert input data to Tensor
const data = inputData[inputName];
// Handle different input data types
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')) {
// Handle image data as Uint8Array (RGB)
const dims = inputData[`${inputName}_dims`] || [1, 3, 224, 224]; // Default image dims if not provided
feeds[inputName] = new ortLib.Tensor('uint8', data, dims);
}
else if (data instanceof Int32Array) {
// Support for Int32Array (often used for token IDs in NLP)
feeds[inputName] = new ortLib.Tensor('int32', data, [data.length]);
}
else if (data instanceof Uint8Array) {
// Generic Uint8Array handling
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); // 75% - Inputs prepared
}
// Run the model
const results = await session.run(feeds);
if (options?.progressCallback) {
options.progressCallback(90); // 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); // 100% - All done
}
return outputs;
}
catch (error) {
throw new OnnxRuntimeError(`ONNX inference failed for model from buffer: ${error instanceof Error ? error.message : String(error)}`);
}
}
/**
* Run inference using a local ONNX model file in Node.js
*/
async function runModelFromNodePath(modelPath, modelConfig, inputData, options) {
// Verify file exists
if (!fs.existsSync(modelPath)) {
throw new OnnxModelError(`Model file not found at path: ${modelPath}`);
}
const config = getConfig();
const modelName = modelConfig.name || path.basename(modelPath, path.extname(modelPath));
// Create cache keys
const resultCacheKey = `onnx:result:${modelName}:${JSON.stringify(inputData)}`;
const modelCacheKey = `onnx:model:${modelName}:${modelPath}`;
// Check for cached results
if (options?.cacheResults !== false) {
const cached = getCached(resultCacheKey);
if (cached) {
return cached;
}
}
// Get or create session
let session = getCached(modelCacheKey);
if (!session) {
try {
// Notify about model loading if configured
if (config.onRequest) {
config.onRequest({
provider: 'onnx',
function: 'runModelFromPath',
input: { model: modelName, path: modelPath }
});
}
// Load ONNX Runtime and create session
const ortLib = await loadOrt();
session = await ortLib.InferenceSession.create(modelPath);
// Cache the model if requested
if (options?.cacheModel !== false) {
setCached(modelCacheKey, session);
}
}
catch (error) {
throw new OnnxModelError(`Failed to load local ONNX model at '${modelPath}': ${error instanceof Error ? error.message : String(error)}`);
}
}
try {
// Prepare inputs
const ortLib = await loadOrt();
const feeds = {};
for (const inputName of modelConfig.inputNames) {
if (!inputData[inputName]) {
throw new OnnxModelError(`Missing required input '${inputName}' for ONNX model '${modelName}'`);
}
// Convert input data to Tensor
const data = inputData[inputName];
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 Int32Array) {
feeds[inputName] = new ortLib.Tensor('int32', data, [data.length]);
}
else if (data instanceof Uint8Array) {
feeds[inputName] = new ortLib.Tensor('uint8', data, [data.length]);
}
else {
throw new OnnxModelError(`Unsupported data type for input '${inputName}'`);
}
}
// Run the model
const results = await session.run(feeds);
// 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);
}
return outputs;
}
catch (error) {
throw new OnnxRuntimeError(`ONNX inference failed for local model '${modelName}': ${error instanceof Error ? error.message : String(error)}`);
}
}
/**
* Run inference using a local ONNX model file in browser
*/
async function runModelFromBrowserFile(modelFile, modelConfig, inputData, options) {
if (!modelFile) {
throw new OnnxModelError('No model file provided');
}
const config = getConfig();
const modelName = modelConfig.name || modelFile.name.replace(/\.[^/.]+$/, ""); // Remove file extension
// Create cache keys
const fileId = `${modelFile.name}-${modelFile.lastModified}`;
const resultCacheKey = `onnx:result:${modelName}:${fileId}:${JSON.stringify(inputData)}`;
const modelCacheKey = `onnx:model:${modelName}:${fileId}`;
// Progress reporting if callback provided
if (options?.progressCallback) {
options.progressCallback(5); // 5% - Starting
}
// Check for cached results
if (options?.cacheResults !== false) {
const cached = getCached(resultCacheKey);
if (cached) {
if (options?.progressCallback) {
options.progressCallback(100); // Complete - Using cached result
}
return cached;
}
}
// Get or create session
let session = getCached(modelCacheKey);
if (!session) {
try {
// Load ONNX Runtime Web from CDN
const ortLib = await loadOrt();
if (options?.progressCallback) {
options.progressCallback(20); // 20% - ONNX Runtime loaded
}
// Convert File to ArrayBuffer
const arrayBuffer = await modelFile.arrayBuffer();
if (options?.progressCallback) {
options.progressCallback(40); // 40% - File read
}
// Notify about model loading if configured
if (config.onRequest) {
config.onRequest({
provider: 'onnx',
function: 'runModelFromFile',
input: { model: modelName, fileSize: modelFile.size }
});
}
// Create ONNX Session from array buffer using CDN-loaded library
session = await ortLib.InferenceSession.create(arrayBuffer, {
executionProviders: ['webgpu', 'wasm'], // Try WebGPU first, then fall back to WASM
graphOptimizationLevel: 'all'
});
if (options?.progressCallback) {
options.progressCallback(60); // 60% - Model loaded
}
// Cache the model if requested
if (options?.cacheModel !== false) {
setCached(modelCacheKey, session);
}
}
catch (error) {
throw new OnnxModelError(`Failed to load local ONNX model '${modelName}': ${error instanceof Error ? error.message : String(error)}`);
}
}
else if (options?.progressCallback) {
options.progressCallback(50); // 50% - Using cached model
}
try {
// Load ONNX Runtime if not already loaded
const ortLib = await loadOrt();
// Prepare inputs
const feeds = {};
for (const inputName of modelConfig.inputNames) {
if (!inputData[inputName]) {
throw new OnnxModelError(`Missing required input '${inputName}' for ONNX model '${modelName}'`);
}
// Convert input data to Tensor
const data = inputData[inputName];
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')) {
// Handle image data as Uint8Array (RGB)
const dims = inputData[`${inputName}_dims`] || [1, 3, 224, 224]; // Default image dims if not provided
feeds[inputName] = new ortLib.Tensor('uint8', data, dims);
}
else if (data instanceof Int32Array) {
// Support for Int32Array (often used for token IDs in NLP)
feeds[inputName] = new ortLib.Tensor('int32', data, [data.length]);
}
else {
throw new OnnxModelError(`Unsupported data type for input '${inputName}'`);
}
}
if (options?.progressCallback) {
options.progressCallback(75); // 75% - Inputs prepared
}
// Run the model
const results = await session.run(feeds);
if (options?.progressCallback) {
options.progressCallback(90); // 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); // 100% - All done
}
return outputs;
}
catch (error) {
throw new OnnxRuntimeError(`ONNX inference failed for local model '${modelName}': ${error instanceof Error ? error.message : String(error)}`);
}
}
// Main function to run inference with an ONNX model
export async function runOnnxInference(modelConfig, inputData, options) {
const config = getConfig();
// Try to get cached results if enabled
if (options?.cacheResults !== false) {
const cacheKey = `onnx:result:${modelConfig.name}:${JSON.stringify(inputData)}`;
const cached = getCached(cacheKey);
if (cached) {
if (options?.progressCallback) {
options.progressCallback(100); // Complete - Using cached result
}
return cached;
}
}
// Check if model is already loaded and cached
const modelCacheKey = `onnx:model:${modelConfig.name}`;
let session = getCached(modelCacheKey);
if (!session) {
try {
// Progress reporting if callback provided
if (options?.progressCallback) {
options.progressCallback(10); // 10% - Starting model loading
}
// Load ONNX Runtime
const ortLib = await loadOrt();
// Create ONNX Session with appropriate options
const sessionOptions = isBrowser() ?
{
executionProviders: ['webgpu', 'wasm'], // Try WebGPU first, then fall back to WASM
graphOptimizationLevel: 'all'
} : undefined;
session = await ortLib.InferenceSession.create(modelConfig.modelUrl, sessionOptions);
if (options?.progressCallback) {
options.progressCallback(50); // 50% - Model loaded
}
// Cache the model if requested
if (options?.cacheModel !== false) {
setCached(modelCacheKey, session);
}
}
catch (error) {
throw new OnnxModelError(`Failed to load ONNX model '${modelConfig.name}': ${error instanceof Error ? error.message : String(error)}`);
}
}
else if (options?.progressCallback) {
options.progressCallback(40); // 40% - Using cached model
}
// Notify about usage if configured
if (config.onRequest) {
config.onRequest({
provider: 'onnx',
function: 'runOnnxInference',
input: { model: modelConfig.name, inputKeys: Object.keys(inputData) }
});
}
try {
// Load ONNX Runtime if not already loaded
const ortLib = await loadOrt();
// Prepare inputs
const feeds = {};
for (const inputName of modelConfig.inputNames) {
if (!inputData[inputName]) {
throw new OnnxModelError(`Missing required input '${inputName}' for ONNX model '${modelConfig.name}'`);
}
// Convert input data to Tensor
const data = inputData[inputName];
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')) {
// Handle image data as Uint8Array (RGB)
const dims = inputData[`${inputName}_dims`] || [1, 3, 224, 224]; // Default image dims if not provided
feeds[inputName] = new ortLib.Tensor('uint8', data, dims);
}
else if (data instanceof Int32Array) {
feeds[inputName] = new ortLib.Tensor('int32', data, [data.length]);
}
else {
throw new OnnxModelError(`Unsupported data type for input '${inputName}'`);
}
}
if (options?.progressCallback) {
options.progressCallback(70); // 70% - Inputs prepared
}
// Run the model
const results = await session.run(feeds);
if (options?.progressCallback) {
options.progressCallback(90); // 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) {
const cacheKey = `onnx:result:${modelConfig.name}:${JSON.stringify(inputData)}`;
setCached(cacheKey, outputs);
}
if (options?.progressCallback) {
options.progressCallback(100); // 100% - All done
}
return outputs;
}
catch (error) {
throw new OnnxRuntimeError(`ONNX inference failed for model '${modelConfig.name}': ${error instanceof Error ? error.message : String(error)}`);
}
}
// Function to unload a model from cache
export function unloadOnnxModel(modelName) {
const modelCacheKey = `onnx:model:${modelName}`;
const session = getCached(modelCacheKey);
if (session) {
try {
// Remove from cache
setCached(modelCacheKey, null);
return true;
}
catch (error) {
console.error(`Error unloading ONNX model '${modelName}':`, error);
return false;
}
}
return false;
}
// Helper to check if WebAssembly is supported in the current browser
export function isWasmSupported() {
return (typeof WebAssembly === 'object' &&
WebAssembly !== null &&
typeof WebAssembly.instantiate === 'function');
}
// Helper to check if WebGPU is supported in the current browser
export function isWebGPUSupported() {
return typeof navigator !== 'undefined' &&
navigator.gpu !== undefined;
}
/**
* Download an ONNX model from a URL and save it for offline use
*
* @param modelUrl URL to download the model from
* @param options Download options
* @returns Path to the downloaded model or the model buffer (browser)
*/
export async function downloadModelForOfflineUse(modelUrl, options) {
const modelName = options?.modelName || `model-${new Date().getTime()}`;
if (isBrowser()) {
try {
// Report progress
if (options?.progressCallback) {
options.progressCallback(10); // 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(modelUrl, { 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); // 85% - Download complete
}
// Store in IndexedDB if requested
if (options?.saveToIndexedDB) {
try {
const dbName = 'onnx_models_db';
const storeName = 'models';
const request = indexedDB.open(dbName, 1);
// Create object store if needed
request.onupgradeneeded = (event) => {
const db = event.target.result;
if (!db.objectStoreNames.contains(storeName)) {
db.createObjectStore(storeName);
}
};
// Save the model to IndexedDB
request.onsuccess = async (event) => {
const db = event.target.result;
const transaction = db.transaction(storeName, 'readwrite');
const store = transaction.objectStore(storeName);
store.put(modelBuffer.buffer, modelName);
if (options?.progressCallback) {
options.progressCallback(95); // 95% - Saved to IndexedDB
}
};
request.onerror = () => {
console.warn('Failed to save model to IndexedDB');
};
}
catch (err) {
console.warn('IndexedDB storage failed:', err);
}
}
// Store in cache if requested
if (options?.saveToCache) {
try {
const cache = await caches.open('onnx-models');
const response = new Response(modelBuffer.buffer);
await cache.put(modelUrl, response);
if (options?.progressCallback) {
options.progressCallback(95); // 95% - Saved to cache
}
}
catch (err) {
console.warn('Cache storage failed:', err);
}
}
if (options?.progressCallback) {
options.progressCallback(100); // 100% - All done
}
return modelBuffer.buffer;
}
catch (error) {
throw new OnnxModelError(`Failed to download model for offline use: ${error instanceof Error ? error.message : String(error)}`);
}
}
else {
// Node.js environment - download and save to file
const MODEL_CACHE_DIR = path.join(process.cwd(), '.model-cache');
try {
// Create cache directory if it doesn't exist
if (!fs.existsSync(MODEL_CACHE_DIR)) {
fs.mkdirSync(MODEL_CACHE_DIR, { recursive: true });
}
const modelFileName = `${modelName}.onnx`;
const modelPath = path.join(MODEL_CACHE_DIR, modelFileName);
// Check if model already exists in cache
if (fs.existsSync(modelPath)) {
return modelPath;
}
// Model doesn't exist, we need to download it
console.log(`Downloading model ${modelName} from ${modelUrl}`);
if (options?.progressCallback) {
options.progressCallback(10); // 10% - Starting download
}
// Download the model with progress tracking
const { default: fetch } = await import('node-fetch');
const response = await fetch(modelUrl);
if (!response.ok) {
throw new Error(`Failed to download model: ${response.status} ${response.statusText}`);
}
// Get the content as an array buffer
const buffer = Buffer.from(await response.arrayBuffer());
if (options?.progressCallback) {
options.progressCallback(80); // 80% - Download complete
}
// Save to file
fs.writeFileSync(modelPath, buffer);
if (options?.progressCallback) {
options.progressCallback(100); // 100% - Saved to file
}
console.log(`Model ${modelName} downloaded and saved to ${modelPath}`);
return modelPath;
}
catch (error) {
throw new OnnxModelError(`Failed to download model: ${error instanceof Error ? error.message : String(error)}`);
}
}
}
/**
* Load a saved ONNX model from IndexedDB (browser only)
*/
export async function loadModelFromIndexedDB(modelName) {
if (!isBrowser()) {
throw new OnnxModelError('IndexedDB is only available in browser environments');
}
return new Promise((resolve, reject) => {
try {
const dbName = 'onnx_models_db';
const storeName = 'models';
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(modelName);
getRequest.onsuccess = () => {
resolve(getRequest.result || null);
};
getRequest.onerror = () => {
resolve(null);
};
}
catch (err) {
resolve(null);
}
};
request.onerror = () => {
reject(new OnnxModelError('Failed to open IndexedDB'));
};
}
catch (err) {
reject(new OnnxModelError(`IndexedDB error: ${err instanceof Error ? err.message : String(err)}`));
}
});
}
/**
* Check if a model exists in the cache
*/
export function isModelInCache(modelName) {
const modelCacheKey = `onnx:model:${modelName}`;
return getCached(modelCacheKey) !== null;
}
/**
* Get available execution providers for the current environment
*/
export function getAvailableExecutionProviders() {
if (isBrowser()) {
const providers = ['wasm'];
// Check for WebGPU support
if (isWebGPUSupported()) {
providers.unshift('webgpu'); // Add as preferred provider
}
return providers;
}
else {
// In Node.js, defer to the available providers in onnxruntime-node
try {
const ortLib = require('onnxruntime-node');
return ortLib.getAvailableProviders ? ortLib.getAvailableProviders() : ['cpu'];
}
catch (err) {
return ['cpu']; // Default fallback
}
}
}
export { validateInputData };