inference-server 1.0.0-beta.19

package/dist/store.js

@@ -0,0 +1,203 @@
+import { promises as fs, existsSync } from 'node:fs';
+import PQueue from 'p-queue';
+import prettyMilliseconds from 'pretty-ms';
+import prettyBytes from 'pretty-bytes';
+import { LogLevels, createSublogger, } from './lib/logger.js';
+import { formatBytesPerSecond, mergeAbortSignals } from './lib/util.js';
+export class ModelStore {
+    prepareQueue;
+    models = {};
+    engines;
+    prepareController;
+    modelsCachePath;
+    log;
+    constructor(options) {
+        this.prepareController = new AbortController();
+        this.log = createSublogger(options.log);
+        this.prepareQueue = new PQueue({
+            concurrency: options.prepareConcurrency ?? 2,
+        });
+        this.modelsCachePath = options.modelsCachePath;
+        this.models = Object.fromEntries(Object.entries(options.models).map(([modelId, model]) => [
+            modelId,
+            {
+                ...model,
+                status: 'unloaded',
+            },
+        ]));
+    }
+    async init(engines) {
+        this.engines = engines;
+        if (!existsSync(this.modelsCachePath)) {
+            await fs.mkdir(this.modelsCachePath, { recursive: true });
+        }
+        const blockingPromises = [];
+        for (const modelId in this.models) {
+            const model = this.models[modelId];
+            if (model.prepare === 'blocking' || model.minInstances > 0) {
+                blockingPromises.push(this.prepareModel(modelId));
+            }
+            else if (model.prepare === 'async') {
+                this.prepareModel(modelId);
+            }
+        }
+        await Promise.all(blockingPromises);
+    }
+    dispose() {
+        this.prepareController.abort();
+    }
+    onDownloadProgress(modelId, progress) {
+        const model = this.models[modelId];
+        if (!model.downloads) {
+            model.downloads = new Map();
+        }
+        if (model.downloads.has(progress.file)) {
+            const tracker = model.downloads.get(progress.file);
+            tracker.pushProgress(progress);
+        }
+        else {
+            const tracker = new DownloadTracker(5000);
+            tracker.pushProgress(progress);
+            model.downloads.set(progress.file, tracker);
+        }
+    }
+    // makes sure all required files for the model exist and are valid
+    // checking model checksums and reading metadata is model + engine specific and can be slow
+    async prepareModel(modelId, signal) {
+        const model = this.models[modelId];
+        if (!this.engines) {
+            throw new Error('No engines available - did you call init()?');
+        }
+        model.status = 'preparing';
+        const engine = this.engines[model.engine];
+        this.log(LogLevels.info, 'Preparing model', {
+            model: modelId,
+            task: model.task,
+        });
+        await this.prepareQueue.add(async () => {
+            if (!('prepareModel' in engine)) {
+                model.status = 'ready';
+                return model;
+            }
+            const logProgressInterval = setInterval(() => {
+                const progress = Array.from(model.downloads?.values() ?? [])
+                    .map((tracker) => tracker.getStatus())
+                    .reduce((acc, status) => {
+                    acc.loadedBytes += status?.loadedBytes || 0;
+                    acc.totalBytes += status?.totalBytes || 0;
+                    acc.speed += status?.speed || 0;
+                    return acc;
+                }, { loadedBytes: 0, totalBytes: 0, speed: 0 });
+                if (progress.totalBytes) {
+                    const percent = (progress.loadedBytes / progress.totalBytes) * 100;
+                    const formattedTotalBytes = prettyBytes(progress.totalBytes, { space: false });
+                    const formattedLoadedBytes = prettyBytes(progress.loadedBytes, { space: false });
+                    this.log(LogLevels.info, `Downloading at ${formatBytesPerSecond(progress.speed)} ${percent.toFixed(1)}% - ${formattedLoadedBytes} of ${formattedTotalBytes}`, {
+                        model: modelId,
+                    });
+                }
+            }, 10000);
+            try {
+                const modelMeta = await engine.prepareModel({ config: model, log: this.log }, (progress) => {
+                    this.onDownloadProgress(model.id, progress);
+                }, mergeAbortSignals([signal, this.prepareController.signal]));
+                model.downloads = undefined;
+                model.meta = modelMeta;
+                model.status = 'ready';
+                this.log(LogLevels.info, 'Model ready', {
+                    model: modelId,
+                    task: model.task,
+                });
+            }
+            catch (error) {
+                this.log(LogLevels.error, 'Error preparing model', {
+                    model: modelId,
+                    error: error,
+                });
+                model.status = 'error';
+            }
+            finally {
+                clearInterval(logProgressInterval);
+            }
+            return model;
+        });
+    }
+    getStatus() {
+        const formatFloat = (num) => parseFloat(num?.toFixed(2) || '0');
+        const storeStatusInfo = Object.fromEntries(Object.entries(this.models).map(([modelId, model]) => {
+            let downloads = undefined;
+            if (model.downloads) {
+                downloads = [...model.downloads].reduce((acc, [key, download]) => {
+                    const status = download.getStatus();
+                    const latestState = download.progressBuffer[download.progressBuffer.length - 1];
+                    const totalBytes = latestState?.totalBytes ?? 0;
+                    const loadedBytes = latestState?.loadedBytes ?? 0;
+                    const etaSeconds = status?.etaSeconds ?? 0;
+                    const formattedEta = prettyMilliseconds(etaSeconds * 1000);
+                    const formattedTotalBytes = prettyBytes(totalBytes);
+                    const formattedLoadedBytes = prettyBytes(loadedBytes);
+                    acc.push({
+                        file: key,
+                        loadedBytes,
+                        formattedLoadedBytes,
+                        totalBytes,
+                        formattedTotalBytes,
+                        percent: formatFloat(status?.percent),
+                        speed: formatFloat(status?.speed),
+                        etaSeconds: formatFloat(etaSeconds),
+                        formattedEta,
+                    });
+                    return acc;
+                }, []);
+            }
+            return [
+                modelId,
+                {
+                    engine: model.engine,
+                    device: model.device,
+                    minInstances: model.minInstances,
+                    maxInstances: model.maxInstances,
+                    status: model.status,
+                    downloads,
+                },
+            ];
+        }));
+        return storeStatusInfo;
+    }
+}
+class DownloadTracker {
+    progressBuffer = [];
+    timeWindow;
+    constructor(timeWindow = 1000) {
+        this.timeWindow = timeWindow;
+    }
+    pushProgress({ loadedBytes, totalBytes }) {
+        const timestamp = Date.now();
+        this.progressBuffer.push({ loadedBytes, totalBytes, timestamp });
+        this.cleanup();
+    }
+    cleanup() {
+        const cutoffTime = Date.now() - this.timeWindow;
+        this.progressBuffer = this.progressBuffer.filter((item) => item.timestamp >= cutoffTime);
+    }
+    getStatus() {
+        if (this.progressBuffer.length < 2) {
+            return null; // Not enough data to calculate speed and ETA
+        }
+        const latestState = this.progressBuffer[this.progressBuffer.length - 1];
+        const previousState = this.progressBuffer[0]; // oldest state within the time window
+        const bytesLoaded = latestState.loadedBytes - previousState.loadedBytes;
+        const timeElapsed = latestState.timestamp - previousState.timestamp; // in milliseconds
+        const speed = bytesLoaded / (timeElapsed / 1000); // bytes per second
+        const remainingBytes = latestState.totalBytes - latestState.loadedBytes;
+        const eta = speed > 0 ? remainingBytes / speed : 0;
+        return {
+            speed,
+            etaSeconds: eta,
+            percent: latestState.loadedBytes / latestState.totalBytes,
+            loadedBytes: latestState.loadedBytes,
+            totalBytes: latestState.totalBytes,
+        };
+    }
+}
+//# sourceMappingURL=store.js.map

