Merge pull request #2667 from SillyTavern/dataroot-cache Move transformers.js model cache under the data root

fc02898a97897df3e1285989d603ea76f4760a37

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed
5 files changed, +48 -16Ignore whitespace
index.d.ts+5 -0
@@ -9,6 +9,11 @@ declare global {
9 };9 };
10 }10 }
11 }11 }
12
13 /**
14 * The root directory for user data.
15 */
16 var DATA_ROOT: string;
12}17}
1318
14declare module 'express-session' {19declare module 'express-session' {
package-lock.json+1 -1
@@ -42,7 +42,7 @@
42 "rate-limiter-flexible": "^5.0.0",42 "rate-limiter-flexible": "^5.0.0",
43 "response-time": "^2.3.2",43 "response-time": "^2.3.2",
44 "sanitize-filename": "^1.6.3",44 "sanitize-filename": "^1.6.3",
45 "sillytavern-transformers": "^2.14.6",45 "sillytavern-transformers": "2.14.6",
46 "simple-git": "^3.19.1",46 "simple-git": "^3.19.1",
47 "tiktoken": "^1.0.15",47 "tiktoken": "^1.0.15",
48 "vectra": "^0.2.2",48 "vectra": "^0.2.2",
package.json+1 -1
@@ -32,7 +32,7 @@
32 "rate-limiter-flexible": "^5.0.0",32 "rate-limiter-flexible": "^5.0.0",
33 "response-time": "^2.3.2",33 "response-time": "^2.3.2",
34 "sanitize-filename": "^1.6.3",34 "sanitize-filename": "^1.6.3",
35 "sillytavern-transformers": "^2.14.6",35 "sillytavern-transformers": "2.14.6",
36 "simple-git": "^3.19.1",36 "simple-git": "^3.19.1",
37 "tiktoken": "^1.0.15",37 "tiktoken": "^1.0.15",
38 "vectra": "^0.2.2",38 "vectra": "^0.2.2",
src/transformers.mjs+36 -3
@@ -1,6 +1,7 @@
1import { pipeline, env, RawImage, Pipeline } from 'sillytavern-transformers';1import { pipeline, env, RawImage, Pipeline } from 'sillytavern-transformers';
2import { getConfigValue } from './util.js';2import { getConfigValue } from './util.js';
3import path from 'path';3import path from 'path';
4import fs from 'fs';
45
5configureTransformers();6configureTransformers();
67
@@ -48,7 +49,7 @@ const tasks = {
48 configField: 'extras.textToSpeechModel',49 configField: 'extras.textToSpeechModel',
49 quantized: false,50 quantized: false,
50 },51 },
51}52};
5253
53/**54/**
54 * Gets a RawImage object from a base64-encoded image.55 * Gets a RawImage object from a base64-encoded image.
@@ -85,6 +86,36 @@ function getModelForTask(task) {
85 }86 }
86}87}
8788
89async function migrateCacheToDataDir() {
90 const oldCacheDir = path.join(process.cwd(), 'cache');
91 const newCacheDir = path.join(global.DATA_ROOT, '_cache');
92
93 if (!fs.existsSync(newCacheDir)) {
94 fs.mkdirSync(newCacheDir, { recursive: true });
95 }
96
97 if (fs.existsSync(oldCacheDir) && fs.statSync(oldCacheDir).isDirectory()) {
98 const files = fs.readdirSync(oldCacheDir);
99
100 if (files.length === 0) {
101 return;
102 }
103
104 console.log('Migrating model cache files to data directory. Please wait...');
105
106 for (const file of files) {
107 try {
108 const oldPath = path.join(oldCacheDir, file);
109 const newPath = path.join(newCacheDir, file);
110 fs.cpSync(oldPath, newPath, { recursive: true, force: true });
111 fs.rmSync(oldPath, { recursive: true, force: true });
112 } catch (error) {
113 console.warn('Failed to migrate cache file. The model will be re-downloaded.', error);
114 }
115 }
116 }
117}
118
88/**119/**
89 * Gets the transformers.js pipeline for a given task.120 * Gets the transformers.js pipeline for a given task.
90 * @param {import('sillytavern-transformers').PipelineType} task The task to get the pipeline for121 * @param {import('sillytavern-transformers').PipelineType} task The task to get the pipeline for
@@ -92,6 +123,8 @@ function getModelForTask(task) {
92 * @returns {Promise<Pipeline>} Pipeline for the task123 * @returns {Promise<Pipeline>} Pipeline for the task
93 */124 */
94async function getPipeline(task, forceModel = '') {125async function getPipeline(task, forceModel = '') {
126 await migrateCacheToDataDir();
127
95 if (tasks[task].pipeline) {128 if (tasks[task].pipeline) {
96 if (forceModel === '' || tasks[task].currentModel === forceModel) {129 if (forceModel === '' || tasks[task].currentModel === forceModel) {
97 return tasks[task].pipeline;130 return tasks[task].pipeline;
@@ -100,11 +133,11 @@ async function getPipeline(task, forceModel = '') {
100 await tasks[task].pipeline.dispose();133 await tasks[task].pipeline.dispose();
101 }134 }
102135
103 const cache_dir = path.join(process.cwd(), 'cache');136 const cacheDir = path.join(global.DATA_ROOT, '_cache');
104 const model = forceModel || getModelForTask(task);137 const model = forceModel || getModelForTask(task);
105 const localOnly = getConfigValue('extras.disableAutoDownload', false);138 const localOnly = getConfigValue('extras.disableAutoDownload', false);
106 console.log('Initializing transformers.js pipeline for task', task, 'with model', model);139 console.log('Initializing transformers.js pipeline for task', task, 'with model', model);
107 const instance = await pipeline(task, model, { cache_dir, quantized: tasks[task].quantized ?? true, local_files_only: localOnly });140 const instance = await pipeline(task, model, { cache_dir: cacheDir, quantized: tasks[task].quantized ?? true, local_files_only: localOnly });
108 tasks[task].pipeline = instance;141 tasks[task].pipeline = instance;
109 tasks[task].currentModel = model;142 tasks[task].currentModel = model;
110 return instance;143 return instance;
src/users.js+5 -11
@@ -20,12 +20,6 @@ const ENABLE_ACCOUNTS = getConfigValue('enableUserAccounts', false);
20const ANON_CSRF_SECRET = crypto.randomBytes(64).toString('base64');20const ANON_CSRF_SECRET = crypto.randomBytes(64).toString('base64');
2121
22/**22/**
23 * The root directory for user data.
24 * @type {string}
25 */
26let DATA_ROOT = './data';
27
28/**
29 * Cache for user directories.23 * Cache for user directories.
30 * @type {Map<string, UserDirectoryList>}24 * @type {Map<string, UserDirectoryList>}
31 */25 */
@@ -138,7 +132,7 @@ async function migrateUserData() {
138132
139 console.log();133 console.log();
140 console.log(color.magenta('Preparing to migrate user data...'));134 console.log(color.magenta('Preparing to migrate user data...'));
141 console.log(`All public data will be moved to the ${DATA_ROOT} directory.`);135 console.log(`All public data will be moved to the ${global.DATA_ROOT} directory.`);
142 console.log('This process may take a while depending on the amount of data to move.');136 console.log('This process may take a while depending on the amount of data to move.');
143 console.log(`Backups will be placed in the ${PUBLIC_DIRECTORIES.backups} directory.`);137 console.log(`Backups will be placed in the ${PUBLIC_DIRECTORIES.backups} directory.`);
144 console.log(`The process will start in ${TIMEOUT} seconds. Press Ctrl+C to cancel.`);138 console.log(`The process will start in ${TIMEOUT} seconds. Press Ctrl+C to cancel.`);
@@ -352,11 +346,11 @@ function toAvatarKey(handle) {
352 * @returns {Promise<void>}346 * @returns {Promise<void>}
353 */347 */
354async function initUserStorage(dataRoot) {348async function initUserStorage(dataRoot) {
355 DATA_ROOT = dataRoot;349 global.DATA_ROOT = dataRoot;
356 console.log('Using data root:', color.green(DATA_ROOT));350 console.log('Using data root:', color.green(global.DATA_ROOT));
357 console.log();351 console.log();
358 await storage.init({352 await storage.init({
359 dir: path.join(DATA_ROOT, '_storage'),353 dir: path.join(global.DATA_ROOT, '_storage'),
360 ttl: false, // Never expire354 ttl: false, // Never expire
361 });355 });
362356
@@ -457,7 +451,7 @@ function getUserDirectories(handle) {
457451
458 const directories = structuredClone(USER_DIRECTORY_TEMPLATE);452 const directories = structuredClone(USER_DIRECTORY_TEMPLATE);
459 for (const key in directories) {453 for (const key in directories) {
460 directories[key] = path.join(DATA_ROOT, handle, USER_DIRECTORY_TEMPLATE[key]);454 directories[key] = path.join(global.DATA_ROOT, handle, USER_DIRECTORY_TEMPLATE[key]);
461 }455 }
462 DIRECTORIES_CACHE.set(handle, directories);456 DIRECTORIES_CACHE.set(handle, directories);
463 return directories;457 return directories;