1
import path from 'node:path';
2
import fs from 'node:fs';
3
import process from 'node:process';
4
import { Buffer } from 'node:buffer';
5

6
import { pipeline, env, RawImage } from 'sillytavern-transformers';
7
import { getConfigValue } from './util.js';
8
import { serverDirectory } from './server-directory.js';
9

10
configureTransformers();
11

12
function configureTransformers() {
13
    // Limit the number of threads to 1 to avoid issues on Android
14
    env.backends.onnx.wasm.numThreads = 1;
15
    // Use WASM from a local folder to avoid CDN connections
16
    env.backends.onnx.wasm.wasmPaths = path.join(serverDirectory, 'node_modules', 'sillytavern-transformers', 'dist') + path.sep;
17
}
18

19
const tasks = {
20
    'text-classification': {
21
        defaultModel: 'Cohee/distilbert-base-uncased-go-emotions-onnx',
22
        pipeline: null,
23
        configField: 'extensions.models.classification',
24
        quantized: true,
25
    },
26
    'image-to-text': {
27
        defaultModel: 'Xenova/vit-gpt2-image-captioning',
28
        pipeline: null,
29
        configField: 'extensions.models.captioning',
30
        quantized: true,
31
    },
32
    'feature-extraction': {
33
        defaultModel: 'Xenova/all-mpnet-base-v2',
34
        pipeline: null,
35
        configField: 'extensions.models.embedding',
36
        quantized: true,
37
    },
38
    'automatic-speech-recognition': {
39
        defaultModel: 'Xenova/whisper-small',
40
        pipeline: null,
41
        configField: 'extensions.models.speechToText',
42
        quantized: true,
43
    },
44
    'text-to-speech': {
45
        defaultModel: 'Xenova/speecht5_tts',
46
        pipeline: null,
47
        configField: 'extensions.models.textToSpeech',
48
        quantized: false,
49
    },
50
};
51

52
/**
53
 * Gets a RawImage object from a base64-encoded image.
54
 * @param {string} image Base64-encoded image
55
 * @returns {Promise<RawImage|null>} Object representing the image
56
 */
57
export async function getRawImage(image) {
58
    try {
59
        const buffer = Buffer.from(image, 'base64');
60
        const byteArray = new Uint8Array(buffer);
61
        const blob = new Blob([byteArray]);
62

63
        const rawImage = await RawImage.fromBlob(blob);
64
        return rawImage;
65
    } catch {
66
        return null;
67
    }
68
}
69

70
/**
71
 * Gets the model to use for a given transformers.js task.
72
 * @param {string} task The task to get the model for
73
 * @returns {string} The model to use for the given task
74
 */
75
function getModelForTask(task) {
76
    const defaultModel = tasks[task].defaultModel;
77

78
    try {
79
        const model = getConfigValue(tasks[task].configField, null);
80
        return model || defaultModel;
81
    } catch (error) {
82
        console.warn('Failed to read config.yaml, using default classification model.');
83
        return defaultModel;
84
    }
85
}
86

87
async function migrateCacheToDataDir() {
88
    const oldCacheDir = path.join(process.cwd(), 'cache');
89
    const newCacheDir = path.join(globalThis.DATA_ROOT, '_cache');
90

91
    if (!fs.existsSync(newCacheDir)) {
92
        fs.mkdirSync(newCacheDir, { recursive: true });
93
    }
94

95
    if (fs.existsSync(oldCacheDir) && fs.statSync(oldCacheDir).isDirectory()) {
96
        const files = fs.readdirSync(oldCacheDir);
97

98
        if (files.length === 0) {
99
            return;
100
        }
101

102
        console.log('Migrating model cache files to data directory. Please wait...');
103

104
        for (const file of files) {
105
            try {
106
                const oldPath = path.join(oldCacheDir, file);
107
                const newPath = path.join(newCacheDir, file);
108
                fs.cpSync(oldPath, newPath, { recursive: true, force: true });
109
                fs.rmSync(oldPath, { recursive: true, force: true });
110
            } catch (error) {
111
                console.warn('Failed to migrate cache file. The model will be re-downloaded.', error);
112
            }
113
        }
114
    }
115
}
116

117
/**
118
 * Gets the transformers.js pipeline for a given task.
119
 * @param {import('sillytavern-transformers').PipelineType} task The task to get the pipeline for
120
 * @param {string} forceModel The model to use for the pipeline, if any
121
 * @returns {Promise<import('sillytavern-transformers').Pipeline>} The transformers.js pipeline
122
 */
123
export async function getPipeline(task, forceModel = '') {
124
    await migrateCacheToDataDir();
125

126
    if (tasks[task].pipeline) {
127
        if (forceModel === '' || tasks[task].currentModel === forceModel) {
128
            return tasks[task].pipeline;
129
        }
130
        console.log('Disposing transformers.js pipeline for for task', task, 'with model', tasks[task].currentModel);
131
        await tasks[task].pipeline.dispose();
132
    }
133

134
    const cacheDir = path.join(globalThis.DATA_ROOT, '_cache');
135
    const model = forceModel || getModelForTask(task);
136
    const localOnly = !getConfigValue('extensions.models.autoDownload', true, 'boolean');
137
    console.log('Initializing transformers.js pipeline for task', task, 'with model', model);
138
    const instance = await pipeline(task, model, { cache_dir: cacheDir, quantized: tasks[task].quantized ?? true, local_files_only: localOnly });
139
    tasks[task].pipeline = instance;
140
    tasks[task].currentModel = model;
141
    // @ts-ignore
142
    return instance;
143
}
144

145
export default {
146
    getRawImage,
147
    getPipeline,
148
};