diff --git a/index.d.ts b/index.d.ts index 2769df508..015d8e353 100644 --- a/index.d.ts +++ b/index.d.ts @@ -9,6 +9,11 @@ declare global { }; } } + + /** + * The root directory for user data. + */ + var DATA_ROOT: string; } declare module 'express-session' { diff --git a/package-lock.json b/package-lock.json index f27b2d3e3..9c640abdd 100644 --- a/package-lock.json +++ b/package-lock.json @@ -42,7 +42,7 @@ "rate-limiter-flexible": "^5.0.0", "response-time": "^2.3.2", "sanitize-filename": "^1.6.3", - "sillytavern-transformers": "^2.14.6", + "sillytavern-transformers": "2.14.6", "simple-git": "^3.19.1", "tiktoken": "^1.0.15", "vectra": "^0.2.2", diff --git a/package.json b/package.json index ac113a7e6..6a282736c 100644 --- a/package.json +++ b/package.json @@ -32,7 +32,7 @@ "rate-limiter-flexible": "^5.0.0", "response-time": "^2.3.2", "sanitize-filename": "^1.6.3", - "sillytavern-transformers": "^2.14.6", + "sillytavern-transformers": "2.14.6", "simple-git": "^3.19.1", "tiktoken": "^1.0.15", "vectra": "^0.2.2", diff --git a/src/transformers.mjs b/src/transformers.mjs index ab382f58a..f3c7389a4 100644 --- a/src/transformers.mjs +++ b/src/transformers.mjs @@ -1,6 +1,7 @@ import { pipeline, env, RawImage, Pipeline } from 'sillytavern-transformers'; import { getConfigValue } from './util.js'; import path from 'path'; +import fs from 'fs'; configureTransformers(); @@ -48,7 +49,7 @@ const tasks = { configField: 'extras.textToSpeechModel', quantized: false, }, -} +}; /** * Gets a RawImage object from a base64-encoded image. @@ -85,6 +86,36 @@ function getModelForTask(task) { } } +async function migrateCacheToDataDir() { + const oldCacheDir = path.join(process.cwd(), 'cache'); + const newCacheDir = path.join(global.DATA_ROOT, '_cache'); + + if (!fs.existsSync(newCacheDir)) { + fs.mkdirSync(newCacheDir, { recursive: true }); + } + + if (fs.existsSync(oldCacheDir) && fs.statSync(oldCacheDir).isDirectory()) { + const files = fs.readdirSync(oldCacheDir); + + if (files.length === 0) { + return; + } + + console.log('Migrating model cache files to data directory. Please wait...'); + + for (const file of files) { + try { + const oldPath = path.join(oldCacheDir, file); + const newPath = path.join(newCacheDir, file); + fs.cpSync(oldPath, newPath, { recursive: true, force: true }); + fs.rmSync(oldPath, { recursive: true, force: true }); + } catch (error) { + console.warn('Failed to migrate cache file. The model will be re-downloaded.', error); + } + } + } +} + /** * Gets the transformers.js pipeline for a given task. * @param {import('sillytavern-transformers').PipelineType} task The task to get the pipeline for @@ -92,6 +123,8 @@ function getModelForTask(task) { * @returns {Promise} Pipeline for the task */ async function getPipeline(task, forceModel = '') { + await migrateCacheToDataDir(); + if (tasks[task].pipeline) { if (forceModel === '' || tasks[task].currentModel === forceModel) { return tasks[task].pipeline; @@ -100,11 +133,11 @@ async function getPipeline(task, forceModel = '') { await tasks[task].pipeline.dispose(); } - const cache_dir = path.join(process.cwd(), 'cache'); + const cacheDir = path.join(global.DATA_ROOT, '_cache'); const model = forceModel || getModelForTask(task); const localOnly = getConfigValue('extras.disableAutoDownload', false); console.log('Initializing transformers.js pipeline for task', task, 'with model', model); - const instance = await pipeline(task, model, { cache_dir, quantized: tasks[task].quantized ?? true, local_files_only: localOnly }); + const instance = await pipeline(task, model, { cache_dir: cacheDir, quantized: tasks[task].quantized ?? true, local_files_only: localOnly }); tasks[task].pipeline = instance; tasks[task].currentModel = model; return instance; diff --git a/src/users.js b/src/users.js index dcf3718a9..ab0352294 100644 --- a/src/users.js +++ b/src/users.js @@ -19,12 +19,6 @@ const AVATAR_PREFIX = 'avatar:'; const ENABLE_ACCOUNTS = getConfigValue('enableUserAccounts', false); const ANON_CSRF_SECRET = crypto.randomBytes(64).toString('base64'); -/** - * The root directory for user data. - * @type {string} - */ -let DATA_ROOT = './data'; - /** * Cache for user directories. * @type {Map} @@ -138,7 +132,7 @@ async function migrateUserData() { console.log(); console.log(color.magenta('Preparing to migrate user data...')); - console.log(`All public data will be moved to the ${DATA_ROOT} directory.`); + console.log(`All public data will be moved to the ${global.DATA_ROOT} directory.`); console.log('This process may take a while depending on the amount of data to move.'); console.log(`Backups will be placed in the ${PUBLIC_DIRECTORIES.backups} directory.`); console.log(`The process will start in ${TIMEOUT} seconds. Press Ctrl+C to cancel.`); @@ -352,11 +346,11 @@ function toAvatarKey(handle) { * @returns {Promise} */ async function initUserStorage(dataRoot) { - DATA_ROOT = dataRoot; - console.log('Using data root:', color.green(DATA_ROOT)); + global.DATA_ROOT = dataRoot; + console.log('Using data root:', color.green(global.DATA_ROOT)); console.log(); await storage.init({ - dir: path.join(DATA_ROOT, '_storage'), + dir: path.join(global.DATA_ROOT, '_storage'), ttl: false, // Never expire }); @@ -457,7 +451,7 @@ function getUserDirectories(handle) { const directories = structuredClone(USER_DIRECTORY_TEMPLATE); for (const key in directories) { - directories[key] = path.join(DATA_ROOT, handle, USER_DIRECTORY_TEMPLATE[key]); + directories[key] = path.join(global.DATA_ROOT, handle, USER_DIRECTORY_TEMPLATE[key]); } DIRECTORIES_CACHE.set(handle, directories); return directories;