package/dist/store.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"store.js","sourceRoot":"","sources":["../src/store.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,QAAQ,IAAI,EAAE,EAAE,UAAU,EAAE,MAAM,SAAS,CAAA;AACpD,OAAO,MAAM,MAAM,SAAS,CAAA;AAC5B,OAAO,kBAAkB,MAAM,WAAW,CAAA;AAC1C,OAAO,WAAW,MAAM,cAAc,CAAA;AAMtC,OAAO,EAEN,SAAS,EAET,eAAe,GACf,MAAM,wBAAwB,CAAA;AAC/B,OAAO,EAAE,oBAAoB,EAAE,iBAAiB,EAAE,MAAM,sBAAsB,CAAA;AAmB9E,MAAM,OAAO,UAAU;IACtB,YAAY,CAAQ;IACpB,MAAM,GAAgC,EAAE,CAAA;IACxC,OAAO,CAA8B;IAC7B,iBAAiB,CAAiB;IAClC,eAAe,CAAQ;IACvB,GAAG,CAAQ;IAEnB,YAAY,OAA0B;QACrC,IAAI,CAAC,iBAAiB,GAAG,IAAI,eAAe,EAAE,CAAA;QAC9C,IAAI,CAAC,GAAG,GAAG,eAAe,CAAC,OAAO,CAAC,GAAG,CAAC,CAAA;QACvC,IAAI,CAAC,YAAY,GAAG,IAAI,MAAM,CAAC;YAC9B,WAAW,EAAE,OAAO,CAAC,kBAAkB,IAAI,CAAC;SAC5C,CAAC,CAAA;QACF,IAAI,CAAC,eAAe,GAAG,OAAO,CAAC,eAAe,CAAA;QAC9C,IAAI,CAAC,MAAM,GAAG,MAAM,CAAC,WAAW,CAC/B,MAAM,CAAC,OAAO,CAAC,OAAO,CAAC,MAAM,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,OAAO,EAAE,KAAK,CAAC,EAAE,EAAE,CAAC;YACxD,OAAO;YACP;gBACC,GAAG,KAAK;gBACR,MAAM,EAAE,UAAU;aAClB;SACD,CAAC,CACF,CAAA;IACF,CAAC;IAED,KAAK,CAAC,IAAI,CAAC,OAAoC;QAC9C,IAAI,CAAC,OAAO,GAAG,OAAO,CAAA;QACtB,IAAI,CAAC,UAAU,CAAC,IAAI,CAAC,eAAe,CAAC,EAAE,CAAC;YACvC,MAAM,EAAE,CAAC,KAAK,CAAC,IAAI,CAAC,eAAe,EAAE,EAAE,SAAS,EAAE,IAAI,EAAE,CAAC,CAAA;QAC1D,CAAC;QAED,MAAM,gBAAgB,GAAG,EAAE,CAAA;QAC3B,KAAK,MAAM,OAAO,IAAI,IAAI,CAAC,MAAM,EAAE,CAAC;YACnC,MAAM,KAAK,GAAG,IAAI,CAAC,MAAM,CAAC,OAAO,CAAC,CAAA;YAClC,IAAI,KAAK,CAAC,OAAO,KAAK,UAAU,IAAI,KAAK,CAAC,YAAY,GAAG,CAAC,EAAE,CAAC;gBAC5D,gBAAgB,CAAC,IAAI,CAAC,IAAI,CAAC,YAAY,CAAC,OAAO,CAAC,CAAC,CAAA;YAClD,CAAC;iBAAM,IAAI,KAAK,CAAC,OAAO,KAAK,OAAO,EAAE,CAAC;gBACtC,IAAI,CAAC,YAAY,CAAC,OAAO,CAAC,CAAA;YAC3B,CAAC;QACF,CAAC;QACD,MAAM,OAAO,CAAC,GAAG,CAAC,gBAAgB,CAAC,CAAA;IACpC,CAAC;IAED,OAAO;QACN,IAAI,CAAC,iBAAiB,CAAC,KAAK,EAAE,CAAA;IAC/B,CAAC;IAEO,kBAAkB,CACzB,OAAe,EACf,QAAmE;QAEnE,MAAM,KAAK,GAAG,IAAI,CAAC,MAAM,CAAC,OAAO,CAAC,CAAA;QAClC,IAAI,CAAC,KAAK,CAAC,SAAS,EAAE,CAAC;YACtB,KAAK,CAAC,SAAS,GAAG,IAAI,GAAG,EAAE,CAAA;QAC5B,CAAC;QAED,IAAI,KAAK,CAAC,SAAS,CAAC,GAAG,CAAC,QAAQ,CAAC,IAAI,CAAC,EAAE,CAAC;YACxC,MAAM,OAAO,GAAG,KAAK,CAAC,SAAS,CAAC,GAAG,CAAC,QAAQ,CAAC,IAAI,CAAE,CAAA;YACnD,OAAO,CAAC,YAAY,CAAC,QAAQ,CAAC,CAAA;QAC/B,CAAC;aAAM,CAAC;YACP,MAAM,OAAO,GAAG,IAAI,eAAe,CAAC,IAAI,CAAC,CAAA;YACzC,OAAO,CAAC,YAAY,CAAC,QAAQ,CAAC,CAAA;YAC9B,KAAK,CAAC,SAAS,CAAC,GAAG,CAAC,QAAQ,CAAC,IAAI,EAAE,OAAO,CAAC,CAAA;QAC5C,CAAC;IACF,CAAC;IAED,kEAAkE;IAClE,2FAA2F;IAC3F,KAAK,CAAC,YAAY,CAAC,OAAe,EAAE,MAAoB;QACvD,MAAM,KAAK,GAAG,IAAI,CAAC,MAAM,CAAC,OAAO,CAAC,CAAA;QAClC,IAAI,CAAC,IAAI,CAAC,OAAO,EAAE,CAAC;YACnB,MAAM,IAAI,KAAK,CAAC,6CAA6C,CAAC,CAAA;QAC/D,CAAC;QACD,KAAK,CAAC,MAAM,GAAG,WAAW,CAAA;QAC1B,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC,KAAK,CAAC,MAAM,CAAC,CAAA;QACzC,IAAI,CAAC,GAAG,CAAC,SAAS,CAAC,IAAI,EAAE,iBAAiB,EAAE;YAC3C,KAAK,EAAE,OAAO;YACd,IAAI,EAAE,KAAK,CAAC,IAAI;SAChB,CAAC,CAAA;QAEF,MAAM,IAAI,CAAC,YAAY,CAAC,GAAG,CAAC,KAAK,IAAI,EAAE;YACtC,IAAI,CAAC,CAAC,cAAc,IAAI,MAAM,CAAC,EAAE,CAAC;gBACjC,KAAK,CAAC,MAAM,GAAG,OAAO,CAAA;gBACtB,OAAO,KAAK,CAAA;YACb,CAAC;YACD,MAAM,mBAAmB,GAAG,WAAW,CAAC,GAAG,EAAE;gBAC5C,MAAM,QAAQ,GAAG,KAAK,CAAC,IAAI,CAAC,KAAK,CAAC,SAAS,EAAE,MAAM,EAAE,IAAI,EAAE,CAAC;qBAC1D,GAAG,CAAC,CAAC,OAAO,EAAE,EAAE,CAAC,OAAO,CAAC,SAAS,EAAE,CAAC;qBACrC,MAAM,CACN,CAAC,GAAG,EAAE,MAAM,EAAE,EAAE;oBACf,GAAG,CAAC,WAAW,IAAI,MAAM,EAAE,WAAW,IAAI,CAAC,CAAA;oBAC3C,GAAG,CAAC,UAAU,IAAI,MAAM,EAAE,UAAU,IAAI,CAAC,CAAA;oBACzC,GAAG,CAAC,KAAK,IAAI,MAAM,EAAE,KAAK,IAAI,CAAC,CAAA;oBAC/B,OAAO,GAAG,CAAA;gBACX,CAAC,EACD,EAAE,WAAW,EAAE,CAAC,EAAE,UAAU,EAAE,CAAC,EAAE,KAAK,EAAE,CAAC,EAAE,CAC3C,CAAA;gBACF,IAAI,QAAQ,CAAC,UAAU,EAAE,CAAC;oBACzB,MAAM,OAAO,GAAG,CAAC,QAAQ,CAAC,WAAW,GAAG,QAAQ,CAAC,UAAU,CAAC,GAAG,GAAG,CAAA;oBAClE,MAAM,mBAAmB,GAAG,WAAW,CAAC,QAAQ,CAAC,UAAU,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,CAAA;oBAC9E,MAAM,oBAAoB,GAAG,WAAW,CAAC,QAAQ,CAAC,WAAW,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,CAAA;oBAChF,IAAI,CAAC,GAAG,CAAC,SAAS,CAAC,IAAI,EAAE,kBAAkB,oBAAoB,CAAC,QAAQ,CAAC,KAAK,CAAC,IAAI,OAAO,CAAC,OAAO,CAAC,CAAC,CAAC,OAAO,oBAAoB,OAAO,mBAAmB,EAAE,EAAE;wBAC7J,KAAK,EAAE,OAAO;qBACd,CAAC,CAAA;gBACH,CAAC;YACF,CAAC,EAAE,KAAK,CAAC,CAAA;YACT,IAAI,CAAC;gBACJ,MAAM,SAAS,GAAG,MAAM,MAAM,CAAC,YAAY,CAC1C,EAAE,MAAM,EAAE,KAAK,EAAE,GAAG,EAAE,IAAI,CAAC,GAAG,EAAE,EAChC,CAAC,QAAQ,EAAE,EAAE;oBACZ,IAAI,CAAC,kBAAkB,CAAC,KAAK,CAAC,EAAE,EAAE,QAAQ,CAAC,CAAA;gBAC5C,CAAC,EACD,iBAAiB,CAAC,CAAC,MAAM,EAAE,IAAI,CAAC,iBAAiB,CAAC,MAAM,CAAC,CAAC,CAC1D,CAAA;gBACD,KAAK,CAAC,SAAS,GAAG,SAAS,CAAA;gBAC3B,KAAK,CAAC,IAAI,GAAG,SAAS,CAAA;gBACtB,KAAK,CAAC,MAAM,GAAG,OAAO,CAAA;gBACtB,IAAI,CAAC,GAAG,CAAC,SAAS,CAAC,IAAI,EAAE,aAAa,EAAE;oBACvC,KAAK,EAAE,OAAO;oBACd,IAAI,EAAE,KAAK,CAAC,IAAI;iBAChB,CAAC,CAAA;YACH,CAAC;YAAC,OAAO,KAAK,EAAE,CAAC;gBAChB,IAAI,CAAC,GAAG,CAAC,SAAS,CAAC,KAAK,EAAE,uBAAuB,EAAE;oBAClD,KAAK,EAAE,OAAO;oBACd,KAAK,EAAE,KAAK;iBACZ,CAAC,CAAA;gBACF,KAAK,CAAC,MAAM,GAAG,OAAO,CAAA;YACvB,CAAC;oBAAS,CAAC;gBACV,aAAa,CAAC,mBAAmB,CAAC,CAAA;YACnC,CAAC;YACD,OAAO,KAAK,CAAA;QACb,CAAC,CAAC,CAAA;IACH,CAAC;IAED,SAAS;QACR,MAAM,WAAW,GAAG,CAAC,GAAY,EAAE,EAAE,CAAC,UAAU,CAAC,GAAG,EAAE,OAAO,CAAC,CAAC,CAAC,IAAI,GAAG,CAAC,CAAA;QACxE,MAAM,eAAe,GAAG,MAAM,CAAC,WAAW,CACzC,MAAM,CAAC,OAAO,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,OAAO,EAAE,KAAK,CAAC,EAAE,EAAE;YACpD,IAAI,SAAS,GAAQ,SAAS,CAAA;YAC9B,IAAI,KAAK,CAAC,SAAS,EAAE,CAAC;gBACrB,SAAS,GAAG,CAAC,GAAG,KAAK,CAAC,SAAS,CAAC,CAAC,MAAM,CACtC,CAAC,GAAG,EAAE,CAAC,GAAG,EAAE,QAAQ,CAAC,EAAE,EAAE;oBACxB,MAAM,MAAM,GAAG,QAAQ,CAAC,SAAS,EAAE,CAAA;oBACnC,MAAM,WAAW,GAChB,QAAQ,CAAC,cAAc,CAAC,QAAQ,CAAC,cAAc,CAAC,MAAM,GAAG,CAAC,CAAC,CAAA;oBAC5D,MAAM,UAAU,GAAG,WAAW,EAAE,UAAU,IAAI,CAAC,CAAA;oBAC/C,MAAM,WAAW,GAAG,WAAW,EAAE,WAAW,IAAI,CAAC,CAAA;oBACjD,MAAM,UAAU,GAAG,MAAM,EAAE,UAAU,IAAI,CAAC,CAAA;oBAC1C,MAAM,YAAY,GAAG,kBAAkB,CAAC,UAAU,GAAG,IAAI,CAAC,CAAA;oBAC1D,MAAM,mBAAmB,GAAG,WAAW,CAAC,UAAU,CAAC,CAAA;oBACnD,MAAM,oBAAoB,GAAG,WAAW,CAAC,WAAW,CAAC,CAAA;oBACrD,GAAG,CAAC,IAAI,CAAC;wBACR,IAAI,EAAE,GAAG;wBACT,WAAW;wBACX,oBAAoB;wBACpB,UAAU;wBACV,mBAAmB;wBACnB,OAAO,EAAE,WAAW,CAAC,MAAM,EAAE,OAAO,CAAC;wBACrC,KAAK,EAAE,WAAW,CAAC,MAAM,EAAE,KAAK,CAAC;wBACjC,UAAU,EAAE,WAAW,CAAC,UAAU,CAAC;wBACnC,YAAY;qBACZ,CAAC,CAAA;oBACF,OAAO,GAAG,CAAA;gBACX,CAAC,EACD,EAAE,CACF,CAAA;YACF,CAAC;YACD,OAAO;gBACN,OAAO;gBACP;oBACC,MAAM,EAAE,KAAK,CAAC,MAAM;oBACpB,MAAM,EAAE,KAAK,CAAC,MAAM;oBACpB,YAAY,EAAE,KAAK,CAAC,YAAY;oBAChC,YAAY,EAAE,KAAK,CAAC,YAAY;oBAChC,MAAM,EAAE,KAAK,CAAC,MAAM;oBACpB,SAAS;iBACT;aACD,CAAA;QACF,CAAC,CAAC,CACF,CAAA;QACD,OAAO,eAAe,CAAA;IACvB,CAAC;CACD;AAgBD,MAAM,eAAe;IACpB,cAAc,GAAoB,EAAE,CAAA;IAC5B,UAAU,CAAQ;IAE1B,YAAY,aAAqB,IAAI;QACpC,IAAI,CAAC,UAAU,GAAG,UAAU,CAAA;IAC7B,CAAC;IAED,YAAY,CAAC,EAAE,WAAW,EAAE,UAAU,EAAwB;QAC7D,MAAM,SAAS,GAAG,IAAI,CAAC,GAAG,EAAE,CAAA;QAC5B,IAAI,CAAC,cAAc,CAAC,IAAI,CAAC,EAAE,WAAW,EAAE,UAAU,EAAE,SAAS,EAAE,CAAC,CAAA;QAChE,IAAI,CAAC,OAAO,EAAE,CAAA;IACf,CAAC;IAEO,OAAO;QACd,MAAM,UAAU,GAAG,IAAI,CAAC,GAAG,EAAE,GAAG,IAAI,CAAC,UAAU,CAAA;QAC/C,IAAI,CAAC,cAAc,GAAG,IAAI,CAAC,cAAc,CAAC,MAAM,CAC/C,CAAC,IAAI,EAAE,EAAE,CAAC,IAAI,CAAC,SAAS,IAAI,UAAU,CACtC,CAAA;IACF,CAAC;IAED,SAAS;QACR,IAAI,IAAI,CAAC,cAAc,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;YACpC,OAAO,IAAI,CAAA,CAAC,6CAA6C;QAC1D,CAAC;QAED,MAAM,WAAW,GAAG,IAAI,CAAC,cAAc,CAAC,IAAI,CAAC,cAAc,CAAC,MAAM,GAAG,CAAC,CAAC,CAAA;QACvE,MAAM,aAAa,GAAG,IAAI,CAAC,cAAc,CAAC,CAAC,CAAC,CAAA,CAAC,sCAAsC;QAEnF,MAAM,WAAW,GAAG,WAAW,CAAC,WAAW,GAAG,aAAa,CAAC,WAAW,CAAA;QACvE,MAAM,WAAW,GAAG,WAAW,CAAC,SAAS,GAAG,aAAa,CAAC,SAAS,CAAA,CAAC,kBAAkB;QAEtF,MAAM,KAAK,GAAG,WAAW,GAAG,CAAC,WAAW,GAAG,IAAI,CAAC,CAAA,CAAC,mBAAmB;QACpE,MAAM,cAAc,GAAG,WAAW,CAAC,UAAU,GAAG,WAAW,CAAC,WAAW,CAAA;QACvE,MAAM,GAAG,GAAG,KAAK,GAAG,CAAC,CAAC,CAAC,CAAC,cAAc,GAAG,KAAK,CAAC,CAAC,CAAC,CAAC,CAAA;QAElD,OAAO;YACN,KAAK;YACL,UAAU,EAAE,GAAG;YACf,OAAO,EAAE,WAAW,CAAC,WAAW,GAAG,WAAW,CAAC,UAAU;YACzD,WAAW,EAAE,WAAW,CAAC,WAAW;YACpC,UAAU,EAAE,WAAW,CAAC,UAAU;SAClC,CAAA;IACF,CAAC;CACD"}

package/dist/types/completions.d.ts

@@ -0,0 +1,57 @@
+import type { SomeJSONSchema } from 'ajv/dist/types/json-schema';
+import { Image } from './index.js';
+export type CompletionFinishReason = 'maxTokens' | 'toolCalls' | 'eogToken' | 'stopTrigger' | 'abort' | 'cancel' | 'timeout';
+export interface AssistantToolCall {
+    id: string;
+    name: string;
+    parameters?: Record<string, any>;
+}
+export type ChatMessage = UserMessage | SystemMessage | AssistantMessage | ToolCallResultMessage;
+export interface MessageTextContentPart {
+    type: 'text';
+    text: string;
+}
+export interface MessageImageContentPart {
+    type: 'image';
+    image: Image;
+}
+export type MessageContentPart = MessageTextContentPart | MessageImageContentPart;
+export interface UserMessage {
+    role: 'user';
+    content: string | MessageContentPart[];
+}
+export interface SystemMessage {
+    role: 'system';
+    content: string | MessageContentPart[];
+}
+export interface AssistantMessage {
+    role: 'assistant';
+    content: string;
+    toolCalls?: AssistantToolCall[];
+}
+export interface ToolCallResultMessage {
+    role: 'tool';
+    content: string | MessageContentPart[];
+    callId: string;
+}
+export type ToolDefinitionParams<TParams> = SomeJSONSchema;
+export interface ToolDefinition<TParams extends Record<string, any> = any> {
+    description?: string;
+    parameters?: ToolDefinitionParams<TParams>;
+    handler?: (params: TParams) => Promise<string>;
+}
+export interface TextCompletionParams {
+    temperature?: number;
+    maxTokens?: number;
+    seed?: number;
+    stop?: string[];
+    repeatPenalty?: number;
+    repeatPenaltyNum?: number;
+    frequencyPenalty?: number;
+    presencePenalty?: number;
+    grammar?: string;
+    topP?: number;
+    minP?: number;
+    topK?: number;
+    tokenBias?: Record<string, number>;
+}

package/dist/types/completions.js

@@ -0,0 +1,2 @@
+export {};
+//# sourceMappingURL=completions.js.map

package/dist/types/completions.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"completions.js","sourceRoot":"","sources":["../../src/types/completions.ts"],"names":[],"mappings":""}

package/dist/types/index.d.ts

@@ -0,0 +1,326 @@
+import type { SomeJSONSchema } from 'ajv/dist/types/json-schema';
+import type { Sharp } from 'sharp';
+import type { BuiltInEngineName } from '../engines/index.js';
+import type { Logger } from '../lib/logger.js';
+import type { ModelPool } from '../pool.js';
+import type { ModelStore } from '../store.js';
+import { AssistantMessage, ChatMessage, CompletionFinishReason, TextCompletionParams, ToolDefinition } from '../types/completions.js';
+import type { ContextShiftStrategy } from '../engines/node-llama-cpp/types.js';
+import type { StableDiffusionWeightType, StableDiffusionSamplingMethod, StableDiffusionSchedule } from '../engines/stable-diffusion-cpp/types.js';
+import type { TransformersJsModelClass, TransformersJsTokenizerClass, TransformersJsProcessorClass, TransformersJsDataType } from '../engines/transformers-js/types.js';
+export * from '../types/completions.js';
+export type ModelTaskType = 'text-completion' | 'embedding' | 'image-to-text' | 'image-to-image' | 'text-to-image' | 'speech-to-text';
+export interface ModelOptionsBase {
+    engine: BuiltInEngineName | (string & {});
+    task: ModelTaskType | (string & {});
+    prepare?: 'blocking' | 'async' | 'on-demand';
+    minInstances?: number;
+    maxInstances?: number;
+    location?: string;
+}
+export interface BuiltInModelOptionsBase extends ModelOptionsBase {
+    engine: BuiltInEngineName;
+    task: ModelTaskType;
+    url?: string;
+    location?: string;
+}
+export interface ModelConfigBase extends ModelOptionsBase {
+    id: string;
+    minInstances: number;
+    maxInstances: number;
+    modelsCachePath: string;
+}
+export interface ModelConfig extends ModelConfigBase {
+    url?: string;
+    location?: string;
+    task: ModelTaskType | (string & {});
+    engine: BuiltInEngineName | (string & {});
+    ttl?: number;
+    prefix?: string;
+    initialMessages?: ChatMessage[];
+    device?: {
+        gpu?: boolean | 'auto' | (string & {});
+    };
+}
+export interface CompletionChunk {
+    tokens: number[];
+    text: string;
+}
+export interface ProcessingOptions {
+    timeout?: number;
+    signal?: AbortSignal;
+}
+export interface Image {
+    handle: Sharp;
+    width: number;
+    height: number;
+    channels: 1 | 2 | 3 | 4;
+}
+export interface CompletionProcessingOptions extends ProcessingOptions {
+    onChunk?: (chunk: CompletionChunk) => void;
+}
+export interface SpeechToTextProcessingOptions extends ProcessingOptions {
+    onChunk?: (chunk: {
+        text: string;
+    }) => void;
+}
+export interface EngineContext<TModelConfig = ModelConfig, TModelMeta = unknown> {
+    config: TModelConfig;
+    meta?: TModelMeta;
+    log: Logger;
+}
+export interface TextCompletionRequestBase extends TextCompletionParams {
+    model: string;
+    stream?: boolean;
+}
+export interface TextCompletionRequest extends TextCompletionRequestBase {
+    prompt?: string;
+}
+export interface ChatCompletionRequest extends TextCompletionRequestBase {
+    messages: ChatMessage[];
+    grammar?: string;
+    tools?: Record<string, ToolDefinition>;
+}
+export interface TextEmbeddingInput {
+    type: 'text';
+    content: string;
+}
+export interface ImageEmbeddingInput {
+    type: 'image';
+    content: Image;
+}
+export type EmbeddingInput = TextEmbeddingInput | ImageEmbeddingInput | string;
+export interface EmbeddingRequest {
+    model: string;
+    input: EmbeddingInput | EmbeddingInput[];
+    dimensions?: number;
+    pooling?: 'cls' | 'mean';
+}
+export interface ImageToTextRequest {
+    model: string;
+    image: Image;
+    prompt?: string;
+    maxTokens?: number;
+}
+export interface StableDiffusionRequest {
+    negativePrompt?: string;
+    guidance?: number;
+    styleRatio?: number;
+    strength?: number;
+    sampleSteps?: number;
+    batchCount?: number;
+    samplingMethod?: StableDiffusionSamplingMethod;
+    cfgScale?: number;
+    controlStrength?: number;
+}
+export interface TextToImageRequest extends StableDiffusionRequest {
+    model: string;
+    prompt: string;
+    width?: number;
+    height?: number;
+    seed?: number;
+}
+export interface ImageToImageRequest extends StableDiffusionRequest {
+    model: string;
+    image: Image;
+    prompt: string;
+    width?: number;
+    height?: number;
+    seed?: number;
+}
+export interface SpeechToTextRequest {
+    model: string;
+    url?: string;
+    file?: string;
+    language?: string;
+    prompt?: string;
+    maxTokens?: number;
+}
+export interface ModelRequestMeta {
+    sequence: number;
+    abortController: AbortController;
+}
+export type IncomingRequest = TextCompletionRequest | ChatCompletionRequest | EmbeddingRequest | ImageToTextRequest | SpeechToTextRequest;
+export type ModelInstanceRequest = ModelRequestMeta & IncomingRequest;
+export interface EngineTextCompletionArgs<TModelConfig = unknown, TModelMeta = unknown> extends EngineContext<TModelConfig, TModelMeta> {
+    onChunk?: (chunk: CompletionChunk) => void;
+    resetContext?: boolean;
+    request: TextCompletionRequest;
+}
+export interface EngineChatCompletionArgs<TModelConfig = unknown, TModelMeta = unknown> extends EngineContext<TModelConfig, TModelMeta> {
+    onChunk?: (chunk: CompletionChunk) => void;
+    resetContext?: boolean;
+    request: ChatCompletionRequest;
+}
+export interface EngineEmbeddingArgs<TModelConfig = unknown, TModelMeta = unknown> extends EngineContext<TModelConfig, TModelMeta> {
+    request: EmbeddingRequest;
+}
+export interface EngineImageToTextArgs<TModelConfig = unknown, TModelMeta = unknown> extends EngineContext<TModelConfig, TModelMeta> {
+    request: ImageToTextRequest;
+}
+export interface EngineTextToImageArgs<TModelConfig = unknown, TModelMeta = unknown> extends EngineContext<TModelConfig, TModelMeta> {
+    request: TextToImageRequest;
+}
+export interface EngineImageToImageArgs<TModelConfig = unknown, TModelMeta = unknown> extends EngineContext<TModelConfig, TModelMeta> {
+    request: ImageToImageRequest;
+}
+export interface EngineSpeechToTextArgs<TModelConfig = unknown, TModelMeta = unknown> extends EngineContext<TModelConfig, TModelMeta> {
+    request: SpeechToTextRequest;
+    onChunk?: (chunk: {
+        text: string;
+    }) => void;
+}
+export interface FileDownloadProgress {
+    file: string;
+    loadedBytes: number;
+    totalBytes: number;
+}
+export interface EngineStartContext {
+    pool: ModelPool;
+    store: ModelStore;
+}
+export interface ModelEngine<TInstance = unknown, TModelConfig extends ModelConfig = ModelConfig, TModelMeta = unknown> {
+    autoGpu?: boolean;
+    start?: (ctx: EngineStartContext) => Promise<void>;
+    prepareModel: (ctx: EngineContext<TModelConfig, TModelMeta>, onProgress?: (progress: FileDownloadProgress) => void, signal?: AbortSignal) => Promise<TModelMeta>;
+    createInstance: (ctx: EngineContext<TModelConfig, TModelMeta>, signal?: AbortSignal) => Promise<TInstance>;
+    disposeInstance: (instance: TInstance) => Promise<void>;
+    processChatCompletionTask?: (args: EngineChatCompletionArgs<TModelConfig, TModelMeta>, instance: TInstance, signal?: AbortSignal) => Promise<EngineChatCompletionResult>;
+    processTextCompletionTask?: (args: EngineTextCompletionArgs<TModelConfig, TModelMeta>, instance: TInstance, signal?: AbortSignal) => Promise<EngineTextCompletionResult>;
+    processEmbeddingTask?: (args: EngineEmbeddingArgs<TModelConfig, TModelMeta>, instance: TInstance, signal?: AbortSignal) => Promise<EngineEmbeddingResult>;
+    processImageToTextTask?: (args: EngineImageToTextArgs<TModelConfig, TModelMeta>, instance: TInstance, signal?: AbortSignal) => Promise<EngineImageToTextResult>;
+    processSpeechToTextTask?: (args: EngineSpeechToTextArgs<TModelConfig, TModelMeta>, instance: TInstance, signal?: AbortSignal) => Promise<EngineSpeechToTextResult>;
+    processTextToImageTask?: (args: EngineTextToImageArgs<TModelConfig, TModelMeta>, instance: TInstance, signal?: AbortSignal) => Promise<EngineTextToImageResult>;
+    processImageToImageTask?: (args: EngineImageToImageArgs<TModelConfig, TModelMeta>, instance: TInstance, signal?: AbortSignal) => Promise<EngineImageToImageResult>;
+}
+interface EmbeddingModelOptions {
+    task: 'embedding';
+}
+export type TextCompletionGrammar = string | SomeJSONSchema;
+interface TextCompletionModelOptions {
+    task: 'text-completion';
+    contextSize?: number;
+    grammars?: Record<string, TextCompletionGrammar>;
+    completionDefaults?: TextCompletionParams;
+    initialMessages?: ChatMessage[];
+    prefix?: string;
+    batchSize?: number;
+}
+interface LlamaCppModelOptionsBase extends BuiltInModelOptionsBase {
+    engine: 'node-llama-cpp';
+    task: 'text-completion' | 'embedding';
+    sha256?: string;
+    batchSize?: number;
+    contextShiftStrategy?: ContextShiftStrategy;
+    tools?: {
+        definitions: Record<string, ToolDefinition>;
+        includeParamsDocumentation?: boolean;
+        parallelism?: number;
+    };
+    device?: {
+        gpu?: boolean | 'auto' | (string & {});
+        gpuLayers?: number;
+        cpuThreads?: number;
+        memLock?: boolean;
+    };
+}
+interface LlamaCppEmbeddingModelOptions extends LlamaCppModelOptionsBase, EmbeddingModelOptions {
+    task: 'embedding';
+}
+export interface LlamaCppTextCompletionModelOptions extends LlamaCppModelOptionsBase, TextCompletionModelOptions {
+    task: 'text-completion';
+}
+interface GPT4AllModelOptions extends BuiltInModelOptionsBase {
+    engine: 'gpt4all';
+    task: 'text-completion' | 'embedding';
+    md5?: string;
+    device?: {
+        gpu?: boolean | 'auto' | (string & {});
+        gpuLayers?: number;
+        cpuThreads?: number;
+    };
+}
+type GPT4AllTextCompletionModelOptions = TextCompletionModelOptions & GPT4AllModelOptions;
+type GPT4AllEmbeddingModelOptions = GPT4AllModelOptions & EmbeddingModelOptions;
+export interface TransformersJsModel {
+    processor?: {
+        url?: string;
+        file?: string;
+    };
+    processorClass?: TransformersJsProcessorClass;
+    tokenizerClass?: TransformersJsTokenizerClass;
+    modelClass?: TransformersJsModelClass;
+    dtype?: Record<string, TransformersJsDataType> | TransformersJsDataType;
+}
+interface TransformersJsModelOptions extends BuiltInModelOptionsBase {
+    engine: 'transformers-js';
+    task: 'image-to-text' | 'speech-to-text' | 'text-completion' | 'embedding';
+    textModel?: TransformersJsModel;
+    visionModel?: TransformersJsModel;
+    speechModel?: TransformersJsModel;
+    device?: {
+        gpu?: boolean | 'auto' | (string & {});
+    };
+}
+export interface ModelFileSource {
+    url?: string;
+    file?: string;
+    sha256?: string;
+}
+interface StableDiffusionModelOptions extends BuiltInModelOptionsBase {
+    engine: 'stable-diffusion-cpp';
+    task: 'image-to-text' | 'text-to-image' | 'image-to-image';
+    sha256?: string;
+    url?: string;
+    diffusionModel?: boolean;
+    vae?: ModelFileSource;
+    clipL?: ModelFileSource;
+    clipG?: ModelFileSource;
+    t5xxl?: ModelFileSource;
+    taesd?: ModelFileSource;
+    controlNet?: ModelFileSource;
+    samplingMethod?: StableDiffusionSamplingMethod;
+    weightType?: StableDiffusionWeightType;
+    schedule?: StableDiffusionSchedule;
+    loras?: ModelFileSource[];
+}
+export interface CustomEngineModelOptions extends ModelOptionsBase {
+}
+export type BuiltInModelOptions = LlamaCppTextCompletionModelOptions | LlamaCppEmbeddingModelOptions | GPT4AllTextCompletionModelOptions | GPT4AllEmbeddingModelOptions | TransformersJsModelOptions | StableDiffusionModelOptions;
+export type ModelOptions = BuiltInModelOptions | CustomEngineModelOptions;
+export interface EngineEmbeddingResult {
+    embeddings: Float32Array[];
+    inputTokens: number;
+}
+export interface ChatCompletionResult extends EngineChatCompletionResult {
+    id: string;
+    model: string;
+}
+export interface EngineChatCompletionResult {
+    message: AssistantMessage;
+    finishReason: CompletionFinishReason;
+    promptTokens: number;
+    completionTokens: number;
+    contextTokens: number;
+}
+export interface EngineTextCompletionResult {
+    text: string;
+    finishReason?: CompletionFinishReason;
+    promptTokens: number;
+    completionTokens: number;
+    contextTokens: number;
+}
+export interface EngineImageToTextResult {
+    text: string;
+}
+export interface EngineTextToImageResult {
+    images: Image[];
+    seed: number;
+}
+export interface EngineImageToImageResult {
+    images: Image[];
+    seed: number;
+}
+export interface EngineSpeechToTextResult {
+    text: string;
+}

package/dist/types/index.js

@@ -0,0 +1,2 @@
+export * from '../types/completions.js';
+//# sourceMappingURL=index.js.map

package/dist/types/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../../src/types/index.ts"],"names":[],"mappings":"AAyBA,cAAc,+BAA+B,CAAA"}

package/docs/engines.md

@@ -0,0 +1,28 @@
+
+### Engines
+
+Currently the built-in inference engines are `node-llama-cpp`, `gpt4all` and `transformers-js` (highly experimental). Install the corresponding peer dependency before using an engine.
+
+#### node-llama-cpp
+
+Can be used for `text-completion` and `embedding` tasks. See the [node-llama-cpp docs](https://node-llama-cpp.withcat.ai/) for more information.
+
+Find available GGUF models on [huggingface.co](https://huggingface.co/models?library=gguf).
+
+#### gpt4all
+
+Can be used for `text-completion` and `embedding` tasks. You can find parameter docs [here](https://github.com/nomic-ai/gpt4all/blob/c73f0e5c8c25ede56e3eeb28ff9dd37f09212994/gpt4all-bindings/typescript/src/gpt4all.d.ts#L615).
+
+You can find available models [here](https://github.com/nomic-ai/gpt4all/blob/c73f0e5c8c25ede56e3eeb28ff9dd37f09212994/gpt4all-chat/metadata/models3.json)
+
+#### transformers-js
+
+Currently supporting `speech-to-text` and `image-to-text` tasks. See [tests](./tests/engines/transformers.test.ts).
+
+#### node-stable-diffusion-cpp
+
+WIP. See [tests](./tests/engines/stable-diffusion.test.ts).
+
+#### Custom Engines
+
+You can also write your own engine implementation. See [./src/engines](./src/engines) for how the built-in engines are implemented and [here](./tests/engines/experiments.test.ts) for examples of how to utilize custom engines to combine models and add multimodality to your chat completion endpoint. (Or to any other consumer of the ModelServer class.) Multiple ModelServers are allowed and can also be nested to create more complex pipelines.